


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 _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			str									=    UnCLIPImageVariationPipeline
       lowercase_				:			str									=    IMAGE_VARIATION_PARAMS - {"""height""", """width""", """guidance_scale"""}
       lowercase_				:			Any									=    IMAGE_VARIATION_BATCH_PARAMS

       lowercase_				:			Optional[Any]									=    [
           """generator""",
           """return_dict""",
           """decoder_num_inference_steps""",
           """super_res_num_inference_steps""",
       ]
       lowercase_				:			Union[str, Any]									=    False
       @property
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

           return 32
       @property
       def     UpperCamelCase  (      self) ->     Tuple:




           """simple docstring"""

           return 32
       @property
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           return self.time_input_dim
       @property
       def     UpperCamelCase  (      self) ->     str:




           """simple docstring"""

           return self.time_input_dim * 4
       @property
       def     UpperCamelCase  (      self) ->     Optional[Any]:




           """simple docstring"""

           return 1_00
       @property
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           _lowercase :							Union[str, Any]   	=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')
           return tokenizer
       @property
       def     UpperCamelCase  (      self) ->     Any:




           """simple docstring"""

           torch.manual_seed(0)
           _lowercase :							List[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=10_00,   )
           return CLIPTextModelWithProjection(lowerCamelCase)
       @property
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           torch.manual_seed(0)
           _lowercase :							List[Any]   	=	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(lowerCamelCase)
       @property
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

           torch.manual_seed(0)

           _lowercase :							List[str]   	=	{
               'clip_embeddings_dim': self.text_embedder_hidden_size,
               'time_embed_dim': self.time_embed_dim,
               'cross_attention_dim': self.cross_attention_dim,
           }

           _lowercase :							Tuple   	=	UnCLIPTextProjModel(**lowerCamelCase)
           return model
       @property
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           torch.manual_seed(0)

           _lowercase :							str   	=	{
               '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',
           }

           _lowercase :							Optional[Any]   	=	UNetaDConditionModel(**lowerCamelCase)
           return model
       @property
       def     UpperCamelCase  (      self) ->     List[Any]:




           """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) ->     List[Any]:




           """simple docstring"""

           torch.manual_seed(0)

           _lowercase :							Any   	=	UNetaDModel(**self.dummy_super_res_kwargs)
           return model
       @property
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

           torch.manual_seed(1)

           _lowercase :							int   	=	UNetaDModel(**self.dummy_super_res_kwargs)
           return model
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           _lowercase :							List[Any]   	=	self.dummy_decoder
           _lowercase :							List[str]   	=	self.dummy_text_proj
           _lowercase :							str   	=	self.dummy_text_encoder
           _lowercase :							List[str]   	=	self.dummy_tokenizer
           _lowercase :							Dict   	=	self.dummy_super_res_first
           _lowercase :							Union[str, Any]   	=	self.dummy_super_res_last

           _lowercase :							Optional[Any]   	=	UnCLIPScheduler(
               variance_type='learned_range',   prediction_type='epsilon',   num_train_timesteps=10_00,   )

           _lowercase :							Optional[int]   	=	UnCLIPScheduler(
               variance_type='fixed_small_log',   prediction_type='epsilon',   num_train_timesteps=10_00,   )

           _lowercase :							List[Any]   	=	CLIPImageProcessor(crop_size=32,   size=32)

           _lowercase :							Optional[int]   	=	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     UpperCamelCase  (      self,   lowerCamelCase,   lowerCamelCase=0,   lowerCamelCase=True) ->     List[Any]:




           """simple docstring"""

           _lowercase :							Tuple   	=	floats_tensor((1, 3, 32, 32),   rng=random.Random(lowerCamelCase)).to(lowerCamelCase)
           if str(lowerCamelCase).startswith('mps'):
               _lowercase :							int   	=	torch.manual_seed(lowerCamelCase)
           else:
               _lowercase :							str   	=	torch.Generator(device=lowerCamelCase).manual_seed(lowerCamelCase)

           if pil_image:
               _lowercase :							Any   	=	input_image * 0.5 + 0.5
               _lowercase :							Optional[Any]   	=	input_image.clamp(0,   1)
               _lowercase :							List[Any]   	=	input_image.cpu().permute(0,   2,   3,   1).float().numpy()
               _lowercase :							str   	=	DiffusionPipeline.numpy_to_pil(lowerCamelCase)[0]

           return {
               "image": input_image,
               "generator": generator,
               "decoder_num_inference_steps": 2,
               "super_res_num_inference_steps": 2,
               "output_type": "np",
           }
       def     UpperCamelCase  (      self) ->     Optional[int]:




           """simple docstring"""

           _lowercase :							Any   	=	'cpu'

           _lowercase :							Dict   	=	self.get_dummy_components()

           _lowercase :							Union[str, Any]   	=	self.pipeline_class(**lowerCamelCase)
           _lowercase :							Optional[int]   	=	pipe.to(lowerCamelCase)

           pipe.set_progress_bar_config(disable=lowerCamelCase)

           _lowercase :							int   	=	self.get_dummy_inputs(lowerCamelCase,   pil_image=lowerCamelCase)

           _lowercase :							Dict   	=	pipe(**lowerCamelCase)
           _lowercase :							Tuple   	=	output.images

           _lowercase :							List[str]   	=	self.get_dummy_inputs(lowerCamelCase,   pil_image=lowerCamelCase)

           _lowercase :							str   	=	pipe(
               **lowerCamelCase,   return_dict=lowerCamelCase,   )[0]

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

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

           _lowercase :							Any   	=	np.array(
               [
                   0.9_9_9_7,
                   0.0_0_0_2,
                   0.9_9_9_7,
                   0.9_9_9_7,
                   0.9_9_6_9,
                   0.0_0_2_3,
                   0.9_9_9_7,
                   0.9_9_6_9,
                   0.9_9_7_0,
               ])

           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) ->     str:




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	'cpu'

           _lowercase :							Union[str, Any]   	=	self.get_dummy_components()

           _lowercase :							Union[str, Any]   	=	self.pipeline_class(**lowerCamelCase)
           _lowercase :							Dict   	=	pipe.to(lowerCamelCase)

           pipe.set_progress_bar_config(disable=lowerCamelCase)

           _lowercase :							Dict   	=	self.get_dummy_inputs(lowerCamelCase,   pil_image=lowerCamelCase)

           _lowercase :							Any   	=	pipe(**lowerCamelCase)
           _lowercase :							Tuple   	=	output.images

           _lowercase :							Any   	=	self.get_dummy_inputs(lowerCamelCase,   pil_image=lowerCamelCase)

           _lowercase :							Optional[Any]   	=	pipe(
               **lowerCamelCase,   return_dict=lowerCamelCase,   )[0]

           _lowercase :							int   	=	image[0, -3:, -3:, -1]
           _lowercase :							Dict   	=	image_from_tuple[0, -3:, -3:, -1]

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

           _lowercase :							Tuple   	=	np.array([0.9_9_9_7, 0.0_0_0_3, 0.9_9_9_7, 0.9_9_9_7, 0.9_9_7_0, 0.0_0_2_4, 0.9_9_9_7, 0.9_9_7_1, 0.9_9_7_1])

           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) ->     Any:




           """simple docstring"""

           _lowercase :							List[Any]   	=	'cpu'

           _lowercase :							List[str]   	=	self.get_dummy_components()

           _lowercase :							List[str]   	=	self.pipeline_class(**lowerCamelCase)
           _lowercase :							List[str]   	=	pipe.to(lowerCamelCase)

           pipe.set_progress_bar_config(disable=lowerCamelCase)

           _lowercase :							Union[str, Any]   	=	self.get_dummy_inputs(lowerCamelCase,   pil_image=lowerCamelCase)
           _lowercase :							Union[str, Any]   	=	[
               pipeline_inputs['image'],
               pipeline_inputs['image'],
           ]

           _lowercase :							Union[str, Any]   	=	pipe(**lowerCamelCase)
           _lowercase :							str   	=	output.images

           _lowercase :							List[Any]   	=	self.get_dummy_inputs(lowerCamelCase,   pil_image=lowerCamelCase)
           _lowercase :							List[Any]   	=	[
               tuple_pipeline_inputs['image'],
               tuple_pipeline_inputs['image'],
           ]

           _lowercase :							List[str]   	=	pipe(
               **lowerCamelCase,   return_dict=lowerCamelCase,   )[0]

           _lowercase :							str   	=	image[0, -3:, -3:, -1]
           _lowercase :							str   	=	image_from_tuple[0, -3:, -3:, -1]

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

           _lowercase :							Any   	=	np.array(
               [
                   0.9_9_9_7,
                   0.9_9_8_9,
                   0.0_0_0_8,
                   0.0_0_2_1,
                   0.9_9_6_0,
                   0.0_0_1_8,
                   0.0_0_1_4,
                   0.0_0_0_2,
                   0.9_9_3_3,
               ])

           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) ->     Any:




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	torch.device('cpu')






           class _lowerCamelCase:
                  lowercase_				:			List[Any]									=    1


           _lowercase :							int   	=	self.get_dummy_components()

           _lowercase :							Any   	=	self.pipeline_class(**lowerCamelCase)
           _lowercase :							Union[str, Any]   	=	pipe.to(lowerCamelCase)

           pipe.set_progress_bar_config(disable=lowerCamelCase)

           _lowercase :							Union[str, Any]   	=	torch.Generator(device=lowerCamelCase).manual_seed(0)
           _lowercase :							List[str]   	=	pipe.decoder.dtype
           _lowercase :							str   	=	1

           _lowercase :							List[Any]   	=	(
               batch_size,
               pipe.decoder.config.in_channels,
               pipe.decoder.config.sample_size,
               pipe.decoder.config.sample_size,
           )
           _lowercase :							Optional[Any]   	=	pipe.prepare_latents(
               lowerCamelCase,   dtype=lowerCamelCase,   device=lowerCamelCase,   generator=lowerCamelCase,   latents=lowerCamelCase,   scheduler=DummyScheduler())

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

           _lowercase :							int   	=	self.get_dummy_inputs(lowerCamelCase,   pil_image=lowerCamelCase)

           _lowercase :							List[str]   	=	pipe(
               **lowerCamelCase,   decoder_latents=lowerCamelCase,   super_res_latents=lowerCamelCase).images

           _lowercase :							Any   	=	self.get_dummy_inputs(lowerCamelCase,   pil_image=lowerCamelCase)
           # Don't pass image, instead pass embedding
           _lowercase :							int   	=	pipeline_inputs.pop('image')
           _lowercase :							Optional[int]   	=	pipe.image_encoder(lowerCamelCase).image_embeds

           _lowercase :							str   	=	pipe(
               **lowerCamelCase,   decoder_latents=lowerCamelCase,   super_res_latents=lowerCamelCase,   image_embeddings=lowerCamelCase,   ).images

           # make sure passing text embeddings manually is identical
           assert np.abs(img_out_a - img_out_a).max() < 1E-4
       @skip_mps
       def     UpperCamelCase  (      self) ->     Optional[int]:




           """simple docstring"""

           _lowercase :							List[Any]   	=	torch_device == 'cpu'

           # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
           _lowercase :							int   	=	1E-2

           self._test_attention_slicing_forward_pass(
               test_max_difference=lowerCamelCase,   expected_max_diff=lowerCamelCase)
       @skip_mps
       def     UpperCamelCase  (      self) ->     Any:




           """simple docstring"""

           _lowercase :							Tuple   	=	torch_device == 'cpu'
           _lowercase :							Tuple   	=	True
           _lowercase :							Optional[int]   	=	[
               '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     UpperCamelCase  (      self) ->     Optional[int]:




           """simple docstring"""

           _lowercase :							List[Any]   	=	[
               'decoder_num_inference_steps',
               'super_res_num_inference_steps',
           ]

           if torch_device == "mps":
               # TODO: MPS errors with larger batch sizes
               _lowercase :							List[str]   	=	[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     UpperCamelCase  (      self) ->     Tuple:




           """simple docstring"""

           return super().test_dict_tuple_outputs_equivalent()
       @skip_mps
       def     UpperCamelCase  (      self) ->     Any:




           """simple docstring"""

           return super().test_save_load_local()




       @skip_mps
       def     UpperCamelCase  (      self) ->     List[str]:




           """simple docstring"""

           return super().test_save_load_optional_components()






@slow
@require_torch_gpu
class _lowerCamelCase(		unittest.TestCase     ):
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

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




       def     UpperCamelCase  (      self) ->     Dict:




           """simple docstring"""

           _lowercase :							List[Any]   	=	load_image(
               'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png')
           _lowercase :							int   	=	load_numpy(
               'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
               '/unclip/karlo_v1_alpha_cat_variation_fp16.npy')

           _lowercase :							List[str]   	=	UnCLIPImageVariationPipeline.from_pretrained(
               'kakaobrain/karlo-v1-alpha-image-variations',   torch_dtype=torch.floataa)
           _lowercase :							Optional[Any]   	=	pipeline.to(lowerCamelCase)
           pipeline.set_progress_bar_config(disable=lowerCamelCase)

           _lowercase :							List[Any]   	=	torch.Generator(device='cpu').manual_seed(0)
           _lowercase :							List[Any]   	=	pipeline(
               lowerCamelCase,   generator=lowerCamelCase,   output_type='np',   )

           _lowercase :							Optional[Any]   	=	output.images[0]

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

           assert_mean_pixel_difference(lowerCamelCase,   lowerCamelCase,   15)






import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding






class _lowerCamelCase(		_a     ):
       lowercase_				:			int									=    ["""image_processor""", """tokenizer"""]
       lowercase_				:			List[str]									=    """CLIPImageProcessor"""
       lowercase_				:			Union[str, Any]									=    ("""CLIPTokenizer""", """CLIPTokenizerFast""")
       def __init__(      self,   lowerCamelCase=None,   lowerCamelCase=None,   **lowerCamelCase) ->     Dict:




           """simple docstring"""

           _lowercase :							List[Any]   	=	None
           if "feature_extractor" in kwargs:
               warnings.warn(
                   'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
                   ' instead.',   lowerCamelCase,   )
               _lowercase :							Dict   	=	kwargs.pop('feature_extractor')

           _lowercase :							int   	=	image_processor if image_processor is not None else feature_extractor
           if image_processor is None:
               raise ValueError('You need to specify an `image_processor`.')
           if tokenizer is None:
               raise ValueError('You need to specify a `tokenizer`.')

           super().__init__(lowerCamelCase,   lowerCamelCase)
       def __call__(      self,   lowerCamelCase=None,   lowerCamelCase=None,   lowerCamelCase=None,   **lowerCamelCase) ->     Tuple:




           """simple docstring"""

           if text is None and images is None:
               raise ValueError('You have to specify either text or images. Both cannot be none.')

           if text is not None:
               _lowercase :							Optional[Any]   	=	self.tokenizer(lowerCamelCase,   return_tensors=lowerCamelCase,   **lowerCamelCase)

           if images is not None:
               _lowercase :							List[Any]   	=	self.image_processor(lowerCamelCase,   return_tensors=lowerCamelCase,   **lowerCamelCase)

           if text is not None and images is not None:
               _lowercase :							int   	=	image_features.pixel_values
               return encoding
           elif text is not None:
               return encoding
           else:
               return BatchEncoding(data=dict(**lowerCamelCase),   tensor_type=lowerCamelCase)
       def     UpperCamelCase  (      self,   *lowerCamelCase,   **lowerCamelCase) ->     Union[str, Any]:




           """simple docstring"""

           return self.tokenizer.batch_decode(*lowerCamelCase,   **lowerCamelCase)
       def     UpperCamelCase  (      self,   *lowerCamelCase,   **lowerCamelCase) ->     int:




           """simple docstring"""

           return self.tokenizer.decode(*lowerCamelCase,   **lowerCamelCase)
       @property
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           _lowercase :							int   	=	self.tokenizer.model_input_names
           _lowercase :							Dict   	=	self.image_processor.model_input_names
           return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
       @property
       def     UpperCamelCase  (      self) ->     List[str]:




           """simple docstring"""

           warnings.warn(
               '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.',   lowerCamelCase,   )
           return self.image_processor_class




       @property
       def     UpperCamelCase  (      self) ->     Optional[Any]:




           """simple docstring"""

           warnings.warn(
               '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.',   lowerCamelCase,   )
           return self.image_processor


