




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


lowercase__  		=		logging.get_logger(__name__)

enable_full_determinism()






class lowerCAmelCase__			(  lowercase,			lowercase,			unittest.TestCase		):






			'''simple docstring'''



			lowerCamelCase__   =    UNetaDModel
			lowerCamelCase__   =    """sample"""
			@property
			def       A_					(						self						):
							_lowerCamelCase  :    List[str]					=					4
							_lowerCamelCase  :    Optional[int]					=					3
							_lowerCamelCase  :    List[Any]					=					(32, 32)

							_lowerCamelCase  :    Optional[Any]					=					floats_tensor((batch_size, num_channels) + sizes						).to(lowercase						)
							_lowerCamelCase  :    Dict					=					torch.tensor([10]						).to(lowercase						)

							return {"sample": noise, "timestep": time_step}
			@property
			def       A_					(						self						):
							return (3, 32, 32)
			@property
			def       A_					(						self						):
							return (3, 32, 32)






			def       A_					(						self						):
							_lowerCamelCase  :    int					=					{
							    '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  :    List[str]					=					self.dummy_input
							return init_dict, inputs_dict






class lowerCAmelCase__			(  lowercase,			lowercase,			unittest.TestCase		):






			'''simple docstring'''



			lowerCamelCase__   =    UNetaDModel
			lowerCamelCase__   =    """sample"""
			@property
			def       A_					(						self						):
							_lowerCamelCase  :    str					=					4
							_lowerCamelCase  :    Optional[int]					=					4
							_lowerCamelCase  :    Any					=					(32, 32)

							_lowerCamelCase  :    Dict					=					floats_tensor((batch_size, num_channels) + sizes						).to(lowercase						)
							_lowerCamelCase  :    Tuple					=					torch.tensor([10]						).to(lowercase						)

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

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

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

							_lowerCamelCase  :    Union[str, 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  :    Any					=					noise.to(lowercase						)
							_lowerCamelCase  :    str					=					torch.tensor([10] * noise.shape[0]						).to(lowercase						)

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

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






			def       A_					(						self						):
							_lowerCamelCase  :    Optional[int]					=					UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'						)
							model.eval()
							model.to(lowercase						)

							_lowerCamelCase  :    Optional[int]					=					torch.randn(
							    1						,   model.config.in_channels						,   model.config.sample_size						,   model.config.sample_size						,   generator=torch.manual_seed(0						)						,   )
							_lowerCamelCase  :    List[str]					=					noise.to(lowercase						)
							_lowerCamelCase  :    List[Any]					=					torch.tensor([10] * noise.shape[0]						).to(lowercase						)

							with torch.no_grad():
											_lowerCamelCase  :    Tuple					=					model(lowercase						,   lowercase						).sample

							_lowerCamelCase  :    Any					=					output[0, -1, -3:, -3:].flatten().cpu()
							# fmt: off
							_lowerCamelCase  :    Dict					=					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(lowercase						,   lowercase						,   rtol=1E-3						)						)






class lowerCAmelCase__			(  lowercase,			lowercase,			unittest.TestCase		):






			'''simple docstring'''



			lowerCamelCase__   =    UNetaDModel
			lowerCamelCase__   =    """sample"""
			@property
			def       A_					(						self						,   lowercase=(32, 32)						):
							_lowerCamelCase  :    Union[str, Any]					=					4
							_lowerCamelCase  :    str					=					3

							_lowerCamelCase  :    Optional[Any]					=					floats_tensor((batch_size, num_channels) + sizes						).to(lowercase						)
							_lowerCamelCase  :    Tuple					=					torch.tensor(batch_size * [10]						).to(dtype=torch.intaa						,   device=lowercase						)

							return {"sample": noise, "timestep": time_step}
			@property
			def       A_					(						self						):
							return (3, 32, 32)
			@property
			def       A_					(						self						):
							return (3, 32, 32)
			def       A_					(						self						):
							_lowerCamelCase  :    str					=					{
							    '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       A_					(						self						):
							_lowerCamelCase,				_lowerCamelCase  :    Union[str, Any]					=					UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'						,   output_loading_info=lowercase						)
							self.assertIsNotNone(lowercase						)
							self.assertEqual(len(loading_info['missing_keys']						)						,   0						)

							model.to(lowercase						)
							_lowerCamelCase  :    str					=					self.dummy_input
							_lowerCamelCase  :    Dict					=					floats_tensor((4, 3) + (256, 256)						).to(lowercase						)
							_lowerCamelCase  :    Tuple					=					noise
							_lowerCamelCase  :    Optional[Any]					=					model(**lowercase						)

							assert image is not None, "Make sure output is not None"
			@slow
			def       A_					(						self						):
							_lowerCamelCase  :    Optional[Any]					=					UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'						)
							model.to(lowercase						)

							_lowerCamelCase  :    Optional[int]					=					4
							_lowerCamelCase  :    Optional[int]					=					3
							_lowerCamelCase  :    Any					=					(256, 256)

							_lowerCamelCase  :    Optional[int]					=					torch.ones((batch_size, num_channels) + sizes						).to(lowercase						)
							_lowerCamelCase  :    Tuple					=					torch.tensor(batch_size * [1E-4]						).to(lowercase						)

							with torch.no_grad():
											_lowerCamelCase  :    Tuple					=					model(lowercase						,   lowercase						).sample

							_lowerCamelCase  :    Dict					=					output[0, -3:, -3:, -1].flatten().cpu()
							# fmt: off
							_lowerCamelCase  :    Optional[int]					=					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(lowercase						,   lowercase						,   rtol=1E-2						)						)
			def       A_					(						self						):
							_lowerCamelCase  :    Optional[Any]					=					UNetaDModel.from_pretrained('fusing/ncsnpp-ffhq-ve-dummy-update'						)
							model.to(lowercase						)

							_lowerCamelCase  :    List[str]					=					4
							_lowerCamelCase  :    str					=					3
							_lowerCamelCase  :    Any					=					(32, 32)

							_lowerCamelCase  :    Dict					=					torch.ones((batch_size, num_channels) + sizes						).to(lowercase						)
							_lowerCamelCase  :    Any					=					torch.tensor(batch_size * [1E-4]						).to(lowercase						)

							with torch.no_grad():
											_lowerCamelCase  :    Any					=					model(lowercase						,   lowercase						).sample

							_lowerCamelCase  :    Dict					=					output[0, -3:, -3:, -1].flatten().cpu()
							# fmt: off
							_lowerCamelCase  :    int					=					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(lowercase						,   lowercase						,   rtol=1E-2						)						)






			def       A_					(						self						):
							# not required for this model
							pass





"""simple docstring"""






def   _snake_case    (       lowercase__	):
				return "".join(chr(ord(lowercase__	) - 32	) if 'a' <= char <= 'z' else char for char in word	)


if __name__ == "__main__":
			from doctest import testmod

			testmod()