"""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


SCREAMING_SNAKE_CASE__   =					logging.get_logger(__name__)

enable_full_determinism()



class 					lowerCAmelCase_      (     lowerCAmelCase       ,      lowerCAmelCase       ,      unittest.TestCase   ):



				"""simple docstring"""
				_lowerCAmelCase       :      Optional[int]     					=  UNetaDModel
				_lowerCAmelCase       :      int     					=  """sample"""

				@property
				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     4
										snake_case        =     3
										snake_case        =     (32, 32)

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

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

				@property
				def 						snake_case	(			self      ):






										"""simple docstring"""

										return (3, 32, 32)

				@property
				def 						snake_case	(			self      ):






										"""simple docstring"""

										return (3, 32, 32)

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     {
										    '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,
										}
										snake_case        =     self.dummy_input
										return init_dict, inputs_dict





class 					lowerCAmelCase_      (     lowerCAmelCase       ,      lowerCAmelCase       ,      unittest.TestCase   ):



				"""simple docstring"""
				_lowerCAmelCase       :      List[Any]     					=  UNetaDModel
				_lowerCAmelCase       :      Union[str, Any]     					=  """sample"""

				@property
				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     4
										snake_case        =     4
										snake_case        =     (32, 32)

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

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

				@property
				def 						snake_case	(			self      ):






										"""simple docstring"""

										return (4, 32, 32)

				@property
				def 						snake_case	(			self      ):






										"""simple docstring"""

										return (4, 32, 32)

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     {
										    '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'),
										}
										snake_case        =     self.dummy_input
										return init_dict, inputs_dict

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case   ,snake_case        =     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      )
										snake_case        =     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 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case   ,snake_case        =     UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'						,      output_loading_info=lowerCAmelCase      )
										model.to(lowerCAmelCase      )
										snake_case        =     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 						snake_case	(			self      ):






										"""simple docstring"""

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

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

										snake_case        =     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()

										snake_case   ,snake_case        =     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()
										snake_case        =     model_normal_load(lowerCAmelCase						,      lowerCAmelCase      )['sample']

										assert torch_all_close(lowerCAmelCase						,      lowerCAmelCase						,      rtol=1E-3      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

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

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

										with torch.no_grad():
																snake_case        =     model(lowerCAmelCase						,      lowerCAmelCase      ).sample

										snake_case        =     output[0, -1, -3:, -3:].flatten().cpu()
										# fmt: off
										snake_case        =     torch.tensor([-13.32_58, -20.11_00, -15.98_73, -17.66_17, -23.05_96, -17.94_19, -13.36_75, -16.18_89, -12.38_00]      )
										# fmt: on

										self.assertTrue(torch_all_close(lowerCAmelCase						,      lowerCAmelCase						,      rtol=1E-3      )      )





class 					lowerCAmelCase_      (     lowerCAmelCase       ,      lowerCAmelCase       ,      unittest.TestCase   ):



				"""simple docstring"""
				_lowerCAmelCase       :      int     					=  UNetaDModel
				_lowerCAmelCase       :      Dict     					=  """sample"""

				@property
				def 						snake_case	(			self						,      lowerCAmelCase=(32, 32)      ):






										"""simple docstring"""

										snake_case        =     4
										snake_case        =     3

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

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

				@property
				def 						snake_case	(			self      ):






										"""simple docstring"""

										return (3, 32, 32)

				@property
				def 						snake_case	(			self      ):






										"""simple docstring"""

										return (3, 32, 32)

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     {
										    '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',
										    ],
										}
										snake_case        =     self.dummy_input
										return init_dict, inputs_dict

				@slow
				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case   ,snake_case        =     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      )
										snake_case        =     self.dummy_input
										snake_case        =     floats_tensor((4, 3) + (2_56, 2_56)      ).to(lowerCAmelCase      )
										snake_case        =     noise
										snake_case        =     model(**lowerCAmelCase      )

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

				@slow
				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'      )
										model.to(lowerCAmelCase      )

										snake_case        =     4
										snake_case        =     3
										snake_case        =     (2_56, 2_56)

										snake_case        =     torch.ones((batch_size, num_channels) + sizes      ).to(lowerCAmelCase      )
										snake_case        =     torch.tensor(batch_size * [1E-4]      ).to(lowerCAmelCase      )

										with torch.no_grad():
																snake_case        =     model(lowerCAmelCase						,      lowerCAmelCase      ).sample

										snake_case        =     output[0, -3:, -3:, -1].flatten().cpu()
										# fmt: off
										snake_case        =     torch.tensor([-48_42.86_91, -64_99.66_31, -38_00.19_53, -79_78.26_86, -1_09_80.71_29, -2_00_28.85_35, 81_48.28_22, 23_42.29_05, 5_67.76_08]      )
										# fmt: on

										self.assertTrue(torch_all_close(lowerCAmelCase						,      lowerCAmelCase						,      rtol=1E-2      )      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

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

										snake_case        =     4
										snake_case        =     3
										snake_case        =     (32, 32)

										snake_case        =     torch.ones((batch_size, num_channels) + sizes      ).to(lowerCAmelCase      )
										snake_case        =     torch.tensor(batch_size * [1E-4]      ).to(lowerCAmelCase      )

										with torch.no_grad():
																snake_case        =     model(lowerCAmelCase						,      lowerCAmelCase      ).sample

										snake_case        =     output[0, -3:, -3:, -1].flatten().cpu()
										# fmt: off
										snake_case        =     torch.tensor([-0.03_25, -0.09_00, -0.08_69, -0.03_32, -0.07_25, -0.02_70, -0.01_01, 0.02_27, 0.02_56]      )
										# fmt: on

										self.assertTrue(torch_all_close(lowerCAmelCase						,      lowerCAmelCase						,      rtol=1E-2      )      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										pass



"""simple docstring"""
import os
import sys
import unittest


SCREAMING_SNAKE_CASE__   =					os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, "utils"))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
SCREAMING_SNAKE_CASE__   =					os.path.join(git_repo_path, "src", "transformers")

SCREAMING_SNAKE_CASE__   =					"\n{0} = None\n"

SCREAMING_SNAKE_CASE__   =					"\nclass {0}(metaclass=DummyObject):\n    _backends = {1}\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, {1})\n"


SCREAMING_SNAKE_CASE__   =					"\ndef {0}(*args, **kwargs):\n    requires_backends({0}, {1})\n"



class 					lowerCAmelCase_      (     unittest.TestCase   ):



				"""simple docstring"""

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     find_backend('    _import_structure["models.albert"].append("AlbertTokenizerFast")'      )
										self.assertIsNone(lowerCAmelCase      )

										snake_case        =     find_backend('    if not is_tokenizers_available():'      )
										self.assertEqual(lowerCAmelCase						,      'tokenizers'      )

										snake_case        =     find_backend('    if not is_tensorflow_text_available():'      )
										self.assertEqual(lowerCAmelCase						,      'tensorflow_text'      )

										snake_case        =     find_backend('    if not (is_sentencepiece_available() and is_tokenizers_available()):'      )
										self.assertEqual(lowerCAmelCase						,      'sentencepiece_and_tokenizers'      )

										snake_case        =     find_backend(
										    '    if not (is_sentencepiece_available() and is_tensorflow_text_available()):'      )
										self.assertEqual(lowerCAmelCase						,      'sentencepiece_and_tensorflow_text'      )

										snake_case        =     find_backend(
										    '    if not (is_sentencepiece_available() and is_tokenizers_available() and is_vision_available()):'      )
										self.assertEqual(lowerCAmelCase						,      'sentencepiece_and_tokenizers_and_vision'      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     read_init()
										# We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
										self.assertIn('torch'						,      lowerCAmelCase      )
										self.assertIn('tensorflow_text'						,      lowerCAmelCase      )
										self.assertIn('sentencepiece_and_tokenizers'						,      lowerCAmelCase      )

										# Likewise, we can't assert on the exact content of a key
										self.assertIn('BertModel'						,      objects['torch']      )
										self.assertIn('TFBertModel'						,      objects['tf']      )
										self.assertIn('FlaxBertModel'						,      objects['flax']      )
										self.assertIn('BertModel'						,      objects['torch']      )
										self.assertIn('TFBertTokenizer'						,      objects['tensorflow_text']      )
										self.assertIn('convert_slow_tokenizer'						,      objects['sentencepiece_and_tokenizers']      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     create_dummy_object('CONSTANT'						,      '\'torch\''      )
										self.assertEqual(lowerCAmelCase						,      '\nCONSTANT = None\n'      )

										snake_case        =     create_dummy_object('function'						,      '\'torch\''      )
										self.assertEqual(
										    lowerCAmelCase						,      '\ndef function(*args, **kwargs):\n    requires_backends(function, \'torch\')\n'      )

										snake_case        =     '\nclass FakeClass(metaclass=DummyObject):\n    _backends = \'torch\'\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, \'torch\')\n'
										snake_case        =     create_dummy_object('FakeClass'						,      '\'torch\''      )
										self.assertEqual(lowerCAmelCase						,      lowerCAmelCase      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     '# This file is autogenerated by the command `make fix-copies`, do not edit.\nfrom ..utils import DummyObject, requires_backends\n\n\nCONSTANT = None\n\n\ndef function(*args, **kwargs):\n    requires_backends(function, ["torch"])\n\n\nclass FakeClass(metaclass=DummyObject):\n    _backends = ["torch"]\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, ["torch"])\n'
										snake_case        =     create_dummy_files({'torch': ['CONSTANT', 'function', 'FakeClass']}      )
										self.assertEqual(dummy_files['torch']						,      lowerCAmelCase      )


