






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







		'''simple docstring'''



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



		def       __UpperCAmelCase    (							self       ,	_UpperCamelCase=0.99_99     )			->							int:
							UpperCAmelCase_					:   Optional[int]   =		UNetaDConditionModel.from_pretrained(self.model_id       ,	subfolder='unet'     )
							UpperCAmelCase_					:   Optional[int]   =		unet.to(_UpperCamelCase     )
							UpperCAmelCase_					:   Dict   =		EMAModel(unet.parameters()       ,	decay=_UpperCamelCase       ,	model_cls=_UpperCamelCase       ,	model_config=unet.config     )
							return unet, ema_unet



		def       __UpperCAmelCase    (							self     )			->							Optional[Any]:
							UpperCAmelCase_					:   List[str]   =		torch.randn(
							    self.batch_size       ,	self.num_in_channels       ,	self.latent_height       ,	self.latent_width       ,	generator=self.generator     ).to(_UpperCamelCase     )
							UpperCAmelCase_					:   Dict   =		torch.randint(0       ,	1_0_0_0       ,	size=(self.batch_size,)       ,	generator=self.generator     ).to(_UpperCamelCase     )
							UpperCAmelCase_					:   Union[str, Any]   =		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     )			->							Dict:
							UpperCAmelCase_					:   Dict   =		{}
							for k, param in unet.state_dict().items():
												UpperCAmelCase_					:   int   =		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     )			->							List[str]:
							UpperCAmelCase_ ,						UpperCAmelCase_					:   Tuple   =		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     )			->							Union[str, Any]:
							UpperCAmelCase_ ,						UpperCAmelCase_					:   int   =		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_					:   str   =		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     )			->							int:
							UpperCAmelCase_ ,						UpperCAmelCase_					:   List[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.
							UpperCAmelCase_					:   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.
							UpperCAmelCase_					:   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       __UpperCAmelCase    (							self     )			->							Optional[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_					:   int   =		self.get_models()

							# First backprop + EMA
							UpperCAmelCase_					:   str   =		self.simulate_backprop(_UpperCamelCase     )
							ema_unet.step(unet_step_one.parameters()     )
							UpperCAmelCase_					:   Any   =		ema_unet.shadow_params

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

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



		def       __UpperCAmelCase    (							self     )			->							List[str]:
							# 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_					:   Optional[Any]   =		self.get_models(decay=0.0     )
							UpperCAmelCase_					:   Optional[int]   =		self.simulate_backprop(_UpperCamelCase     )
							ema_unet.step(unet_step_one.parameters()     )
							UpperCAmelCase_					:   Union[str, Any]   =		ema_unet.shadow_params

							UpperCAmelCase_					:   Any   =		self.simulate_backprop(_UpperCamelCase     )
							ema_unet.step(unet_step_two.parameters()     )
							UpperCAmelCase_					:   Tuple   =		ema_unet.shadow_params

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



		@skip_mps
		def       __UpperCAmelCase    (							self     )			->							Union[str, Any]:
							UpperCAmelCase_ ,						UpperCAmelCase_					:   int   =		self.get_models()
							UpperCAmelCase_ ,						UpperCAmelCase_ ,						UpperCAmelCase_					:   List[Any]   =		self.get_dummy_inputs()

							with tempfile.TemporaryDirectory() as tmpdir:
												ema_unet.save_pretrained(_UpperCamelCase     )
												UpperCAmelCase_					:   Optional[Any]   =		UNetaDConditionModel.from_pretrained(_UpperCamelCase       ,	model_cls=_UpperCamelCase     )
												UpperCAmelCase_					:   Optional[int]   =		loaded_unet.to(unet.device     )

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

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














import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
		import transformer_engine.pytorch as te




def   lowercase__   (  __snake_case					:							List[str]  , __snake_case					:							List[Any]=True  , __snake_case					:							str=True  , __snake_case					:							Tuple=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(__snake_case  , 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
															UpperCAmelCase_					:   Optional[int]   =		module.bias is not None
															UpperCAmelCase_					:   Optional[int]   =		te.Linear(
															    module.in_features  , module.out_features  , bias=__snake_case  , params_dtype=module.weight.dtype						)
															UpperCAmelCase_					:   List[str]   =		module.weight.data.clone()
															if has_bias:
																				UpperCAmelCase_					:   List[Any]   =		module.bias.data.clone()

															setattr(__snake_case  , __snake_case  , __snake_case						)
										elif isinstance(__snake_case  , nn.LayerNorm						) and to_transformer_engine and _convert_ln:
															UpperCAmelCase_					:   int   =		te.LayerNorm(module.normalized_shape[0]  , eps=module.eps  , params_dtype=module.weight.dtype						)
															UpperCAmelCase_					:   Any   =		module.weight.data.clone()
															UpperCAmelCase_					:   Any   =		module.bias.data.clone()

															setattr(__snake_case  , __snake_case  , __snake_case						)
										elif isinstance(__snake_case  , te.Linear						) and not to_transformer_engine and _convert_linear:
															UpperCAmelCase_					:   int   =		module.bias is not None
															UpperCAmelCase_					:   Tuple   =		nn.Linear(
															    module.in_features  , module.out_features  , bias=__snake_case  , params_dtype=module.weight.dtype						)
															UpperCAmelCase_					:   str   =		module.weight.data.clone()
															if has_bias:
																				UpperCAmelCase_					:   Dict   =		module.bias.data.clone()

															setattr(__snake_case  , __snake_case  , __snake_case						)
										elif isinstance(__snake_case  , te.LayerNorm						) and not to_transformer_engine and _convert_ln:
															UpperCAmelCase_					:   List[Any]   =		nn.LayerNorm(module.normalized_shape[0]  , eps=module.eps  , params_dtype=module.weight.dtype						)
															UpperCAmelCase_					:   List[str]   =		module.weight.data.clone()
															UpperCAmelCase_					:   Tuple   =		module.bias.data.clone()

															setattr(__snake_case  , __snake_case  , __snake_case						)
										else:
															convert_model(
															    __snake_case  , to_transformer_engine=__snake_case  , _convert_linear=_convert_linear  , _convert_ln=_convert_ln  , )






def   lowercase__   (  __snake_case					:							List[Any]						):


					'''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(__snake_case  , (te.LayerNorm, te.Linear)						):
															return True
					return False






