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     snake_case	(			__UpperCAmelCase				,   __UpperCAmelCase				,   unittest.TestCase    ):




					"""simple docstring"""

					snake_case__       =			UNetaDModel
					snake_case__       =			"sample"







					@property
					def 		__lowerCAmelCase						(    self       :     Optional[int]      ):
									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 		__lowerCAmelCase						(    self       :     str      ):
									return (4, 14, 16)







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







					def 		__lowerCAmelCase						(    self       :     List[str]      ):
									pass







					def 		__lowerCAmelCase						(    self       :     int      ):
									pass







					@unittest.skipIf(torch_device == 'mps'   ,'mish op not supported in MPS'      )
					def 		__lowerCAmelCase						(    self       :     Any      ):
									super().test_determinism()







					@unittest.skipIf(torch_device == 'mps'   ,'mish op not supported in MPS'      )
					def 		__lowerCAmelCase						(    self       :     List[Any]      ):
									super().test_outputs_equivalence()







					@unittest.skipIf(torch_device == 'mps'   ,'mish op not supported in MPS'      )
					def 		__lowerCAmelCase						(    self       :     str      ):
									super().test_from_save_pretrained()







					@unittest.skipIf(torch_device == 'mps'   ,'mish op not supported in MPS'      )
					def 		__lowerCAmelCase						(    self       :     Dict      ):
									super().test_from_save_pretrained_variant()







					@unittest.skipIf(torch_device == 'mps'   ,'mish op not supported in MPS'      )
					def 		__lowerCAmelCase						(    self       :     Dict      ):
									super().test_model_from_pretrained()







					@unittest.skipIf(torch_device == 'mps'   ,'mish op not supported in MPS'      )
					def 		__lowerCAmelCase						(    self       :     Tuple      ):
									super().test_output()







					def 		__lowerCAmelCase						(    self       :     List[str]      ):
									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 		__lowerCAmelCase						(    self       :     Union[str, 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 		__lowerCAmelCase						(    self       :     Union[str, 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 		__lowerCAmelCase						(    self       :     Optional[Any]      ):
									# Not implemented yet for this UNet
									pass







					@slow
					def 		__lowerCAmelCase						(    self       :     List[Any]      ):
									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 - 224.0_896).abs() < 0.5
									assert (output_max - 0.0_6_0_7).abs() < 4e-4






class     snake_case	(			__UpperCAmelCase				,   __UpperCAmelCase				,   unittest.TestCase    ):




					"""simple docstring"""

					snake_case__       =			UNetaDModel
					snake_case__       =			"sample"







					@property
					def 		__lowerCAmelCase						(    self       :     Dict      ):
									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 		__lowerCAmelCase						(    self       :     Tuple      ):
									return (4, 14, 16)







					@property
					def 		__lowerCAmelCase						(    self       :     List[str]      ):
									return (4, 14, 1)







					@unittest.skipIf(torch_device == 'mps'   ,'mish op not supported in MPS'      )
					def 		__lowerCAmelCase						(    self       :     Dict      ):
									super().test_determinism()







					@unittest.skipIf(torch_device == 'mps'   ,'mish op not supported in MPS'      )
					def 		__lowerCAmelCase						(    self       :     int      ):
									super().test_outputs_equivalence()







					@unittest.skipIf(torch_device == 'mps'   ,'mish op not supported in MPS'      )
					def 		__lowerCAmelCase						(    self       :     List[Any]      ):
									super().test_from_save_pretrained()







					@unittest.skipIf(torch_device == 'mps'   ,'mish op not supported in MPS'      )
					def 		__lowerCAmelCase						(    self       :     int      ):
									super().test_from_save_pretrained_variant()







					@unittest.skipIf(torch_device == 'mps'   ,'mish op not supported in MPS'      )
					def 		__lowerCAmelCase						(    self       :     List[str]      ):
									super().test_model_from_pretrained()







					@unittest.skipIf(torch_device == 'mps'   ,'mish op not supported in MPS'      )
					def 		__lowerCAmelCase						(    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 		__lowerCAmelCase						(    self       :     List[Any]      ):
									pass







					def 		__lowerCAmelCase						(    self       :     int      ):
									pass







					def 		__lowerCAmelCase						(    self       :     Optional[Any]      ):
									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 		__lowerCAmelCase						(    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 		__lowerCAmelCase						(    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([165.25] * seq_len      )
									# fmt: on
									self.assertTrue(torch.allclose(lowerCamelCase__   ,lowerCamelCase__   ,rtol=1e-3      )      )







					def 		__lowerCAmelCase						(    self       :     str      ):
									# Not implemented yet for this UNet
									pass




"""simple docstring"""






import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()
class     snake_case	(			__UpperCAmelCase				,   unittest.TestCase    ):




					"""simple docstring"""

					snake_case__       =			KandinskyPriorPipeline
					snake_case__       =			["prompt"]
					snake_case__       =			["prompt", "negative_prompt"]
					snake_case__       =			[
					    "num_images_per_prompt",
					    "generator",
					    "num_inference_steps",
					    "latents",
					    "negative_prompt",
					    "guidance_scale",
					    "output_type",
					    "return_dict",
					]
					snake_case__       =			False







					@property
					def 		__lowerCAmelCase						(    self       :     Optional[int]      ):
									return 32







					@property
					def 		__lowerCAmelCase						(    self       :     int      ):
									return 32







					@property
					def 		__lowerCAmelCase						(    self       :     Tuple      ):
									return self.time_input_dim







					@property
					def 		__lowerCAmelCase						(    self       :     Tuple      ):
									return self.time_input_dim * 4







					@property
					def 		__lowerCAmelCase						(    self       :     Optional[int]      ):
									return 100







					@property
					def 		__lowerCAmelCase						(    self       :     List[str]      ):
									UpperCAmelCase__					=							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )
									return tokenizer







					@property
					def 		__lowerCAmelCase						(    self       :     int      ):
									torch.manual_seed(0      )
									UpperCAmelCase__					=							CLIPTextConfig(
									    bos_token_id=0   ,eos_token_id=2   ,hidden_size=self.text_embedder_hidden_size   ,projection_dim=self.text_embedder_hidden_size   ,intermediate_size=37   ,layer_norm_eps=1e-05   ,num_attention_heads=4   ,num_hidden_layers=5   ,pad_token_id=1   ,vocab_size=1_000   ,)
									return CLIPTextModelWithProjection(lowerCamelCase__      )







					@property
					def 		__lowerCAmelCase						(    self       :     Tuple      ):
									torch.manual_seed(0      )

									UpperCAmelCase__					=							{
									    'num_attention_heads': 2,
									    'attention_head_dim': 12,
									    'embedding_dim': self.text_embedder_hidden_size,
									    'num_layers': 1,
									}

									UpperCAmelCase__					=							PriorTransformer(**lowerCamelCase__      )
									# clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
									UpperCAmelCase__					=							nn.Parameter(torch.ones(model.clip_std.shape      )      )
									return model







					@property
					def 		__lowerCAmelCase						(    self       :     Dict      ):
									torch.manual_seed(0      )
									UpperCAmelCase__					=							CLIPVisionConfig(
									    hidden_size=self.text_embedder_hidden_size   ,image_size=224   ,projection_dim=self.text_embedder_hidden_size   ,intermediate_size=37   ,num_attention_heads=4   ,num_channels=3   ,num_hidden_layers=5   ,patch_size=14   ,)

									UpperCAmelCase__					=							CLIPVisionModelWithProjection(lowerCamelCase__      )
									return model







					@property
					def 		__lowerCAmelCase						(    self       :     Dict      ):
									UpperCAmelCase__					=							CLIPImageProcessor(
									    crop_size=224   ,do_center_crop=lowerCamelCase__   ,do_normalize=lowerCamelCase__   ,do_resize=lowerCamelCase__   ,image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3]   ,image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1]   ,resample=3   ,size=224   ,)

									return image_processor







					def 		__lowerCAmelCase						(    self       :     List[Any]      ):
									UpperCAmelCase__					=							self.dummy_prior
									UpperCAmelCase__					=							self.dummy_image_encoder
									UpperCAmelCase__					=							self.dummy_text_encoder
									UpperCAmelCase__					=							self.dummy_tokenizer
									UpperCAmelCase__					=							self.dummy_image_processor

									UpperCAmelCase__					=							UnCLIPScheduler(
									    variance_type='fixed_small_log'   ,prediction_type='sample'   ,num_train_timesteps=1_000   ,clip_sample=lowerCamelCase__   ,clip_sample_range=1_0.0   ,)

									UpperCAmelCase__					=							{
									    'prior': prior,
									    'image_encoder': image_encoder,
									    'text_encoder': text_encoder,
									    'tokenizer': tokenizer,
									    'scheduler': scheduler,
									    'image_processor': image_processor,
									}

									return components







					def 		__lowerCAmelCase						(    self       :     Dict   ,lowerCamelCase__       :     Union[str, Any]   ,lowerCamelCase__       :     Optional[int]=0      ):
									if str(lowerCamelCase__      ).startswith('mps'      ):
													UpperCAmelCase__					=							torch.manual_seed(lowerCamelCase__      )
									else:
													UpperCAmelCase__					=							torch.Generator(device=lowerCamelCase__      ).manual_seed(lowerCamelCase__      )
									UpperCAmelCase__					=							{
									    'prompt': 'horse',
									    'generator': generator,
									    'guidance_scale': 4.0,
									    'num_inference_steps': 2,
									    'output_type': 'np',
									}
									return inputs







					def 		__lowerCAmelCase						(    self       :     List[str]      ):
									UpperCAmelCase__					=							'cpu'

									UpperCAmelCase__					=							self.get_dummy_components()

									UpperCAmelCase__					=							self.pipeline_class(**lowerCamelCase__      )
									UpperCAmelCase__					=							pipe.to(lowerCamelCase__      )

									pipe.set_progress_bar_config(disable=lowerCamelCase__      )

									UpperCAmelCase__					=							pipe(**self.get_dummy_inputs(lowerCamelCase__      )      )
									UpperCAmelCase__					=							output.image_embeds

									UpperCAmelCase__					=							pipe(
									    **self.get_dummy_inputs(lowerCamelCase__      )   ,return_dict=lowerCamelCase__   ,)[0]

									UpperCAmelCase__					=							image[0, -10:]
									UpperCAmelCase__					=							image_from_tuple[0, -10:]

									assert image.shape == (1, 32)

									UpperCAmelCase__					=							np.array(
									    [-0.0_5_3_2, 1.7_1_2_0, 0.3_6_5_6, -1.0_8_5_2, -0.8_9_4_6, -1.1_7_5_6, 0.4_3_4_8, 0.2_4_8_2, 0.5_1_4_6, -0.1_1_5_6]      )

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







					@skip_mps
					def 		__lowerCAmelCase						(    self       :     Optional[int]      ):
									UpperCAmelCase__					=							torch_device == 'cpu'
									UpperCAmelCase__					=							True
									UpperCAmelCase__					=							False

									self._test_inference_batch_single_identical(
									    test_max_difference=lowerCamelCase__   ,relax_max_difference=lowerCamelCase__   ,test_mean_pixel_difference=lowerCamelCase__   ,)







					@skip_mps
					def 		__lowerCAmelCase						(    self       :     Optional[int]      ):
									UpperCAmelCase__					=							torch_device == 'cpu'
									UpperCAmelCase__					=							False

									self._test_attention_slicing_forward_pass(
									    test_max_difference=lowerCamelCase__   ,test_mean_pixel_difference=lowerCamelCase__   ,)



