



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


UpperCAmelCase     		=   logging.get_logger(__name__)

enable_full_determinism()

class __snake_case(				_lowerCAmelCase		,							_lowerCAmelCase		,							unittest.TestCase ):


						'''simple docstring'''



						UpperCAmelCase							:   int									=		UNetaDModel
						UpperCAmelCase							:   Optional[int]									=		"sample"


						@property
						def     __snake_case   ( self				)     ->     List[str]:
							lowerCAmelCase   =     4
							lowerCAmelCase   =     3
							lowerCAmelCase   =     (32, 32)

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

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


						@property
						def     __snake_case   ( self				)     ->     List[str]:
							return (3, 32, 32)


						@property
						def     __snake_case   ( self				)     ->     Dict:
							return (3, 32, 32)




						def     __snake_case   ( self				)     ->     str:
							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 __snake_case(				_lowerCAmelCase		,							_lowerCAmelCase		,							unittest.TestCase ):


						'''simple docstring'''



						UpperCAmelCase							:   str									=		UNetaDModel
						UpperCAmelCase							:   Optional[int]									=		"sample"


						@property
						def     __snake_case   ( self				)     ->     Union[str, Any]:
							lowerCAmelCase   =     4
							lowerCAmelCase   =     4
							lowerCAmelCase   =     (32, 32)

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

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


						@property
						def     __snake_case   ( self				)     ->     int:
							return (4, 32, 32)


						@property
						def     __snake_case   ( self				)     ->     List[str]:
							return (4, 32, 32)


						def     __snake_case   ( self				)     ->     List[Any]:
							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     __snake_case   ( self				)     ->     List[str]:
							lowerCAmelCase,       lowerCAmelCase   =     UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""     ,	output_loading_info=A_				)

							self.assertIsNotNone(A_				)
							self.assertEqual(len(loading_info["""missing_keys"""]				)     ,	0				)

							model.to(A_				)
							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     __snake_case   ( self				)     ->     Dict:
							lowerCAmelCase,       lowerCAmelCase   =     UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""     ,	output_loading_info=A_				)
							model.to(A_				)
							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     __snake_case   ( self				)     ->     Any:
							# by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
							lowerCAmelCase,       lowerCAmelCase   =     UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""     ,	output_loading_info=A_				)
							model_accelerate.to(A_				)
							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(A_				)
							lowerCAmelCase   =     torch.tensor([10] * noise.shape[0]				).to(A_				)

							lowerCAmelCase   =     model_accelerate(A_     ,	A_				)["""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   =     UNetaDModel.from_pretrained(
							    """fusing/unet-ldm-dummy-update"""     ,	output_loading_info=A_     ,	low_cpu_mem_usage=A_				)
							model_normal_load.to(A_				)
							model_normal_load.eval()
							lowerCAmelCase   =     model_normal_load(A_     ,	A_				)["""sample"""]

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




						def     __snake_case   ( self				)     ->     Any:
							lowerCAmelCase   =     UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""				)
							model.eval()
							model.to(A_				)

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

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

							lowerCAmelCase   =     output[0, -1, -3:, -3:].flatten().cpu()
							# fmt: off
							lowerCAmelCase   =     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(A_     ,	A_     ,	rtol=1e-3				)				)







class __snake_case(				_lowerCAmelCase		,							_lowerCAmelCase		,							unittest.TestCase ):


						'''simple docstring'''



						UpperCAmelCase							:   List[Any]									=		UNetaDModel
						UpperCAmelCase							:   Optional[Any]									=		"sample"


						@property
						def     __snake_case   ( self     ,	A_=(32, 32)				)     ->     int:
							lowerCAmelCase   =     4
							lowerCAmelCase   =     3

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

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


						@property
						def     __snake_case   ( self				)     ->     Union[str, Any]:
							return (3, 32, 32)


						@property
						def     __snake_case   ( self				)     ->     str:
							return (3, 32, 32)


						def     __snake_case   ( self				)     ->     Dict:
							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     __snake_case   ( self				)     ->     Optional[Any]:
							lowerCAmelCase,       lowerCAmelCase   =     UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""     ,	output_loading_info=A_				)
							self.assertIsNotNone(A_				)
							self.assertEqual(len(loading_info["""missing_keys"""]				)     ,	0				)

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

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


						@slow
						def     __snake_case   ( self				)     ->     Optional[int]:
							lowerCAmelCase   =     UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""				)
							model.to(A_				)

							lowerCAmelCase   =     4
							lowerCAmelCase   =     3
							lowerCAmelCase   =     (256, 256)

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

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

							lowerCAmelCase   =     output[0, -3:, -3:, -1].flatten().cpu()
							# fmt: off
							lowerCAmelCase   =     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(A_     ,	A_     ,	rtol=1e-2				)				)


						def     __snake_case   ( self				)     ->     List[Any]:
							lowerCAmelCase   =     UNetaDModel.from_pretrained("""fusing/ncsnpp-ffhq-ve-dummy-update"""				)
							model.to(A_				)

							lowerCAmelCase   =     4
							lowerCAmelCase   =     3
							lowerCAmelCase   =     (32, 32)

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

							with torch.no_grad():
								lowerCAmelCase   =     model(A_     ,	A_				).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(A_     ,	A_     ,	rtol=1e-2				)				)




						def     __snake_case   ( self				)     ->     int:
							# not required for this model
							pass




'''simple docstring'''

from torch import nn






def 		_snake_case		( _SCREAMING_SNAKE_CASE					:						Optional[int]							)      ->						Union[str, Any]:







	"""simple docstring"""
	if act_fn in ["swish", "silu"]:
		return nn.SiLU()
	elif act_fn == "mish":
		return nn.Mish()
	elif act_fn == "gelu":
		return nn.GELU()
	else:
		raise ValueError(f'Unsupported activation function: {act_fn}'							)