

import gc
import unittest

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

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()



class    A__(a_,       unittest.TestCase    ):
						"""simple docstring"""






						_A  :							Optional[int]   			=	UnCLIPPipeline
						_A  :							Tuple   			=	TEXT_TO_IMAGE_PARAMS - {
						    '''negative_prompt''',
						    '''height''',
						    '''width''',
						    '''negative_prompt_embeds''',
						    '''guidance_scale''',
						    '''prompt_embeds''',
						    '''cross_attention_kwargs''',
						}
						_A  :							Optional[int]   			=	TEXT_TO_IMAGE_BATCH_PARAMS
						_A  :							str   			=	[
						    '''generator''',
						    '''return_dict''',
						    '''prior_num_inference_steps''',
						    '''decoder_num_inference_steps''',
						    '''super_res_num_inference_steps''',
						]
						_A  :							Optional[int]   			=	False





						@property
						def 	UpperCamelCase__			(  self  )      ->				Any:
									return 32





						@property
						def 	UpperCamelCase__			(  self  )      ->				str:
									return 32





						@property
						def 	UpperCamelCase__			(  self  )      ->				List[Any]:
									return self.time_input_dim





						@property
						def 	UpperCamelCase__			(  self  )      ->				Any:
									return self.time_input_dim * 4





						@property
						def 	UpperCamelCase__			(  self  )      ->				str:
									return 100





						@property
						def 	UpperCamelCase__			(  self  )      ->				List[str]:
									a_		:		Optional[int]				    =				CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  )
									return tokenizer





						@property
						def 	UpperCamelCase__			(  self  )      ->				int:
									torch.manual_seed(0  )
									a_		:		str				    =				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(_lowercase  )





						@property
						def 	UpperCamelCase__			(  self  )      ->				Optional[int]:
									torch.manual_seed(0  )

									a_		:		Optional[int]				    =				{
									    """num_attention_heads""": 2,
									    """attention_head_dim""": 12,
									    """embedding_dim""": self.text_embedder_hidden_size,
									    """num_layers""": 1,
									}

									a_		:		Any				    =				PriorTransformer(**_lowercase  )
									return model





						@property
						def 	UpperCamelCase__			(  self  )      ->				Dict:
									torch.manual_seed(0  )

									a_		:		Union[str, Any]				    =				{
									    """clip_embeddings_dim""": self.text_embedder_hidden_size,
									    """time_embed_dim""": self.time_embed_dim,
									    """cross_attention_dim""": self.cross_attention_dim,
									}

									a_		:		List[Any]				    =				UnCLIPTextProjModel(**_lowercase  )
									return model





						@property
						def 	UpperCamelCase__			(  self  )      ->				Any:
									torch.manual_seed(0  )

									a_		:		Tuple				    =				{
									    """sample_size""": 32,
									    # RGB in channels
									    """in_channels""": 3,
									    # Out channels is double in channels because predicts mean and variance
									    """out_channels""": 6,
									    """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
									    """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
									    """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
									    """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
									    """layers_per_block""": 1,
									    """cross_attention_dim""": self.cross_attention_dim,
									    """attention_head_dim""": 4,
									    """resnet_time_scale_shift""": """scale_shift""",
									    """class_embed_type""": """identity""",
									}

									a_		:		Any				    =				UNetaDConditionModel(**_lowercase  )
									return model





						@property
						def 	UpperCamelCase__			(  self  )      ->				List[Any]:
									return {
									    "sample_size": 64,
									    "layers_per_block": 1,
									    "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
									    "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
									    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
									    "in_channels": 6,
									    "out_channels": 3,
									}





						@property
						def 	UpperCamelCase__			(  self  )      ->				Dict:
									torch.manual_seed(0  )

									a_		:		List[str]				    =				UNetaDModel(**self.dummy_super_res_kwargs  )
									return model





						@property
						def 	UpperCamelCase__			(  self  )      ->				Dict:
									# seeded differently to get different unet than `self.dummy_super_res_first`
									torch.manual_seed(1  )

									a_		:		str				    =				UNetaDModel(**self.dummy_super_res_kwargs  )
									return model





						def 	UpperCamelCase__			(  self  )      ->				Optional[int]:
									a_		:		Union[str, Any]				    =				self.dummy_prior
									a_		:		Any				    =				self.dummy_decoder
									a_		:		Tuple				    =				self.dummy_text_proj
									a_		:		int				    =				self.dummy_text_encoder
									a_		:		int				    =				self.dummy_tokenizer
									a_		:		Optional[int]				    =				self.dummy_super_res_first
									a_		:		Tuple				    =				self.dummy_super_res_last

									a_		:		Union[str, Any]				    =				UnCLIPScheduler(
									    variance_type="""fixed_small_log"""       ,     prediction_type="""sample"""       ,     num_train_timesteps=1_000       ,     clip_sample_range=5.0       ,     )

									a_		:		Any				    =				UnCLIPScheduler(
									    variance_type="""learned_range"""       ,     prediction_type="""epsilon"""       ,     num_train_timesteps=1_000       ,     )

									a_		:		List[str]				    =				UnCLIPScheduler(
									    variance_type="""fixed_small_log"""       ,     prediction_type="""epsilon"""       ,     num_train_timesteps=1_000       ,     )

									a_		:		str				    =				{
									    """prior""": prior,
									    """decoder""": decoder,
									    """text_proj""": text_proj,
									    """text_encoder""": text_encoder,
									    """tokenizer""": tokenizer,
									    """super_res_first""": super_res_first,
									    """super_res_last""": super_res_last,
									    """prior_scheduler""": prior_scheduler,
									    """decoder_scheduler""": decoder_scheduler,
									    """super_res_scheduler""": super_res_scheduler,
									}

									return components





						def 	UpperCamelCase__			(  self       ,     _lowercase       ,     _lowercase=0  )      ->				List[str]:
									if str(_lowercase  ).startswith("""mps"""  ):
												a_		:		str				    =				torch.manual_seed(_lowercase  )
									else:
												a_		:		Any				    =				torch.Generator(device=_lowercase  ).manual_seed(_lowercase  )
									a_		:		Optional[Any]				    =				{
									    """prompt""": """horse""",
									    """generator""": generator,
									    """prior_num_inference_steps""": 2,
									    """decoder_num_inference_steps""": 2,
									    """super_res_num_inference_steps""": 2,
									    """output_type""": """numpy""",
									}
									return inputs





						def 	UpperCamelCase__			(  self  )      ->				Optional[int]:
									a_		:		Optional[Any]				    =				"""cpu"""

									a_		:		Dict				    =				self.get_dummy_components()

									a_		:		Optional[int]				    =				self.pipeline_class(**_lowercase  )
									a_		:		List[Any]				    =				pipe.to(_lowercase  )

									pipe.set_progress_bar_config(disable=_lowercase  )

									a_		:		Any				    =				pipe(**self.get_dummy_inputs(_lowercase  )  )
									a_		:		List[str]				    =				output.images

									a_		:		List[str]				    =				pipe(
									    **self.get_dummy_inputs(_lowercase  )       ,     return_dict=_lowercase       ,     )[0]

									a_		:		Any				    =				image[0, -3:, -3:, -1]
									a_		:		Any				    =				image_from_tuple[0, -3:, -3:, -1]

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

									a_		:		List[str]				    =				np.array(
									    [
									        0.9_9_9_7,
									        0.9_9_8_8,
									        0.0_0_2_8,
									        0.9_9_9_7,
									        0.9_9_8_4,
									        0.9_9_6_5,
									        0.0_0_2_9,
									        0.9_9_8_6,
									        0.0_0_2_5,
									    ]  )

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





						def 	UpperCamelCase__			(  self  )      ->				int:
									a_		:		Dict				    =				torch.device("""cpu"""  )



									class    A__:
															"""simple docstring"""






															_A  :							str   			=	1







									a_		:		str				    =				self.get_dummy_components()

									a_		:		List[str]				    =				self.pipeline_class(**_lowercase  )
									a_		:		str				    =				pipe.to(_lowercase  )

									a_		:		Any				    =				components["""prior"""]
									a_		:		Union[str, Any]				    =				components["""decoder"""]
									a_		:		List[Any]				    =				components["""super_res_first"""]
									a_		:		str				    =				components["""tokenizer"""]
									a_		:		Any				    =				components["""text_encoder"""]

									a_		:		str				    =				torch.Generator(device=_lowercase  ).manual_seed(0  )
									a_		:		Optional[int]				    =				prior.dtype
									a_		:		Optional[Any]				    =				1

									a_		:		Union[str, Any]				    =				(batch_size, prior.config.embedding_dim)
									a_		:		int				    =				pipe.prepare_latents(
									    _lowercase       ,     dtype=_lowercase       ,     device=_lowercase       ,     generator=_lowercase       ,     latents=_lowercase       ,     scheduler=DummyScheduler()  )
									a_		:		Optional[int]				    =				(batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
									a_		:		int				    =				pipe.prepare_latents(
									    _lowercase       ,     dtype=_lowercase       ,     device=_lowercase       ,     generator=_lowercase       ,     latents=_lowercase       ,     scheduler=DummyScheduler()  )

									a_		:		Tuple				    =				(
									    batch_size,
									    super_res_first.config.in_channels // 2,
									    super_res_first.config.sample_size,
									    super_res_first.config.sample_size,
									)
									a_		:		Any				    =				pipe.prepare_latents(
									    _lowercase       ,     dtype=_lowercase       ,     device=_lowercase       ,     generator=_lowercase       ,     latents=_lowercase       ,     scheduler=DummyScheduler()  )

									pipe.set_progress_bar_config(disable=_lowercase  )

									a_		:		Optional[int]				    =				"""this is a prompt example"""

									a_		:		Optional[Any]				    =				torch.Generator(device=_lowercase  ).manual_seed(0  )
									a_		:		Optional[Any]				    =				pipe(
									    [prompt]       ,     generator=_lowercase       ,     prior_num_inference_steps=2       ,     decoder_num_inference_steps=2       ,     super_res_num_inference_steps=2       ,     prior_latents=_lowercase       ,     decoder_latents=_lowercase       ,     super_res_latents=_lowercase       ,     output_type="""np"""       ,     )
									a_		:		str				    =				output.images

									a_		:		Optional[int]				    =				tokenizer(
									    _lowercase       ,     padding="""max_length"""       ,     max_length=tokenizer.model_max_length       ,     return_tensors="""pt"""       ,     )
									a_		:		Dict				    =				text_encoder(text_inputs.input_ids  )
									a_		:		Optional[Any]				    =				text_inputs.attention_mask

									a_		:		str				    =				torch.Generator(device=_lowercase  ).manual_seed(0  )
									a_		:		Tuple				    =				pipe(
									    generator=_lowercase       ,     prior_num_inference_steps=2       ,     decoder_num_inference_steps=2       ,     super_res_num_inference_steps=2       ,     prior_latents=_lowercase       ,     decoder_latents=_lowercase       ,     super_res_latents=_lowercase       ,     text_model_output=_lowercase       ,     text_attention_mask=_lowercase       ,     output_type="""np"""       ,     )[0]

									# make sure passing text embeddings manually is identical
									assert np.abs(image - image_from_text  ).max() < 1e-4





						@skip_mps
						def 	UpperCamelCase__			(  self  )      ->				Tuple:
									a_		:		Any				    =				torch_device == """cpu"""

									self._test_attention_slicing_forward_pass(test_max_difference=_lowercase       ,     expected_max_diff=0.0_1  )





						@skip_mps
						def 	UpperCamelCase__			(  self  )      ->				Union[str, Any]:
									a_		:		Union[str, Any]				    =				torch_device == """cpu"""
									a_		:		Optional[int]				    =				True
									a_		:		Dict				    =				[
									    """prior_num_inference_steps""",
									    """decoder_num_inference_steps""",
									    """super_res_num_inference_steps""",
									]

									self._test_inference_batch_single_identical(
									    test_max_difference=_lowercase       ,     relax_max_difference=_lowercase       ,     additional_params_copy_to_batched_inputs=_lowercase       ,     )





						def 	UpperCamelCase__			(  self  )      ->				Any:
									a_		:		Optional[int]				    =				[
									    """prior_num_inference_steps""",
									    """decoder_num_inference_steps""",
									    """super_res_num_inference_steps""",
									]

									if torch_device == "mps":
												# TODO: MPS errors with larger batch sizes
												a_		:		Any				    =				[2, 3]
												self._test_inference_batch_consistent(
												    batch_sizes=_lowercase       ,     additional_params_copy_to_batched_inputs=_lowercase       ,     )
									else:
												self._test_inference_batch_consistent(
												    additional_params_copy_to_batched_inputs=_lowercase  )





						@skip_mps
						def 	UpperCamelCase__			(  self  )      ->				List[Any]:
									return super().test_dict_tuple_outputs_equivalent()





						@skip_mps
						def 	UpperCamelCase__			(  self  )      ->				Union[str, Any]:
									return super().test_save_load_local()





						@skip_mps
						def 	UpperCamelCase__			(  self  )      ->				List[str]:
									return super().test_save_load_optional_components()







@nightly
class    A__(unittest.TestCase    ):
						"""simple docstring"""






						def 	UpperCamelCase__			(  self  )      ->				Union[str, Any]:
									# clean up the VRAM after each test
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()





						def 	UpperCamelCase__			(  self  )      ->				int:
									a_		:		int				    =				load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
									    """/unclip/karlo_v1_alpha_horse_cpu.npy"""  )

									a_		:		Optional[Any]				    =				UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""  )
									pipeline.set_progress_bar_config(disable=_lowercase  )

									a_		:		Optional[int]				    =				torch.manual_seed(0  )
									a_		:		Union[str, Any]				    =				pipeline(
									    """horse"""       ,     num_images_per_prompt=1       ,     generator=_lowercase       ,     output_type="""np"""       ,     )

									a_		:		Dict				    =				output.images[0]

									assert image.shape == (256, 256, 3)
									assert np.abs(expected_image - image  ).max() < 1e-1







@slow
@require_torch_gpu
class    A__(unittest.TestCase    ):
						"""simple docstring"""






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





						def 	UpperCamelCase__			(  self  )      ->				Optional[int]:
									a_		:		List[str]				    =				load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
									    """/unclip/karlo_v1_alpha_horse_fp16.npy"""  )

									a_		:		Optional[int]				    =				UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""       ,     torch_dtype=torch.floataa  )
									a_		:		List[str]				    =				pipeline.to(_lowercase  )
									pipeline.set_progress_bar_config(disable=_lowercase  )

									a_		:		List[str]				    =				torch.Generator(device="""cpu"""  ).manual_seed(0  )
									a_		:		str				    =				pipeline(
									    """horse"""       ,     generator=_lowercase       ,     output_type="""np"""       ,     )

									a_		:		Any				    =				output.images[0]

									assert image.shape == (256, 256, 3)

									assert_mean_pixel_difference(_lowercase       ,     _lowercase  )





						def 	UpperCamelCase__			(  self  )      ->				str:
									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()
									torch.cuda.reset_peak_memory_stats()

									a_		:		List[str]				    =				UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""       ,     torch_dtype=torch.floataa  )
									a_		:		int				    =				pipe.to(_lowercase  )
									pipe.set_progress_bar_config(disable=_lowercase  )
									pipe.enable_attention_slicing()
									pipe.enable_sequential_cpu_offload()

									a_		:		str				    =				pipe(
									    """horse"""       ,     num_images_per_prompt=1       ,     prior_num_inference_steps=2       ,     decoder_num_inference_steps=2       ,     super_res_num_inference_steps=2       ,     output_type="""np"""       ,     )

									a_		:		Dict				    =				torch.cuda.max_memory_allocated()
									# make sure that less than 7 GB is allocated
									assert mem_bytes < 7 * 10**9





from collections import namedtuple

__snake_case							:     Optional[Any] 						=	namedtuple("""from_to""", """from_ to""")

__snake_case							:     Tuple 						=	{
    """millimeter""": """mm""",
    """centimeter""": """cm""",
    """meter""": """m""",
    """kilometer""": """km""",
    """inch""": """in""",
    """inche""": """in""",  # Trailing 's' has been stripped off
    """feet""": """ft""",
    """foot""": """ft""",
    """yard""": """yd""",
    """mile""": """mi""",
}

__snake_case							:     Optional[int] 						=	{
    """mm""": from_to(0.001, 10_00),
    """cm""": from_to(0.01, 1_00),
    """m""": from_to(1, 1),
    """km""": from_to(10_00, 0.001),
    """in""": from_to(0.0254, 39.3701),
    """ft""": from_to(0.3048, 3.2_8084),
    """yd""": from_to(0.9144, 1.0_9361),
    """mi""": from_to(1609.34, 0.0_0062_1371),
}


def 			_UpperCAmelCase							(		a__       ,					a__       ,					a__):

			'''simple docstring'''



			a_		:		List[str]				    =				from_type.lower().rstrip("""s""")
			a_		:		Union[str, Any]				    =				TYPE_CONVERSION.get(a__       ,					a__)
			a_		:		Optional[int]				    =				to_type.lower().rstrip("""s""")
			a_		:		Optional[int]				    =				TYPE_CONVERSION.get(a__       ,					a__)
			if new_from not in METRIC_CONVERSION:
						a_		:		Any				    =				(
						    f'''Invalid \'from_type\' value: {from_type!r}.\n'''
						    f'''Conversion abbreviations are: {", ".join(a__)}'''
						)
						raise ValueError(a__)
			if new_to not in METRIC_CONVERSION:
						a_		:		Union[str, Any]				    =				(
						    f'''Invalid \'to_type\' value: {to_type!r}.\n'''
						    f'''Conversion abbreviations are: {", ".join(a__)}'''
						)
						raise ValueError(a__)
			return value * METRIC_CONVERSION[new_from].from_ * METRIC_CONVERSION[new_to].to


if __name__ == "__main__":
	import doctest

	doctest.testmod()


