





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


				"""simple docstring"""


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



				@property
				def  UpperCAmelCase_		(      self					):
									return 32



				@property
				def  UpperCAmelCase_		(      self					):
									return 32



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



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



				@property
				def  UpperCAmelCase_		(      self					):
									return 100



				@property
				def  UpperCAmelCase_		(      self					):
									__A     :						str          =    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)
									return tokenizer



				@property
				def  UpperCAmelCase_		(      self					):
									torch.manual_seed(0					)
									__A     :						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-0_5					, num_attention_heads=4					, num_hidden_layers=5					, pad_token_id=1					, vocab_size=1000					, )
									return CLIPTextModelWithProjection(lowercase_					)



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

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

									__A     :						Dict          =    PriorTransformer(**lowercase_					)
									return model



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

									__A     :						Dict          =    {
									    'clip_embeddings_dim': self.text_embedder_hidden_size,
									    'time_embed_dim': self.time_embed_dim,
									    'cross_attention_dim': self.cross_attention_dim,
									}

									__A     :						Union[str, Any]          =    UnCLIPTextProjModel(**lowercase_					)
									return model



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

									__A     :						List[Any]          =    {
									    '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					):
									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					):
									torch.manual_seed(0					)

									__A     :						Optional[Any]          =    UNetaDModel(**self.dummy_super_res_kwargs					)
									return model



				@property
				def  UpperCAmelCase_		(      self					):
									torch.manual_seed(1					)

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



				def  UpperCAmelCase_		(      self					):
									__A     :						str          =    self.dummy_prior
									__A     :						Union[str, Any]          =    self.dummy_decoder
									__A     :						Union[str, Any]          =    self.dummy_text_proj
									__A     :						Dict          =    self.dummy_text_encoder
									__A     :						Tuple          =    self.dummy_tokenizer
									__A     :						Optional[int]          =    self.dummy_super_res_first
									__A     :						Dict          =    self.dummy_super_res_last

									__A     :						Any          =    UnCLIPScheduler(
									    variance_type='fixed_small_log'					, prediction_type='sample'					, num_train_timesteps=1000					, clip_sample_range=5.0					, )

									__A     :						Union[str, Any]          =    UnCLIPScheduler(
									    variance_type='learned_range'					, prediction_type='epsilon'					, num_train_timesteps=1000					, )

									__A     :						Any          =    UnCLIPScheduler(
									    variance_type='fixed_small_log'					, prediction_type='epsilon'					, num_train_timesteps=1000					, )

									__A     :						Any          =    {
									    '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					, _A					, _A=0					):






									if str(lowercase_					).startswith('mps'					):
														__A     :						Union[str, Any]          =    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					):
									__A     :						str          =    'cpu'

									__A     :						Tuple          =    self.get_dummy_components()

									__A     :						int          =    self.pipeline_class(**lowercase_					)
									__A     :						Tuple          =    pipe.to(lowercase_					)

									pipe.set_progress_bar_config(disable=lowercase_					)

									__A     :						Optional[int]          =    pipe(**self.get_dummy_inputs(lowercase_					)					)
									__A     :						str          =    output.images

									__A     :						Optional[Any]          =    pipe(
									    **self.get_dummy_inputs(lowercase_					)					, return_dict=lowercase_					, )[0]

									__A     :						Dict          =    image[0, -3:, -3:, -1]
									__A     :						List[str]          =    image_from_tuple[0, -3:, -3:, -1]

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

									__A     :						Any          =    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					):
									__A     :						Any          =    torch.device('cpu'					)





									class 						_A:


													"""simple docstring"""


													UpperCamelCase   :  str    =     1




									__A     :						Dict          =    self.get_dummy_components()

									__A     :						int          =    self.pipeline_class(**lowercase_					)
									__A     :						List[str]          =    pipe.to(lowercase_					)

									__A     :						int          =    components['prior']
									__A     :						Optional[int]          =    components['decoder']
									__A     :						List[str]          =    components['super_res_first']
									__A     :						Tuple          =    components['tokenizer']
									__A     :						Any          =    components['text_encoder']

									__A     :						List[Any]          =    torch.Generator(device=lowercase_					).manual_seed(0					)
									__A     :						Dict          =    prior.dtype
									__A     :						List[str]          =    1

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

									__A     :						int          =    (
									    batch_size,
									    super_res_first.config.in_channels // 2,
									    super_res_first.config.sample_size,
									    super_res_first.config.sample_size,
									)
									__A     :						Optional[int]          =    pipe.prepare_latents(
									    lowercase_					, dtype=lowercase_					, device=lowercase_					, generator=lowercase_					, latents=lowercase_					, scheduler=DummyScheduler()					)

									pipe.set_progress_bar_config(disable=lowercase_					)

									__A     :						Dict          =    'this is a prompt example'

									__A     :						Optional[int]          =    torch.Generator(device=lowercase_					).manual_seed(0					)
									__A     :						Union[str, 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     :						int          =    output.images

									__A     :						str          =    tokenizer(
									    lowercase_					, padding='max_length'					, max_length=tokenizer.model_max_length					, return_tensors='pt'					, )
									__A     :						Tuple          =    text_encoder(text_inputs.input_ids					)
									__A     :						Tuple          =    text_inputs.attention_mask

									__A     :						Optional[Any]          =    torch.Generator(device=lowercase_					).manual_seed(0					)
									__A     :						List[Any]          =    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					):
									__A     :						Union[str, 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					):
									__A     :						Any          =    torch_device == 'cpu'
									__A     :						int          =    True
									__A     :						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=lowercase_					, relax_max_difference=lowercase_					, additional_params_copy_to_batched_inputs=lowercase_					, )



				def  UpperCAmelCase_		(      self					):
									__A     :						Union[str, 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
														__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					):
									return super().test_dict_tuple_outputs_equivalent()



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





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





@nightly
class 						_A( unittest.TestCase       ):


				"""simple docstring"""


				def  UpperCAmelCase_		(      self					):
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()





				def  UpperCAmelCase_		(      self					):
									__A     :						Tuple          =    load_numpy(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
									    '/unclip/karlo_v1_alpha_horse_cpu.npy'					)

									__A     :						Union[str, Any]          =    UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha'					)
									pipeline.set_progress_bar_config(disable=lowercase_					)

									__A     :						Tuple          =    torch.manual_seed(0					)
									__A     :						Optional[Any]          =    pipeline(
									    'horse'					, num_images_per_prompt=1					, generator=lowercase_					, output_type='np'					, )

									__A     :						List[str]          =    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					):
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()



				def  UpperCAmelCase_		(      self					):
									__A     :						Union[str, Any]          =    load_numpy(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
									    '/unclip/karlo_v1_alpha_horse_fp16.npy'					)

									__A     :						Optional[Any]          =    UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha'					, torch_dtype=torch.floataa					)
									__A     :						Dict          =    pipeline.to(lowercase_					)
									pipeline.set_progress_bar_config(disable=lowercase_					)

									__A     :						List[Any]          =    torch.Generator(device='cpu'					).manual_seed(0					)
									__A     :						Optional[Any]          =    pipeline(
									    'horse'					, generator=lowercase_					, output_type='np'					, )

									__A     :						Union[str, Any]          =    output.images[0]

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

									assert_mean_pixel_difference(lowercase_					, lowercase_					)





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

									__A     :						Union[str, Any]          =    UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha'					, torch_dtype=torch.floataa					)
									__A     :						Union[str, Any]          =    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     :						List[str]          =    torch.cuda.max_memory_allocated()
									# make sure that less than 7 GB is allocated
									assert mem_bytes < 7 * 10**9











from .hash_table import HashTable





class 						_A( snake_case__       ):


				"""simple docstring"""


				def __init__(      self					, *_A					, **_A					):
									super().__init__(*_A					, **_A					)





				def  UpperCAmelCase_		(      self					, _A					, _A=None					):
									__A     :						Optional[int]          =    1
									__A     :						Union[str, Any]          =    self.hash_function(key + i * i					)

									while self.values[new_key] is not None and self.values[new_key] != key:
														i += 1
														__A     :						Tuple          =    (
														    self.hash_function(key + i * i					)
														    if not self.balanced_factor() >= self.lim_charge
														    else None
														)

														if new_key is None:
																			break

									return new_key




