




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    __lowerCAmelCase			(	UpperCAmelCase__ ,		unittest.TestCase			):
			snake_case_				:			Any								=					UnCLIPPipeline
			snake_case_				:			int								=					TEXT_TO_IMAGE_PARAMS - {
			    "negative_prompt",
			    "height",
			    "width",
			    "negative_prompt_embeds",
			    "guidance_scale",
			    "prompt_embeds",
			    "cross_attention_kwargs",
			}
			snake_case_				:			Union[str, Any]								=					TEXT_TO_IMAGE_BATCH_PARAMS
			snake_case_				:			Any								=					[
			    "generator",
			    "return_dict",
			    "prior_num_inference_steps",
			    "decoder_num_inference_steps",
			    "super_res_num_inference_steps",
			]
			snake_case_				:			Optional[Any]								=					False







			@property
			def UpperCamelCase      (	self   :      Optional[int]				):




									"""simple docstring"""


									return 32







			@property
			def UpperCamelCase      (	self   :      Tuple				):




									"""simple docstring"""


									return 32







			@property
			def UpperCamelCase      (	self   :      List[Any]				):




									"""simple docstring"""


									return self.time_input_dim







			@property
			def UpperCamelCase      (	self   :      Tuple				):




									"""simple docstring"""


									return self.time_input_dim * 4







			@property
			def UpperCamelCase      (	self   :      Optional[int]				):




									"""simple docstring"""


									return 100







			@property
			def UpperCamelCase      (	self   :      Tuple				):




									"""simple docstring"""


									_UpperCAmelCase													=					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				)
									return tokenizer







			@property
			def UpperCamelCase      (	self   :      int				):




									"""simple docstring"""


									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(snake_case__				)







			@property
			def UpperCamelCase      (	self   :      Tuple				):




									"""simple docstring"""


									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(**snake_case__				)
									return model







			@property
			def UpperCamelCase      (	self   :      str				):




									"""simple docstring"""


									torch.manual_seed(0				)

									_UpperCAmelCase													=					{
									    "clip_embeddings_dim": self.text_embedder_hidden_size,
									    "time_embed_dim": self.time_embed_dim,
									    "cross_attention_dim": self.cross_attention_dim,
									}

									_UpperCAmelCase													=					UnCLIPTextProjModel(**snake_case__				)
									return model







			@property
			def UpperCamelCase      (	self   :      Union[str, Any]				):




									"""simple docstring"""


									torch.manual_seed(0				)

									_UpperCAmelCase													=					{
									    "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",
									}

									_UpperCAmelCase													=					UNetaDConditionModel(**snake_case__				)
									return model







			@property
			def UpperCamelCase      (	self   :      int				):




									"""simple docstring"""


									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   :      Optional[int]				):




									"""simple docstring"""


									torch.manual_seed(0				)

									_UpperCAmelCase													=					UNetaDModel(**self.dummy_super_res_kwargs				)
									return model







			@property
			def UpperCamelCase      (	self   :      str				):




									"""simple docstring"""


									torch.manual_seed(1				)

									_UpperCAmelCase													=					UNetaDModel(**self.dummy_super_res_kwargs				)
									return model







			def UpperCamelCase      (	self   :      Optional[Any]				):




									"""simple docstring"""


									_UpperCAmelCase													=					self.dummy_prior
									_UpperCAmelCase													=					self.dummy_decoder
									_UpperCAmelCase													=					self.dummy_text_proj
									_UpperCAmelCase													=					self.dummy_text_encoder
									_UpperCAmelCase													=					self.dummy_tokenizer
									_UpperCAmelCase													=					self.dummy_super_res_first
									_UpperCAmelCase													=					self.dummy_super_res_last

									_UpperCAmelCase													=					UnCLIPScheduler(
									    variance_type="fixed_small_log"   ,   prediction_type="sample"   ,   num_train_timesteps=1_000   ,   clip_sample_range=5.0   ,   )

									_UpperCAmelCase													=					UnCLIPScheduler(
									    variance_type="learned_range"   ,   prediction_type="epsilon"   ,   num_train_timesteps=1_000   ,   )

									_UpperCAmelCase													=					UnCLIPScheduler(
									    variance_type="fixed_small_log"   ,   prediction_type="epsilon"   ,   num_train_timesteps=1_000   ,   )

									_UpperCAmelCase													=					{
									    "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   :      Tuple   ,   snake_case__   :      Any   ,   snake_case__   :      Optional[Any]=0				):




									"""simple docstring"""


									if str(snake_case__				).startswith("mps"				):
															_UpperCAmelCase													=					torch.manual_seed(snake_case__				)
									else:
															_UpperCAmelCase													=					torch.Generator(device=snake_case__				).manual_seed(snake_case__				)
									_UpperCAmelCase													=					{
									    "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   :      int				):




									"""simple docstring"""


									_UpperCAmelCase													=					"cpu"

									_UpperCAmelCase													=					self.get_dummy_components()

									_UpperCAmelCase													=					self.pipeline_class(**snake_case__				)
									_UpperCAmelCase													=					pipe.to(snake_case__				)

									pipe.set_progress_bar_config(disable=snake_case__				)

									_UpperCAmelCase													=					pipe(**self.get_dummy_inputs(snake_case__				)				)
									_UpperCAmelCase													=					output.images

									_UpperCAmelCase													=					pipe(
									    **self.get_dummy_inputs(snake_case__				)   ,   return_dict=snake_case__   ,   )[0]

									_UpperCAmelCase													=					image[0, -3:, -3:, -1]
									_UpperCAmelCase													=					image_from_tuple[0, -3:, -3:, -1]

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

									_UpperCAmelCase													=					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 UpperCamelCase      (	self   :      Tuple				):




									"""simple docstring"""


									_UpperCAmelCase													=					torch.device("cpu"				)






									class    __lowerCAmelCase			:
												snake_case_				:			List[str]								=					1




									_UpperCAmelCase													=					self.get_dummy_components()

									_UpperCAmelCase													=					self.pipeline_class(**snake_case__				)
									_UpperCAmelCase													=					pipe.to(snake_case__				)

									_UpperCAmelCase													=					components["prior"]
									_UpperCAmelCase													=					components["decoder"]
									_UpperCAmelCase													=					components["super_res_first"]
									_UpperCAmelCase													=					components["tokenizer"]
									_UpperCAmelCase													=					components["text_encoder"]

									_UpperCAmelCase													=					torch.Generator(device=snake_case__				).manual_seed(0				)
									_UpperCAmelCase													=					prior.dtype
									_UpperCAmelCase													=					1

									_UpperCAmelCase													=					(batch_size, prior.config.embedding_dim)
									_UpperCAmelCase													=					pipe.prepare_latents(
									    snake_case__   ,   dtype=snake_case__   ,   device=snake_case__   ,   generator=snake_case__   ,   latents=snake_case__   ,   scheduler=DummyScheduler()				)
									_UpperCAmelCase													=					(batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
									_UpperCAmelCase													=					pipe.prepare_latents(
									    snake_case__   ,   dtype=snake_case__   ,   device=snake_case__   ,   generator=snake_case__   ,   latents=snake_case__   ,   scheduler=DummyScheduler()				)

									_UpperCAmelCase													=					(
									    batch_size,
									    super_res_first.config.in_channels // 2,
									    super_res_first.config.sample_size,
									    super_res_first.config.sample_size,
									)
									_UpperCAmelCase													=					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__				)

									_UpperCAmelCase													=					"this is a prompt example"

									_UpperCAmelCase													=					torch.Generator(device=snake_case__				).manual_seed(0				)
									_UpperCAmelCase													=					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"   ,   )
									_UpperCAmelCase													=					output.images

									_UpperCAmelCase													=					tokenizer(
									    snake_case__   ,   padding="max_length"   ,   max_length=tokenizer.model_max_length   ,   return_tensors="pt"   ,   )
									_UpperCAmelCase													=					text_encoder(text_inputs.input_ids				)
									_UpperCAmelCase													=					text_inputs.attention_mask

									_UpperCAmelCase													=					torch.Generator(device=snake_case__				).manual_seed(0				)
									_UpperCAmelCase													=					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 UpperCamelCase      (	self   :      Dict				):




									"""simple docstring"""


									_UpperCAmelCase													=					torch_device == "cpu"

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







			@skip_mps
			def UpperCamelCase      (	self   :      str				):




									"""simple docstring"""


									_UpperCAmelCase													=					torch_device == "cpu"
									_UpperCAmelCase													=					True
									_UpperCAmelCase													=					[
									    "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 UpperCamelCase      (	self   :      Optional[Any]				):




									"""simple docstring"""


									_UpperCAmelCase													=					[
									    "prior_num_inference_steps",
									    "decoder_num_inference_steps",
									    "super_res_num_inference_steps",
									]

									if torch_device == "mps":
															# TODO: MPS errors with larger batch sizes
															_UpperCAmelCase													=					[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 UpperCamelCase      (	self   :      Any				):




									"""simple docstring"""


									return super().test_dict_tuple_outputs_equivalent()







			@skip_mps
			def UpperCamelCase      (	self   :      Dict				):




									"""simple docstring"""


									return super().test_save_load_local()







			@skip_mps
			def UpperCamelCase      (	self   :      Tuple				):




									"""simple docstring"""


									return super().test_save_load_optional_components()






@nightly
class    __lowerCAmelCase			(	unittest.TestCase			):







			def UpperCamelCase      (	self   :      Optional[Any]				):




									"""simple docstring"""


									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()







			def UpperCamelCase      (	self   :      Optional[int]				):




									"""simple docstring"""


									_UpperCAmelCase													=					load_numpy(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
									    "/unclip/karlo_v1_alpha_horse_cpu.npy"				)

									_UpperCAmelCase													=					UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"				)
									pipeline.set_progress_bar_config(disable=snake_case__				)

									_UpperCAmelCase													=					torch.manual_seed(0				)
									_UpperCAmelCase													=					pipeline(
									    "horse"   ,   num_images_per_prompt=1   ,   generator=snake_case__   ,   output_type="np"   ,   )

									_UpperCAmelCase													=					output.images[0]

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






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







			def UpperCamelCase      (	self   :      Optional[int]				):




									"""simple docstring"""


									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()







			def UpperCamelCase      (	self   :      Any				):




									"""simple docstring"""


									_UpperCAmelCase													=					load_numpy(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
									    "/unclip/karlo_v1_alpha_horse_fp16.npy"				)

									_UpperCAmelCase													=					UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"   ,   torch_dtype=torch.floataa				)
									_UpperCAmelCase													=					pipeline.to(snake_case__				)
									pipeline.set_progress_bar_config(disable=snake_case__				)

									_UpperCAmelCase													=					torch.Generator(device="cpu"				).manual_seed(0				)
									_UpperCAmelCase													=					pipeline(
									    "horse"   ,   generator=snake_case__   ,   output_type="np"   ,   )

									_UpperCAmelCase													=					output.images[0]

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

									assert_mean_pixel_difference(snake_case__   ,   snake_case__				)







			def UpperCamelCase      (	self   :      str				):




									"""simple docstring"""


									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()
									torch.cuda.reset_peak_memory_stats()

									_UpperCAmelCase													=					UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"   ,   torch_dtype=torch.floataa				)
									_UpperCAmelCase													=					pipe.to(snake_case__				)
									pipe.set_progress_bar_config(disable=snake_case__				)
									pipe.enable_attention_slicing()
									pipe.enable_sequential_cpu_offload()

									_UpperCAmelCase													=					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"   ,   )

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










from __future__ import annotations


def 							__SCREAMING_SNAKE_CASE  (      snake_case_     ,   snake_case_						):

						'''simple docstring'''


						_UpperCAmelCase													=					{}
						for index, val in enumerate(snake_case_						):
												_UpperCAmelCase													=					target - val
												if compl in chk_map:
																		return [chk_map[compl], index]
												_UpperCAmelCase													=					index
						return []


if __name__ == "__main__":
				import doctest

				doctest.testmod()
				print(f"""{two_sum([2, 7, 11, 15], 9) = }""")




