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   snake_case__					(			unittest.TestCase				):
  
  
  
  
  
  
  
  		"""simple docstring"""
  
  		lowerCamelCase           =	"""hf-internal-testing/tiny-stable-diffusion-pipe"""
  		lowerCamelCase           =	1
  		lowerCamelCase           =	77
  		lowerCamelCase           =	32
  		lowerCamelCase           =	4
  		lowerCamelCase           =	lowerCamelCase           =	64
  		lowerCamelCase           =	torch.manual_seed(0				)
  
  
  
  		def 			lowerCAmelCase							(			self							: Union[str, Any]     ,  UpperCamelCase__							: List[Any]=0.9_999			)	-> Dict:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					List[str]     =	UNetaDConditionModel.from_pretrained(self.model_id     ,  subfolder='''unet'''			)
  				snake_case     :					Tuple     =	unet.to(UpperCamelCase__			)
  				snake_case     :					Any     =	EMAModel(unet.parameters()     ,  decay=UpperCamelCase__     ,  model_cls=UpperCamelCase__     ,  model_config=unet.config			)
  				return unet, ema_unet
  
  
  
  		def 			lowerCAmelCase							(			self							: Tuple			)	-> List[Any]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					Dict     =	torch.randn(
  				    self.batch_size     ,  self.num_in_channels     ,  self.latent_height     ,  self.latent_width     ,  generator=self.generator			).to(UpperCamelCase__			)
  				snake_case     :					List[Any]     =	torch.randint(0     ,  1000     ,  size=(self.batch_size,)     ,  generator=self.generator			).to(UpperCamelCase__			)
  				snake_case     :					Optional[int]     =	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 			lowerCAmelCase							(			self							: str     ,  UpperCamelCase__							: List[str]			)	-> Tuple:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					List[Any]     =	{}
  				for k, param in unet.state_dict().items():
  						snake_case     :					Dict     =	torch.randn_like(UpperCamelCase__			) + (param * torch.randn_like(UpperCamelCase__			))
  						updated_state_dict.update({k: updated_param}			)
  				unet.load_state_dict(UpperCamelCase__			)
  				return unet
  
  
  
  		def 			lowerCAmelCase							(			self							: List[str]			)	-> List[Any]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case ,snake_case     :					Optional[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 			lowerCAmelCase							(			self							: Tuple			)	-> Dict:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case ,snake_case     :					Optional[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()			)
  				snake_case     :					Optional[Any]     =	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 			lowerCAmelCase							(			self							: List[Any]			)	-> List[str]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case ,snake_case     :					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.
  				snake_case     :					List[str]     =	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.
  				snake_case     :					str     =	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 			lowerCAmelCase							(			self							: int			)	-> Any:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case ,snake_case     :					Union[str, Any]     =	self.get_models()
  
  				# First backprop + EMA
  				snake_case     :					Any     =	self.simulate_backprop(UpperCamelCase__			)
  				ema_unet.step(unet_step_one.parameters()			)
  				snake_case     :					Optional[int]     =	ema_unet.shadow_params
  
  				# Second backprop + EMA
  				snake_case     :					Optional[int]     =	self.simulate_backprop(UpperCamelCase__			)
  				ema_unet.step(unet_step_two.parameters()			)
  				snake_case     :					int     =	ema_unet.shadow_params
  
  				for step_one, step_two in zip(UpperCamelCase__     ,  UpperCamelCase__			):
  						assert ~torch.allclose(UpperCamelCase__     ,  UpperCamelCase__			)
  
  
  
  		def 			lowerCAmelCase							(			self							: str			)	-> List[Any]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case ,snake_case     :					int     =	self.get_models(decay=0.0			)
  				snake_case     :					Optional[Any]     =	self.simulate_backprop(UpperCamelCase__			)
  				ema_unet.step(unet_step_one.parameters()			)
  				snake_case     :					int     =	ema_unet.shadow_params
  
  				snake_case     :					Optional[int]     =	self.simulate_backprop(UpperCamelCase__			)
  				ema_unet.step(unet_step_two.parameters()			)
  				snake_case     :					int     =	ema_unet.shadow_params
  
  				for step_one, step_two in zip(UpperCamelCase__     ,  UpperCamelCase__			):
  						assert torch.allclose(UpperCamelCase__     ,  UpperCamelCase__			)
  
  
  
  		@skip_mps
  		def 			lowerCAmelCase							(			self							: int			)	-> Union[str, Any]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case ,snake_case     :					List[str]     =	self.get_models()
  				snake_case ,snake_case ,snake_case     :					Optional[int]     =	self.get_dummy_inputs()
  
  				with tempfile.TemporaryDirectory() as tmpdir:
  						ema_unet.save_pretrained(UpperCamelCase__			)
  						snake_case     :					int     =	UNetaDConditionModel.from_pretrained(UpperCamelCase__     ,  model_cls=UpperCamelCase__			)
  						snake_case     :					List[str]     =	loaded_unet.to(unet.device			)
  
  				# Since no EMA step has been performed the outputs should match.
  				snake_case     :					Optional[Any]     =	unet(UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__			).sample
  				snake_case     :					Optional[Any]     =	loaded_unet(UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__			).sample
  
  				assert torch.allclose(UpperCamelCase__     ,  UpperCamelCase__     ,  atol=1e-4			)
  
  
  
style_context: 
  '''simple docstring'''
  def       _UpperCamelCase	(					SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__ )     ->       int:
  
  
  		'''simple docstring'''
  
  
  		snake_case     :					Optional[Any]     =	0
  		for ele in arr:
  				total_sum += product_sum(SCREAMING_SNAKE_CASE__   ,					depth + 1 ) if isinstance(SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__ ) else ele
  		return total_sum * depth
  
  
  
  
  def       _UpperCamelCase	(					SCREAMING_SNAKE_CASE__ )     ->       int:
  
  
  		'''simple docstring'''
  
  
  		return product_sum(SCREAMING_SNAKE_CASE__   ,					1 )
  
  
  if __name__ == "__main__":
  						import doctest
  
  						doctest.testmod()
  
  
  