





'''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    UpperCAmelCase	(       a_			,       unittest.TestCase				):
 __lowercase   		=				UnCLIPImageVariationPipeline
 __lowercase   		=				IMAGE_VARIATION_PARAMS - {"""height""", """width""", """guidance_scale"""}
 __lowercase   		=				IMAGE_VARIATION_BATCH_PARAMS

 __lowercase   		=				[
     """generator""",
     """return_dict""",
     """decoder_num_inference_steps""",
     """super_res_num_inference_steps""",
 ]
 __lowercase   		=				False
 @property
 def      UpperCAmelCase_							(				self      :Optional[int]							)->							List[str]:
     return 32
 @property
 def      UpperCAmelCase_							(				self      :List[Any]							)->							Dict:
     return 32
 @property
 def      UpperCAmelCase_							(				self      :Optional[Any]							)->							str:
     return self.time_input_dim
 @property
 def      UpperCAmelCase_							(				self      :Optional[int]							)->							List[Any]:
     return self.time_input_dim * 4
 @property
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							int:
     return 1_00
 @property
 def      UpperCAmelCase_							(				self      :Tuple							)->							int:
     A__				=			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"							)
     return tokenizer
 @property
 def      UpperCAmelCase_							(				self      :Tuple							)->							List[str]:
     torch.manual_seed(0							)
     A__				=			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(_a							)
 @property
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							Optional[Any]:
     torch.manual_seed(0							)
     A__				=			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(_a							)
 @property
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							str:
     torch.manual_seed(0							)

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

     A__				=			UnCLIPTextProjModel(**_a							)
     return model
 @property
 def      UpperCAmelCase_							(				self      :Any							)->							str:
     torch.manual_seed(0							)

     A__				=			{
         "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__				=			UNetaDConditionModel(**_a							)
     return model
 @property
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							Optional[int]:
     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      :Union[str, Any]							)->							Union[str, Any]:
     torch.manual_seed(0							)

     A__				=			UNetaDModel(**self.dummy_super_res_kwargs							)
     return model
 @property
 def      UpperCAmelCase_							(				self      :List[Any]							)->							Optional[int]:
     # seeded differently to get different unet than `self.dummy_super_res_first`
     torch.manual_seed(1							)

     A__				=			UNetaDModel(**self.dummy_super_res_kwargs							)
     return model
 def      UpperCAmelCase_							(				self      :str							)->							str:
     A__				=			self.dummy_decoder
     A__				=			self.dummy_text_proj
     A__				=			self.dummy_text_encoder
     A__				=			self.dummy_tokenizer
     A__				=			self.dummy_super_res_first
     A__				=			self.dummy_super_res_last

     A__				=			UnCLIPScheduler(
         variance_type="learned_range"						,       prediction_type="epsilon"						,       num_train_timesteps=10_00						,       )

     A__				=			UnCLIPScheduler(
         variance_type="fixed_small_log"						,       prediction_type="epsilon"						,       num_train_timesteps=10_00						,       )

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

     A__				=			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      :Union[str, Any]						,       lowercase_      :Union[str, Any]						,       lowercase_      :Optional[Any]=0						,       lowercase_      :List[str]=True							)->							List[Any]:
     A__				=			floats_tensor((1, 3, 32, 32)						,       rng=random.Random(_a							)							).to(_a							)
     if str(_a							).startswith("mps"							):
         A__				=			torch.manual_seed(_a							)
     else:
         A__				=			torch.Generator(device=_a							).manual_seed(_a							)

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

     return {
         "image": input_image,
         "generator": generator,
         "decoder_num_inference_steps": 2,
         "super_res_num_inference_steps": 2,
         "output_type": "np",
     }
 def      UpperCAmelCase_							(				self      :Any							)->							int:
     A__				=			"cpu"

     A__				=			self.get_dummy_components()

     A__				=			self.pipeline_class(**_a							)
     A__				=			pipe.to(_a							)

     pipe.set_progress_bar_config(disable=_a							)

     A__				=			self.get_dummy_inputs(_a						,       pil_image=_a							)

     A__				=			pipe(**_a							)
     A__				=			output.images

     A__				=			self.get_dummy_inputs(_a						,       pil_image=_a							)

     A__				=			pipe(
         **_a						,       return_dict=_a						,       )[0]

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

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

     A__				=			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      :int							)->							Any:
     A__				=			"cpu"

     A__				=			self.get_dummy_components()

     A__				=			self.pipeline_class(**_a							)
     A__				=			pipe.to(_a							)

     pipe.set_progress_bar_config(disable=_a							)

     A__				=			self.get_dummy_inputs(_a						,       pil_image=_a							)

     A__				=			pipe(**_a							)
     A__				=			output.images

     A__				=			self.get_dummy_inputs(_a						,       pil_image=_a							)

     A__				=			pipe(
         **_a						,       return_dict=_a						,       )[0]

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

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

     A__				=			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      :Optional[Any]							)->							str:
     A__				=			"cpu"

     A__				=			self.get_dummy_components()

     A__				=			self.pipeline_class(**_a							)
     A__				=			pipe.to(_a							)

     pipe.set_progress_bar_config(disable=_a							)

     A__				=			self.get_dummy_inputs(_a						,       pil_image=_a							)
     A__				=			[
         pipeline_inputs["image"],
         pipeline_inputs["image"],
     ]

     A__				=			pipe(**_a							)
     A__				=			output.images

     A__				=			self.get_dummy_inputs(_a						,       pil_image=_a							)
     A__				=			[
         tuple_pipeline_inputs["image"],
         tuple_pipeline_inputs["image"],
     ]

     A__				=			pipe(
         **_a						,       return_dict=_a						,       )[0]

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

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

     A__				=			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							)->							Union[str, Any]:
     A__				=			torch.device("cpu"							)






     class    UpperCAmelCase	:
      __lowercase   		=				1

     A__				=			self.get_dummy_components()

     A__				=			self.pipeline_class(**_a							)
     A__				=			pipe.to(_a							)

     pipe.set_progress_bar_config(disable=_a							)

     A__				=			torch.Generator(device=_a							).manual_seed(0							)
     A__				=			pipe.decoder.dtype
     A__				=			1

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

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

     A__				=			self.get_dummy_inputs(_a						,       pil_image=_a							)

     A__				=			pipe(
         **_a						,       decoder_latents=_a						,       super_res_latents=_a							).images

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

     A__				=			pipe(
         **_a						,       decoder_latents=_a						,       super_res_latents=_a						,       image_embeddings=_a						,       ).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[Any]							)->							List[str]:
     A__				=			torch_device == "cpu"

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

     self._test_attention_slicing_forward_pass(
         test_max_difference=_a						,       expected_max_diff=_a							)
 @skip_mps
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							Any:
     A__				=			torch_device == "cpu"
     A__				=			True
     A__				=			[
         "decoder_num_inference_steps",
         "super_res_num_inference_steps",
     ]

     self._test_inference_batch_single_identical(
         test_max_difference=_a						,       relax_max_difference=_a						,       additional_params_copy_to_batched_inputs=_a						,       )
 def      UpperCAmelCase_							(				self      :Dict							)->							List[Any]:
     A__				=			[
         "decoder_num_inference_steps",
         "super_res_num_inference_steps",
     ]

     if torch_device == "mps":
         # TODO: MPS errors with larger batch sizes
         A__				=			[2, 3]
         self._test_inference_batch_consistent(
             batch_sizes=_a						,       additional_params_copy_to_batched_inputs=_a						,       )
     else:
         self._test_inference_batch_consistent(
             additional_params_copy_to_batched_inputs=_a							)
 @skip_mps
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							Tuple:
     return super().test_dict_tuple_outputs_equivalent()
 @skip_mps
 def      UpperCAmelCase_							(				self      :Optional[int]							)->							str:
     return super().test_save_load_local()



 @skip_mps
 def      UpperCAmelCase_							(				self      :Dict							)->							str:
     return super().test_save_load_optional_components()






@slow
@require_torch_gpu
class    UpperCAmelCase	(       unittest.TestCase				):
 def      UpperCAmelCase_							(				self      :Dict							)->							List[str]:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()



 def      UpperCAmelCase_							(				self      :List[Any]							)->							Union[str, Any]:
     A__				=			load_image(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png"							)
     A__				=			load_numpy(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
         "/unclip/karlo_v1_alpha_cat_variation_fp16.npy"							)

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

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

     A__				=			output.images[0]

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

     assert_mean_pixel_difference(_a						,       _a						,       15							)





"""simple docstring"""




from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    flip_channel_order,
    get_resize_output_image_size,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, is_vision_available, logging


if is_vision_available():
   import PIL

if is_torch_available():
   import torch


_A			:       List[Any]   		=       logging.get_logger(__name__)

class 		a__						(      a_						):
   __lowerCAmelCase            =				["""pixel_values"""]
   def __init__(			self	,    _a = True	,    _a = None	,    _a = PILImageResampling.BILINEAR	,    _a = True	,    _a = 1 / 255	,    _a = True	,    _a = None	,    _a = True	,    **_a	,    ):
          super().__init__(**_a				)
          lowercase    :				Optional[Any]       		=							size if size is not None else {"shortest_edge": 224}
          lowercase    :				List[Any]       		=							get_size_dict(_a	,    default_to_square=_a				)
          lowercase    :				str       		=							crop_size if crop_size is not None else {"height": 256, "width": 256}
          lowercase    :				List[str]       		=							get_size_dict(_a	,    param_name="crop_size"				)

          lowercase    :				int       		=							do_resize
          lowercase    :				Optional[int]       		=							size
          lowercase    :				str       		=							resample
          lowercase    :				List[Any]       		=							do_rescale
          lowercase    :				Union[str, Any]       		=							rescale_factor
          lowercase    :				Optional[int]       		=							do_center_crop
          lowercase    :				Union[str, Any]       		=							crop_size
          lowercase    :				Optional[Any]       		=							do_flip_channel_order
   def 			__magic_name__			(			self	,    _a	,    _a	,    _a = PIL.Image.BILINEAR	,    _a = None	,    **_a	,    ):
          lowercase    :				List[Any]       		=							get_size_dict(_a	,    default_to_square=_a				)
          if "shortest_edge" not in size:
                 raise ValueError(f"""The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}"""				)
          lowercase    :				Union[str, Any]       		=							get_resize_output_image_size(_a	,    size=size["shortest_edge"]	,    default_to_square=_a				)
          return resize(_a	,    size=_a	,    resample=_a	,    data_format=_a	,    **_a				)
   def 			__magic_name__			(			self	,    _a	,    _a	,    _a = None	,    **_a	,    ):
          lowercase    :				str       		=							get_size_dict(_a				)
          if "height" not in size or "width" not in size:
                 raise ValueError(f"""The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}"""				)
          return center_crop(_a	,    size=(size["height"], size["width"])	,    data_format=_a	,    **_a				)
   def 			__magic_name__			(			self	,    _a	,    _a	,    _a = None	,    **_a	,    ):
          return rescale(_a	,    scale=_a	,    data_format=_a	,    **_a				)
   def 			__magic_name__			(			self	,    _a	,    _a = None				):
          return flip_channel_order(_a	,    data_format=_a				)
   def 			__magic_name__			(			self	,    _a	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = ChannelDimension.FIRST	,    **_a	,    ):
          lowercase    :				Optional[Any]       		=							do_resize if do_resize is not None else self.do_resize
          lowercase    :				Tuple       		=							resample if resample is not None else self.resample
          lowercase    :				List[str]       		=							do_rescale if do_rescale is not None else self.do_rescale
          lowercase    :				List[Any]       		=							rescale_factor if rescale_factor is not None else self.rescale_factor
          lowercase    :				Union[str, Any]       		=							do_center_crop if do_center_crop is not None else self.do_center_crop
          lowercase    :				Optional[int]       		=							(
              do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order
          )

          lowercase    :				str       		=							size if size is not None else self.size
          lowercase    :				Any       		=							get_size_dict(_a	,    default_to_square=_a				)
          lowercase    :				int       		=							crop_size if crop_size is not None else self.crop_size
          lowercase    :				Any       		=							get_size_dict(_a	,    param_name="crop_size"				)

          lowercase    :				int       		=							make_list_of_images(_a				)

          if not valid_images(_a				):
                 raise ValueError(
                     "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
                     "torch.Tensor, tf.Tensor or jax.ndarray."				)

          if do_resize and size is None:
                 raise ValueError("Size must be specified if do_resize is True."				)

          if do_rescale and rescale_factor is None:
                 raise ValueError("Rescale factor must be specified if do_rescale is True."				)

          if do_center_crop and crop_size is None:
                 raise ValueError("Crop size must be specified if do_center_crop is True."				)

          # All transformations expect numpy arrays.
          lowercase    :				Any       		=							[to_numpy_array(_a				) for image in images]

          if do_resize:
                 lowercase    :				Optional[int]       		=							[self.resize(image=_a	,    size=_a	,    resample=_a				) for image in images]

          if do_center_crop:
                 lowercase    :				str       		=							[self.center_crop(image=_a	,    size=_a				) for image in images]

          if do_rescale:
                 lowercase    :				Union[str, Any]       		=							[self.rescale(image=_a	,    scale=_a				) for image in images]

          # the pretrained checkpoints assume images are BGR, not RGB
          if do_flip_channel_order:
                 lowercase    :				int       		=							[self.flip_channel_order(image=_a				) for image in images]

          lowercase    :				int       		=							[to_channel_dimension_format(_a	,    _a				) for image in images]

          lowercase    :				Optional[Any]       		=							{"pixel_values": images}
          return BatchFeature(data=_a	,    tensor_type=_a				)

   def 			__magic_name__			(			self	,    _a	,    _a = None				):
          lowercase    :				Optional[int]       		=							outputs.logits

          # Resize logits and compute semantic segmentation maps
          if target_sizes is not None:
                 if len(_a				) != len(_a				):
                        raise ValueError(
                            "Make sure that you pass in as many target sizes as the batch dimension of the logits"				)

                 if is_torch_tensor(_a				):
                        lowercase    :				Tuple       		=							target_sizes.numpy()

                 lowercase    :				List[Any]       		=							[]

                 for idx in range(len(_a				)				):
                        lowercase    :				Dict       		=							torch.nn.functional.interpolate(
                            logits[idx].unsqueeze(dim=0				)	,    size=target_sizes[idx]	,    mode="bilinear"	,    align_corners=_a				)
                        lowercase    :				Tuple       		=							resized_logits[0].argmax(dim=0				)
                        semantic_segmentation.append(_a				)
          else:
                 lowercase    :				str       		=							logits.argmax(dim=1				)
                 lowercase    :				Dict       		=							[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]				)]

          return semantic_segmentation


