
'''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 		snake_case_       (						unittest.TestCase						):



   """simple docstring"""





   __lowerCAmelCase	:							List[Any]								='''hf-internal-testing/tiny-stable-diffusion-pipe'''
   __lowerCAmelCase	:							Optional[int]								=1
   __lowerCAmelCase	:							List[str]								=7_7
   __lowerCAmelCase	:							List[str]								=3_2
   __lowerCAmelCase	:							Union[str, Any]								=4
   __lowerCAmelCase	:							Tuple								=6_4
   __lowerCAmelCase	:							Optional[int]								=torch.manual_seed(0						)







   def     __UpperCAmelCase      ( self			,   UpperCamelCase=0.9_9_9_9):
          lowerCamelCase__				   =      UNetaDConditionModel.from_pretrained(self.model_id			,   subfolder="unet")
          lowerCamelCase__				   =      unet.to(UpperCamelCase)
          lowerCamelCase__				   =      EMAModel(unet.parameters()			,   decay=UpperCamelCase			,   model_cls=UpperCamelCase			,   model_config=unet.config)
          return unet, ema_unet







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      torch.randn(
              self.batch_size			,   self.num_in_channels			,   self.latent_height			,   self.latent_width			,   generator=self.generator).to(UpperCamelCase)
          lowerCamelCase__				   =      torch.randint(0			,   10_00			,   size=(self.batch_size,)			,   generator=self.generator).to(UpperCamelCase)
          lowerCamelCase__				   =      torch.randn(
              self.batch_size			,   self.prompt_length			,   self.text_encoder_hidden_dim			,   generator=self.generator).to(UpperCamelCase)
          return noisy_latents, timesteps, encoder_hidden_states







   def     __UpperCAmelCase      ( self			,   UpperCamelCase):
          lowerCamelCase__				   =      {}
          for k, param in unet.state_dict().items():
                 lowerCamelCase__				   =      torch.randn_like(UpperCamelCase) + (param * torch.randn_like(UpperCamelCase))
                 updated_state_dict.update({k: updated_param})
          unet.load_state_dict(UpperCamelCase)
          return unet







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__   ,		lowerCamelCase__				   =      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     __UpperCAmelCase      ( self):
          lowerCamelCase__   ,		lowerCamelCase__				   =      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())
          lowerCamelCase__				   =      list(unet.parameters())
          for s_param, param in zip(ema_unet.shadow_params			,   UpperCamelCase):
                 assert torch.allclose(UpperCamelCase			,   UpperCamelCase)

          # 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			,   UpperCamelCase):
                 assert torch.allclose(UpperCamelCase			,   UpperCamelCase)







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__   ,		lowerCamelCase__				   =      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.
          lowerCamelCase__				   =      self.simulate_backprop(UpperCamelCase)

          # 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.
          lowerCamelCase__				   =      list(unet_pseudo_updated_step_one.parameters())
          for s_param, param in zip(ema_unet.shadow_params			,   UpperCamelCase):
                 assert ~torch.allclose(UpperCamelCase			,   UpperCamelCase)

          # 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			,   UpperCamelCase):
                 assert ~torch.allclose(UpperCamelCase			,   UpperCamelCase)







   def     __UpperCAmelCase      ( self):
          # If we call EMA step after a backpropagation consecutively for two times,
          # the shadow params from those two steps should be different.
          lowerCamelCase__   ,		lowerCamelCase__				   =      self.get_models()

          # First backprop + EMA
          lowerCamelCase__				   =      self.simulate_backprop(UpperCamelCase)
          ema_unet.step(unet_step_one.parameters())
          lowerCamelCase__				   =      ema_unet.shadow_params

          # Second backprop + EMA
          lowerCamelCase__				   =      self.simulate_backprop(UpperCamelCase)
          ema_unet.step(unet_step_two.parameters())
          lowerCamelCase__				   =      ema_unet.shadow_params

          for step_one, step_two in zip(UpperCamelCase			,   UpperCamelCase):
                 assert ~torch.allclose(UpperCamelCase			,   UpperCamelCase)







   def     __UpperCAmelCase      ( 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.
          lowerCamelCase__   ,		lowerCamelCase__				   =      self.get_models(decay=0.0)
          lowerCamelCase__				   =      self.simulate_backprop(UpperCamelCase)
          ema_unet.step(unet_step_one.parameters())
          lowerCamelCase__				   =      ema_unet.shadow_params

          lowerCamelCase__				   =      self.simulate_backprop(UpperCamelCase)
          ema_unet.step(unet_step_two.parameters())
          lowerCamelCase__				   =      ema_unet.shadow_params

          for step_one, step_two in zip(UpperCamelCase			,   UpperCamelCase):
                 assert torch.allclose(UpperCamelCase			,   UpperCamelCase)







   @skip_mps
   def     __UpperCAmelCase      ( self):
          lowerCamelCase__   ,		lowerCamelCase__				   =      self.get_models()
          lowerCamelCase__   ,		lowerCamelCase__   ,		lowerCamelCase__				   =      self.get_dummy_inputs()

          with tempfile.TemporaryDirectory() as tmpdir:
                 ema_unet.save_pretrained(UpperCamelCase)
                 lowerCamelCase__				   =      UNetaDConditionModel.from_pretrained(UpperCamelCase			,   model_cls=UpperCamelCase)
                 lowerCamelCase__				   =      loaded_unet.to(unet.device)

          # Since no EMA step has been performed the outputs should match.
          lowerCamelCase__				   =      unet(UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase).sample
          lowerCamelCase__				   =      loaded_unet(UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase).sample

          assert torch.allclose(UpperCamelCase			,   UpperCamelCase			,   atol=1E-4)



'''simple docstring'''







from ..utils import DummyObject, requires_backends







class 		snake_case_       (						metaclass=A__						):



   """simple docstring"""





   __lowerCAmelCase	:							List[str]								=['''torch''', '''librosa''']







   def __init__( self			,   *UpperCamelCase			,   **UpperCamelCase):
          requires_backends(self			,   ["torch", "librosa"])







   @classmethod
   def     __UpperCAmelCase      ( cls			,   *UpperCamelCase			,   **UpperCamelCase):
          requires_backends(cls			,   ["torch", "librosa"])







   @classmethod
   def     __UpperCAmelCase      ( cls			,   *UpperCamelCase			,   **UpperCamelCase):
          requires_backends(cls			,   ["torch", "librosa"])







class 		snake_case_       (						metaclass=A__						):



   """simple docstring"""





   __lowerCAmelCase	:							Union[str, Any]								=['''torch''', '''librosa''']







   def __init__( self			,   *UpperCamelCase			,   **UpperCamelCase):
          requires_backends(self			,   ["torch", "librosa"])







   @classmethod
   def     __UpperCAmelCase      ( cls			,   *UpperCamelCase			,   **UpperCamelCase):
          requires_backends(cls			,   ["torch", "librosa"])







   @classmethod
   def     __UpperCAmelCase      ( cls			,   *UpperCamelCase			,   **UpperCamelCase):
          requires_backends(cls			,   ["torch", "librosa"])

