






"""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 						__UpperCamelCase     (     UpperCamelCase   ,     UpperCamelCase   ,     unittest.TestCase						):
					lowercase_    :				List[Any]      =  UNetaDModel
					lowercase_    :				str      =  """sample"""
					@property
					def        UpperCAmelCase__ (							self							:	List[str]   )	->       Union[str, Any]:
											lowerCAmelCase   :Dict						=			4
											lowerCAmelCase   :str						=			3
											lowerCAmelCase   :Optional[Any]						=			(32, 32)

											lowerCAmelCase   :Optional[int]						=			floats_tensor((batch_size, num_channels) + sizes   ).to(UpperCAmelCase   )
											lowerCAmelCase   :Optional[int]						=			torch.tensor([10]   ).to(UpperCAmelCase   )

											return {"sample": noise, "timestep": time_step}
					@property
					def        UpperCAmelCase__ (							self							:	Any   )	->       str:
											return (3, 32, 32)
					@property
					def        UpperCAmelCase__ (							self							:	Optional[Any]   )	->       int:
											return (3, 32, 32)
					def        UpperCAmelCase__ (							self							:	Tuple   )	->       Tuple:
											lowerCAmelCase   :Tuple						=			{
											    '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   :Optional[Any]						=			self.dummy_input
											return init_dict, inputs_dict




class 						__UpperCamelCase     (     UpperCamelCase   ,     UpperCamelCase   ,     unittest.TestCase						):
					lowercase_    :				int      =  UNetaDModel
					lowercase_    :				int      =  """sample"""
					@property
					def        UpperCAmelCase__ (							self							:	str   )	->       int:
											lowerCAmelCase   :Tuple						=			4
											lowerCAmelCase   :int						=			4
											lowerCAmelCase   :Optional[int]						=			(32, 32)

											lowerCAmelCase   :Union[str, Any]						=			floats_tensor((batch_size, num_channels) + sizes   ).to(UpperCAmelCase   )
											lowerCAmelCase   :Any						=			torch.tensor([10]   ).to(UpperCAmelCase   )

											return {"sample": noise, "timestep": time_step}
					@property
					def        UpperCAmelCase__ (							self							:	Any   )	->       List[str]:
											return (4, 32, 32)
					@property
					def        UpperCAmelCase__ (							self							:	Optional[Any]   )	->       str:
											return (4, 32, 32)
					def        UpperCAmelCase__ (							self							:	Optional[Any]   )	->       Dict:
											lowerCAmelCase   :Optional[int]						=			{
											    '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   :List[Any]						=			self.dummy_input
											return init_dict, inputs_dict
					def        UpperCAmelCase__ (							self							:	int   )	->       int:
											lowerCAmelCase	, lowerCAmelCase   :Dict						=			UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'     , output_loading_info=UpperCAmelCase   )

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

											model.to(UpperCAmelCase   )
											lowerCAmelCase   :Any						=			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        UpperCAmelCase__ (							self							:	str   )	->       int:
											lowerCAmelCase	, lowerCAmelCase   :List[Any]						=			UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'     , output_loading_info=UpperCAmelCase   )
											model.to(UpperCAmelCase   )
											lowerCAmelCase   :Any						=			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        UpperCAmelCase__ (							self							:	Union[str, Any]   )	->       List[str]:
											# by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
											lowerCAmelCase	, lowerCAmelCase   :Dict						=			UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'     , output_loading_info=UpperCAmelCase   )
											model_accelerate.to(UpperCAmelCase   )
											model_accelerate.eval()

											lowerCAmelCase   :Any						=			torch.randn(
											    1     , model_accelerate.config.in_channels     , model_accelerate.config.sample_size     , model_accelerate.config.sample_size     , generator=torch.manual_seed(0   )     , )
											lowerCAmelCase   :Dict						=			noise.to(UpperCAmelCase   )
											lowerCAmelCase   :Optional[int]						=			torch.tensor([10] * noise.shape[0]   ).to(UpperCAmelCase   )

											lowerCAmelCase   :List[str]						=			model_accelerate(UpperCAmelCase     , UpperCAmelCase   )['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	, lowerCAmelCase   :int						=			UNetaDModel.from_pretrained(
											    'fusing/unet-ldm-dummy-update'     , output_loading_info=UpperCAmelCase     , low_cpu_mem_usage=UpperCAmelCase   )
											model_normal_load.to(UpperCAmelCase   )
											model_normal_load.eval()
											lowerCAmelCase   :Optional[Any]						=			model_normal_load(UpperCAmelCase     , UpperCAmelCase   )['sample']

											assert torch_all_close(UpperCAmelCase     , UpperCAmelCase     , rtol=1e-3   )
					def        UpperCAmelCase__ (							self							:	Optional[int]   )	->       Optional[int]:
											lowerCAmelCase   :Union[str, Any]						=			UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'   )
											model.eval()
											model.to(UpperCAmelCase   )

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

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

											lowerCAmelCase   :List[Any]						=			output[0, -1, -3:, -3:].flatten().cpu()
											# fmt: off
											lowerCAmelCase   :Optional[Any]						=			torch.tensor([-1_3.3_2_5_8, -2_0.1_1_0_0, -1_5.9_8_7_3, -1_7.6_6_1_7, -2_3.0_5_9_6, -1_7.9_4_1_9, -1_3.3_6_7_5, -1_6.1_8_8_9, -1_2.3_8_0_0]   )
											# fmt: on

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




class 						__UpperCamelCase     (     UpperCamelCase   ,     UpperCamelCase   ,     unittest.TestCase						):
					lowercase_    :				Union[str, Any]      =  UNetaDModel
					lowercase_    :				Any      =  """sample"""
					@property
					def        UpperCAmelCase__ (							self							:	List[str]     , UpperCAmelCase							:	List[Any]=(32, 32)   )	->       int:
											lowerCAmelCase   :List[str]						=			4
											lowerCAmelCase   :Optional[Any]						=			3

											lowerCAmelCase   :int						=			floats_tensor((batch_size, num_channels) + sizes   ).to(UpperCAmelCase   )
											lowerCAmelCase   :Optional[Any]						=			torch.tensor(batch_size * [10]   ).to(dtype=torch.intaa     , device=UpperCAmelCase   )

											return {"sample": noise, "timestep": time_step}
					@property
					def        UpperCAmelCase__ (							self							:	str   )	->       Tuple:
											return (3, 32, 32)
					@property
					def        UpperCAmelCase__ (							self							:	Tuple   )	->       Union[str, Any]:
											return (3, 32, 32)
					def        UpperCAmelCase__ (							self							:	int   )	->       Optional[Any]:
											lowerCAmelCase   :List[Any]						=			{
											    '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   :Dict						=			self.dummy_input
											return init_dict, inputs_dict
					@slow
					def        UpperCAmelCase__ (							self							:	Tuple   )	->       str:
											lowerCAmelCase	, lowerCAmelCase   :Optional[Any]						=			UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'     , output_loading_info=UpperCAmelCase   )
											self.assertIsNotNone(UpperCAmelCase   )
											self.assertEqual(len(loading_info['missing_keys']   )     , 0   )

											model.to(UpperCAmelCase   )
											lowerCAmelCase   :Dict						=			self.dummy_input
											lowerCAmelCase   :int						=			floats_tensor((4, 3) + (256, 256)   ).to(UpperCAmelCase   )
											lowerCAmelCase   :Optional[int]						=			noise
											lowerCAmelCase   :Tuple						=			model(**UpperCAmelCase   )

											assert image is not None, "Make sure output is not None"
					@slow
					def        UpperCAmelCase__ (							self							:	Tuple   )	->       Dict:
											lowerCAmelCase   :str						=			UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'   )
											model.to(UpperCAmelCase   )

											lowerCAmelCase   :List[str]						=			4
											lowerCAmelCase   :Optional[Any]						=			3
											lowerCAmelCase   :Any						=			(256, 256)

											lowerCAmelCase   :int						=			torch.ones((batch_size, num_channels) + sizes   ).to(UpperCAmelCase   )
											lowerCAmelCase   :List[Any]						=			torch.tensor(batch_size * [1e-4]   ).to(UpperCAmelCase   )

											with torch.no_grad():
																	lowerCAmelCase   :List[Any]						=			model(UpperCAmelCase     , UpperCAmelCase   ).sample

											lowerCAmelCase   :List[Any]						=			output[0, -3:, -3:, -1].flatten().cpu()
											# fmt: off
											lowerCAmelCase   :Union[str, Any]						=			torch.tensor([-4_8_4_2.8_6_9_1, -6_4_9_9.6_6_3_1, -3_8_0_0.1_9_5_3, -7_9_7_8.2_6_8_6, -1_0_9_8_0.7_1_2_9, -2_0_0_2_8.8_5_3_5, 8_1_4_8.2_8_2_2, 2_3_4_2.2_9_0_5, 5_6_7.7_6_0_8]   )
											# fmt: on

											self.assertTrue(torch_all_close(UpperCAmelCase     , UpperCAmelCase     , rtol=1e-2   )   )
					def        UpperCAmelCase__ (							self							:	Union[str, Any]   )	->       Union[str, Any]:
											lowerCAmelCase   :Tuple						=			UNetaDModel.from_pretrained('fusing/ncsnpp-ffhq-ve-dummy-update'   )
											model.to(UpperCAmelCase   )

											lowerCAmelCase   :Dict						=			4
											lowerCAmelCase   :Optional[int]						=			3
											lowerCAmelCase   :Union[str, Any]						=			(32, 32)

											lowerCAmelCase   :Optional[int]						=			torch.ones((batch_size, num_channels) + sizes   ).to(UpperCAmelCase   )
											lowerCAmelCase   :Tuple						=			torch.tensor(batch_size * [1e-4]   ).to(UpperCAmelCase   )

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

											lowerCAmelCase   :Optional[Any]						=			output[0, -3:, -3:, -1].flatten().cpu()
											# fmt: off
											lowerCAmelCase   :int						=			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(UpperCAmelCase     , UpperCAmelCase     , rtol=1e-2   )   )
					def        UpperCAmelCase__ (							self							:	Tuple   )	->       Optional[Any]:
											# not required for this model
											pass







"""simple docstring"""


from abc import ABC, abstractmethod
from argparse import ArgumentParser




class 						__UpperCamelCase     (     UpperCamelCase						):
					@staticmethod
					@abstractmethod
					def        UpperCAmelCase__ (							UpperCAmelCase							:	ArgumentParser   )	->       Optional[Any]:
											raise NotImplementedError()
					@abstractmethod
					def        UpperCAmelCase__ (							self							:	str   )	->       List[str]:
											raise NotImplementedError()