code: 
  
  
  '''simple docstring'''
  
  
  
  
  
  import tempfile
  import unittest
  
  import torch
  
  from diffusers import UNetaDConditionModel
  from diffusers.training_utils import EMAModel
  from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device
  
  
  enable_full_determinism()
  
  
  class 		_SCREAMING_SNAKE_CASE    (	unittest.TestCase    ):
       __SCREAMING_SNAKE_CASE    				=						'''hf-internal-testing/tiny-stable-diffusion-pipe'''
       __SCREAMING_SNAKE_CASE    				=						1
       __SCREAMING_SNAKE_CASE    				=						7_7
       __SCREAMING_SNAKE_CASE    				=						3_2
       __SCREAMING_SNAKE_CASE    				=						4
       __SCREAMING_SNAKE_CASE    				=						__SCREAMING_SNAKE_CASE    				=						6_4
       __SCREAMING_SNAKE_CASE    				=						torch.manual_seed(0    )
  
  
       def 					__snake_case(     self       ,				A_=0.9_9_9_9						):
         _UpperCAmelCase		:       Optional[int]       =    UNetaDConditionModel.from_pretrained(self.model_id       ,				subfolder="""unet"""						)
         _UpperCAmelCase		:       List[str]       =    unet.to(A_						)
         _UpperCAmelCase		:       Dict       =    EMAModel(unet.parameters()       ,				decay=A_       ,				model_cls=A_       ,				model_config=unet.config						)
         return unet, ema_unet
  
  
       def 					__snake_case(     self						):
         _UpperCAmelCase		:       Union[str, Any]       =    torch.randn(
             self.batch_size       ,				self.num_in_channels       ,				self.latent_height       ,				self.latent_width       ,				generator=self.generator						).to(A_						)
         _UpperCAmelCase		:       Any       =    torch.randint(0       ,				10_00       ,				size=(self.batch_size,)       ,				generator=self.generator						).to(A_						)
         _UpperCAmelCase		:       Union[str, Any]       =    torch.randn(
             self.batch_size       ,				self.prompt_length       ,				self.text_encoder_hidden_dim       ,				generator=self.generator						).to(A_						)
         return noisy_latents, timesteps, encoder_hidden_states
  
  
       def 					__snake_case(     self       ,				A_						):
         _UpperCAmelCase		:       List[Any]       =    {}
         for k, param in unet.state_dict().items():
           _UpperCAmelCase		:       Tuple       =    torch.randn_like(A_						) + (param * torch.randn_like(A_						))
           updated_state_dict.update({k: updated_param}						)
         unet.load_state_dict(A_						)
         return unet
  
  
       def 					__snake_case(     self						):
         _UpperCAmelCase		:       List[str]       =    self.get_models()
         # Take the first (hypothetical) EMA step.
         ema_unet.step(unet.parameters()						)
         assert ema_unet.optimization_step == 1
  
         # Take two more.
         for _ in range(2						):
           ema_unet.step(unet.parameters()						)
         assert ema_unet.optimization_step == 3
  
  
       def 					__snake_case(     self						):
         _UpperCAmelCase		:       Tuple       =    self.get_models()
         # Since the `unet` is not being updated (i.e., backprop'd)
         # there won't be any difference between the `params` of `unet`
         # and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
         ema_unet.step(unet.parameters()						)
         _UpperCAmelCase		:       Any       =    list(unet.parameters()						)
         for s_param, param in zip(ema_unet.shadow_params       ,				A_						):
           assert torch.allclose(A_       ,				A_						)
  
         # The above holds true even if we call `ema.step()` multiple times since
         # `unet` params are still not being updated.
         for _ in range(4						):
           ema_unet.step(unet.parameters()						)
         for s_param, param in zip(ema_unet.shadow_params       ,				A_						):
           assert torch.allclose(A_       ,				A_						)
  
  
       def 					__snake_case(     self						):
         _UpperCAmelCase		:       Tuple       =    self.get_models()
         # Here we simulate the parameter updates for `unet`. Since there might
         # be some parameters which are initialized to zero we take extra care to
         # initialize their values to something non-zero before the multiplication.
         _UpperCAmelCase		:       Optional[Any]       =    self.simulate_backprop(A_						)
  
         # Take the EMA step.
         ema_unet.step(unet_pseudo_updated_step_one.parameters()						)
  
         # Now the EMA'd parameters won't be equal to the original model parameters.
         _UpperCAmelCase		:       str       =    list(unet_pseudo_updated_step_one.parameters()						)
         for s_param, param in zip(ema_unet.shadow_params       ,				A_						):
           assert ~torch.allclose(A_       ,				A_						)
  
         # Ensure this is the case when we take multiple EMA steps.
         for _ in range(4						):
           ema_unet.step(unet.parameters()						)
         for s_param, param in zip(ema_unet.shadow_params       ,				A_						):
           assert ~torch.allclose(A_       ,				A_						)
  
  
       def 					__snake_case(     self						):
         # If we call EMA step after a backpropagation consecutively for two times,
         # the shadow params from those two steps should be different.
         _UpperCAmelCase		:       Dict       =    self.get_models()
  
         # First backprop + EMA
         _UpperCAmelCase		:       Optional[int]       =    self.simulate_backprop(A_						)
         ema_unet.step(unet_step_one.parameters()						)
         _UpperCAmelCase		:       Tuple       =    ema_unet.shadow_params
  
         # Second backprop + EMA
         _UpperCAmelCase		:       Union[str, Any]       =    self.simulate_backprop(A_						)
         ema_unet.step(unet_step_two.parameters()						)
         _UpperCAmelCase		:       List[str]       =    ema_unet.shadow_params
  
         for step_one, step_two in zip(A_       ,				A_						):
           assert ~torch.allclose(A_       ,				A_						)
  
  
       def 					__snake_case(     self						):
         # If there's no decay even if there are backprops, EMA steps
         # won't take any effect i.e., the shadow params would remain the
         # same.
         _UpperCAmelCase		:       str       =    self.get_models(decay=0.0						)
         _UpperCAmelCase		:       Union[str, Any]       =    self.simulate_backprop(A_						)
         ema_unet.step(unet_step_one.parameters()						)
         _UpperCAmelCase		:       str       =    ema_unet.shadow_params
  
         _UpperCAmelCase		:       int       =    self.simulate_backprop(A_						)
         ema_unet.step(unet_step_two.parameters()						)
         _UpperCAmelCase		:       Dict       =    ema_unet.shadow_params
  
         for step_one, step_two in zip(A_       ,				A_						):
           assert torch.allclose(A_       ,				A_						)
  
  
  
  
       @skip_mps
       def 					__snake_case(     self						):
         _UpperCAmelCase		:       Tuple       =    self.get_models()
         _UpperCAmelCase		:       Any       =    self.get_dummy_inputs()
  
         with tempfile.TemporaryDirectory() as tmpdir:
           ema_unet.save_pretrained(A_						)
           _UpperCAmelCase		:       str       =    UNetaDConditionModel.from_pretrained(A_       ,				model_cls=A_						)
           _UpperCAmelCase		:       Dict       =    loaded_unet.to(unet.device						)
  
         # Since no EMA step has been performed the outputs should match.
         _UpperCAmelCase		:       Tuple       =    unet(A_       ,				A_       ,				A_						).sample
         _UpperCAmelCase		:       int       =    loaded_unet(A_       ,				A_       ,				A_						).sample
  
         assert torch.allclose(A_       ,				A_       ,				atol=1e-4						)
  
  
  
  
  
  
style_context: 
  
  
  from dataclasses import dataclass
  from enum import Enum
  from typing import List, Optional, Union
  
  import numpy as np
  import PIL
  from PIL import Image
  
  from ...utils import BaseOutput, is_torch_available, is_transformers_available
  
  
  @dataclass
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
       __SCREAMING_SNAKE_CASE    				=						{
           '''sld_warmup_steps''': 1_5,
           '''sld_guidance_scale''': 2_0,
           '''sld_threshold''': 0.0,
           '''sld_momentum_scale''': 0.0,
           '''sld_mom_beta''': 0.0,
       }
       __SCREAMING_SNAKE_CASE    				=						{
           '''sld_warmup_steps''': 1_0,
           '''sld_guidance_scale''': 1_0_0_0,
           '''sld_threshold''': 0.01,
           '''sld_momentum_scale''': 0.3,
           '''sld_mom_beta''': 0.4,
       }
       __SCREAMING_SNAKE_CASE    				=						{
           '''sld_warmup_steps''': 7,
           '''sld_guidance_scale''': 2_0_0_0,
           '''sld_threshold''': 0.025,
           '''sld_momentum_scale''': 0.5,
           '''sld_mom_beta''': 0.7,
       }
       __SCREAMING_SNAKE_CASE    				=						{
           '''sld_warmup_steps''': 0,
           '''sld_guidance_scale''': 5_0_0_0,
           '''sld_threshold''': 1.0,
           '''sld_momentum_scale''': 0.5,
           '''sld_mom_beta''': 0.7,
       }
  
  
  @dataclass
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
       __SCREAMING_SNAKE_CASE    				=						42
       __SCREAMING_SNAKE_CASE    				=						42
       __SCREAMING_SNAKE_CASE    				=						42
       __SCREAMING_SNAKE_CASE    				=						42
  
  
  if is_transformers_available() and is_torch_available():
   from .pipeline_stable_diffusion_safe import StableDiffusionPipelineSafe
   from .safety_checker import SafeStableDiffusionSafetyChecker
  
  
  
  
  
  