'''simple docstring'''

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 		SCREAMING_SNAKE_CASE_					(_UpperCAmelCase							,				unittest.TestCase      ):



			"""simple docstring"""






			__magic_name__ :				int   							=       UnCLIPPipeline
			__magic_name__ :				List[Any]   							=       TEXT_TO_IMAGE_PARAMS - {
			    """negative_prompt""",
			    """height""",
			    """width""",
			    """negative_prompt_embeds""",
			    """guidance_scale""",
			    """prompt_embeds""",
			    """cross_attention_kwargs""",
			}
			__magic_name__ :				Any   							=       TEXT_TO_IMAGE_BATCH_PARAMS
			__magic_name__ :				List[str]   							=       [
			    """generator""",
			    """return_dict""",
			    """prior_num_inference_steps""",
			    """decoder_num_inference_steps""",
			    """super_res_num_inference_steps""",
			]
			__magic_name__ :				str   							=       False



			@property
			def   lowerCamelCase (							self					):
									return 32



			@property
			def   lowerCamelCase (							self					):
									return 32



			@property
			def   lowerCamelCase (							self					):
									return self.time_input_dim



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



			@property
			def   lowerCamelCase (							self					):
									return 100



			@property
			def   lowerCamelCase (							self					):
									snake_case :       Any				=    CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)
									return tokenizer



			@property
			def   lowerCamelCase (							self					):
									torch.manual_seed(0					)
									snake_case :       Optional[Any]				=    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(snake_case_					)



			@property
			def   lowerCamelCase (							self					):
									torch.manual_seed(0					)

									snake_case :       Any				=    {
									    "num_attention_heads": 2,
									    "attention_head_dim": 12,
									    "embedding_dim": self.text_embedder_hidden_size,
									    "num_layers": 1,
									}

									snake_case :       List[str]				=    PriorTransformer(**snake_case_					)
									return model



			@property
			def   lowerCamelCase (							self					):
									torch.manual_seed(0					)

									snake_case :       Optional[int]				=    {
									    "clip_embeddings_dim": self.text_embedder_hidden_size,
									    "time_embed_dim": self.time_embed_dim,
									    "cross_attention_dim": self.cross_attention_dim,
									}

									snake_case :       List[str]				=    UnCLIPTextProjModel(**snake_case_					)
									return model



			@property
			def   lowerCamelCase (							self					):
									torch.manual_seed(0					)

									snake_case :       Dict				=    {
									    "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",
									}

									snake_case :       str				=    UNetaDConditionModel(**snake_case_					)
									return model



			@property
			def   lowerCamelCase (							self					):
									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   lowerCamelCase (							self					):
									torch.manual_seed(0					)

									snake_case :       int				=    UNetaDModel(**self.dummy_super_res_kwargs					)
									return model



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

									snake_case :       Any				=    UNetaDModel(**self.dummy_super_res_kwargs					)
									return model



			def   lowerCamelCase (							self					):
									snake_case :       Union[str, Any]				=    self.dummy_prior
									snake_case :       Union[str, Any]				=    self.dummy_decoder
									snake_case :       Dict				=    self.dummy_text_proj
									snake_case :       Optional[Any]				=    self.dummy_text_encoder
									snake_case :       Dict				=    self.dummy_tokenizer
									snake_case :       str				=    self.dummy_super_res_first
									snake_case :       str				=    self.dummy_super_res_last

									snake_case :       str				=    UnCLIPScheduler(
									    variance_type="fixed_small_log"				,			prediction_type="sample"				,			num_train_timesteps=1_000				,			clip_sample_range=5.0				,			)

									snake_case :       int				=    UnCLIPScheduler(
									    variance_type="learned_range"				,			prediction_type="epsilon"				,			num_train_timesteps=1_000				,			)

									snake_case :       Optional[Any]				=    UnCLIPScheduler(
									    variance_type="fixed_small_log"				,			prediction_type="epsilon"				,			num_train_timesteps=1_000				,			)

									snake_case :       Optional[int]				=    {
									    "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   lowerCamelCase (							self				,			snake_case_				,			snake_case_=0					):
									if str(snake_case_					).startswith("mps"					):
															snake_case :       List[str]				=    torch.manual_seed(snake_case_					)
									else:
															snake_case :       Union[str, Any]				=    torch.Generator(device=snake_case_					).manual_seed(snake_case_					)
									snake_case :       Union[str, 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   lowerCamelCase (							self					):
									snake_case :       Any				=    "cpu"

									snake_case :       Dict				=    self.get_dummy_components()

									snake_case :       Union[str, Any]				=    self.pipeline_class(**snake_case_					)
									snake_case :       Tuple				=    pipe.to(snake_case_					)

									pipe.set_progress_bar_config(disable=snake_case_					)

									snake_case :       Optional[Any]				=    pipe(**self.get_dummy_inputs(snake_case_					)					)
									snake_case :       Optional[int]				=    output.images

									snake_case :       Optional[int]				=    pipe(
									    **self.get_dummy_inputs(snake_case_					)				,			return_dict=snake_case_				,			)[0]

									snake_case :       Dict				=    image[0, -3:, -3:, -1]
									snake_case :       List[Any]				=    image_from_tuple[0, -3:, -3:, -1]

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

									snake_case :       Dict				=    np.array(
									    [
									        0.9_997,
									        0.9_988,
									        0.0_028,
									        0.9_997,
									        0.9_984,
									        0.9_965,
									        0.0_029,
									        0.9_986,
									        0.0_025,
									    ]					)

									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   lowerCamelCase (							self					):
									snake_case :       Tuple				=    torch.device("cpu"					)




									class 		SCREAMING_SNAKE_CASE_					:



												"""simple docstring"""






												__magic_name__ :				str   							=       1







									snake_case :       str				=    self.get_dummy_components()

									snake_case :       List[str]				=    self.pipeline_class(**snake_case_					)
									snake_case :       Tuple				=    pipe.to(snake_case_					)

									snake_case :       int				=    components["prior"]
									snake_case :       Dict				=    components["decoder"]
									snake_case :       List[Any]				=    components["super_res_first"]
									snake_case :       List[str]				=    components["tokenizer"]
									snake_case :       Any				=    components["text_encoder"]

									snake_case :       Optional[int]				=    torch.Generator(device=snake_case_					).manual_seed(0					)
									snake_case :       Union[str, Any]				=    prior.dtype
									snake_case :       List[str]				=    1

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

									snake_case :       str				=    (
									    batch_size,
									    super_res_first.config.in_channels // 2,
									    super_res_first.config.sample_size,
									    super_res_first.config.sample_size,
									)
									snake_case :       List[str]				=    pipe.prepare_latents(
									    snake_case_				,			dtype=snake_case_				,			device=snake_case_				,			generator=snake_case_				,			latents=snake_case_				,			scheduler=DummyScheduler()					)

									pipe.set_progress_bar_config(disable=snake_case_					)

									snake_case :       Any				=    "this is a prompt example"

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

									snake_case :       Optional[Any]				=    tokenizer(
									    snake_case_				,			padding="max_length"				,			max_length=tokenizer.model_max_length				,			return_tensors="pt"				,			)
									snake_case :       Any				=    text_encoder(text_inputs.input_ids					)
									snake_case :       List[Any]				=    text_inputs.attention_mask

									snake_case :       int				=    torch.Generator(device=snake_case_					).manual_seed(0					)
									snake_case :       Union[str, Any]				=    pipe(
									    generator=snake_case_				,			prior_num_inference_steps=2				,			decoder_num_inference_steps=2				,			super_res_num_inference_steps=2				,			prior_latents=snake_case_				,			decoder_latents=snake_case_				,			super_res_latents=snake_case_				,			text_model_output=snake_case_				,			text_attention_mask=snake_case_				,			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   lowerCamelCase (							self					):
									snake_case :       Dict				=    torch_device == "cpu"

									self._test_attention_slicing_forward_pass(test_max_difference=snake_case_				,			expected_max_diff=0.01					)



			@skip_mps
			def   lowerCamelCase (							self					):
									snake_case :       Any				=    torch_device == "cpu"
									snake_case :       str				=    True
									snake_case :       Union[str, Any]				=    [
									    "prior_num_inference_steps",
									    "decoder_num_inference_steps",
									    "super_res_num_inference_steps",
									]

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



			def   lowerCamelCase (							self					):
									snake_case :       Optional[Any]				=    [
									    "prior_num_inference_steps",
									    "decoder_num_inference_steps",
									    "super_res_num_inference_steps",
									]

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



			@skip_mps
			def   lowerCamelCase (							self					):
									return super().test_dict_tuple_outputs_equivalent()



			@skip_mps
			def   lowerCamelCase (							self					):
									return super().test_save_load_local()






			@skip_mps
			def   lowerCamelCase (							self					):
									return super().test_save_load_optional_components()







@nightly
class 		SCREAMING_SNAKE_CASE_					(unittest.TestCase      ):



			"""simple docstring"""






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






			def   lowerCamelCase (							self					):
									snake_case :       Optional[int]				=    load_numpy(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
									    "/unclip/karlo_v1_alpha_horse_cpu.npy"					)

									snake_case :       Tuple				=    UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"					)
									pipeline.set_progress_bar_config(disable=snake_case_					)

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

									snake_case :       Any				=    output.images[0]

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







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



			"""simple docstring"""






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



			def   lowerCamelCase (							self					):
									snake_case :       str				=    load_numpy(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
									    "/unclip/karlo_v1_alpha_horse_fp16.npy"					)

									snake_case :       Union[str, Any]				=    UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"				,			torch_dtype=torch.floataa					)
									snake_case :       int				=    pipeline.to(snake_case_					)
									pipeline.set_progress_bar_config(disable=snake_case_					)

									snake_case :       Union[str, Any]				=    torch.Generator(device="cpu"					).manual_seed(0					)
									snake_case :       Optional[Any]				=    pipeline(
									    "horse"				,			generator=snake_case_				,			output_type="np"				,			)

									snake_case :       int				=    output.images[0]

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

									assert_mean_pixel_difference(snake_case_				,			snake_case_					)






			def   lowerCamelCase (							self					):
									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()
									torch.cuda.reset_peak_memory_stats()

									snake_case :       Dict				=    UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"				,			torch_dtype=torch.floataa					)
									snake_case :       Dict				=    pipe.to(snake_case_					)
									pipe.set_progress_bar_config(disable=snake_case_					)
									pipe.enable_attention_slicing()
									pipe.enable_sequential_cpu_offload()

									snake_case :       Tuple				=    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"				,			)

									snake_case :       Optional[int]				=    torch.cuda.max_memory_allocated()
									# make sure that less than 7 GB is allocated
									assert mem_bytes < 7 * 10**9






'''simple docstring'''

from __future__ import annotations


def        UpperCAmelCase__   (				_lowerCAmelCase    :			list[int | float])    ->		int | float:



						"""simple docstring"""

						if len(_lowerCAmelCase) == 0:
												raise ValueError("find_max() arg is an empty sequence")
						snake_case :       Any				=    nums[0]
						for x in nums:
												if x > max_num:
																		snake_case :       int				=    x
						return max_num


if __name__ == "__main__":
						import doctest

						doctest.testmod(verbose=True)





