


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


				"""simple docstring"""







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






				@property
				def        lowerCAmelCase__			(				self     :    Any					):
								return 32






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






				@property
				def        lowerCAmelCase__			(				self     :    str					):
								return self.time_input_dim






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






				@property
				def        lowerCAmelCase__			(				self     :    Any					):
								return 100






				@property
				def        lowerCAmelCase__			(				self     :    List[Any]					):
								_lowerCAmelCase  						=CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""					)
								return tokenizer






				@property
				def        lowerCAmelCase__			(				self     :    int					):
								torch.manual_seed(0					)
								_lowerCAmelCase  						=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=1000	,							)
								return CLIPTextModelWithProjection(lowerCamelCase_					)






				@property
				def        lowerCAmelCase__			(				self     :    Any					):
								torch.manual_seed(0					)

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

								_lowerCAmelCase  						=PriorTransformer(**lowerCamelCase_					)
								return model






				@property
				def        lowerCAmelCase__			(				self     :    Optional[int]					):
								torch.manual_seed(0					)

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

								_lowerCAmelCase  						=UnCLIPTextProjModel(**lowerCamelCase_					)
								return model






				@property
				def        lowerCAmelCase__			(				self     :    str					):
								torch.manual_seed(0					)

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

								_lowerCAmelCase  						=UNetaDConditionModel(**lowerCamelCase_					)
								return model






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

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






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

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






				def        lowerCAmelCase__			(				self     :    Optional[int]					):
								_lowerCAmelCase  						=self.dummy_prior
								_lowerCAmelCase  						=self.dummy_decoder
								_lowerCAmelCase  						=self.dummy_text_proj
								_lowerCAmelCase  						=self.dummy_text_encoder
								_lowerCAmelCase  						=self.dummy_tokenizer
								_lowerCAmelCase  						=self.dummy_super_res_first
								_lowerCAmelCase  						=self.dummy_super_res_last

								_lowerCAmelCase  						=UnCLIPScheduler(
								    variance_type="""fixed_small_log"""	,							prediction_type="""sample"""	,							num_train_timesteps=1000	,							clip_sample_range=5.0	,							)

								_lowerCAmelCase  						=UnCLIPScheduler(
								    variance_type="""learned_range"""	,							prediction_type="""epsilon"""	,							num_train_timesteps=1000	,							)

								_lowerCAmelCase  						=UnCLIPScheduler(
								    variance_type="""fixed_small_log"""	,							prediction_type="""epsilon"""	,							num_train_timesteps=1000	,							)

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

								_lowerCAmelCase  						=self.get_dummy_components()

								_lowerCAmelCase  						=self.pipeline_class(**lowerCamelCase_					)
								_lowerCAmelCase  						=pipe.to(lowerCamelCase_					)

								pipe.set_progress_bar_config(disable=lowerCamelCase_					)

								_lowerCAmelCase  						=pipe(**self.get_dummy_inputs(lowerCamelCase_					)					)
								_lowerCAmelCase  						=output.images

								_lowerCAmelCase  						=pipe(
								    **self.get_dummy_inputs(lowerCamelCase_					)	,							return_dict=lowerCamelCase_	,							)[0]

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

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

								_lowerCAmelCase  						=np.array(
								    [
								        0.9997,
								        0.9988,
								        0.0028,
								        0.9997,
								        0.9984,
								        0.9965,
								        0.0029,
								        0.9986,
								        0.0025,
								    ]					)

								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     :    List[str]					):
								_lowerCAmelCase  						=torch.device("""cpu"""					)


								class        __lowerCamelCase      :


												"""simple docstring"""







												a_:		int					      =	1
								_lowerCAmelCase  						=self.get_dummy_components()

								_lowerCAmelCase  						=self.pipeline_class(**lowerCamelCase_					)
								_lowerCAmelCase  						=pipe.to(lowerCamelCase_					)

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

								_lowerCAmelCase  						=torch.Generator(device=lowerCamelCase_					).manual_seed(0					)
								_lowerCAmelCase  						=prior.dtype
								_lowerCAmelCase  						=1

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

								_lowerCAmelCase  						=(
								    batch_size,
								    super_res_first.config.in_channels // 2,
								    super_res_first.config.sample_size,
								    super_res_first.config.sample_size,
								)
								_lowerCAmelCase  						=pipe.prepare_latents(
								    lowerCamelCase_	,							dtype=lowerCamelCase_	,							device=lowerCamelCase_	,							generator=lowerCamelCase_	,							latents=lowerCamelCase_	,							scheduler=DummyScheduler()					)

								pipe.set_progress_bar_config(disable=lowerCamelCase_					)

								_lowerCAmelCase  						="""this is a prompt example"""

								_lowerCAmelCase  						=torch.Generator(device=lowerCamelCase_					).manual_seed(0					)
								_lowerCAmelCase  						=pipe(
								    [prompt]	,							generator=lowerCamelCase_	,							prior_num_inference_steps=2	,							decoder_num_inference_steps=2	,							super_res_num_inference_steps=2	,							prior_latents=lowerCamelCase_	,							decoder_latents=lowerCamelCase_	,							super_res_latents=lowerCamelCase_	,							output_type="""np"""	,							)
								_lowerCAmelCase  						=output.images

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

								_lowerCAmelCase  						=torch.Generator(device=lowerCamelCase_					).manual_seed(0					)
								_lowerCAmelCase  						=pipe(
								    generator=lowerCamelCase_	,							prior_num_inference_steps=2	,							decoder_num_inference_steps=2	,							super_res_num_inference_steps=2	,							prior_latents=lowerCamelCase_	,							decoder_latents=lowerCamelCase_	,							super_res_latents=lowerCamelCase_	,							text_model_output=lowerCamelCase_	,							text_attention_mask=lowerCamelCase_	,							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     :    int					):
								_lowerCAmelCase  						=torch_device == """cpu"""

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






				@skip_mps
				def        lowerCAmelCase__			(				self     :    Optional[Any]					):
								_lowerCAmelCase  						=torch_device == """cpu"""
								_lowerCAmelCase  						=True
								_lowerCAmelCase  						=[
								    """prior_num_inference_steps""",
								    """decoder_num_inference_steps""",
								    """super_res_num_inference_steps""",
								]

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






				def        lowerCAmelCase__			(				self     :    Tuple					):
								_lowerCAmelCase  						=[
								    """prior_num_inference_steps""",
								    """decoder_num_inference_steps""",
								    """super_res_num_inference_steps""",
								]

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






				@skip_mps
				def        lowerCAmelCase__			(				self     :    str					):
								return super().test_dict_tuple_outputs_equivalent()






				@skip_mps
				def        lowerCAmelCase__			(				self     :    Dict					):
								return super().test_save_load_local()






				@skip_mps
				def        lowerCAmelCase__			(				self     :    List[str]					):
								return super().test_save_load_optional_components()


@nightly
class        __lowerCamelCase      (				unittest.TestCase ):


				"""simple docstring"""







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






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

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

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

								_lowerCAmelCase  						=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 ):


				"""simple docstring"""







				def        lowerCAmelCase__			(				self     :    List[str]					):
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()






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

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

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

								_lowerCAmelCase  						=output.images[0]

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

								assert_mean_pixel_difference(lowerCamelCase_	,							lowerCamelCase_					)






				def        lowerCAmelCase__			(				self     :    str					):
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

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

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

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





def 		snake_case_   (				lowercase__     :    str     ):






				'''simple docstring'''






				_lowerCAmelCase  						=53_81
				for x in s:
								_lowerCAmelCase  						=((hash_value << 5) + hash_value) + ord(lowercase__     )
				return hash_value & 0Xffff_ffff

