


import logging
import os
import sys
import threading
from logging import (
    CRITICAL,  # NOQA
    DEBUG,  # NOQA
    ERROR,  # NOQA
    FATAL,  # NOQA
    INFO,  # NOQA
    NOTSET,  # NOQA
    WARN,  # NOQA
    WARNING,  # NOQA
)
from typing import Optional

from tqdm import auto as tqdm_lib


snake_case__  				=			threading.Lock()
snake_case__  				=			None

snake_case__  				=			{
    'debug': logging.DEBUG,
    'info': logging.INFO,
    'warning': logging.WARNING,
    'error': logging.ERROR,
    'critical': logging.CRITICAL,
}

snake_case__  				=			logging.WARNING

snake_case__  				=			True
def 			__SCREAMING_SNAKE_CASE							(					)							->	str:
						__magic_name__					:				str				=   os.getenv('DIFFUSERS_VERBOSITY'   ,    lowercase__					)
						if env_level_str:
												if env_level_str in log_levels:
																		return log_levels[env_level_str]
												else:
																		logging.getLogger().warning(
																		    f"Unknown option DIFFUSERS_VERBOSITY={env_level_str}, "
																		    f"has to be one of: { ', '.join(log_levels.keys()					) }"					)
						return _default_log_level
def 			__SCREAMING_SNAKE_CASE							(					)							->	str:
						return __name__.split('.'					)[0]
def 			__SCREAMING_SNAKE_CASE							(					)							->	logging.Logger:
						return logging.getLogger(_get_library_name()					)
def 			__SCREAMING_SNAKE_CASE							(					)							->	None:
						global _default_handler

						with _lock:
												if _default_handler:
																		# This library has already configured the library root logger.
																		return
												__magic_name__					:				Any				=   logging.StreamHandler()  # Set sys.stderr as stream.
												__magic_name__					:				Dict				=   sys.stderr.flush

												# Apply our default configuration to the library root logger.
												__magic_name__					:				Union[str, Any]				=   _get_library_root_logger()
												library_root_logger.addHandler(_default_handler					)
												library_root_logger.setLevel(_get_default_logging_level()					)
												__magic_name__					:				str				=   False
def 			__SCREAMING_SNAKE_CASE							(					)							->	None:
						global _default_handler

						with _lock:
												if not _default_handler:
																		return

												__magic_name__					:				Dict				=   _get_library_root_logger()
												library_root_logger.removeHandler(_default_handler					)
												library_root_logger.setLevel(logging.NOTSET					)
												__magic_name__					:				Dict				=   None
def 			__SCREAMING_SNAKE_CASE							(					)							->	str:
						return log_levels
def 			__SCREAMING_SNAKE_CASE							(					lowercase__	:Optional[str] = None					)							->	logging.Logger:

						if name is None:
												__magic_name__					:				Union[str, Any]				=   _get_library_name()

						_configure_library_root_logger()
						return logging.getLogger(lowercase__					)
def 			__SCREAMING_SNAKE_CASE							(					)							->	int:
						_configure_library_root_logger()
						return _get_library_root_logger().getEffectiveLevel()
def 			__SCREAMING_SNAKE_CASE							(					lowercase__	:int					)							->	None:
						_configure_library_root_logger()
						_get_library_root_logger().setLevel(lowercase__					)
def 			__SCREAMING_SNAKE_CASE							(					)							->	Tuple:
						return set_verbosity(lowercase__					)
def 			__SCREAMING_SNAKE_CASE							(					)							->	Tuple:
						return set_verbosity(lowercase__					)
def 			__SCREAMING_SNAKE_CASE							(					)							->	Tuple:
						return set_verbosity(lowercase__					)
def 			__SCREAMING_SNAKE_CASE							(					)							->	Dict:
						return set_verbosity(lowercase__					)
def 			__SCREAMING_SNAKE_CASE							(					)							->	None:
						_configure_library_root_logger()

						assert _default_handler is not None
						_get_library_root_logger().removeHandler(_default_handler					)
def 			__SCREAMING_SNAKE_CASE							(					)							->	None:
						_configure_library_root_logger()

						assert _default_handler is not None
						_get_library_root_logger().addHandler(_default_handler					)
def 			__SCREAMING_SNAKE_CASE							(					lowercase__	:logging.Handler					)							->	None:
						_configure_library_root_logger()

						assert handler is not None
						_get_library_root_logger().addHandler(lowercase__					)
def 			__SCREAMING_SNAKE_CASE							(					lowercase__	:logging.Handler					)							->	None:
						_configure_library_root_logger()

						assert handler is not None and handler not in _get_library_root_logger().handlers
						_get_library_root_logger().removeHandler(lowercase__					)
def 			__SCREAMING_SNAKE_CASE							(					)							->	None:
						_configure_library_root_logger()
						__magic_name__					:				int				=   False
def 			__SCREAMING_SNAKE_CASE							(					)							->	None:
						_configure_library_root_logger()
						__magic_name__					:				List[str]				=   True
def 			__SCREAMING_SNAKE_CASE							(					)							->	None:
						__magic_name__					:				Optional[Any]				=   _get_library_root_logger().handlers

						for handler in handlers:
												__magic_name__					:				Optional[Any]				=   logging.Formatter('[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s'					)
												handler.setFormatter(lowercase__					)
def 			__SCREAMING_SNAKE_CASE							(					)							->	None:
						__magic_name__					:				Dict				=   _get_library_root_logger().handlers

						for handler in handlers:
												handler.setFormatter(lowercase__					)
def 			__SCREAMING_SNAKE_CASE							(					self	:Union[str, Any]   ,    *lowercase__	:List[str]   ,    **lowercase__	:Tuple					)							->	int:
						__magic_name__					:				List[str]				=   os.getenv('DIFFUSERS_NO_ADVISORY_WARNINGS'   ,    lowercase__					)
						if no_advisory_warnings:
												return
						self.warning(*lowercase__   ,    **lowercase__					)


snake_case__  				=			warning_advice



class      snake_case__	:
				def __init__(	self  ,   *_UpperCAmelCase  ,   **_UpperCAmelCase   ):  # pylint: disable=unused-argument




										"""simple docstring"""




										__magic_name__					:				Tuple				=   args[0] if args else None
				def __iter__(	self   ):




										"""simple docstring"""




										return iter(self._iterator   )
				def __getattr__(	self  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										def empty_fn(*_UpperCAmelCase  ,   **_UpperCAmelCase   ):  # pylint: disable=unused-argument
																return

										return empty_fn
				def __enter__(	self   ):




										"""simple docstring"""




										return self





				def __exit__(	self  ,   _UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										return






class      snake_case__	:
				def __call__(	self  ,   *_UpperCAmelCase  ,   **_UpperCAmelCase   ):




										"""simple docstring"""




										if _tqdm_active:
																return tqdm_lib.tqdm(*_UpperCAmelCase  ,   **_UpperCAmelCase   )
										else:
																return EmptyTqdm(*_UpperCAmelCase  ,   **_UpperCAmelCase   )
				def 			UpperCamelCase_    (	self  ,   *_UpperCAmelCase  ,   **_UpperCAmelCase   ):




										"""simple docstring"""




										__magic_name__					:				int				=   None
										if _tqdm_active:
																return tqdm_lib.tqdm.set_lock(*_UpperCAmelCase  ,   **_UpperCAmelCase   )





				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										if _tqdm_active:
																return tqdm_lib.tqdm.get_lock()






snake_case__  				=			_tqdm_cls()
def 			__SCREAMING_SNAKE_CASE							(					)							->	bool:
						global _tqdm_active
						return bool(_tqdm_active					)
def 			__SCREAMING_SNAKE_CASE							(					)							->	int:
						global _tqdm_active
						__magic_name__					:				int				=   True

def 			__SCREAMING_SNAKE_CASE							(					)							->	Optional[int]:
						global _tqdm_active
						__magic_name__					:				Optional[Any]				=   False






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      snake_case__	(   unittest.TestCase	):
				__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 			UpperCamelCase_    (	self  ,   _UpperCAmelCase=0.9999   ):




										"""simple docstring"""




										__magic_name__					:				List[Any]				=   UNetaDConditionModel.from_pretrained(self.model_id  ,   subfolder='unet'   )
										__magic_name__					:				int				=   unet.to(_UpperCAmelCase   )
										__magic_name__					:				Optional[int]				=   EMAModel(unet.parameters()  ,   decay=_UpperCAmelCase  ,   model_cls=_UpperCAmelCase  ,   model_config=unet.config   )
										return unet, ema_unet
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				List[str]				=   torch.randn(
										    self.batch_size  ,   self.num_in_channels  ,   self.latent_height  ,   self.latent_width  ,   generator=self.generator   ).to(_UpperCAmelCase   )
										__magic_name__					:				List[str]				=   torch.randint(0  ,   1000  ,   size=(self.batch_size,)  ,   generator=self.generator   ).to(_UpperCAmelCase   )
										__magic_name__					:				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   ):




										"""simple docstring"""




										__magic_name__					:				Any				=   {}
										for k, param in unet.state_dict().items():
																__magic_name__					:				str				=   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   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				int				=   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   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Union[str, 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()   )
										__magic_name__					:				Any				=   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   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Optional[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.
										__magic_name__					:				Any				=   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.
										__magic_name__					:				Any				=   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   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				List[str]				=   self.get_models()

										# First backprop + EMA
										__magic_name__					:				int				=   self.simulate_backprop(_UpperCAmelCase   )
										ema_unet.step(unet_step_one.parameters()   )
										__magic_name__					:				Any				=   ema_unet.shadow_params

										# Second backprop + EMA
										__magic_name__					:				Any				=   self.simulate_backprop(_UpperCAmelCase   )
										ema_unet.step(unet_step_two.parameters()   )
										__magic_name__					:				Union[str, Any]				=   ema_unet.shadow_params

										for step_one, step_two in zip(_UpperCAmelCase  ,   _UpperCAmelCase   ):
																assert ~torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Optional[Any]				=   self.get_models(decay=0.0   )
										__magic_name__					:				str				=   self.simulate_backprop(_UpperCAmelCase   )
										ema_unet.step(unet_step_one.parameters()   )
										__magic_name__					:				List[Any]				=   ema_unet.shadow_params

										__magic_name__					:				Optional[Any]				=   self.simulate_backprop(_UpperCAmelCase   )
										ema_unet.step(unet_step_two.parameters()   )
										__magic_name__					:				Optional[Any]				=   ema_unet.shadow_params

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





				@skip_mps
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Optional[int]				=   self.get_models()
										__magic_name__		,			__magic_name__		,			__magic_name__					:				str				=   self.get_dummy_inputs()

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

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

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


