






import contextlib
import copy
import random
from typing import Any, Dict, Iterable, Optional, Union

import numpy as np
import torch

from .utils import deprecate, is_transformers_available


if is_transformers_available():
	import transformers
def       snake_case_		(     lowerCAmelCase_       :						int      ):
							random.seed(lowerCAmelCase_      )
							np.random.seed(lowerCAmelCase_      )
							torch.manual_seed(lowerCAmelCase_      )
							torch.cuda.manual_seed_all(lowerCAmelCase_      )
							# ^^ safe to call this function even if cuda is not available



class 		lowerCAmelCase     :

					'''simple docstring'''

					def __init__(  self				:	Dict		,     __a				:	Iterable[torch.nn.Parameter]		,     __a				:	float = 0.9999		,     __a				:	float = 0.0		,     __a				:	int = 0		,     __a				:	bool = False		,     __a				:	Union[float, int] = 1.0		,     __a				:	Union[float, int] = 2 / 3		,     __a				:	Optional[Any] = None		,     __a				:	Dict[str, Any] = None		,     **__a				:	Optional[Any]		,     )			->  int:


												"""simple docstring"""

												if isinstance(__a		,     torch.nn.Module							):
																			__lowercase      :							Tuple				=     (
																			    """Passing a `torch.nn.Module` to `ExponentialMovingAverage` is deprecated. """
																			    """Please pass the parameters of the module instead."""
																			)
																			deprecate(
																			    """passing a `torch.nn.Module` to `ExponentialMovingAverage`"""		,     """1.0.0"""		,     __a		,     standard_warn=__a		,     )
																			__lowercase      :							Dict				=     parameters.parameters()

																			# set use_ema_warmup to True if a torch.nn.Module is passed for backwards compatibility
																			__lowercase      :							str				=     True

												if kwargs.get("""max_value"""		,     __a							) is not None:
																			__lowercase      :							Optional[int]				=     """The `max_value` argument is deprecated. Please use `decay` instead."""
																			deprecate("""max_value"""		,     """1.0.0"""		,     __a		,     standard_warn=__a							)
																			__lowercase      :							Any				=     kwargs["""max_value"""]

												if kwargs.get("""min_value"""		,     __a							) is not None:
																			__lowercase      :							List[str]				=     """The `min_value` argument is deprecated. Please use `min_decay` instead."""
																			deprecate("""min_value"""		,     """1.0.0"""		,     __a		,     standard_warn=__a							)
																			__lowercase      :							Any				=     kwargs["""min_value"""]

												__lowercase      :							Any				=     list(__a							)
												__lowercase      :							Any				=     [p.clone().detach() for p in parameters]

												if kwargs.get("""device"""		,     __a							) is not None:
																			__lowercase      :							Tuple				=     """The `device` argument is deprecated. Please use `to` instead."""
																			deprecate("""device"""		,     """1.0.0"""		,     __a		,     standard_warn=__a							)
																			self.to(device=kwargs["""device"""]							)

												__lowercase      :							Optional[Any]				=     None

												__lowercase      :							str				=     decay
												__lowercase      :							Union[str, Any]				=     min_decay
												__lowercase      :							int				=     update_after_step
												__lowercase      :							Any				=     use_ema_warmup
												__lowercase      :							List[str]				=     inv_gamma
												__lowercase      :							Dict				=     power
												__lowercase      :							str				=     0
												__lowercase      :							Optional[int]				=     None  # set in `step()`

												__lowercase      :							Any				=     model_cls
												__lowercase      :							Optional[Any]				=     model_config



					@classmethod
					def       lowerCAmelCase				(  cls				:	List[str]		,     __a				:	int		,     __a				:	Any							)			->  "EMAModel":


												"""simple docstring"""

												__lowercase			,			__lowercase      :							str				=     model_cls.load_config(__a		,     return_unused_kwargs=__a							)
												__lowercase      :							List[Any]				=     model_cls.from_pretrained(__a							)

												__lowercase      :							Union[str, Any]				=     cls(model.parameters()		,     model_cls=__a		,     model_config=model.config							)

												ema_model.load_state_dict(__a							)
												return ema_model



					def       lowerCAmelCase				(  self				:	Optional[int]		,     __a				:	str							)			->  str:


												"""simple docstring"""

												if self.model_cls is None:
																			raise ValueError("""`save_pretrained` can only be used if `model_cls` was defined at __init__."""							)

												if self.model_config is None:
																			raise ValueError("""`save_pretrained` can only be used if `model_config` was defined at __init__."""							)

												__lowercase      :							List[Any]				=     self.model_cls.from_config(self.model_config							)
												__lowercase      :							Optional[int]				=     self.state_dict()
												state_dict.pop("""shadow_params"""		,     __a							)

												model.register_to_config(**__a							)
												self.copy_to(model.parameters()							)
												model.save_pretrained(__a							)



					def       lowerCAmelCase				(  self				:	Union[str, Any]		,     __a				:	int							)			->  float:


												"""simple docstring"""

												__lowercase      :							Optional[Any]				=     max(0		,     optimization_step - self.update_after_step - 1							)

												if step <= 0:
																			return 0.0

												if self.use_ema_warmup:
																			__lowercase      :							Tuple				=     1 - (1 + step / self.inv_gamma) ** -self.power
												else:
																			__lowercase      :							Optional[int]				=     (1 + step) / (10 + step)

												__lowercase      :							int				=     min(__a		,     self.decay							)
												# make sure decay is not smaller than min_decay
												__lowercase      :							Optional[Any]				=     max(__a		,     self.min_decay							)
												return cur_decay_value



					@torch.no_grad()
					def       lowerCAmelCase				(  self				:	List[Any]		,     __a				:	Iterable[torch.nn.Parameter]							)			->  List[Any]:


												"""simple docstring"""

												if isinstance(__a		,     torch.nn.Module							):
																			__lowercase      :							Tuple				=     (
																			    """Passing a `torch.nn.Module` to `ExponentialMovingAverage.step` is deprecated. """
																			    """Please pass the parameters of the module instead."""
																			)
																			deprecate(
																			    """passing a `torch.nn.Module` to `ExponentialMovingAverage.step`"""		,     """1.0.0"""		,     __a		,     standard_warn=__a		,     )
																			__lowercase      :							Union[str, Any]				=     parameters.parameters()

												__lowercase      :							Any				=     list(__a							)

												self.optimization_step += 1

												# Compute the decay factor for the exponential moving average.
												__lowercase      :							Union[str, Any]				=     self.get_decay(self.optimization_step							)
												__lowercase      :							List[Any]				=     decay
												__lowercase      :							List[str]				=     1 - decay

												__lowercase      :							Any				=     contextlib.nullcontext
												if is_transformers_available() and transformers.deepspeed.is_deepspeed_zeroa_enabled():
																			import deepspeed

												for s_param, param in zip(self.shadow_params		,     __a							):
																			if is_transformers_available() and transformers.deepspeed.is_deepspeed_zeroa_enabled():
																										__lowercase      :							Optional[int]				=     deepspeed.zero.GatheredParameters(__a		,     modifier_rank=__a							)

																			with context_manager():
																										if param.requires_grad:
																																	s_param.sub_(one_minus_decay * (s_param - param)							)
																										else:
																																	s_param.copy_(__a							)



					def       lowerCAmelCase				(  self				:	Union[str, Any]		,     __a				:	Iterable[torch.nn.Parameter]							)			->  None:


												"""simple docstring"""

												__lowercase      :							Tuple				=     list(__a							)
												for s_param, param in zip(self.shadow_params		,     __a							):
																			param.data.copy_(s_param.to(param.device							).data							)



					def       lowerCAmelCase				(  self				:	Any		,     __a				:	List[Any]=None		,     __a				:	List[str]=None							)			->  None:


												"""simple docstring"""

												__lowercase      :							List[str]				=     [
												    p.to(device=__a		,     dtype=__a							) if p.is_floating_point() else p.to(device=__a							)
												    for p in self.shadow_params
												]



					def       lowerCAmelCase				(  self				:	Tuple							)			->  dict:


												"""simple docstring"""

												return {
												    "decay": self.decay,
												    "min_decay": self.min_decay,
												    "optimization_step": self.optimization_step,
												    "update_after_step": self.update_after_step,
												    "use_ema_warmup": self.use_ema_warmup,
												    "inv_gamma": self.inv_gamma,
												    "power": self.power,
												    "shadow_params": self.shadow_params,
												}



					def       lowerCAmelCase				(  self				:	Union[str, Any]		,     __a				:	Iterable[torch.nn.Parameter]							)			->  None:


												"""simple docstring"""

												__lowercase      :							int				=     [param.detach().cpu().clone() for param in parameters]



					def       lowerCAmelCase				(  self				:	List[Any]		,     __a				:	Iterable[torch.nn.Parameter]							)			->  None:


												"""simple docstring"""

												if self.temp_stored_params is None:
																			raise RuntimeError("""This ExponentialMovingAverage has no `store()`ed weights """ """to `restore()`"""							)
												for c_param, param in zip(self.temp_stored_params		,     __a							):
																			param.data.copy_(c_param.data							)

												# Better memory-wise.
												__lowercase      :							Union[str, Any]				=     None



					def       lowerCAmelCase				(  self				:	List[Any]		,     __a				:	dict							)			->  None:


												"""simple docstring"""

												__lowercase      :							Union[str, Any]				=     copy.deepcopy(__a							)

												__lowercase      :							List[str]				=     state_dict.get("""decay"""		,     self.decay							)
												if self.decay < 0.0 or self.decay > 1.0:
																			raise ValueError("""Decay must be between 0 and 1"""							)

												__lowercase      :							int				=     state_dict.get("""min_decay"""		,     self.min_decay							)
												if not isinstance(self.min_decay		,     __a							):
																			raise ValueError("""Invalid min_decay"""							)

												__lowercase      :							Tuple				=     state_dict.get("""optimization_step"""		,     self.optimization_step							)
												if not isinstance(self.optimization_step		,     __a							):
																			raise ValueError("""Invalid optimization_step"""							)

												__lowercase      :							int				=     state_dict.get("""update_after_step"""		,     self.update_after_step							)
												if not isinstance(self.update_after_step		,     __a							):
																			raise ValueError("""Invalid update_after_step"""							)

												__lowercase      :							Tuple				=     state_dict.get("""use_ema_warmup"""		,     self.use_ema_warmup							)
												if not isinstance(self.use_ema_warmup		,     __a							):
																			raise ValueError("""Invalid use_ema_warmup"""							)

												__lowercase      :							Dict				=     state_dict.get("""inv_gamma"""		,     self.inv_gamma							)
												if not isinstance(self.inv_gamma		,     (float, int)							):
																			raise ValueError("""Invalid inv_gamma"""							)

												__lowercase      :							str				=     state_dict.get("""power"""		,     self.power							)
												if not isinstance(self.power		,     (float, int)							):
																			raise ValueError("""Invalid power"""							)

												__lowercase      :							Dict				=     state_dict.get("""shadow_params"""		,     __a							)
												if shadow_params is not None:
																			__lowercase      :							int				=     shadow_params
																			if not isinstance(self.shadow_params		,     __a							):
																										raise ValueError("""shadow_params must be a list"""							)
																			if not all(isinstance(__a		,     torch.Tensor							) for p in self.shadow_params							):
																										raise ValueError("""shadow_params must all be Tensors"""							)







lowerCamelCase : List[str]						   =       '''0.18.2'''

from .configuration_utils import ConfigMixin
from .utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_inflect_available,
    is_invisible_watermark_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_librosa_available,
    is_note_seq_available,
    is_onnx_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
    is_transformers_available,
    is_transformers_version,
    is_unidecode_available,
    logging,
)


try:
	if not is_onnx_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_onnx_objects import *  # noqa F403
else:
	from .pipelines import OnnxRuntimeModel

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_pt_objects import *  # noqa F403
else:
	from .models import (
	    AutoencoderKL,
	    ControlNetModel,
	    ModelMixin,
	    PriorTransformer,
	    TaFilmDecoder,
	    TransformeraDModel,
	    UNetaDModel,
	    UNetaDConditionModel,
	    UNetaDModel,
	    UNetaDConditionModel,
	    VQModel,
	)
	from .optimization import (
	    get_constant_schedule,
	    get_constant_schedule_with_warmup,
	    get_cosine_schedule_with_warmup,
	    get_cosine_with_hard_restarts_schedule_with_warmup,
	    get_linear_schedule_with_warmup,
	    get_polynomial_decay_schedule_with_warmup,
	    get_scheduler,
	)
	from .pipelines import (
	    AudioPipelineOutput,
	    ConsistencyModelPipeline,
	    DanceDiffusionPipeline,
	    DDIMPipeline,
	    DDPMPipeline,
	    DiffusionPipeline,
	    DiTPipeline,
	    ImagePipelineOutput,
	    KarrasVePipeline,
	    LDMPipeline,
	    LDMSuperResolutionPipeline,
	    PNDMPipeline,
	    RePaintPipeline,
	    ScoreSdeVePipeline,
	)
	from .schedulers import (
	    CMStochasticIterativeScheduler,
	    DDIMInverseScheduler,
	    DDIMParallelScheduler,
	    DDIMScheduler,
	    DDPMParallelScheduler,
	    DDPMScheduler,
	    DEISMultistepScheduler,
	    DPMSolverMultistepInverseScheduler,
	    DPMSolverMultistepScheduler,
	    DPMSolverSinglestepScheduler,
	    EulerAncestralDiscreteScheduler,
	    EulerDiscreteScheduler,
	    HeunDiscreteScheduler,
	    IPNDMScheduler,
	    KarrasVeScheduler,
	    KDPMaAncestralDiscreteScheduler,
	    KDPMaDiscreteScheduler,
	    PNDMScheduler,
	    RePaintScheduler,
	    SchedulerMixin,
	    ScoreSdeVeScheduler,
	    UnCLIPScheduler,
	    UniPCMultistepScheduler,
	    VQDiffusionScheduler,
	)
	from .training_utils import EMAModel

try:
	if not (is_torch_available() and is_scipy_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_torch_and_scipy_objects import *  # noqa F403
else:
	from .schedulers import LMSDiscreteScheduler

try:
	if not (is_torch_available() and is_torchsde_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_torch_and_torchsde_objects import *  # noqa F403
else:
	from .schedulers import DPMSolverSDEScheduler

try:
	if not (is_torch_available() and is_transformers_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
	from .pipelines import (
	    AltDiffusionImgaImgPipeline,
	    AltDiffusionPipeline,
	    AudioLDMPipeline,
	    CycleDiffusionPipeline,
	    IFImgaImgPipeline,
	    IFImgaImgSuperResolutionPipeline,
	    IFInpaintingPipeline,
	    IFInpaintingSuperResolutionPipeline,
	    IFPipeline,
	    IFSuperResolutionPipeline,
	    ImageTextPipelineOutput,
	    KandinskyImgaImgPipeline,
	    KandinskyInpaintPipeline,
	    KandinskyPipeline,
	    KandinskyPriorPipeline,
	    KandinskyVaaControlnetImgaImgPipeline,
	    KandinskyVaaControlnetPipeline,
	    KandinskyVaaImgaImgPipeline,
	    KandinskyVaaInpaintPipeline,
	    KandinskyVaaPipeline,
	    KandinskyVaaPriorEmbaEmbPipeline,
	    KandinskyVaaPriorPipeline,
	    LDMTextToImagePipeline,
	    PaintByExamplePipeline,
	    SemanticStableDiffusionPipeline,
	    ShapEImgaImgPipeline,
	    ShapEPipeline,
	    StableDiffusionAttendAndExcitePipeline,
	    StableDiffusionControlNetImgaImgPipeline,
	    StableDiffusionControlNetInpaintPipeline,
	    StableDiffusionControlNetPipeline,
	    StableDiffusionDepthaImgPipeline,
	    StableDiffusionDiffEditPipeline,
	    StableDiffusionImageVariationPipeline,
	    StableDiffusionImgaImgPipeline,
	    StableDiffusionInpaintPipeline,
	    StableDiffusionInpaintPipelineLegacy,
	    StableDiffusionInstructPixaPixPipeline,
	    StableDiffusionLatentUpscalePipeline,
	    StableDiffusionLDMaDPipeline,
	    StableDiffusionModelEditingPipeline,
	    StableDiffusionPanoramaPipeline,
	    StableDiffusionParadigmsPipeline,
	    StableDiffusionPipeline,
	    StableDiffusionPipelineSafe,
	    StableDiffusionPixaPixZeroPipeline,
	    StableDiffusionSAGPipeline,
	    StableDiffusionUpscalePipeline,
	    StableUnCLIPImgaImgPipeline,
	    StableUnCLIPPipeline,
	    TextToVideoSDPipeline,
	    TextToVideoZeroPipeline,
	    UnCLIPImageVariationPipeline,
	    UnCLIPPipeline,
	    UniDiffuserModel,
	    UniDiffuserPipeline,
	    UniDiffuserTextDecoder,
	    VersatileDiffusionDualGuidedPipeline,
	    VersatileDiffusionImageVariationPipeline,
	    VersatileDiffusionPipeline,
	    VersatileDiffusionTextToImagePipeline,
	    VideoToVideoSDPipeline,
	    VQDiffusionPipeline,
	)

try:
	if not (is_torch_available() and is_transformers_available() and is_invisible_watermark_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_torch_and_transformers_and_invisible_watermark_objects import *  # noqa F403
else:
	from .pipelines import StableDiffusionXLImgaImgPipeline, StableDiffusionXLPipeline

try:
	if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
else:
	from .pipelines import StableDiffusionKDiffusionPipeline

try:
	if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_torch_and_transformers_and_onnx_objects import *  # noqa F403
else:
	from .pipelines import (
	    OnnxStableDiffusionImgaImgPipeline,
	    OnnxStableDiffusionInpaintPipeline,
	    OnnxStableDiffusionInpaintPipelineLegacy,
	    OnnxStableDiffusionPipeline,
	    OnnxStableDiffusionUpscalePipeline,
	    StableDiffusionOnnxPipeline,
	)

try:
	if not (is_torch_available() and is_librosa_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_torch_and_librosa_objects import *  # noqa F403
else:
	from .pipelines import AudioDiffusionPipeline, Mel

try:
	if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
else:
	from .pipelines import SpectrogramDiffusionPipeline

try:
	if not is_flax_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_flax_objects import *  # noqa F403
else:
	from .models.controlnet_flax import FlaxControlNetModel
	from .models.modeling_flax_utils import FlaxModelMixin
	from .models.unet_ad_condition_flax import FlaxUNetaDConditionModel
	from .models.vae_flax import FlaxAutoencoderKL
	from .pipelines import FlaxDiffusionPipeline
	from .schedulers import (
	    FlaxDDIMScheduler,
	    FlaxDDPMScheduler,
	    FlaxDPMSolverMultistepScheduler,
	    FlaxKarrasVeScheduler,
	    FlaxLMSDiscreteScheduler,
	    FlaxPNDMScheduler,
	    FlaxSchedulerMixin,
	    FlaxScoreSdeVeScheduler,
	)


try:
	if not (is_flax_available() and is_transformers_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
else:
	from .pipelines import (
	    FlaxStableDiffusionControlNetPipeline,
	    FlaxStableDiffusionImgaImgPipeline,
	    FlaxStableDiffusionInpaintPipeline,
	    FlaxStableDiffusionPipeline,
	)

try:
	if not (is_note_seq_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_note_seq_objects import *  # noqa F403
else:
	from .pipelines import MidiProcessor