



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  __snake_case      (   SCREAMING_SNAKE_CASE			,unittest.TestCase):







							'''simple docstring'''






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


							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		TransformeraDModel(
											    sample_size=16  ,  num_layers=2  ,  patch_size=4  ,  attention_head_dim=8  ,  num_attention_heads=2  ,  in_channels=4  ,  out_channels=8  ,  attention_bias=a_  ,  activation_fn="""gelu-approximate"""  ,  num_embeds_ada_norm=1_000  ,  norm_type="""ada_norm_zero"""  ,  norm_elementwise_affine=a_  ,  )
											a__				      =		AutoencoderKL()
											a__				      =		DDIMScheduler()
											a__				      =		{"""transformer""": transformer.eval(), """vae""": vae.eval(), """scheduler""": scheduler}
											return components


							def _a      (							self  ,  a_  ,  a_=0		):
											if str(a_		).startswith("""mps"""		):
															a__				      =		torch.manual_seed(a_		)
											else:
															a__				      =		torch.Generator(device=a_		).manual_seed(a_		)
											a__				      =		{
											    """class_labels""": [1],
											    """generator""": generator,
											    """num_inference_steps""": 2,
											    """output_type""": """numpy""",
											}
											return inputs


							def _a      (							self		):
											a__				      =		"""cpu"""

											a__				      =		self.get_dummy_components()
											a__				      =		self.pipeline_class(**a_		)
											pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

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

											self.assertEqual(image.shape  ,  (1, 16, 16, 3)		)
											a__				      =		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]		)
											a__				      =		np.abs(image_slice.flatten() - expected_slice		).max()
											self.assertLessEqual(a_  ,  1E-3		)


							def _a      (							self		):
											self._test_inference_batch_single_identical(relax_max_difference=a_  ,  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 _a      (							self		):
											self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3		)




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







							'''simple docstring'''






							def _a      (							self		):
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()


							def _a      (							self		):
											a__				      =		torch.manual_seed(0		)

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

											a__				      =		["""vase""", """umbrella""", """white shark""", """white wolf"""]
											a__				      =		pipe.get_label_ids(a_		)

											a__				      =		pipe(a_  ,  generator=a_  ,  num_inference_steps=40  ,  output_type="""np"""		).images

											for word, image in zip(a_  ,  a_		):
															a__				      =		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 _a      (							self		):
											a__				      =		DiTPipeline.from_pretrained("""facebook/DiT-XL-2-512"""		)
											a__				      =		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
											pipe.to("""cuda"""		)

											a__				      =		["""vase""", """umbrella"""]
											a__				      =		pipe.get_label_ids(a_		)

											a__				      =		torch.manual_seed(0		)
											a__				      =		pipe(a_  ,  generator=a_  ,  num_inference_steps=25  ,  output_type="""np"""		).images

											for word, image in zip(a_  ,  a_		):
															a__				      =		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






import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

class  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							def _a      (							self		):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		UNetaDModel(
											    sample_size=(32, 64)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")  ,  )
											return model


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		UNetaDConditionModel(
											    sample_size=(64, 32)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")  ,  cross_attention_dim=10  ,  )
											return model


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		AutoencoderKL(
											    sample_size=(128, 64)  ,  in_channels=1  ,  out_channels=1  ,  latent_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""DownEncoderBlock2D""", """DownEncoderBlock2D""")  ,  up_block_types=("""UpDecoderBlock2D""", """UpDecoderBlock2D""")  ,  )
											a__				      =		UNetaDModel(
											    sample_size=(64, 32)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")  ,  )
											return vqvae, unet







							@slow
							def _a      (							self		):
											a__				      =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
											a__				      =		Mel(
											    x_res=self.dummy_unet.config.sample_size[1]  ,  y_res=self.dummy_unet.config.sample_size[0]  ,  )

											a__				      =		DDPMScheduler()
											a__				      =		AudioDiffusionPipeline(vqvae=a_  ,  unet=self.dummy_unet  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_  ,  steps=4		)
											a__				      =		output.audios[0]
											a__				      =		output.images[0]

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_  ,  steps=4  ,  return_dict=a_		)
											a__				      =		output[0][0]

											assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
											assert (
											    image.height == self.dummy_unet.config.sample_size[0]
											    and image.width == self.dummy_unet.config.sample_size[1]
											)
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.frombuffer(image_from_tuple.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0
											assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() == 0

											a__				      =		Mel(
											    x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]  ,  y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]  ,  )

											a__				      =		DDIMScheduler()
											a__				      =		self.dummy_vqvae_and_unet
											a__				      =		AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]  ,  unet=dummy_vqvae_and_unet[1]  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											np.random.seed(0		)
											a__				      =		np.random.uniform(-1  ,  1  ,  ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)		)
											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(raw_audio=a_  ,  generator=a_  ,  start_step=5  ,  steps=10		)
											a__				      =		output.images[0]

											assert (
											    image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
											    and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
											)
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0

											a__				      =		self.dummy_unet_condition
											a__				      =		AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]  ,  unet=a_  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											np.random.seed(0		)
											a__				      =		torch.rand((1, 1, 10)		)
											a__				      =		pipe(generator=a_  ,  encoding=a_		)
											a__				      =		output.images[0]
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0




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







							'''simple docstring'''






							def _a      (							self		):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()







							def _a      (							self		):
											a__				      =		torch_device

											a__				      =		DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_		)
											a__				      =		output.audios[0]
											a__				      =		output.images[0]

											assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
											assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0

