
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					):
     lowercase__				  =			"""hf-internal-testing/tiny-stable-diffusion-pipe"""
     lowercase__				  =			1
     lowercase__				  =			77
     lowercase__				  =			32
     lowercase__				  =			4
     lowercase__				  =			lowercase__				  =			64
     lowercase__				  =			torch.manual_seed(0					)

     def 						__lowerCamelCase				(  self							,       SCREAMING_SNAKE_CASE_=0.9_9_9_9    ):
            __snake_case						       =   UNetaDConditionModel.from_pretrained(self.model_id							,       subfolder='unet'    )
            __snake_case						       =   unet.to(lowerCamelCase_    )
            __snake_case						       =   EMAModel(unet.parameters()							,       decay=lowerCamelCase_							,       model_cls=lowerCamelCase_							,       model_config=unet.config    )
            return unet, ema_unet

     def 						__lowerCamelCase				(  self    ):
            __snake_case						       =   torch.randn(
                self.batch_size							,       self.num_in_channels							,       self.latent_height							,       self.latent_width							,       generator=self.generator    ).to(lowerCamelCase_    )
            __snake_case						       =   torch.randint(0							,       1_000							,       size=(self.batch_size,)							,       generator=self.generator    ).to(lowerCamelCase_    )
            __snake_case						       =   torch.randn(
                self.batch_size							,       self.prompt_length							,       self.text_encoder_hidden_dim							,       generator=self.generator    ).to(lowerCamelCase_    )
            return noisy_latents, timesteps, encoder_hidden_states

     def 						__lowerCamelCase				(  self							,       SCREAMING_SNAKE_CASE_    ):
            __snake_case						       =   {}
            for k, param in unet.state_dict().items():
                   __snake_case						       =   torch.randn_like(lowerCamelCase_    ) + (param * torch.randn_like(lowerCamelCase_    ))
                   updated_state_dict.update({k: updated_param}    )
            unet.load_state_dict(lowerCamelCase_    )
            return unet

     def 						__lowerCamelCase				(  self    ):
            __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 						__lowerCamelCase				(  self    ):
            __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()    )
            __snake_case						       =   list(unet.parameters()    )
            for s_param, param in zip(ema_unet.shadow_params							,       lowerCamelCase_    ):
                   assert torch.allclose(lowerCamelCase_							,       lowerCamelCase_    )

            # 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							,       lowerCamelCase_    ):
                   assert torch.allclose(lowerCamelCase_							,       lowerCamelCase_    )

     def 						__lowerCamelCase				(  self    ):
            __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.
            __snake_case						       =   self.simulate_backprop(lowerCamelCase_    )

            # 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						       =   list(unet_pseudo_updated_step_one.parameters()    )
            for s_param, param in zip(ema_unet.shadow_params							,       lowerCamelCase_    ):
                   assert ~torch.allclose(lowerCamelCase_							,       lowerCamelCase_    )

            # 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							,       lowerCamelCase_    ):
                   assert ~torch.allclose(lowerCamelCase_							,       lowerCamelCase_    )

     def 						__lowerCamelCase				(  self    ):
            # If we call EMA step after a backpropagation consecutively for two times,
            # the shadow params from those two steps should be different.
            __snake_case						       =   self.get_models()

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

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

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

     def 						__lowerCamelCase				(  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.
            __snake_case						       =   self.get_models(decay=0.0    )
            __snake_case						       =   self.simulate_backprop(lowerCamelCase_    )
            ema_unet.step(unet_step_one.parameters()    )
            __snake_case						       =   ema_unet.shadow_params

            __snake_case						       =   self.simulate_backprop(lowerCamelCase_    )
            ema_unet.step(unet_step_two.parameters()    )
            __snake_case						       =   ema_unet.shadow_params

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

     @skip_mps
     def 						__lowerCamelCase				(  self    ):
            __snake_case						       =   self.get_models()
            __snake_case						       =   self.get_dummy_inputs()

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

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

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






from typing import TYPE_CHECKING

from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_tf_available, is_torch_available


__lowerCamelCase :    Tuple      						=      {
    '''configuration_utils''': ['''GenerationConfig'''],
    '''streamers''': ['''TextIteratorStreamer''', '''TextStreamer'''],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase :    Any      						=      [
	    '''Constraint''',
	    '''ConstraintListState''',
	    '''DisjunctiveConstraint''',
	    '''PhrasalConstraint''',
	]
	__lowerCamelCase :    int      						=      [
	    '''BeamHypotheses''',
	    '''BeamScorer''',
	    '''BeamSearchScorer''',
	    '''ConstrainedBeamSearchScorer''',
	]
	__lowerCamelCase :    Union[str, Any]      						=      [
	    '''EpsilonLogitsWarper''',
	    '''EtaLogitsWarper''',
	    '''ForcedBOSTokenLogitsProcessor''',
	    '''ForcedEOSTokenLogitsProcessor''',
	    '''HammingDiversityLogitsProcessor''',
	    '''InfNanRemoveLogitsProcessor''',
	    '''LogitsProcessor''',
	    '''LogitsProcessorList''',
	    '''LogitsWarper''',
	    '''MinLengthLogitsProcessor''',
	    '''MinNewTokensLengthLogitsProcessor''',
	    '''NoBadWordsLogitsProcessor''',
	    '''NoRepeatNGramLogitsProcessor''',
	    '''PrefixConstrainedLogitsProcessor''',
	    '''RepetitionPenaltyLogitsProcessor''',
	    '''SequenceBiasLogitsProcessor''',
	    '''EncoderRepetitionPenaltyLogitsProcessor''',
	    '''TemperatureLogitsWarper''',
	    '''TopKLogitsWarper''',
	    '''TopPLogitsWarper''',
	    '''TypicalLogitsWarper''',
	    '''EncoderNoRepeatNGramLogitsProcessor''',
	    '''ExponentialDecayLengthPenalty''',
	    '''LogitNormalization''',
	]
	__lowerCamelCase :    Union[str, Any]      						=      [
	    '''MaxNewTokensCriteria''',
	    '''MaxLengthCriteria''',
	    '''MaxTimeCriteria''',
	    '''StoppingCriteria''',
	    '''StoppingCriteriaList''',
	    '''validate_stopping_criteria''',
	]
	__lowerCamelCase :    Union[str, Any]      						=      [
	    '''GenerationMixin''',
	    '''top_k_top_p_filtering''',
	    '''GreedySearchEncoderDecoderOutput''',
	    '''GreedySearchDecoderOnlyOutput''',
	    '''SampleEncoderDecoderOutput''',
	    '''SampleDecoderOnlyOutput''',
	    '''BeamSearchEncoderDecoderOutput''',
	    '''BeamSearchDecoderOnlyOutput''',
	    '''BeamSampleEncoderDecoderOutput''',
	    '''BeamSampleDecoderOnlyOutput''',
	    '''ContrastiveSearchEncoderDecoderOutput''',
	    '''ContrastiveSearchDecoderOnlyOutput''',
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase :    Tuple      						=      [
	    '''TFForcedBOSTokenLogitsProcessor''',
	    '''TFForcedEOSTokenLogitsProcessor''',
	    '''TFLogitsProcessor''',
	    '''TFLogitsProcessorList''',
	    '''TFLogitsWarper''',
	    '''TFMinLengthLogitsProcessor''',
	    '''TFNoBadWordsLogitsProcessor''',
	    '''TFNoRepeatNGramLogitsProcessor''',
	    '''TFRepetitionPenaltyLogitsProcessor''',
	    '''TFTemperatureLogitsWarper''',
	    '''TFTopKLogitsWarper''',
	    '''TFTopPLogitsWarper''',
	    '''TFForceTokensLogitsProcessor''',
	    '''TFSuppressTokensAtBeginLogitsProcessor''',
	    '''TFSuppressTokensLogitsProcessor''',
	]
	__lowerCamelCase :    Any      						=      [
	    '''TFGenerationMixin''',
	    '''tf_top_k_top_p_filtering''',
	    '''TFGreedySearchDecoderOnlyOutput''',
	    '''TFGreedySearchEncoderDecoderOutput''',
	    '''TFSampleEncoderDecoderOutput''',
	    '''TFSampleDecoderOnlyOutput''',
	    '''TFBeamSearchEncoderDecoderOutput''',
	    '''TFBeamSearchDecoderOnlyOutput''',
	    '''TFBeamSampleEncoderDecoderOutput''',
	    '''TFBeamSampleDecoderOnlyOutput''',
	    '''TFContrastiveSearchEncoderDecoderOutput''',
	    '''TFContrastiveSearchDecoderOnlyOutput''',
	]

try:
	if not is_flax_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase :    Optional[int]      						=      [
	    '''FlaxForcedBOSTokenLogitsProcessor''',
	    '''FlaxForcedEOSTokenLogitsProcessor''',
	    '''FlaxLogitsProcessor''',
	    '''FlaxLogitsProcessorList''',
	    '''FlaxLogitsWarper''',
	    '''FlaxMinLengthLogitsProcessor''',
	    '''FlaxTemperatureLogitsWarper''',
	    '''FlaxTopKLogitsWarper''',
	    '''FlaxTopPLogitsWarper''',
	]
	__lowerCamelCase :    Union[str, Any]      						=      [
	    '''FlaxGenerationMixin''',
	    '''FlaxGreedySearchOutput''',
	    '''FlaxSampleOutput''',
	    '''FlaxBeamSearchOutput''',
	]

if TYPE_CHECKING:
	from .configuration_utils import GenerationConfig
	from .streamers import TextIteratorStreamer, TextStreamer

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint
		from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
		from .logits_process import (
		    EncoderNoRepeatNGramLogitsProcessor,
		    EncoderRepetitionPenaltyLogitsProcessor,
		    EpsilonLogitsWarper,
		    EtaLogitsWarper,
		    ExponentialDecayLengthPenalty,
		    ForcedBOSTokenLogitsProcessor,
		    ForcedEOSTokenLogitsProcessor,
		    HammingDiversityLogitsProcessor,
		    InfNanRemoveLogitsProcessor,
		    LogitNormalization,
		    LogitsProcessor,
		    LogitsProcessorList,
		    LogitsWarper,
		    MinLengthLogitsProcessor,
		    MinNewTokensLengthLogitsProcessor,
		    NoBadWordsLogitsProcessor,
		    NoRepeatNGramLogitsProcessor,
		    PrefixConstrainedLogitsProcessor,
		    RepetitionPenaltyLogitsProcessor,
		    SequenceBiasLogitsProcessor,
		    TemperatureLogitsWarper,
		    TopKLogitsWarper,
		    TopPLogitsWarper,
		    TypicalLogitsWarper,
		)
		from .stopping_criteria import (
		    MaxLengthCriteria,
		    MaxNewTokensCriteria,
		    MaxTimeCriteria,
		    StoppingCriteria,
		    StoppingCriteriaList,
		    validate_stopping_criteria,
		)
		from .utils import (
		    BeamSampleDecoderOnlyOutput,
		    BeamSampleEncoderDecoderOutput,
		    BeamSearchDecoderOnlyOutput,
		    BeamSearchEncoderDecoderOutput,
		    ContrastiveSearchDecoderOnlyOutput,
		    ContrastiveSearchEncoderDecoderOutput,
		    GenerationMixin,
		    GreedySearchDecoderOnlyOutput,
		    GreedySearchEncoderDecoderOutput,
		    SampleDecoderOnlyOutput,
		    SampleEncoderDecoderOutput,
		    top_k_top_p_filtering,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tf_logits_process import (
		    TFForcedBOSTokenLogitsProcessor,
		    TFForcedEOSTokenLogitsProcessor,
		    TFForceTokensLogitsProcessor,
		    TFLogitsProcessor,
		    TFLogitsProcessorList,
		    TFLogitsWarper,
		    TFMinLengthLogitsProcessor,
		    TFNoBadWordsLogitsProcessor,
		    TFNoRepeatNGramLogitsProcessor,
		    TFRepetitionPenaltyLogitsProcessor,
		    TFSuppressTokensAtBeginLogitsProcessor,
		    TFSuppressTokensLogitsProcessor,
		    TFTemperatureLogitsWarper,
		    TFTopKLogitsWarper,
		    TFTopPLogitsWarper,
		)
		from .tf_utils import (
		    TFBeamSampleDecoderOnlyOutput,
		    TFBeamSampleEncoderDecoderOutput,
		    TFBeamSearchDecoderOnlyOutput,
		    TFBeamSearchEncoderDecoderOutput,
		    TFContrastiveSearchDecoderOnlyOutput,
		    TFContrastiveSearchEncoderDecoderOutput,
		    TFGenerationMixin,
		    TFGreedySearchDecoderOnlyOutput,
		    TFGreedySearchEncoderDecoderOutput,
		    TFSampleDecoderOnlyOutput,
		    TFSampleEncoderDecoderOutput,
		    tf_top_k_top_p_filtering,
		)

	try:
		if not is_flax_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .flax_logits_process import (
		    FlaxForcedBOSTokenLogitsProcessor,
		    FlaxForcedEOSTokenLogitsProcessor,
		    FlaxLogitsProcessor,
		    FlaxLogitsProcessorList,
		    FlaxLogitsWarper,
		    FlaxMinLengthLogitsProcessor,
		    FlaxTemperatureLogitsWarper,
		    FlaxTopKLogitsWarper,
		    FlaxTopPLogitsWarper,
		)
		from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput
else:
	import sys

	__lowerCamelCase :    int      						=      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



