

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    __UpperCAmelCase	(			unittest.TestCase		):
   UpperCamelCase  	=	"""hf-internal-testing/tiny-stable-diffusion-pipe"""
   UpperCamelCase  	=	1
   UpperCamelCase  	=	7_7
   UpperCamelCase  	=	3_2
   UpperCamelCase  	=	4
   UpperCamelCase  	=	UpperCamelCase  	=	6_4
   UpperCamelCase  	=	torch.manual_seed(0		)



   def 			__magic_name__       (				self				:	Optional[int],						__A				:	Optional[int]=0.9_9_9_9		):
     UpperCAmelCase			:				int										=					UNetaDConditionModel.from_pretrained(self.model_id,						subfolder='''unet'''		)
     UpperCAmelCase			:				Optional[Any]										=					unet.to(__A		)
     UpperCAmelCase			:				Union[str, Any]										=					EMAModel(unet.parameters(),						decay=__A,						model_cls=__A,						model_config=unet.config		)
     return unet, ema_unet



   def 			__magic_name__       (				self				:	Dict		):
     UpperCAmelCase			:				List[str]										=					torch.randn(
         self.batch_size,						self.num_in_channels,						self.latent_height,						self.latent_width,						generator=self.generator		).to(__A		)
     UpperCAmelCase			:				List[Any]										=					torch.randint(0,						1_0_0_0,						size=(self.batch_size,),						generator=self.generator		).to(__A		)
     UpperCAmelCase			:				Optional[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 			__magic_name__       (				self				:	Dict,						__A				:	Optional[Any]		):
     UpperCAmelCase			:				List[str]										=					{}
     for k, param in unet.state_dict().items():
       UpperCAmelCase			:				str										=					torch.randn_like(__A		) + (param * torch.randn_like(__A		))
       updated_state_dict.update({k: updated_param}		)
     unet.load_state_dict(__A		)
     return unet



   def 			__magic_name__       (				self				:	Tuple		):
     UpperCAmelCase					,		UpperCAmelCase			:				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 			__magic_name__       (				self				:	List[Any]		):
     UpperCAmelCase					,		UpperCAmelCase			:				str										=					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 			__magic_name__       (				self				:	Optional[Any]		):
     UpperCAmelCase					,		UpperCAmelCase			:				Optional[int]										=					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			:				List[str]										=					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			:				Dict										=					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 			__magic_name__       (				self				:	List[Any]		):
     # If we call EMA step after a backpropagation consecutively for two times,
     # the shadow params from those two steps should be different.
     UpperCAmelCase					,		UpperCAmelCase			:				Tuple										=					self.get_models()

     # First backprop + EMA
     UpperCAmelCase			:				List[Any]										=					self.simulate_backprop(__A		)
     ema_unet.step(unet_step_one.parameters()		)
     UpperCAmelCase			:				Any										=					ema_unet.shadow_params

     # Second backprop + EMA
     UpperCAmelCase			:				List[str]										=					self.simulate_backprop(__A		)
     ema_unet.step(unet_step_two.parameters()		)
     UpperCAmelCase			:				Tuple										=					ema_unet.shadow_params

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



   def 			__magic_name__       (				self				:	Union[str, Any]		):
     # 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					,		UpperCAmelCase			:				List[Any]										=					self.get_models(decay=0.0		)
     UpperCAmelCase			:				str										=					self.simulate_backprop(__A		)
     ema_unet.step(unet_step_one.parameters()		)
     UpperCAmelCase			:				str										=					ema_unet.shadow_params

     UpperCAmelCase			:				Optional[int]										=					self.simulate_backprop(__A		)
     ema_unet.step(unet_step_two.parameters()		)
     UpperCAmelCase			:				int										=					ema_unet.shadow_params

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





   @skip_mps
   def 			__magic_name__       (				self				:	List[str]		):
     UpperCAmelCase					,		UpperCAmelCase			:				Optional[int]										=					self.get_models()
     UpperCAmelCase					,		UpperCAmelCase					,		UpperCAmelCase			:				Union[str, Any]										=					self.get_dummy_inputs()

     with tempfile.TemporaryDirectory() as tmpdir:
       ema_unet.save_pretrained(__A		)
       UpperCAmelCase			:				List[str]										=					UNetaDConditionModel.from_pretrained(__A,						model_cls=__A		)
       UpperCAmelCase			:				List[Any]										=					loaded_unet.to(unet.device		)

     # Since no EMA step has been performed the outputs should match.
     UpperCAmelCase			:				List[str]										=					unet(__A,						__A,						__A		).sample
     UpperCAmelCase			:				Tuple										=					loaded_unet(__A,						__A,						__A		).sample

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








def      a__	(						UpperCAmelCase				:     list     )    ->		list:
  UpperCAmelCase			:				Tuple										=					len(UpperCAmelCase     )

  # Set the gap value to be decreased by a factor of 1.3
  # after each iteration
  UpperCAmelCase			:				Optional[Any]										=					1.3

  # Continue sorting until the gap is 1
  while gap > 1:
    # Decrease the gap value
    UpperCAmelCase			:				int										=					int(gap / shrink     )

    # Sort the elements using insertion sort
    for i in range(UpperCAmelCase							,					len(UpperCAmelCase     )     ):
      UpperCAmelCase			:				Union[str, Any]										=					collection[i]
      UpperCAmelCase			:				List[Any]										=					i
      while j >= gap and collection[j - gap] > temp:
        UpperCAmelCase			:				int										=					collection[j - gap]
        j -= gap
      UpperCAmelCase			:				List[Any]										=					temp

  return collection


if __name__ == "__main__":
     import doctest

     doctest.testmod()





