


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


a__       :							Dict        =  logging.get_logger(__name__)

enable_full_determinism()




class 				UpperCamelCase_			(				UpperCamelCase    ,						UpperCamelCase    ,						unittest.TestCase):

	"""simple docstring"""





	snake_case__       :						int					      =    UNetaDModel
	snake_case__       :						List[str]					      =    "sample"




	@property
	def 						UpperCAmelCase_   (				self    :       str     ) ->			str:
						__SCREAMING_SNAKE_CASE		     =       4
						__SCREAMING_SNAKE_CASE		     =       3
						__SCREAMING_SNAKE_CASE		     =       (3_2, 3_2)

						__SCREAMING_SNAKE_CASE		     =       floats_tensor((batch_size, num_channels) + sizes     ).to(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       torch.tensor([1_0]     ).to(UpperCAmelCase__     )

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




	@property
	def 						UpperCAmelCase_   (				self    :       str     ) ->			Any:
						return (3, 3_2, 3_2)




	@property
	def 						UpperCAmelCase_   (				self    :       Tuple     ) ->			Tuple:
						return (3, 3_2, 3_2)




	def 						UpperCAmelCase_   (				self    :       Dict     ) ->			Union[str, Any]:
						__SCREAMING_SNAKE_CASE		     =       {
						    "block_out_channels": (3_2, 6_4),
						    "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": 3_2,
						}
						__SCREAMING_SNAKE_CASE		     =       self.dummy_input
						return init_dict, inputs_dict




class 				UpperCamelCase_			(				UpperCamelCase    ,						UpperCamelCase    ,						unittest.TestCase):

	"""simple docstring"""





	snake_case__       :						Union[str, Any]					      =    UNetaDModel
	snake_case__       :						List[str]					      =    "sample"




	@property
	def 						UpperCAmelCase_   (				self    :       str     ) ->			int:
						__SCREAMING_SNAKE_CASE		     =       4
						__SCREAMING_SNAKE_CASE		     =       4
						__SCREAMING_SNAKE_CASE		     =       (3_2, 3_2)

						__SCREAMING_SNAKE_CASE		     =       floats_tensor((batch_size, num_channels) + sizes     ).to(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       torch.tensor([1_0]     ).to(UpperCAmelCase__     )

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




	@property
	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			Optional[int]:
						return (4, 3_2, 3_2)




	@property
	def 						UpperCAmelCase_   (				self    :       Optional[int]     ) ->			Any:
						return (4, 3_2, 3_2)




	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			List[Any]:
						__SCREAMING_SNAKE_CASE		     =       {
						    "sample_size": 3_2,
						    "in_channels": 4,
						    "out_channels": 4,
						    "layers_per_block": 2,
						    "block_out_channels": (3_2, 6_4),
						    "attention_head_dim": 3_2,
						    "down_block_types": ("DownBlock2D", "DownBlock2D"),
						    "up_block_types": ("UpBlock2D", "UpBlock2D"),
						}
						__SCREAMING_SNAKE_CASE		     =       self.dummy_input
						return init_dict, inputs_dict




	def 						UpperCAmelCase_   (				self    :       Dict     ) ->			List[str]:
						__SCREAMING_SNAKE_CASE     ,						__SCREAMING_SNAKE_CASE		     =       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__     )
						__SCREAMING_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 						UpperCAmelCase_   (				self    :       Tuple     ) ->			Tuple:
						__SCREAMING_SNAKE_CASE     ,						__SCREAMING_SNAKE_CASE		     =       UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"    ,    output_loading_info=UpperCAmelCase__     )
						model.to(UpperCAmelCase__     )
						__SCREAMING_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 						UpperCAmelCase_   (				self    :       List[str]     ) ->			Optional[Any]:
						# by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
						__SCREAMING_SNAKE_CASE     ,						__SCREAMING_SNAKE_CASE		     =       UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"    ,    output_loading_info=UpperCAmelCase__     )
						model_accelerate.to(UpperCAmelCase__     )
						model_accelerate.eval()

						__SCREAMING_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     )    ,    )
						__SCREAMING_SNAKE_CASE		     =       noise.to(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       torch.tensor([1_0] * noise.shape[0]     ).to(UpperCAmelCase__     )

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

						__SCREAMING_SNAKE_CASE     ,						__SCREAMING_SNAKE_CASE		     =       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()
						__SCREAMING_SNAKE_CASE		     =       model_normal_load(UpperCAmelCase__    ,    UpperCAmelCase__     )["sample"]

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




	def 						UpperCAmelCase_   (				self    :       int     ) ->			Tuple:
						__SCREAMING_SNAKE_CASE		     =       UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"     )
						model.eval()
						model.to(UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       torch.randn(
						    1    ,    model.config.in_channels    ,    model.config.sample_size    ,    model.config.sample_size    ,    generator=torch.manual_seed(0     )    ,    )
						__SCREAMING_SNAKE_CASE		     =       noise.to(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       torch.tensor([1_0] * noise.shape[0]     ).to(UpperCAmelCase__     )

						with torch.no_grad():
											__SCREAMING_SNAKE_CASE		     =       model(UpperCAmelCase__    ,    UpperCAmelCase__     ).sample

						__SCREAMING_SNAKE_CASE		     =       output[0, -1, -3:, -3:].flatten().cpu()
						# fmt: off
						__SCREAMING_SNAKE_CASE		     =       torch.tensor([-13.3_258, -20.1_100, -15.9_873, -17.6_617, -23.0_596, -17.9_419, -13.3_675, -16.1_889, -12.3_800]     )
						# fmt: on

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




class 				UpperCamelCase_			(				UpperCamelCase    ,						UpperCamelCase    ,						unittest.TestCase):

	"""simple docstring"""





	snake_case__       :						Optional[int]					      =    UNetaDModel
	snake_case__       :						Any					      =    "sample"




	@property
	def 						UpperCAmelCase_   (				self    :       Optional[int]    ,    UpperCAmelCase__    :       Optional[int]=(3_2, 3_2)     ) ->			Tuple:
						__SCREAMING_SNAKE_CASE		     =       4
						__SCREAMING_SNAKE_CASE		     =       3

						__SCREAMING_SNAKE_CASE		     =       floats_tensor((batch_size, num_channels) + sizes     ).to(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       torch.tensor(batch_size * [1_0]     ).to(dtype=torch.intaa    ,    device=UpperCAmelCase__     )

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




	@property
	def 						UpperCAmelCase_   (				self    :       Optional[int]     ) ->			Optional[Any]:
						return (3, 3_2, 3_2)




	@property
	def 						UpperCAmelCase_   (				self    :       Optional[Any]     ) ->			Dict:
						return (3, 3_2, 3_2)




	def 						UpperCAmelCase_   (				self    :       Any     ) ->			Union[str, Any]:
						__SCREAMING_SNAKE_CASE		     =       {
						    "block_out_channels": [3_2, 6_4, 6_4, 6_4],
						    "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",
						    ],
						}
						__SCREAMING_SNAKE_CASE		     =       self.dummy_input
						return init_dict, inputs_dict




	@slow
	def 						UpperCAmelCase_   (				self    :       Tuple     ) ->			Optional[Any]:
						__SCREAMING_SNAKE_CASE     ,						__SCREAMING_SNAKE_CASE		     =       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__     )
						__SCREAMING_SNAKE_CASE		     =       self.dummy_input
						__SCREAMING_SNAKE_CASE		     =       floats_tensor((4, 3) + (2_5_6, 2_5_6)     ).to(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       noise
						__SCREAMING_SNAKE_CASE		     =       model(**UpperCAmelCase__     )

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




	@slow
	def 						UpperCAmelCase_   (				self    :       int     ) ->			Tuple:
						__SCREAMING_SNAKE_CASE		     =       UNetaDModel.from_pretrained("google/ncsnpp-celebahq-256"     )
						model.to(UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       4
						__SCREAMING_SNAKE_CASE		     =       3
						__SCREAMING_SNAKE_CASE		     =       (2_5_6, 2_5_6)

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

						with torch.no_grad():
											__SCREAMING_SNAKE_CASE		     =       model(UpperCAmelCase__    ,    UpperCAmelCase__     ).sample

						__SCREAMING_SNAKE_CASE		     =       output[0, -3:, -3:, -1].flatten().cpu()
						# fmt: off
						__SCREAMING_SNAKE_CASE		     =       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, 567.7_608]     )
						# fmt: on

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




	def 						UpperCAmelCase_   (				self    :       int     ) ->			int:
						__SCREAMING_SNAKE_CASE		     =       UNetaDModel.from_pretrained("fusing/ncsnpp-ffhq-ve-dummy-update"     )
						model.to(UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       4
						__SCREAMING_SNAKE_CASE		     =       3
						__SCREAMING_SNAKE_CASE		     =       (3_2, 3_2)

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

						with torch.no_grad():
											__SCREAMING_SNAKE_CASE		     =       model(UpperCAmelCase__    ,    UpperCAmelCase__     ).sample

						__SCREAMING_SNAKE_CASE		     =       output[0, -3:, -3:, -1].flatten().cpu()
						# fmt: off
						__SCREAMING_SNAKE_CASE		     =       torch.tensor([-0.0_325, -0.0_900, -0.0_869, -0.0_332, -0.0_725, -0.0_270, -0.0_101, 0.0_227, 0.0_256]     )
						# fmt: on

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




	def 						UpperCAmelCase_   (				self    :       List[Any]     ) ->			List[str]:
						# not required for this model
						pass




"""simple docstring"""







import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class 				UpperCamelCase_			(				UpperCamelCase    ,						unittest.TestCase):

	"""simple docstring"""





	snake_case__       :						Union[str, Any]					      =    DiTPipeline
	snake_case__       :						Union[str, Any]					      =    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
	snake_case__       :						Any					      =    PipelineTesterMixin.required_optional_params - {
	    "latents",
	    "num_images_per_prompt",
	    "callback",
	    "callback_steps",
	}
	snake_case__       :						Optional[Any]					      =    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
	snake_case__       :						Optional[Any]					      =    False




	def 						UpperCAmelCase_   (				self    :       Dict     ) ->			Dict:
						torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       TransformeraDModel(
						    sample_size=1_6    ,    num_layers=2    ,    patch_size=4    ,    attention_head_dim=8    ,    num_attention_heads=2    ,    in_channels=4    ,    out_channels=8    ,    attention_bias=UpperCAmelCase__    ,    activation_fn="gelu-approximate"    ,    num_embeds_ada_norm=1_0_0_0    ,    norm_type="ada_norm_zero"    ,    norm_elementwise_affine=UpperCAmelCase__    ,    )
						__SCREAMING_SNAKE_CASE		     =       AutoencoderKL()
						__SCREAMING_SNAKE_CASE		     =       DDIMScheduler()
						__SCREAMING_SNAKE_CASE		     =       {"transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler}
						return components




	def 						UpperCAmelCase_   (				self    :       Optional[int]    ,    UpperCAmelCase__    :       Tuple    ,    UpperCAmelCase__    :       str=0     ) ->			Optional[int]:
						if str(UpperCAmelCase__     ).startswith("mps"     ):
											__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(UpperCAmelCase__     )
						else:
											__SCREAMING_SNAKE_CASE		     =       torch.Generator(device=UpperCAmelCase__     ).manual_seed(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       {
						    "class_labels": [1],
						    "generator": generator,
						    "num_inference_steps": 2,
						    "output_type": "numpy",
						}
						return inputs




	def 						UpperCAmelCase_   (				self    :       Optional[int]     ) ->			List[str]:
						__SCREAMING_SNAKE_CASE		     =       "cpu"

						__SCREAMING_SNAKE_CASE		     =       self.get_dummy_components()
						__SCREAMING_SNAKE_CASE		     =       self.pipeline_class(**UpperCAmelCase__     )
						pipe.to(UpperCAmelCase__     )
						pipe.set_progress_bar_config(disable=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       self.get_dummy_inputs(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       pipe(**UpperCAmelCase__     ).images
						__SCREAMING_SNAKE_CASE		     =       image[0, -3:, -3:, -1]

						self.assertEqual(image.shape    ,    (1, 1_6, 1_6, 3)     )
						__SCREAMING_SNAKE_CASE		     =       np.array([0.2_946, 0.6_601, 0.4_329, 0.3_296, 0.4_144, 0.5_319, 0.7_273, 0.5_013, 0.4_457]     )
						__SCREAMING_SNAKE_CASE		     =       np.abs(image_slice.flatten() - expected_slice     ).max()
						self.assertLessEqual(UpperCAmelCase__    ,    1E-3     )




	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			Union[str, Any]:
						self._test_inference_batch_single_identical(relax_max_difference=UpperCAmelCase__    ,    expected_max_diff=1E-3     )




	@unittest.skipIf(
	    torch_device != "cuda" or not is_xformers_available()    ,    reason="XFormers attention is only available with CUDA and `xformers` installed"    ,    )
	def 						UpperCAmelCase_   (				self    :       int     ) ->			Dict:
						self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3     )




@require_torch_gpu
@slow
class 				UpperCamelCase_			(				unittest.TestCase):

	"""simple docstring"""





	def 						UpperCAmelCase_   (				self    :       Optional[int]     ) ->			Any:
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()




	def 						UpperCAmelCase_   (				self    :       Union[str, Any]     ) ->			Optional[int]:
						__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(0     )

						__SCREAMING_SNAKE_CASE		     =       DiTPipeline.from_pretrained("facebook/DiT-XL-2-256"     )
						pipe.to("cuda"     )

						__SCREAMING_SNAKE_CASE		     =       ["vase", "umbrella", "white shark", "white wolf"]
						__SCREAMING_SNAKE_CASE		     =       pipe.get_label_ids(UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       pipe(UpperCAmelCase__    ,    generator=UpperCAmelCase__    ,    num_inference_steps=4_0    ,    output_type="np"     ).images

						for word, image in zip(UpperCAmelCase__    ,    UpperCAmelCase__     ):
											__SCREAMING_SNAKE_CASE		     =       load_numpy(
											    F"""https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy"""     )
											assert np.abs((expected_image - image).max()     ) < 1E-2




	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			List[Any]:
						__SCREAMING_SNAKE_CASE		     =       DiTPipeline.from_pretrained("facebook/DiT-XL-2-512"     )
						__SCREAMING_SNAKE_CASE		     =       DPMSolverMultistepScheduler.from_config(pipe.scheduler.config     )
						pipe.to("cuda"     )

						__SCREAMING_SNAKE_CASE		     =       ["vase", "umbrella"]
						__SCREAMING_SNAKE_CASE		     =       pipe.get_label_ids(UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       pipe(UpperCAmelCase__    ,    generator=UpperCAmelCase__    ,    num_inference_steps=2_5    ,    output_type="np"     ).images

						for word, image in zip(UpperCAmelCase__    ,    UpperCAmelCase__     ):
											__SCREAMING_SNAKE_CASE		     =       load_numpy(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    F"""/dit/{word}_512.npy"""     )

											assert np.abs((expected_image - image).max()     ) < 1E-1
