





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class 	_UpperCAmelCase						( _lowercase			,							_lowercase			,							unittest.TestCase       ):
					__A        = StableDiffusionParadigmsPipeline
					__A        = TEXT_TO_IMAGE_PARAMS
					__A        = TEXT_TO_IMAGE_BATCH_PARAMS
					__A        = TEXT_TO_IMAGE_IMAGE_PARAMS
					__A        = TEXT_TO_IMAGE_IMAGE_PARAMS




					def 						snake_case_		(    self							)    ->	Optional[int]:
							torch.manual_seed(0							)
							UpperCamelCase__								=						UNetaDConditionModel(
							    block_out_channels=(32, 64)     ,   layers_per_block=2     ,   sample_size=32     ,   in_channels=4     ,   out_channels=4     ,   down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")     ,   up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")     ,   cross_attention_dim=32     ,   attention_head_dim=(2, 4)     ,   use_linear_projection=__lowerCAmelCase     ,   )
							UpperCamelCase__								=						DDIMParallelScheduler(
							    beta_start=0.0_0_0_8_5     ,   beta_end=0.0_1_2     ,   beta_schedule="""scaled_linear"""     ,   clip_sample=__lowerCAmelCase     ,   set_alpha_to_one=__lowerCAmelCase     ,   )
							torch.manual_seed(0							)
							UpperCamelCase__								=						AutoencoderKL(
							    block_out_channels=[32, 64]     ,   in_channels=3     ,   out_channels=3     ,   down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]     ,   up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]     ,   latent_channels=4     ,   sample_size=128     ,   )
							torch.manual_seed(0							)
							UpperCamelCase__								=						CLIPTextConfig(
							    bos_token_id=0     ,   eos_token_id=2     ,   hidden_size=32     ,   intermediate_size=37     ,   layer_norm_eps=1E-0_5     ,   num_attention_heads=4     ,   num_hidden_layers=5     ,   pad_token_id=1     ,   vocab_size=1_000     ,   hidden_act="""gelu"""     ,   projection_dim=512     ,   )
							UpperCamelCase__								=						CLIPTextModel(__lowerCAmelCase							)
							UpperCamelCase__								=						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)

							UpperCamelCase__								=						{
							    """unet""": unet,
							    """scheduler""": scheduler,
							    """vae""": vae,
							    """text_encoder""": text_encoder,
							    """tokenizer""": tokenizer,
							    """safety_checker""": None,
							    """feature_extractor""": None,
							}
							return components




					def 						snake_case_		(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase=0							)    ->	Union[str, Any]:
							if str(__lowerCAmelCase							).startswith("""mps"""							):
									UpperCamelCase__								=						torch.manual_seed(__lowerCAmelCase							)
							else:
									UpperCamelCase__								=						torch.Generator(device=__lowerCAmelCase							).manual_seed(__lowerCAmelCase							)
							UpperCamelCase__								=						{
							    """prompt""": """a photograph of an astronaut riding a horse""",
							    """generator""": generator,
							    """num_inference_steps""": 10,
							    """guidance_scale""": 6.0,
							    """output_type""": """numpy""",
							    """parallel""": 3,
							    """debug""": True,
							}
							return inputs




					def 						snake_case_		(    self							)    ->	Optional[Any]:
							UpperCamelCase__								=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCamelCase__								=						self.get_dummy_components()
							UpperCamelCase__								=						StableDiffusionParadigmsPipeline(**__lowerCAmelCase							)
							UpperCamelCase__								=						sd_pipe.to(__lowerCAmelCase							)
							sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

							UpperCamelCase__								=						self.get_dummy_inputs(__lowerCAmelCase							)
							UpperCamelCase__								=						sd_pipe(**__lowerCAmelCase							).images
							UpperCamelCase__								=						image[0, -3:, -3:, -1]
							assert image.shape == (1, 64, 64, 3)

							UpperCamelCase__								=						np.array([0.4_7_7_3, 0.5_4_1_7, 0.4_7_2_3, 0.4_9_2_5, 0.5_6_3_1, 0.4_7_5_2, 0.5_2_4_0, 0.4_9_3_5, 0.5_0_2_3]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2




					def 						snake_case_		(    self							)    ->	Optional[int]:
							UpperCamelCase__								=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCamelCase__								=						self.get_dummy_components()
							torch.manual_seed(0							)
							UpperCamelCase__								=						DDPMParallelScheduler()
							torch.manual_seed(0							)
							UpperCamelCase__								=						StableDiffusionParadigmsPipeline(**__lowerCAmelCase							)
							UpperCamelCase__								=						sd_pipe.to(__lowerCAmelCase							)
							sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

							UpperCamelCase__								=						self.get_dummy_inputs(__lowerCAmelCase							)
							UpperCamelCase__								=						sd_pipe(**__lowerCAmelCase							).images
							UpperCamelCase__								=						image[0, -3:, -3:, -1]
							assert image.shape == (1, 64, 64, 3)

							UpperCamelCase__								=						np.array([0.3_5_7_3, 0.4_4_2_0, 0.4_9_6_0, 0.4_7_9_9, 0.3_7_9_6, 0.3_8_7_9, 0.4_8_1_9, 0.4_3_6_5, 0.4_4_6_8]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2




					def 						snake_case_		(    self							)    ->	str:
							super().test_inference_batch_consistent(batch_sizes=[1, 2]							)




					def 						snake_case_		(    self							)    ->	Dict:
							super().test_inference_batch_single_identical(batch_size=2     ,   expected_max_diff=3E-3							)




					def 						snake_case_		(    self							)    ->	Tuple:
							UpperCamelCase__								=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCamelCase__								=						self.get_dummy_components()
							UpperCamelCase__								=						StableDiffusionParadigmsPipeline(**__lowerCAmelCase							)
							UpperCamelCase__								=						sd_pipe.to(__lowerCAmelCase							)
							sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

							UpperCamelCase__								=						self.get_dummy_inputs(__lowerCAmelCase							)
							UpperCamelCase__								=						"""french fries"""
							UpperCamelCase__								=						sd_pipe(**__lowerCAmelCase     ,   negative_prompt=__lowerCAmelCase							)
							UpperCamelCase__								=						output.images
							UpperCamelCase__								=						image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)

							UpperCamelCase__								=						np.array([0.4_7_7_1, 0.5_4_2_0, 0.4_6_8_3, 0.4_9_1_8, 0.5_6_3_6, 0.4_7_2_5, 0.5_2_3_0, 0.4_9_2_3, 0.5_0_1_5]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2




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




					def 						snake_case_		(    self							)    ->	List[Any]:
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()




					def 						snake_case_		(    self     ,   __lowerCAmelCase=0							)    ->	Any:
							UpperCamelCase__								=						torch.Generator(device=__lowerCAmelCase							).manual_seed(__lowerCAmelCase							)
							UpperCamelCase__								=						{
							    """prompt""": """a photograph of an astronaut riding a horse""",
							    """generator""": generator,
							    """num_inference_steps""": 10,
							    """guidance_scale""": 7.5,
							    """output_type""": """numpy""",
							    """parallel""": 3,
							    """debug""": True,
							}
							return inputs




					def 						snake_case_		(    self							)    ->	Any:
							UpperCamelCase__								=						"""stabilityai/stable-diffusion-2-base"""
							UpperCamelCase__								=						DDIMParallelScheduler.from_pretrained(__lowerCAmelCase     ,   subfolder="""scheduler"""							)
							UpperCamelCase__								=						StableDiffusionParadigmsPipeline.from_pretrained(__lowerCAmelCase     ,   scheduler=__lowerCAmelCase     ,   safety_checker=__lowerCAmelCase							)
							pipe.to(__lowerCAmelCase							)
							pipe.set_progress_bar_config(disable=__lowerCAmelCase							)
							pipe.enable_attention_slicing()

							UpperCamelCase__								=						self.get_inputs()
							UpperCamelCase__								=						pipe(**__lowerCAmelCase							).images
							UpperCamelCase__								=						image[0, -3:, -3:, -1].flatten()

							assert image.shape == (1, 512, 512, 3)

							UpperCamelCase__								=						np.array([0.9_6_2_2, 0.9_6_0_2, 0.9_7_4_8, 0.9_5_9_1, 0.9_6_3_0, 0.9_6_9_1, 0.9_6_6_1, 0.9_6_3_1, 0.9_7_4_1]							)

							assert np.abs(expected_slice - image_slice							).max() < 1E-2












import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, slow, torch_device

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin




class 	_UpperCAmelCase						( _lowercase			,							_lowercase			,							unittest.TestCase       ):
					__A        = UNetaDModel
					__A        = "sample"




					@property
					def 						snake_case_		(    self							)    ->	Any:
							UpperCamelCase__								=						4
							UpperCamelCase__								=						14
							UpperCamelCase__								=						16

							UpperCamelCase__								=						floats_tensor((batch_size, num_features, seq_len)							).to(__lowerCAmelCase							)
							UpperCamelCase__								=						torch.tensor([10] * batch_size							).to(__lowerCAmelCase							)

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




					@property
					def 						snake_case_		(    self							)    ->	Optional[int]:
							return (4, 14, 16)




					@property
					def 						snake_case_		(    self							)    ->	Any:
							return (4, 14, 16)




					def 						snake_case_		(    self							)    ->	Any:
							pass




					def 						snake_case_		(    self							)    ->	Tuple:
							pass




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Any:
							super().test_determinism()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	int:
							super().test_outputs_equivalence()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[Any]:
							super().test_from_save_pretrained()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Union[str, Any]:
							super().test_from_save_pretrained_variant()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[int]:
							super().test_model_from_pretrained()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Dict:
							super().test_output()




					def 						snake_case_		(    self							)    ->	int:
							UpperCamelCase__								=						{
							    """block_out_channels""": (32, 64, 128, 256),
							    """in_channels""": 14,
							    """out_channels""": 14,
							    """time_embedding_type""": """positional""",
							    """use_timestep_embedding""": True,
							    """flip_sin_to_cos""": False,
							    """freq_shift""": 1.0,
							    """out_block_type""": """OutConv1DBlock""",
							    """mid_block_type""": """MidResTemporalBlock1D""",
							    """down_block_types""": ("""DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D"""),
							    """up_block_types""": ("""UpResnetBlock1D""", """UpResnetBlock1D""", """UpResnetBlock1D"""),
							    """act_fn""": """mish""",
							}
							UpperCamelCase__								=						self.dummy_input
							return init_dict, inputs_dict




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	List[Any]:
							UpperCamelCase__	,      UpperCamelCase__								=						UNetaDModel.from_pretrained(
							    """bglick13/hopper-medium-v2-value-function-hor32"""     ,   output_loading_info=__lowerCAmelCase     ,   subfolder="""unet"""							)
							self.assertIsNotNone(__lowerCAmelCase							)
							self.assertEqual(len(loading_info["""missing_keys"""]							)     ,   0							)

							model.to(__lowerCAmelCase							)
							UpperCamelCase__								=						model(**self.dummy_input							)

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




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[Any]:
							UpperCamelCase__								=						UNetaDModel.from_pretrained("""bglick13/hopper-medium-v2-value-function-hor32"""     ,   subfolder="""unet"""							)
							torch.manual_seed(0							)
							if torch.cuda.is_available():
									torch.cuda.manual_seed_all(0							)

							UpperCamelCase__								=						model.config.in_channels
							UpperCamelCase__								=						16
							UpperCamelCase__								=						torch.randn((1, seq_len, num_features)							).permute(
							    0     ,   2     ,   1							)  # match original, we can update values and remove
							UpperCamelCase__								=						torch.full((num_features,)     ,   0							)

							with torch.no_grad():
									UpperCamelCase__								=						model(__lowerCAmelCase     ,   __lowerCAmelCase							).sample.permute(0     ,   2     ,   1							)

							UpperCamelCase__								=						output[0, -3:, -3:].flatten()
							# fmt: off
							UpperCamelCase__								=						torch.tensor([-2.1_3_7_1_7_2, 1.1_4_2_6_0_1_6, 0.3_6_8_8_6_8_7, -0.7_6_6_9_2_2, 0.7_3_0_3_1_4_6, 0.1_1_0_3_8_8_6_4, -0.4_7_6_0_6_3_3, 0.1_3_2_7_0_1_7_2, 0.0_2_5_9_1_3_4_8]							)
							# fmt: on
							self.assertTrue(torch.allclose(__lowerCAmelCase     ,   __lowerCAmelCase     ,   rtol=1E-3							)							)




					def 						snake_case_		(    self							)    ->	int:
							# Not implemented yet for this UNet
							pass




					@slow
					def 						snake_case_		(    self							)    ->	Tuple:
							UpperCamelCase__								=						"""harmonai/maestro-150k"""
							UpperCamelCase__								=						UNetaDModel.from_pretrained(__lowerCAmelCase     ,   subfolder="""unet"""							)
							model.to(__lowerCAmelCase							)

							UpperCamelCase__								=						65_536
							UpperCamelCase__								=						torch.sin(torch.arange(__lowerCAmelCase							)[None, None, :].repeat(1     ,   2     ,   1							)							).to(__lowerCAmelCase							)
							UpperCamelCase__								=						torch.tensor([1]							).to(__lowerCAmelCase							)

							with torch.no_grad():
									UpperCamelCase__								=						model(__lowerCAmelCase     ,   __lowerCAmelCase							).sample

							UpperCamelCase__								=						output.abs().sum()
							UpperCamelCase__								=						output.abs().max()

							assert (output_sum - 2_2_4.0_8_9_6).abs() < 0.5
							assert (output_max - 0.0_6_0_7).abs() < 4E-4




class 	_UpperCAmelCase						( _lowercase			,							_lowercase			,							unittest.TestCase       ):
					__A        = UNetaDModel
					__A        = "sample"




					@property
					def 						snake_case_		(    self							)    ->	Any:
							UpperCamelCase__								=						4
							UpperCamelCase__								=						14
							UpperCamelCase__								=						16

							UpperCamelCase__								=						floats_tensor((batch_size, num_features, seq_len)							).to(__lowerCAmelCase							)
							UpperCamelCase__								=						torch.tensor([10] * batch_size							).to(__lowerCAmelCase							)

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




					@property
					def 						snake_case_		(    self							)    ->	Union[str, Any]:
							return (4, 14, 16)




					@property
					def 						snake_case_		(    self							)    ->	List[Any]:
							return (4, 14, 1)




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Dict:
							super().test_determinism()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[int]:
							super().test_outputs_equivalence()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[int]:
							super().test_from_save_pretrained()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[Any]:
							super().test_from_save_pretrained_variant()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[Any]:
							super().test_model_from_pretrained()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	List[str]:
							# UNetRL is a value-function is different output shape
							UpperCamelCase__	,      UpperCamelCase__								=						self.prepare_init_args_and_inputs_for_common()
							UpperCamelCase__								=						self.model_class(**__lowerCAmelCase							)
							model.to(__lowerCAmelCase							)
							model.eval()

							with torch.no_grad():
									UpperCamelCase__								=						model(**__lowerCAmelCase							)

									if isinstance(__lowerCAmelCase     ,   __lowerCAmelCase							):
											UpperCamelCase__								=						output.sample

							self.assertIsNotNone(__lowerCAmelCase							)
							UpperCamelCase__								=						torch.Size((inputs_dict["""sample"""].shape[0], 1)							)
							self.assertEqual(output.shape     ,   __lowerCAmelCase     ,   """Input and output shapes do not match"""							)




					def 						snake_case_		(    self							)    ->	Optional[int]:
							pass




					def 						snake_case_		(    self							)    ->	Union[str, Any]:
							pass




					def 						snake_case_		(    self							)    ->	Dict:
							UpperCamelCase__								=						{
							    """in_channels""": 14,
							    """out_channels""": 14,
							    """down_block_types""": ["""DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D"""],
							    """up_block_types""": [],
							    """out_block_type""": """ValueFunction""",
							    """mid_block_type""": """ValueFunctionMidBlock1D""",
							    """block_out_channels""": [32, 64, 128, 256],
							    """layers_per_block""": 1,
							    """downsample_each_block""": True,
							    """use_timestep_embedding""": True,
							    """freq_shift""": 1.0,
							    """flip_sin_to_cos""": False,
							    """time_embedding_type""": """positional""",
							    """act_fn""": """mish""",
							}
							UpperCamelCase__								=						self.dummy_input
							return init_dict, inputs_dict




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Tuple:
							UpperCamelCase__	,      UpperCamelCase__								=						UNetaDModel.from_pretrained(
							    """bglick13/hopper-medium-v2-value-function-hor32"""     ,   output_loading_info=__lowerCAmelCase     ,   subfolder="""value_function"""							)
							self.assertIsNotNone(__lowerCAmelCase							)
							self.assertEqual(len(vf_loading_info["""missing_keys"""]							)     ,   0							)

							value_function.to(__lowerCAmelCase							)
							UpperCamelCase__								=						value_function(**self.dummy_input							)

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




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	List[Any]:
							UpperCamelCase__	,      UpperCamelCase__								=						UNetaDModel.from_pretrained(
							    """bglick13/hopper-medium-v2-value-function-hor32"""     ,   output_loading_info=__lowerCAmelCase     ,   subfolder="""value_function"""							)
							torch.manual_seed(0							)
							if torch.cuda.is_available():
									torch.cuda.manual_seed_all(0							)

							UpperCamelCase__								=						value_function.config.in_channels
							UpperCamelCase__								=						14
							UpperCamelCase__								=						torch.randn((1, seq_len, num_features)							).permute(
							    0     ,   2     ,   1							)  # match original, we can update values and remove
							UpperCamelCase__								=						torch.full((num_features,)     ,   0							)

							with torch.no_grad():
									UpperCamelCase__								=						value_function(__lowerCAmelCase     ,   __lowerCAmelCase							).sample

							# fmt: off
							UpperCamelCase__								=						torch.tensor([1_6_5.2_5] * seq_len							)
							# fmt: on
							self.assertTrue(torch.allclose(__lowerCAmelCase     ,   __lowerCAmelCase     ,   rtol=1E-3							)							)




					def 						snake_case_		(    self							)    ->	Union[str, Any]:
							# Not implemented yet for this UNet
							pass





