'''simple docstring'''




import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


snake_case	=      logging.get_logger(__name__)

enable_full_determinism()





class 							SCREAMING_SNAKE_CASE (					__a						,    __a						,    unittest.TestCase			):
  """simple docstring"""







  __A  =		UNetaDModel
  __A  =		"""sample"""



  @property
  def 	a (	self :       List[str]       ):


        """simple docstring"""



        _lowerCAmelCase		 = 4
        _lowerCAmelCase		 = 3
        _lowerCAmelCase		 = (32, 32)

        _lowerCAmelCase		 = floats_tensor((batch_size, num_channels) + sizes       ).to(__lowerCAmelCase       )
        _lowerCAmelCase		 = torch.tensor([10]       ).to(__lowerCAmelCase       )

        return {"sample": noise, "timestep": time_step}



  @property
  def 	a (	self :       Any       ):


        """simple docstring"""



        return (3, 32, 32)



  @property
  def 	a (	self :       Optional[Any]       ):


        """simple docstring"""



        return (3, 32, 32)



  def 	a (	self :       Tuple       ):


        """simple docstring"""



        _lowerCAmelCase		 = {
            'block_out_channels': (32, 64),
            'down_block_types': ('DownBlock2D', 'AttnDownBlock2D'),
            'up_block_types': ('AttnUpBlock2D', 'UpBlock2D'),
            'attention_head_dim': 3,
            'out_channels': 3,
            'in_channels': 3,
            'layers_per_block': 2,
            'sample_size': 32,
        }
        _lowerCAmelCase		 = self.dummy_input
        return init_dict, inputs_dict





class 							SCREAMING_SNAKE_CASE (					__a						,    __a						,    unittest.TestCase			):
  """simple docstring"""







  __A  =		UNetaDModel
  __A  =		"""sample"""



  @property
  def 	a (	self :       str       ):


        """simple docstring"""



        _lowerCAmelCase		 = 4
        _lowerCAmelCase		 = 4
        _lowerCAmelCase		 = (32, 32)

        _lowerCAmelCase		 = floats_tensor((batch_size, num_channels) + sizes       ).to(__lowerCAmelCase       )
        _lowerCAmelCase		 = torch.tensor([10]       ).to(__lowerCAmelCase       )

        return {"sample": noise, "timestep": time_step}



  @property
  def 	a (	self :       Any       ):


        """simple docstring"""



        return (4, 32, 32)



  @property
  def 	a (	self :       Optional[Any]       ):


        """simple docstring"""



        return (4, 32, 32)



  def 	a (	self :       Optional[Any]       ):


        """simple docstring"""



        _lowerCAmelCase		 = {
            'sample_size': 32,
            'in_channels': 4,
            'out_channels': 4,
            'layers_per_block': 2,
            'block_out_channels': (32, 64),
            'attention_head_dim': 32,
            'down_block_types': ('DownBlock2D', 'DownBlock2D'),
            'up_block_types': ('UpBlock2D', 'UpBlock2D'),
        }
        _lowerCAmelCase		 = self.dummy_input
        return init_dict, inputs_dict



  def 	a (	self :       int       ):


        """simple docstring"""



        _lowerCAmelCase		 = UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'       ,       output_loading_info=__lowerCAmelCase       )

        self.assertIsNotNone(__lowerCAmelCase       )
        self.assertEqual(len(loading_info['missing_keys']       )       ,       0       )

        model.to(__lowerCAmelCase       )
        _lowerCAmelCase		 = model(**self.dummy_input       ).sample

        assert image is not None, "Make sure output is not None"



  @unittest.skipIf(torch_device != 'cuda'       ,       'This test is supposed to run on GPU'       )
  def 	a (	self :       str       ):


        """simple docstring"""



        _lowerCAmelCase		 = UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'       ,       output_loading_info=__lowerCAmelCase       )
        model.to(__lowerCAmelCase       )
        _lowerCAmelCase		 = model(**self.dummy_input       ).sample

        assert image is not None, "Make sure output is not None"



  @unittest.skipIf(torch_device != 'cuda'       ,       'This test is supposed to run on GPU'       )
  def 	a (	self :       Union[str, Any]       ):


        """simple docstring"""



        _lowerCAmelCase		 = UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'       ,       output_loading_info=__lowerCAmelCase       )
        model_accelerate.to(__lowerCAmelCase       )
        model_accelerate.eval()

        _lowerCAmelCase		 = torch.randn(
            1       ,       model_accelerate.config.in_channels       ,       model_accelerate.config.sample_size       ,       model_accelerate.config.sample_size       ,       generator=torch.manual_seed(0       )       ,       )
        _lowerCAmelCase		 = noise.to(__lowerCAmelCase       )
        _lowerCAmelCase		 = torch.tensor([10] * noise.shape[0]       ).to(__lowerCAmelCase       )

        _lowerCAmelCase		 = model_accelerate(__lowerCAmelCase       ,       __lowerCAmelCase       )['sample']

        # two models don't need to stay in the device at the same time
        del model_accelerate
        torch.cuda.empty_cache()
        gc.collect()

        _lowerCAmelCase		 = UNetaDModel.from_pretrained(
            'fusing/unet-ldm-dummy-update'       ,       output_loading_info=__lowerCAmelCase       ,       low_cpu_mem_usage=__lowerCAmelCase       )
        model_normal_load.to(__lowerCAmelCase       )
        model_normal_load.eval()
        _lowerCAmelCase		 = model_normal_load(__lowerCAmelCase       ,       __lowerCAmelCase       )['sample']

        assert torch_all_close(__lowerCAmelCase       ,       __lowerCAmelCase       ,       rtol=1e-3       )



  def 	a (	self :       Optional[int]       ):


        """simple docstring"""



        _lowerCAmelCase		 = UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'       )
        model.eval()
        model.to(__lowerCAmelCase       )

        _lowerCAmelCase		 = torch.randn(
            1       ,       model.config.in_channels       ,       model.config.sample_size       ,       model.config.sample_size       ,       generator=torch.manual_seed(0       )       ,       )
        _lowerCAmelCase		 = noise.to(__lowerCAmelCase       )
        _lowerCAmelCase		 = torch.tensor([10] * noise.shape[0]       ).to(__lowerCAmelCase       )

        with torch.no_grad():
              _lowerCAmelCase		 = model(__lowerCAmelCase       ,       __lowerCAmelCase       ).sample

        _lowerCAmelCase		 = output[0, -1, -3:, -3:].flatten().cpu()
        # fmt: off
        _lowerCAmelCase		 = torch.tensor([-13.3258, -20.1100, -15.9873, -17.6617, -23.0596, -17.9419, -13.3675, -16.1889, -12.3800]       )
        # fmt: on

        self.assertTrue(torch_all_close(__lowerCAmelCase       ,       __lowerCAmelCase       ,       rtol=1e-3       )       )





class 							SCREAMING_SNAKE_CASE (					__a						,    __a						,    unittest.TestCase			):
  """simple docstring"""







  __A  =		UNetaDModel
  __A  =		"""sample"""



  @property
  def 	a (	self :       List[str]       ,       __lowerCAmelCase :       List[Any]=(32, 32)       ):


        """simple docstring"""



        _lowerCAmelCase		 = 4
        _lowerCAmelCase		 = 3

        _lowerCAmelCase		 = floats_tensor((batch_size, num_channels) + sizes       ).to(__lowerCAmelCase       )
        _lowerCAmelCase		 = torch.tensor(batch_size * [10]       ).to(dtype=torch.intaa       ,       device=__lowerCAmelCase       )

        return {"sample": noise, "timestep": time_step}



  @property
  def 	a (	self :       str       ):


        """simple docstring"""



        return (3, 32, 32)



  @property
  def 	a (	self :       Tuple       ):


        """simple docstring"""



        return (3, 32, 32)



  def 	a (	self :       int       ):


        """simple docstring"""



        _lowerCAmelCase		 = {
            'block_out_channels': [32, 64, 64, 64],
            'in_channels': 3,
            'layers_per_block': 1,
            'out_channels': 3,
            'time_embedding_type': 'fourier',
            'norm_eps': 1e-6,
            'mid_block_scale_factor': math.sqrt(2.0       ),
            'norm_num_groups': None,
            'down_block_types': [
                'SkipDownBlock2D',
                'AttnSkipDownBlock2D',
                'SkipDownBlock2D',
                'SkipDownBlock2D',
            ],
            'up_block_types': [
                'SkipUpBlock2D',
                'SkipUpBlock2D',
                'AttnSkipUpBlock2D',
                'SkipUpBlock2D',
            ],
        }
        _lowerCAmelCase		 = self.dummy_input
        return init_dict, inputs_dict



  @slow
  def 	a (	self :       Tuple       ):


        """simple docstring"""



        _lowerCAmelCase		 = UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'       ,       output_loading_info=__lowerCAmelCase       )
        self.assertIsNotNone(__lowerCAmelCase       )
        self.assertEqual(len(loading_info['missing_keys']       )       ,       0       )

        model.to(__lowerCAmelCase       )
        _lowerCAmelCase		 = self.dummy_input
        _lowerCAmelCase		 = floats_tensor((4, 3) + (256, 256)       ).to(__lowerCAmelCase       )
        _lowerCAmelCase		 = noise
        _lowerCAmelCase		 = model(**__lowerCAmelCase       )

        assert image is not None, "Make sure output is not None"



  @slow
  def 	a (	self :       Tuple       ):


        """simple docstring"""



        _lowerCAmelCase		 = UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'       )
        model.to(__lowerCAmelCase       )

        _lowerCAmelCase		 = 4
        _lowerCAmelCase		 = 3
        _lowerCAmelCase		 = (256, 256)

        _lowerCAmelCase		 = torch.ones((batch_size, num_channels) + sizes       ).to(__lowerCAmelCase       )
        _lowerCAmelCase		 = torch.tensor(batch_size * [1e-4]       ).to(__lowerCAmelCase       )

        with torch.no_grad():
              _lowerCAmelCase		 = model(__lowerCAmelCase       ,       __lowerCAmelCase       ).sample

        _lowerCAmelCase		 = output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
        _lowerCAmelCase		 = torch.tensor([-4842.8691, -6499.6631, -3800.1953, -7978.2686, -1_0980.7129, -2_0028.8535, 8148.2822, 2342.2905, 567.7608]       )
        # fmt: on

        self.assertTrue(torch_all_close(__lowerCAmelCase       ,       __lowerCAmelCase       ,       rtol=1e-2       )       )



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


        """simple docstring"""



        _lowerCAmelCase		 = UNetaDModel.from_pretrained('fusing/ncsnpp-ffhq-ve-dummy-update'       )
        model.to(__lowerCAmelCase       )

        _lowerCAmelCase		 = 4
        _lowerCAmelCase		 = 3
        _lowerCAmelCase		 = (32, 32)

        _lowerCAmelCase		 = torch.ones((batch_size, num_channels) + sizes       ).to(__lowerCAmelCase       )
        _lowerCAmelCase		 = torch.tensor(batch_size * [1e-4]       ).to(__lowerCAmelCase       )

        with torch.no_grad():
              _lowerCAmelCase		 = model(__lowerCAmelCase       ,       __lowerCAmelCase       ).sample

        _lowerCAmelCase		 = output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
        _lowerCAmelCase		 = torch.tensor([-0.0_3_2_5, -0.0_9_0_0, -0.0_8_6_9, -0.0_3_3_2, -0.0_7_2_5, -0.0_2_7_0, -0.0_1_0_1, 0.0_2_2_7, 0.0_2_5_6]       )
        # fmt: on

        self.assertTrue(torch_all_close(__lowerCAmelCase       ,       __lowerCAmelCase       ,       rtol=1e-2       )       )



  def 	a (	self :       Tuple       ):


        """simple docstring"""



        pass













"""simple docstring"""


from typing import TYPE_CHECKING

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


__SCREAMING_SNAKE_CASE   =   {
    'configuration_transfo_xl': ['TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TransfoXLConfig'],
    'tokenization_transfo_xl': ['TransfoXLCorpus', 'TransfoXLTokenizer'],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__SCREAMING_SNAKE_CASE   =   [
		    'TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'AdaptiveEmbedding',
		    'TransfoXLForSequenceClassification',
		    'TransfoXLLMHeadModel',
		    'TransfoXLModel',
		    'TransfoXLPreTrainedModel',
		    'load_tf_weights_in_transfo_xl',
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__SCREAMING_SNAKE_CASE   =   [
		    'TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'TFAdaptiveEmbedding',
		    'TFTransfoXLForSequenceClassification',
		    'TFTransfoXLLMHeadModel',
		    'TFTransfoXLMainLayer',
		    'TFTransfoXLModel',
		    'TFTransfoXLPreTrainedModel',
		]


if TYPE_CHECKING:
		from .configuration_transfo_xl import TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP, TransfoXLConfig
		from .tokenization_transfo_xl import TransfoXLCorpus, TransfoXLTokenizer

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_transfo_xl import (
				    TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
				    AdaptiveEmbedding,
				    TransfoXLForSequenceClassification,
				    TransfoXLLMHeadModel,
				    TransfoXLModel,
				    TransfoXLPreTrainedModel,
				    load_tf_weights_in_transfo_xl,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_transfo_xl import (
				    TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFAdaptiveEmbedding,
				    TFTransfoXLForSequenceClassification,
				    TFTransfoXLLMHeadModel,
				    TFTransfoXLMainLayer,
				    TFTransfoXLModel,
				    TFTransfoXLPreTrainedModel,
				)

else:
		import sys

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