






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    ):


						'''simple docstring'''




						A   :      Tuple						 =    'hf-internal-testing/tiny-stable-diffusion-pipe'
						A   :      int						 =    1
						A   :      Any						 =    77
						A   :      int						 =    32
						A   :      Optional[int]						 =    4
						A   :      Union[str, Any]						 =    64
						A   :      List[Any]						 =    torch.manual_seed(0    )






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE=0.9999 )       -> str:
								snake_case_ :							List[str]         =				UNetaDConditionModel.from_pretrained(self.model_id						, subfolder="unet" )
								snake_case_ :							Any         =				unet.to(_SCREAMING_SNAKE_CASE )
								snake_case_ :							List[str]         =				EMAModel(unet.parameters()						, decay=_SCREAMING_SNAKE_CASE						, model_cls=_SCREAMING_SNAKE_CASE						, model_config=unet.config )
								return unet, ema_unet






						def 					_lowerCAmelCase				(		self )       -> Union[str, Any]:
								snake_case_ :							Any         =				torch.randn(
								    self.batch_size						, self.num_in_channels						, self.latent_height						, self.latent_width						, generator=self.generator ).to(_SCREAMING_SNAKE_CASE )
								snake_case_ :							str         =				torch.randint(0						, 1000						, size=(self.batch_size,)						, generator=self.generator ).to(_SCREAMING_SNAKE_CASE )
								snake_case_ :							int         =				torch.randn(
								    self.batch_size						, self.prompt_length						, self.text_encoder_hidden_dim						, generator=self.generator ).to(_SCREAMING_SNAKE_CASE )
								return noisy_latents, timesteps, encoder_hidden_states






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE )       -> Optional[int]:
								snake_case_ :							List[str]         =				{}
								for k, param in unet.state_dict().items():
										snake_case_ :							List[Any]         =				torch.randn_like(_SCREAMING_SNAKE_CASE ) + (param * torch.randn_like(_SCREAMING_SNAKE_CASE ))
										updated_state_dict.update({k: updated_param} )
								unet.load_state_dict(_SCREAMING_SNAKE_CASE )
								return unet






						def 					_lowerCAmelCase				(		self )       -> Dict:
								snake_case_       ,				snake_case_ :							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 )       -> Optional[Any]:
								snake_case_       ,				snake_case_ :							List[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_ :							List[str]         =				list(unet.parameters() )
								for s_param, param in zip(ema_unet.shadow_params						, _SCREAMING_SNAKE_CASE ):
										assert torch.allclose(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )

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






						def 					_lowerCAmelCase				(		self )       -> Any:
								snake_case_       ,				snake_case_ :							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.
								snake_case_ :							int         =				self.simulate_backprop(_SCREAMING_SNAKE_CASE )

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

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






						def 					_lowerCAmelCase				(		self )       -> Optional[int]:
								# If we call EMA step after a backpropagation consecutively for two times,
								# the shadow params from those two steps should be different.
								snake_case_       ,				snake_case_ :							str         =				self.get_models()

								# First backprop + EMA
								snake_case_ :							List[str]         =				self.simulate_backprop(_SCREAMING_SNAKE_CASE )
								ema_unet.step(unet_step_one.parameters() )
								snake_case_ :							Union[str, Any]         =				ema_unet.shadow_params

								# Second backprop + EMA
								snake_case_ :							Dict         =				self.simulate_backprop(_SCREAMING_SNAKE_CASE )
								ema_unet.step(unet_step_two.parameters() )
								snake_case_ :							str         =				ema_unet.shadow_params

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






						def 					_lowerCAmelCase				(		self )       -> Tuple:
								# 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.
								snake_case_       ,				snake_case_ :							Tuple         =				self.get_models(decay=0.0 )
								snake_case_ :							List[Any]         =				self.simulate_backprop(_SCREAMING_SNAKE_CASE )
								ema_unet.step(unet_step_one.parameters() )
								snake_case_ :							List[Any]         =				ema_unet.shadow_params

								snake_case_ :							Dict         =				self.simulate_backprop(_SCREAMING_SNAKE_CASE )
								ema_unet.step(unet_step_two.parameters() )
								snake_case_ :							int         =				ema_unet.shadow_params

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






						@skip_mps
						def 					_lowerCAmelCase				(		self )       -> List[Any]:
								snake_case_       ,				snake_case_ :							Dict         =				self.get_models()
								snake_case_       ,				snake_case_       ,				snake_case_ :							Union[str, Any]         =				self.get_dummy_inputs()

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

								# Since no EMA step has been performed the outputs should match.
								snake_case_ :							Dict         =				unet(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE ).sample
								snake_case_ :							Union[str, Any]         =				loaded_unet(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE ).sample

								assert torch.allclose(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, atol=1e-4 )








import glob
import os

from get_test_info import get_tester_classes


if __name__ == "__main__":
							lowercase     :							List[Any]	     =			[]

							lowercase     :							Optional[int]	     =			os.path.join('''tests''', '''models''', '''**''', '''test_modeling_*.py''')
							lowercase     :							Optional[Any]	     =			glob.glob(pattern)
							# TODO: deal with TF/Flax too
							lowercase     :							Tuple	     =			[
							    x for x in test_files if not (x.startswith('''test_modeling_tf_''') or x.startswith('''test_modeling_flax_'''))
							]

							for test_file in test_files:
														lowercase     :							Union[str, Any]	     =			get_tester_classes(test_file)
														for tester_class in tester_classes:
																					# A few tester classes don't have `parent` parameter in `__init__`.
																					# TODO: deal this better
																					try:
																												lowercase     :							Union[str, Any]	     =			tester_class(parent=None)
																					except Exception:
																												continue
																					if hasattr(tester, '''get_config'''):
																												lowercase     :							Tuple	     =			tester.get_config()
																												for k, v in config.to_dict().items():
																																			if isinstance(v, int):
																																										lowercase     :							Tuple	     =			None
																																										if k in ["vocab_size"]:
																																																	lowercase     :							str	     =			1_00
																																										elif k in ["max_position_embeddings"]:
																																																	lowercase     :							List[str]	     =			1_28
																																										elif k in ["hidden_size", "d_model"]:
																																																	lowercase     :							Union[str, Any]	     =			40
																																										elif k == ["num_layers", "num_hidden_layers", "num_encoder_layers", "num_decoder_layers"]:
																																																	lowercase     :							List[str]	     =			5
																																										if target is not None and v > target:
																																																	failures.append(
																																																	    F"""{tester_class.__name__} will produce a `config` of type `{config.__class__.__name__}`"""
																																																	    F""" with config[\"{k}\"] = {v} which is too large for testing! Set its value to be smaller"""
																																																	    F""" than {target}."""
																																																	)

							if len(failures) > 0:
														raise Exception(F"""There were {len(failures)} failures:\n""" + '''\n'''.join(failures))
