
'''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       lowerCAmelCase			(					unittest.TestCase      ):
		a				:       Optional[Any]							=       """hf-internal-testing/tiny-stable-diffusion-pipe"""
		a				:       List[str]							=       1
		a				:       int							=       77
		a				:       Dict							=       32
		a				:       Dict							=       4
		a				:       List[Any]							=       64
		a				:       Tuple							=       torch.manual_seed(0      )




		def      lowercase							(					self       ,    UpperCamelCase=0.99_99   ):
								_SCREAMING_SNAKE_CASE					=			UNetaDConditionModel.from_pretrained(self.model_id       ,    subfolder="unet"   )
								_SCREAMING_SNAKE_CASE					=			unet.to(UpperCamelCase   )
								_SCREAMING_SNAKE_CASE					=			EMAModel(unet.parameters()       ,    decay=UpperCamelCase       ,    model_cls=UpperCamelCase       ,    model_config=unet.config   )
								return unet, ema_unet




		def      lowercase							(					self   ):
								_SCREAMING_SNAKE_CASE					=			torch.randn(
								    self.batch_size       ,    self.num_in_channels       ,    self.latent_height       ,    self.latent_width       ,    generator=self.generator   ).to(UpperCamelCase   )
								_SCREAMING_SNAKE_CASE					=			torch.randint(0       ,    1_000       ,    size=(self.batch_size,)       ,    generator=self.generator   ).to(UpperCamelCase   )
								_SCREAMING_SNAKE_CASE					=			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      lowercase							(					self       ,    UpperCamelCase   ):
								_SCREAMING_SNAKE_CASE					=			{}
								for k, param in unet.state_dict().items():
														_SCREAMING_SNAKE_CASE					=			torch.randn_like(UpperCamelCase   ) + (param * torch.randn_like(UpperCamelCase   ))
														updated_state_dict.update({k: updated_param}   )
								unet.load_state_dict(UpperCamelCase   )
								return unet




		def      lowercase							(					self   ):
								_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      lowercase							(					self   ):
								_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       ,    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      lowercase							(					self   ):
								_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(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.
								_SCREAMING_SNAKE_CASE					=			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      lowercase							(					self   ):
								# 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(UpperCamelCase   )
								ema_unet.step(unet_step_one.parameters()   )
								_SCREAMING_SNAKE_CASE					=			ema_unet.shadow_params

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

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




		def      lowercase							(					self   ):
								# 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(UpperCamelCase   )
								ema_unet.step(unet_step_one.parameters()   )
								_SCREAMING_SNAKE_CASE					=			ema_unet.shadow_params

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

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




		@skip_mps
		def      lowercase							(					self   ):
								_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(UpperCamelCase   )
														_SCREAMING_SNAKE_CASE					=			UNetaDConditionModel.from_pretrained(UpperCamelCase       ,    model_cls=UpperCamelCase   )
														_SCREAMING_SNAKE_CASE					=			loaded_unet.to(unet.device   )

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

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

'''simple docstring'''







import os
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
						import torch

						from transformers import GPTaLMHeadModel


@require_torch
@require_sentencepiece
@require_tokenizers
class       lowerCAmelCase			(					unittest.TestCase      ):




		@slow
		@unittest.skip("Model is not available."   )
		def      lowercase							(					self   ):
								_SCREAMING_SNAKE_CASE					=			"nvidia/megatron-gpt2-345m/"
								if "MYDIR" in os.environ:
														_SCREAMING_SNAKE_CASE					=			os.path.join(os.environ["MYDIR"]       ,    UpperCamelCase   )
								_SCREAMING_SNAKE_CASE					=			GPTaLMHeadModel.from_pretrained(UpperCamelCase   )
								model.to(UpperCamelCase   )
								model.half()

								_SCREAMING_SNAKE_CASE					=			torch.tensor(
								    [[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]       ,    device=UpperCamelCase       ,    dtype=torch.long       ,    )

								with torch.no_grad():
														_SCREAMING_SNAKE_CASE					=			model(UpperCamelCase   ).logits

								_SCREAMING_SNAKE_CASE					=			torch.Size((1, 9, 50_257)   )
								self.assertEqual(output.shape       ,    UpperCamelCase   )

								_SCREAMING_SNAKE_CASE					=			torch.tensor(
								    [
								        4.94_14,
								        -0.29_20,
								        -1.21_48,
								        -4.02_73,
								        -0.51_61,
								        -5.21_09,
								        -1.24_12,
								        -1.83_01,
								        -1.77_34,
								        -4.71_48,
								        -0.23_17,
								        -1.08_11,
								        -2.17_77,
								        0.41_41,
								        -3.79_69,
								        -4.05_86,
								        -2.53_32,
								        -3.38_09,
								        4.38_67,
								    ]       ,    device=UpperCamelCase       ,    dtype=torch.half       ,    )

								for i in range(19   ):
														_SCREAMING_SNAKE_CASE,				_SCREAMING_SNAKE_CASE					=			8 * i // 17, 2_792 * i  # along the diagonal
														_SCREAMING_SNAKE_CASE,				_SCREAMING_SNAKE_CASE					=			output[0, r, c], expected_diag[i]
														_SCREAMING_SNAKE_CASE					=			F'row={r} col={c} computed={computed} expected={expected}'
														self.assertAlmostEqual(UpperCamelCase       ,    UpperCamelCase       ,    delta=1e-4       ,    msg=UpperCamelCase   )