




"""simple docstring"""







import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()






class 					SCREAMING_SNAKE_CASE__    (    _a  , unittest.TestCase    ):
						_a			  =  UnCLIPImageVariationPipeline
						_a			  =  IMAGE_VARIATION_PARAMS - {'height', 'width', 'guidance_scale'}
						_a			  =  IMAGE_VARIATION_BATCH_PARAMS

						_a			  =  [
						    'generator',
						    'return_dict',
						    'decoder_num_inference_steps',
						    'super_res_num_inference_steps',
						]
						_a			  =  False


						@property
						def      __lowercase       (							self :	Optional[Any]    ):
										return 32


						@property
						def      __lowercase       (							self :	Optional[Any]    ):
										return 32


						@property
						def      __lowercase       (							self :	Any    ):
										return self.time_input_dim


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


						@property
						def      __lowercase       (							self :	Optional[Any]    ):
										return 100


						@property
						def      __lowercase       (							self :	Dict    ):
										lowerCAmelCase              = CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    )
										return tokenizer


						@property
						def      __lowercase       (							self :	Optional[Any]    ):
										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(SCREAMING_SNAKE_CASE_    )


						@property
						def      __lowercase       (							self :	int    ):
										torch.manual_seed(0    )
										lowerCAmelCase              = CLIPVisionConfig(
										    hidden_size=self.text_embedder_hidden_size  ,  projection_dim=self.text_embedder_hidden_size  ,  num_hidden_layers=5  ,  num_attention_heads=4  ,  image_size=32  ,  intermediate_size=37  ,  patch_size=1  ,  )
										return CLIPVisionModelWithProjection(SCREAMING_SNAKE_CASE_    )


						@property
						def      __lowercase       (							self :	Any    ):
										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(**SCREAMING_SNAKE_CASE_    )
										return model


						@property
						def      __lowercase       (							self :	List[Any]    ):
										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(**SCREAMING_SNAKE_CASE_    )
										return model


						@property
						def      __lowercase       (							self :	Tuple    ):
										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      __lowercase       (							self :	Union[str, Any]    ):
										torch.manual_seed(0    )

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


						@property
						def      __lowercase       (							self :	Dict    ):
										# 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      __lowercase       (							self :	int    ):
										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="""learned_range"""  ,  prediction_type="""epsilon"""  ,  num_train_timesteps=1000  ,  )

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

										lowerCAmelCase              = CLIPImageProcessor(crop_size=32  ,  size=32    )

										lowerCAmelCase              = self.dummy_image_encoder

										return {
										    "decoder": decoder,
										    "text_encoder": text_encoder,
										    "tokenizer": tokenizer,
										    "text_proj": text_proj,
										    "feature_extractor": feature_extractor,
										    "image_encoder": image_encoder,
										    "super_res_first": super_res_first,
										    "super_res_last": super_res_last,
										    "decoder_scheduler": decoder_scheduler,
										    "super_res_scheduler": super_res_scheduler,
										}


						def      __lowercase       (							self :	int  ,  lowerCAmelCase :	str  ,  lowerCAmelCase :	Optional[Any]=0  ,  lowerCAmelCase :	Any=True    ):
										lowerCAmelCase              = floats_tensor((1, 3, 32, 32)  ,  rng=random.Random(SCREAMING_SNAKE_CASE_    )    ).to(SCREAMING_SNAKE_CASE_    )
										if str(SCREAMING_SNAKE_CASE_    ).startswith("""mps"""    ):
														lowerCAmelCase              = torch.manual_seed(SCREAMING_SNAKE_CASE_    )
										else:
														lowerCAmelCase              = torch.Generator(device=SCREAMING_SNAKE_CASE_    ).manual_seed(SCREAMING_SNAKE_CASE_    )

										if pil_image:
														lowerCAmelCase              = input_image * 0.5 + 0.5
														lowerCAmelCase              = input_image.clamp(0  ,  1    )
														lowerCAmelCase              = input_image.cpu().permute(0  ,  2  ,  3  ,  1    ).float().numpy()
														lowerCAmelCase              = DiffusionPipeline.numpy_to_pil(SCREAMING_SNAKE_CASE_    )[0]

										return {
										    "image": input_image,
										    "generator": generator,
										    "decoder_num_inference_steps": 2,
										    "super_res_num_inference_steps": 2,
										    "output_type": "np",
										}


						def      __lowercase       (							self :	Optional[int]    ):
										lowerCAmelCase              = """cpu"""

										lowerCAmelCase              = self.get_dummy_components()

										lowerCAmelCase              = self.pipeline_class(**SCREAMING_SNAKE_CASE_    )
										lowerCAmelCase              = pipe.to(SCREAMING_SNAKE_CASE_    )

										pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_    )

										lowerCAmelCase              = self.get_dummy_inputs(SCREAMING_SNAKE_CASE_  ,  pil_image=SCREAMING_SNAKE_CASE_    )

										lowerCAmelCase              = pipe(**SCREAMING_SNAKE_CASE_    )
										lowerCAmelCase              = output.images

										lowerCAmelCase              = self.get_dummy_inputs(SCREAMING_SNAKE_CASE_  ,  pil_image=SCREAMING_SNAKE_CASE_    )

										lowerCAmelCase              = pipe(
										    **SCREAMING_SNAKE_CASE_  ,  return_dict=SCREAMING_SNAKE_CASE_  ,  )[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.0002,
										        0.9997,
										        0.9997,
										        0.9969,
										        0.0023,
										        0.9997,
										        0.9969,
										        0.9970,
										    ]    )

										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      __lowercase       (							self :	str    ):
										lowerCAmelCase              = """cpu"""

										lowerCAmelCase              = self.get_dummy_components()

										lowerCAmelCase              = self.pipeline_class(**SCREAMING_SNAKE_CASE_    )
										lowerCAmelCase              = pipe.to(SCREAMING_SNAKE_CASE_    )

										pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_    )

										lowerCAmelCase              = self.get_dummy_inputs(SCREAMING_SNAKE_CASE_  ,  pil_image=SCREAMING_SNAKE_CASE_    )

										lowerCAmelCase              = pipe(**SCREAMING_SNAKE_CASE_    )
										lowerCAmelCase              = output.images

										lowerCAmelCase              = self.get_dummy_inputs(SCREAMING_SNAKE_CASE_  ,  pil_image=SCREAMING_SNAKE_CASE_    )

										lowerCAmelCase              = pipe(
										    **SCREAMING_SNAKE_CASE_  ,  return_dict=SCREAMING_SNAKE_CASE_  ,  )[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.0003, 0.9997, 0.9997, 0.9970, 0.0024, 0.9997, 0.9971, 0.9971]    )

										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      __lowercase       (							self :	str    ):
										lowerCAmelCase              = """cpu"""

										lowerCAmelCase              = self.get_dummy_components()

										lowerCAmelCase              = self.pipeline_class(**SCREAMING_SNAKE_CASE_    )
										lowerCAmelCase              = pipe.to(SCREAMING_SNAKE_CASE_    )

										pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_    )

										lowerCAmelCase              = self.get_dummy_inputs(SCREAMING_SNAKE_CASE_  ,  pil_image=SCREAMING_SNAKE_CASE_    )
										lowerCAmelCase              = [
										    pipeline_inputs["""image"""],
										    pipeline_inputs["""image"""],
										]

										lowerCAmelCase              = pipe(**SCREAMING_SNAKE_CASE_    )
										lowerCAmelCase              = output.images

										lowerCAmelCase              = self.get_dummy_inputs(SCREAMING_SNAKE_CASE_  ,  pil_image=SCREAMING_SNAKE_CASE_    )
										lowerCAmelCase              = [
										    tuple_pipeline_inputs["""image"""],
										    tuple_pipeline_inputs["""image"""],
										]

										lowerCAmelCase              = pipe(
										    **SCREAMING_SNAKE_CASE_  ,  return_dict=SCREAMING_SNAKE_CASE_  ,  )[0]

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

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

										lowerCAmelCase              = np.array(
										    [
										        0.9997,
										        0.9989,
										        0.0008,
										        0.0021,
										        0.9960,
										        0.0018,
										        0.0014,
										        0.0002,
										        0.9933,
										    ]    )

										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      __lowercase       (							self :	Any    ):
										lowerCAmelCase              = torch.device("""cpu"""    )






										class 					SCREAMING_SNAKE_CASE__    :
																_a			  =  1





										lowerCAmelCase              = self.get_dummy_components()

										lowerCAmelCase              = self.pipeline_class(**SCREAMING_SNAKE_CASE_    )
										lowerCAmelCase              = pipe.to(SCREAMING_SNAKE_CASE_    )

										pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_    )

										lowerCAmelCase              = torch.Generator(device=SCREAMING_SNAKE_CASE_    ).manual_seed(0    )
										lowerCAmelCase              = pipe.decoder.dtype
										lowerCAmelCase              = 1

										lowerCAmelCase              = (
										    batch_size,
										    pipe.decoder.config.in_channels,
										    pipe.decoder.config.sample_size,
										    pipe.decoder.config.sample_size,
										)
										lowerCAmelCase              = pipe.prepare_latents(
										    SCREAMING_SNAKE_CASE_  ,  dtype=SCREAMING_SNAKE_CASE_  ,  device=SCREAMING_SNAKE_CASE_  ,  generator=SCREAMING_SNAKE_CASE_  ,  latents=SCREAMING_SNAKE_CASE_  ,  scheduler=DummyScheduler()    )

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

										lowerCAmelCase              = self.get_dummy_inputs(SCREAMING_SNAKE_CASE_  ,  pil_image=SCREAMING_SNAKE_CASE_    )

										lowerCAmelCase              = pipe(
										    **SCREAMING_SNAKE_CASE_  ,  decoder_latents=SCREAMING_SNAKE_CASE_  ,  super_res_latents=SCREAMING_SNAKE_CASE_    ).images

										lowerCAmelCase              = self.get_dummy_inputs(SCREAMING_SNAKE_CASE_  ,  pil_image=SCREAMING_SNAKE_CASE_    )
										# Don't pass image, instead pass embedding
										lowerCAmelCase              = pipeline_inputs.pop("""image"""    )
										lowerCAmelCase              = pipe.image_encoder(SCREAMING_SNAKE_CASE_    ).image_embeds

										lowerCAmelCase              = pipe(
										    **SCREAMING_SNAKE_CASE_  ,  decoder_latents=SCREAMING_SNAKE_CASE_  ,  super_res_latents=SCREAMING_SNAKE_CASE_  ,  image_embeddings=SCREAMING_SNAKE_CASE_  ,  ).images

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


						@skip_mps
						def      __lowercase       (							self :	Tuple    ):
										lowerCAmelCase              = torch_device == """cpu"""

										# Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
										lowerCAmelCase              = 1e-2

										self._test_attention_slicing_forward_pass(
										    test_max_difference=SCREAMING_SNAKE_CASE_  ,  expected_max_diff=SCREAMING_SNAKE_CASE_    )


						@skip_mps
						def      __lowercase       (							self :	List[str]    ):
										lowerCAmelCase              = torch_device == """cpu"""
										lowerCAmelCase              = True
										lowerCAmelCase              = [
										    """decoder_num_inference_steps""",
										    """super_res_num_inference_steps""",
										]

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


						def      __lowercase       (							self :	Optional[int]    ):
										lowerCAmelCase              = [
										    """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=SCREAMING_SNAKE_CASE_  ,  additional_params_copy_to_batched_inputs=SCREAMING_SNAKE_CASE_  ,  )
										else:
														self._test_inference_batch_consistent(
														    additional_params_copy_to_batched_inputs=SCREAMING_SNAKE_CASE_    )


						@skip_mps
						def      __lowercase       (							self :	Optional[int]    ):
										return super().test_dict_tuple_outputs_equivalent()


						@skip_mps
						def      __lowercase       (							self :	str    ):
										return super().test_save_load_local()




						@skip_mps
						def      __lowercase       (							self :	str    ):
										return super().test_save_load_optional_components()






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


						def      __lowercase       (							self :	Optional[int]    ):
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




						def      __lowercase       (							self :	Optional[int]    ):
										lowerCAmelCase              = load_image(
										    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png"""    )
										lowerCAmelCase              = load_numpy(
										    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
										    """/unclip/karlo_v1_alpha_cat_variation_fp16.npy"""    )

										lowerCAmelCase              = UnCLIPImageVariationPipeline.from_pretrained(
										    """kakaobrain/karlo-v1-alpha-image-variations"""  ,  torch_dtype=torch.floataa    )
										lowerCAmelCase              = pipeline.to(SCREAMING_SNAKE_CASE_    )
										pipeline.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_    )

										lowerCAmelCase              = torch.Generator(device="""cpu"""    ).manual_seed(0    )
										lowerCAmelCase              = pipeline(
										    SCREAMING_SNAKE_CASE_  ,  generator=SCREAMING_SNAKE_CASE_  ,  output_type="""np"""  ,  )

										lowerCAmelCase              = output.images[0]

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

										assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_  ,  SCREAMING_SNAKE_CASE_  ,  15    )







from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
						if not (is_transformers_available() and is_torch_available() and is_transformers_version('>=', '4.25.0')):
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						from ...utils.dummy_torch_and_transformers_objects import UnCLIPImageVariationPipeline, UnCLIPPipeline
else:
						from .pipeline_unclip import UnCLIPPipeline
						from .pipeline_unclip_image_variation import UnCLIPImageVariationPipeline
						from .text_proj import UnCLIPTextProjModel




