






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





    def 							lowerCamelCase__ (				self      :   Tuple							,						A      :   int=0.99_99     ):


      '''simple docstring'''




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





    def 							lowerCamelCase__ (				self      :   Dict     ):


      '''simple docstring'''




      a	:     List[str]							=  torch.randn(
          self.batch_size							,						self.num_in_channels							,						self.latent_height							,						self.latent_width							,						generator=self.generator     ).to(_UpperCAmelCase     )
      a	:     List[str]							=  torch.randint(0							,						1_0_0_0							,						size=(self.batch_size,)							,						generator=self.generator     ).to(_UpperCAmelCase     )
      a	:     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 							lowerCamelCase__ (				self      :   Dict							,						A      :   Any     ):


      '''simple docstring'''




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


      '''simple docstring'''




      a	:     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 							lowerCamelCase__ (				self      :   List[str]     ):


      '''simple docstring'''




      a	:     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()     )
      a	:     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 							lowerCamelCase__ (				self      :   List[Any]     ):


      '''simple docstring'''




      a	:     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.
      a	:     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.
      a	:     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 							lowerCamelCase__ (				self      :   List[Any]     ):


      '''simple docstring'''




      a	:     List[str]							=  self.get_models()

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

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

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





    def 							lowerCamelCase__ (				self      :   Union[str, Any]     ):


      '''simple docstring'''




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

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

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





    @skip_mps
    def 							lowerCamelCase__ (				self      :   Union[str, Any]     ):


      '''simple docstring'''




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

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

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

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








# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
    is_vision_available,
)


snake_case__  				=			{
    'configuration_tvlt': ['TVLT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TvltConfig'],
    'feature_extraction_tvlt': ['TvltFeatureExtractor'],
    'processing_tvlt': ['TvltProcessor'],
}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			snake_case__  				=			[
			    'TVLT_PRETRAINED_MODEL_ARCHIVE_LIST',
			    'TvltModel',
			    'TvltForPreTraining',
			    'TvltForAudioVisualClassification',
			    'TvltPreTrainedModel',
			]

try:
			if not is_vision_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			snake_case__  				=			['TvltImageProcessor']


if TYPE_CHECKING:
			from .configuration_tvlt import TVLT_PRETRAINED_CONFIG_ARCHIVE_MAP, TvltConfig
			from .processing_tvlt import TvltProcessor
			from .feature_extraction_tvlt import TvltFeatureExtractor

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tvlt import (
						    TVLT_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TvltForAudioVisualClassification,
						    TvltForPreTraining,
						    TvltModel,
						    TvltPreTrainedModel,
						)

			try:
						if not is_vision_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .image_processing_tvlt import TvltImageProcessor


else:
			import sys

			snake_case__  				=			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


