code: 
  
  
  
  
  
  '''simple docstring'''
  
  
  
  
  
  import torch.nn as nn
  
  from .imports import is_fpa_available
  
  
  if is_fpa_available():
       import transformer_engine.pytorch as te
  def   _A (	SCREAMING_SNAKE_CASE  :     str , SCREAMING_SNAKE_CASE  :     List[Any]=True , SCREAMING_SNAKE_CASE  :     Any=True , SCREAMING_SNAKE_CASE  :     List[str]=True		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    if not is_fpa_available():
      raise ImportError("Using `convert_model` requires transformer_engine to be installed."		)
    for name, module in model.named_children():
      if isinstance(__lowerCamelCase , nn.Linear		) and to_transformer_engine and _convert_linear:
        # Return early if the linear layer weights are not multiples of 16
        if any(p % 16 != 0 for p in module.weight.shape		):
          return
        a__			:      int            =module.bias is not None
        a__			:      Optional[Any]            =te.Linear(
            module.in_features , module.out_features , bias=__lowerCamelCase , params_dtype=module.weight.dtype		)
        a__			:      Dict            =module.weight.data.clone()
        if has_bias:
          a__			:      Dict            =module.bias.data.clone()
  
        setattr(__lowerCamelCase , __lowerCamelCase , __lowerCamelCase		)
      elif isinstance(__lowerCamelCase , nn.LayerNorm		) and to_transformer_engine and _convert_ln:
        a__			:      List[str]            =te.LayerNorm(module.normalized_shape[0] , eps=module.eps , params_dtype=module.weight.dtype		)
        a__			:      Optional[Any]            =module.weight.data.clone()
        a__			:      Union[str, Any]            =module.bias.data.clone()
  
        setattr(__lowerCamelCase , __lowerCamelCase , __lowerCamelCase		)
      elif isinstance(__lowerCamelCase , te.Linear		) and not to_transformer_engine and _convert_linear:
        a__			:      Any            =module.bias is not None
        a__			:      Any            =nn.Linear(
            module.in_features , module.out_features , bias=__lowerCamelCase , params_dtype=module.weight.dtype		)
        a__			:      Optional[int]            =module.weight.data.clone()
        if has_bias:
          a__			:      Optional[Any]            =module.bias.data.clone()
  
        setattr(__lowerCamelCase , __lowerCamelCase , __lowerCamelCase		)
      elif isinstance(__lowerCamelCase , te.LayerNorm		) and not to_transformer_engine and _convert_ln:
        a__			:      Optional[Any]            =nn.LayerNorm(module.normalized_shape[0] , eps=module.eps , params_dtype=module.weight.dtype		)
        a__			:      Any            =module.weight.data.clone()
        a__			:      Union[str, Any]            =module.bias.data.clone()
  
        setattr(__lowerCamelCase , __lowerCamelCase , __lowerCamelCase		)
      else:
        convert_model(
            __lowerCamelCase , to_transformer_engine=__lowerCamelCase , _convert_linear=_convert_linear , _convert_ln=_convert_ln , )
  
  
  
  
  
  def   _A (	SCREAMING_SNAKE_CASE  :     int		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    if not is_fpa_available():
      raise ImportError("Using `has_transformer_engine_layers` requires transformer_engine to be installed."		)
    for m in model.modules():
      if isinstance(__lowerCamelCase , (te.LayerNorm, te.Linear)		):
        return True
    return False
  
  
  
  
  
  
  
style_context: 
  
  
  
  
  
  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        __lowerCAmelCase (					unittest.TestCase):
        _lowercase       :      str        =    """hf-internal-testing/tiny-stable-diffusion-pipe"""
        _lowercase       :      Optional[int]        =    1
        _lowercase       :      List[str]        =    77
        _lowercase       :      int        =    32
        _lowercase       :      Dict        =    4
        _lowercase       :      str        =    64
        _lowercase       :      Optional[Any]        =    torch.manual_seed(0)
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__=0.99_99							)     ->		Tuple:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      int            =UNetaDConditionModel.from_pretrained(self.model_id	,   subfolder="unet"							)
          a__			:      Any            =unet.to(lowerCAmelCase__							)
          a__			:      int            =EMAModel(unet.parameters()	,   decay=lowerCAmelCase__	,   model_cls=lowerCAmelCase__	,   model_config=unet.config							)
          return unet, ema_unet
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		Dict:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      str            =torch.randn(
              self.batch_size	,   self.num_in_channels	,   self.latent_height	,   self.latent_width	,   generator=self.generator							).to(lowerCAmelCase__							)
          a__			:      Any            =torch.randint(0	,   1_0_0_0	,   size=(self.batch_size,)	,   generator=self.generator							).to(lowerCAmelCase__							)
          a__			:      Tuple            =torch.randn(
              self.batch_size	,   self.prompt_length	,   self.text_encoder_hidden_dim	,   generator=self.generator							).to(lowerCAmelCase__							)
          return noisy_latents, timesteps, encoder_hidden_states
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__							)     ->		Any:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Any            ={}
          for k, param in unet.state_dict().items():
            a__			:      Tuple            =torch.randn_like(lowerCAmelCase__							) + (param * torch.randn_like(lowerCAmelCase__							))
            updated_state_dict.update({k: updated_param}							)
          unet.load_state_dict(lowerCAmelCase__							)
          return unet
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		Tuple:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__			:      List[Any]            =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      _lowercase						(					self							)     ->		List[str]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__			:      Any            =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()							)
          a__			:      str            =list(unet.parameters()							)
          for s_param, param in zip(ema_unet.shadow_params	,   lowerCAmelCase__							):
            assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)
  
          # 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	,   lowerCAmelCase__							):
            assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		List[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__			:      Optional[Any]            =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.
          a__			:      List[str]            =self.simulate_backprop(lowerCAmelCase__							)
  
          # 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.
          a__			:      Optional[int]            =list(unet_pseudo_updated_step_one.parameters()							)
          for s_param, param in zip(ema_unet.shadow_params	,   lowerCAmelCase__							):
            assert ~torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)
  
          # 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	,   lowerCAmelCase__							):
            assert ~torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		Union[str, Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__			:      Dict            =self.get_models()
  
          # First backprop + EMA
          a__			:      Tuple            =self.simulate_backprop(lowerCAmelCase__							)
          ema_unet.step(unet_step_one.parameters()							)
          a__			:      Dict            =ema_unet.shadow_params
  
          # Second backprop + EMA
          a__			:      str            =self.simulate_backprop(lowerCAmelCase__							)
          ema_unet.step(unet_step_two.parameters()							)
          a__			:      List[str]            =ema_unet.shadow_params
  
          for step_one, step_two in zip(lowerCAmelCase__	,   lowerCAmelCase__							):
            assert ~torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		List[str]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__			:      int            =self.get_models(decay=0.0							)
          a__			:      Optional[int]            =self.simulate_backprop(lowerCAmelCase__							)
          ema_unet.step(unet_step_one.parameters()							)
          a__			:      List[Any]            =ema_unet.shadow_params
  
          a__			:      Tuple            =self.simulate_backprop(lowerCAmelCase__							)
          ema_unet.step(unet_step_two.parameters()							)
          a__			:      List[Any]            =ema_unet.shadow_params
  
          for step_one, step_two in zip(lowerCAmelCase__	,   lowerCAmelCase__							):
            assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)
  
  
  
  
  
  
  
        @skip_mps
        def      _lowercase						(					self							)     ->		Optional[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__			:      Union[str, Any]            =self.get_models()
          a__      ,							a__      ,							a__			:      Tuple            =self.get_dummy_inputs()
  
          with tempfile.TemporaryDirectory() as tmpdir:
            ema_unet.save_pretrained(lowerCAmelCase__							)
            a__			:      List[str]            =UNetaDConditionModel.from_pretrained(lowerCAmelCase__	,   model_cls=lowerCAmelCase__							)
            a__			:      Tuple            =loaded_unet.to(unet.device							)
  
          # Since no EMA step has been performed the outputs should match.
          a__			:      Optional[Any]            =unet(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							).sample
          a__			:      Optional[int]            =loaded_unet(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							).sample
  
          assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__	,   atol=1E-4							)
  
  
  
  
  
  
  