


import inspect
import warnings

import torch

from .state import AcceleratorState, GradientState
from .utils import DistributedType, honor_type, is_tpu_available


if is_tpu_available(check_device=False):
							import torch_xla.core.xla_model as xm


def 				SCREAMING_SNAKE_CASE (     UpperCAmelCase__   ,UpperCAmelCase__							):

				"""simple docstring"""






				if isinstance(UpperCAmelCase__   ,(list, tuple)							):
								return honor_type(UpperCAmelCase__   ,(move_to_device(UpperCAmelCase__   ,UpperCAmelCase__							) for t in state)							)
				elif isinstance(UpperCAmelCase__   ,UpperCAmelCase__							):
								return type(UpperCAmelCase__							)({k: move_to_device(UpperCAmelCase__   ,UpperCAmelCase__							) for k, v in state.items()}							)
				elif isinstance(UpperCAmelCase__   ,torch.Tensor							):
								return state.to(UpperCAmelCase__							)
				return state



class    __lowercase  (       torch.optim.Optimizer							):

							"""simple docstring"""


							def __init__(    self				, A_				, A_=True				, A_=None	)->	List[str]:
											_SCREAMING_SNAKE_CASE      = optimizer
											_SCREAMING_SNAKE_CASE      = scaler
											_SCREAMING_SNAKE_CASE      = AcceleratorState()
											_SCREAMING_SNAKE_CASE      = GradientState()
											_SCREAMING_SNAKE_CASE      = device_placement
											_SCREAMING_SNAKE_CASE      = False
											_SCREAMING_SNAKE_CASE      = None

											# Handle device placement
											if device_placement:
															_SCREAMING_SNAKE_CASE      = self.optimizer.state_dict()
															if self.accelerator_state.distributed_type == DistributedType.TPU:
																			xm.send_cpu_data_to_device(A_				, self.accelerator_state.device	)
															else:
																			_SCREAMING_SNAKE_CASE      = move_to_device(A_				, self.accelerator_state.device	)
															self.optimizer.load_state_dict(A_	)





							@property
							def 		__magic_name__			(    self	)->	Tuple:
											return self.optimizer.state





							@state.setter
							def 		__magic_name__			(    self				, A_	)->	str:
											_SCREAMING_SNAKE_CASE      = state





							@property
							def 		__magic_name__			(    self	)->	Optional[Any]:
											return self.optimizer.param_groups





							@param_groups.setter
							def 		__magic_name__			(    self				, A_	)->	Tuple:
											_SCREAMING_SNAKE_CASE      = param_groups





							@property
							def 		__magic_name__			(    self	)->	Any:
											return self.optimizer.defaults





							@defaults.setter
							def 		__magic_name__			(    self				, A_	)->	Dict:
											_SCREAMING_SNAKE_CASE      = defaults





							def 		__magic_name__			(    self				, A_	)->	Any:
											self.optimizer.add_param_group(A_	)





							def 		__magic_name__			(    self				, A_	)->	int:
											if self.accelerator_state.distributed_type == DistributedType.TPU and self.device_placement:
															xm.send_cpu_data_to_device(A_				, self.accelerator_state.device	)
											self.optimizer.load_state_dict(A_	)





							def 		__magic_name__			(    self	)->	List[str]:
											return self.optimizer.state_dict()





							def 		__magic_name__			(    self				, A_=None	)->	Any:
											if self.gradient_state.sync_gradients:
															_SCREAMING_SNAKE_CASE      = 'set_to_none' in inspect.signature(self.optimizer.zero_grad	).parameters
															if accept_arg:
																			if set_to_none is None:
																							_SCREAMING_SNAKE_CASE      = False
																			self.optimizer.zero_grad(set_to_none=A_	)
															else:
																			if set_to_none is not None:
																							raise ValueError('`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.'	)
																			self.optimizer.zero_grad()





							def 		__magic_name__			(    self				, A_=None	)->	Optional[int]:
											if self.gradient_state.sync_gradients:
															if self.accelerator_state.distributed_type == DistributedType.TPU:
																			_SCREAMING_SNAKE_CASE      = {'closure': closure} if closure is not None else {}
																			xm.optimizer_step(self.optimizer				, optimizer_args=A_	)
															elif self.scaler is not None:
																			_SCREAMING_SNAKE_CASE      = False
																			if self._last_scale is None:
																							# `get_scale` is an async operation requiring full synchronization
																							# on CPU and GPUs before finishing. As a result, we store away
																							# the prior one to reduce the call overhead
																							_SCREAMING_SNAKE_CASE      = self.scaler.get_scale()
																							_SCREAMING_SNAKE_CASE      = True
																			self.scaler.step(self.optimizer				, A_	)
																			self.scaler.update()
																			_SCREAMING_SNAKE_CASE      = self.scaler.get_scale()
																			if not new_scale:
																							# If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
																							_SCREAMING_SNAKE_CASE      = scale_after < self._last_scale
																			_SCREAMING_SNAKE_CASE      = scale_after

															else:
																			self.optimizer.step(A_	)





							def 		__magic_name__			(    self				, A_	)->	Dict:
											for param_group in self.optimizer.param_groups:
															_SCREAMING_SNAKE_CASE      = [parameters_map.get(A_				, A_	) for p in param_group['params']]





							@property
							def 		__magic_name__			(    self	)->	str:
											warnings.warn(
											    'The `is_overflow` property is deprecated and will be removed in version 1.0 of Accelerate use '
											    '`optimizer.step_was_skipped` instead.'				, A_				, )
											return self._is_overflow





							@property
							def 		__magic_name__			(    self	)->	int:
											return self._is_overflow





							def __getstate__(    self	)->	Dict:
											return self.__dict__.copy()





							def __setstate__(    self				, A_	)->	List[str]:
											self.__dict__.update(A_	)




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

							"""simple docstring"""


							SCREAMING_SNAKE_CASE   :					str        =						"hf-internal-testing/tiny-stable-diffusion-pipe"
							SCREAMING_SNAKE_CASE   :					str        =						1
							SCREAMING_SNAKE_CASE   :					Optional[Any]        =						77
							SCREAMING_SNAKE_CASE   :					Dict        =						32
							SCREAMING_SNAKE_CASE   :					Optional[Any]        =						4
							SCREAMING_SNAKE_CASE   :					Optional[int]        =						64
							SCREAMING_SNAKE_CASE   :					Tuple        =						torch.manual_seed(0							)





							def 		__magic_name__			(    self				, A_=0.9999	)->	Optional[int]:
											_SCREAMING_SNAKE_CASE      = UNetaDConditionModel.from_pretrained(self.model_id				, subfolder='unet'	)
											_SCREAMING_SNAKE_CASE      = unet.to(A_	)
											_SCREAMING_SNAKE_CASE      = EMAModel(unet.parameters()				, decay=A_				, model_cls=A_				, model_config=unet.config	)
											return unet, ema_unet





							def 		__magic_name__			(    self	)->	Optional[int]:
											_SCREAMING_SNAKE_CASE      = torch.randn(
											    self.batch_size				, self.num_in_channels				, self.latent_height				, self.latent_width				, generator=self.generator	).to(A_	)
											_SCREAMING_SNAKE_CASE      = torch.randint(0				, 1000				, size=(self.batch_size,)				, generator=self.generator	).to(A_	)
											_SCREAMING_SNAKE_CASE      = 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				, A_	)->	str:
											_SCREAMING_SNAKE_CASE      = {}
											for k, param in unet.state_dict().items():
															_SCREAMING_SNAKE_CASE      = 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	)->	Dict:
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = 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	)->	Optional[int]:
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = 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()	)
											_SCREAMING_SNAKE_CASE      = 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	)->	Union[str, Any]:
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = 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.
											_SCREAMING_SNAKE_CASE      = 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.
											_SCREAMING_SNAKE_CASE      = 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	)->	Optional[int]:
											# If we call EMA step after a backpropagation consecutively for two times,
											# the shadow params from those two steps should be different.
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.get_models()

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

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

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





							def 		__magic_name__			(    self	)->	Dict:
											# 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.
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.get_models(decay=0.0	)
											_SCREAMING_SNAKE_CASE      = self.simulate_backprop(A_	)
											ema_unet.step(unet_step_one.parameters()	)
											_SCREAMING_SNAKE_CASE      = ema_unet.shadow_params

											_SCREAMING_SNAKE_CASE      = self.simulate_backprop(A_	)
											ema_unet.step(unet_step_two.parameters()	)
											_SCREAMING_SNAKE_CASE      = ema_unet.shadow_params

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





							@skip_mps
							def 		__magic_name__			(    self	)->	Optional[Any]:
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.get_models()
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.get_dummy_inputs()

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

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

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