


"""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      _UpperCAmelCase					(				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 							SCREAMING_SNAKE_CASE    (self							,			a_=0.9999   ):





			'''simple docstring'''

			__snake_case      :		Optional[Any]         =							UNetaDConditionModel.from_pretrained(self.model_id							,			subfolder='''unet'''   )
			__snake_case      :		Union[str, Any]         =							unet.to(a_   )
			__snake_case      :		Tuple         =							EMAModel(unet.parameters()							,			decay=a_							,			model_cls=a_							,			model_config=unet.config   )
			return unet, ema_unet



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      :		Union[str, Any]         =							torch.randn(
			    self.batch_size							,			self.num_in_channels							,			self.latent_height							,			self.latent_width							,			generator=self.generator   ).to(a_   )
			__snake_case      :		int         =							torch.randint(0							,			10_00							,			size=(self.batch_size,)							,			generator=self.generator   ).to(a_   )
			__snake_case      :		int         =							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 							SCREAMING_SNAKE_CASE    (self							,			a_   ):





			'''simple docstring'''

			__snake_case      :		List[Any]         =							{}
			for k, param in unet.state_dict().items():
				__snake_case      :		int         =							torch.randn_like(a_   ) + (param * torch.randn_like(a_   ))
				updated_state_dict.update({k: updated_param}   )
			unet.load_state_dict(a_   )
			return unet



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''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 							SCREAMING_SNAKE_CASE    (self   ):





			'''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							,			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 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      ,				__snake_case      :		List[str]         =							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      :		Union[str, Any]         =							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.
			__snake_case      :		Any         =							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 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      ,				__snake_case      :		Dict         =							self.get_models()

			# First backprop + EMA
			__snake_case      :		Optional[int]         =							self.simulate_backprop(a_   )
			ema_unet.step(unet_step_one.parameters()   )
			__snake_case      :		Union[str, Any]         =							ema_unet.shadow_params

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

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



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      ,				__snake_case      :		Optional[int]         =							self.get_models(decay=0.0   )
			__snake_case      :		Any         =							self.simulate_backprop(a_   )
			ema_unet.step(unet_step_one.parameters()   )
			__snake_case      :		Any         =							ema_unet.shadow_params

			__snake_case      :		Optional[Any]         =							self.simulate_backprop(a_   )
			ema_unet.step(unet_step_two.parameters()   )
			__snake_case      :		List[str]         =							ema_unet.shadow_params

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



		@skip_mps
		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      ,				__snake_case      :		Optional[int]         =							self.get_models()
			__snake_case      ,				__snake_case      ,				__snake_case      :		Any         =							self.get_dummy_inputs()

			with tempfile.TemporaryDirectory() as tmpdir:
				ema_unet.save_pretrained(a_   )
				__snake_case      :		int         =							UNetaDConditionModel.from_pretrained(a_							,			model_cls=a_   )
				__snake_case      :		Union[str, Any]         =							loaded_unet.to(unet.device   )

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

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









import os

from packaging import version

from .. import __version__
from .accelerate_utils import apply_forward_hook
from .constants import (
    CONFIG_NAME,
    DEPRECATED_REVISION_ARGS,
    DIFFUSERS_CACHE,
    DIFFUSERS_DYNAMIC_MODULE_NAME,
    FLAX_WEIGHTS_NAME,
    HF_MODULES_CACHE,
    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
    ONNX_EXTERNAL_WEIGHTS_NAME,
    ONNX_WEIGHTS_NAME,
    SAFETENSORS_WEIGHTS_NAME,
    TEXT_ENCODER_ATTN_MODULE,
    WEIGHTS_NAME,
)
from .deprecation_utils import deprecate
from .doc_utils import replace_example_docstring
from .dynamic_modules_utils import get_class_from_dynamic_module
from .hub_utils import (
    HF_HUB_OFFLINE,
    _add_variant,
    _get_model_file,
    extract_commit_hash,
    http_user_agent,
)
from .import_utils import (
    BACKENDS_MAPPING,
    ENV_VARS_TRUE_AND_AUTO_VALUES,
    ENV_VARS_TRUE_VALUES,
    USE_JAX,
    USE_TF,
    USE_TORCH,
    DummyObject,
    OptionalDependencyNotAvailable,
    is_accelerate_available,
    is_accelerate_version,
    is_bsa_available,
    is_flax_available,
    is_ftfy_available,
    is_inflect_available,
    is_invisible_watermark_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_librosa_available,
    is_note_seq_available,
    is_omegaconf_available,
    is_onnx_available,
    is_safetensors_available,
    is_scipy_available,
    is_tensorboard_available,
    is_tf_available,
    is_torch_available,
    is_torch_version,
    is_torchsde_available,
    is_transformers_available,
    is_transformers_version,
    is_unidecode_available,
    is_wandb_available,
    is_xformers_available,
    requires_backends,
)
from .logging import get_logger
from .outputs import BaseOutput
from .pil_utils import PIL_INTERPOLATION, numpy_to_pil, pt_to_pil
from .torch_utils import is_compiled_module, randn_tensor


if is_torch_available():
		from .testing_utils import (
		    floats_tensor,
		    load_hf_numpy,
		    load_image,
		    load_numpy,
		    load_pt,
		    nightly,
		    parse_flag_from_env,
		    print_tensor_test,
		    require_torch_a,
		    require_torch_gpu,
		    skip_mps,
		    slow,
		    torch_all_close,
		    torch_device,
		)
		from .torch_utils import maybe_allow_in_graph

from .testing_utils import export_to_gif, export_to_video


lowerCamelCase		=						get_logger(__name__)






def 	lowerCamelCase_       (			_a			):




		"""simple docstring"""






		if version.parse(_a			) < version.parse(_a			):
				if "dev" in min_version:
						lowerCAmelCase__  :     Tuple       	=   (
						    '''This example requires a source install from HuggingFace diffusers (see '''
						    '''`https://huggingface.co/docs/diffusers/installation#install-from-source`),'''
						)
				else:
						lowerCAmelCase__  :     Dict       	=   f'This example requires a minimum version of {min_version},'
				error_message += f' but the version found is {__version__}.\n'
				raise ImportError(_a			)




