

'''simple docstring'''







from typing import Dict, List, Optional, Union

import numpy as np

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


if is_vision_available():
      import PIL


__lowerCamelCase		       =    logging.get_logger(__name__)


class 					A__      (      _snake_case			):
    lowercase						=		["pixel_values"]



    def __init__(					self		,				UpperCamelCase__ = True		,				UpperCamelCase__ = None		,				UpperCamelCase__ = 0.9		,				UpperCamelCase__ = PILImageResampling.BICUBIC		,				UpperCamelCase__ = True		,				UpperCamelCase__ = None		,				UpperCamelCase__ = 1 / 255		,				UpperCamelCase__ = True		,				UpperCamelCase__ = True		,				UpperCamelCase__ = None		,				UpperCamelCase__ = None		,				**UpperCamelCase__		,				)						->      None:





           '''simple docstring'''
           super().__init__(**UpperCamelCase__       )
           A_						=      size if size is not None else {"""shortest_edge""": 224}
           A_						=      get_size_dict(UpperCamelCase__		,				default_to_square=UpperCamelCase__       )
           A_						=      crop_size if crop_size is not None else {"""height""": 224, """width""": 224}
           A_						=      get_size_dict(UpperCamelCase__		,				param_name="""crop_size"""       )

           A_						=      do_resize
           A_						=      size
           A_						=      crop_pct
           A_						=      resample
           A_						=      do_center_crop
           A_						=      crop_size
           A_						=      do_rescale
           A_						=      rescale_factor
           A_						=      do_normalize
           A_						=      image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
           A_						=      image_std if image_std is not None else IMAGENET_DEFAULT_STD



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__ = None		,				UpperCamelCase__ = PILImageResampling.BICUBIC		,				UpperCamelCase__ = None		,				**UpperCamelCase__		,				)						->      np.ndarray:





           '''simple docstring'''
           A_						=      get_size_dict(UpperCamelCase__		,				default_to_square=UpperCamelCase__       )
           if "shortest_edge" not in size and ("height" not in size or "width" not in size):
                  raise ValueError(f'''size must contain \'height\' and \'width\' or \'shortest_edge\' as keys. Got {size.keys()}'''       )
           if crop_pct is not None:
                  if "shortest_edge" in size:
                         A_						=      int(size["""shortest_edge"""] / crop_pct       )
                  elif "height" in size and "width" in size:
                         if size["height"] == size["width"]:
                                A_						=      int(size["""height"""] / crop_pct       )
                         else:
                                A_						=      (int(size["""height"""] / crop_pct       ), int(size["""width"""] / crop_pct       ))
                  else:
                         raise ValueError("""Invalid size for resize: {}""".format(UpperCamelCase__       )       )

                  A_						=      get_resize_output_image_size(UpperCamelCase__		,				size=UpperCamelCase__		,				default_to_square=UpperCamelCase__       )
           else:
                  if "shortest_edge" in size:
                         A_						=      get_resize_output_image_size(UpperCamelCase__		,				size=size["""shortest_edge"""]		,				default_to_square=UpperCamelCase__       )
                  elif "height" in size and "width" in size:
                         A_						=      (size["""height"""], size["""width"""])
                  else:
                         raise ValueError("""Invalid size for resize: {}""".format(UpperCamelCase__       )       )

           return resize(UpperCamelCase__		,				size=UpperCamelCase__		,				resample=UpperCamelCase__		,				data_format=UpperCamelCase__		,				**UpperCamelCase__       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__ = None		,				**UpperCamelCase__		,				)						->      np.ndarray:





           '''simple docstring'''
           A_						=      get_size_dict(UpperCamelCase__       )
           if "height" not in size or "width" not in size:
                  raise ValueError(f'''size must contain \'height\' and \'width\' as keys. Got {size.keys()}'''       )
           return center_crop(UpperCamelCase__		,				size=(size["""height"""], size["""width"""])		,				data_format=UpperCamelCase__		,				**UpperCamelCase__       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__ = None		,				**UpperCamelCase__		,				)						->      Tuple:





           '''simple docstring'''
           return rescale(UpperCamelCase__		,				scale=UpperCamelCase__		,				data_format=UpperCamelCase__		,				**UpperCamelCase__       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__ = None		,				**UpperCamelCase__		,				)						->      np.ndarray:





           '''simple docstring'''
           return normalize(UpperCamelCase__		,				mean=UpperCamelCase__		,				std=UpperCamelCase__		,				data_format=UpperCamelCase__		,				**UpperCamelCase__       )







    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__ = None		,				UpperCamelCase__ = None		,				UpperCamelCase__ = None		,				UpperCamelCase__ = None		,				UpperCamelCase__ = None		,				UpperCamelCase__ = None		,				UpperCamelCase__ = None		,				UpperCamelCase__ = None		,				UpperCamelCase__ = None		,				UpperCamelCase__ = None		,				UpperCamelCase__ = None		,				UpperCamelCase__ = None		,				UpperCamelCase__ = ChannelDimension.FIRST		,				**UpperCamelCase__		,				)						->      PIL.Image.Image:





           '''simple docstring'''
           A_						=      do_resize if do_resize is not None else self.do_resize
           A_						=      crop_pct if crop_pct is not None else self.crop_pct
           A_						=      resample if resample is not None else self.resample
           A_						=      do_center_crop if do_center_crop is not None else self.do_center_crop
           A_						=      do_rescale if do_rescale is not None else self.do_rescale
           A_						=      rescale_factor if rescale_factor is not None else self.rescale_factor
           A_						=      do_normalize if do_normalize is not None else self.do_normalize
           A_						=      image_mean if image_mean is not None else self.image_mean
           A_						=      image_std if image_std is not None else self.image_std

           A_						=      size if size is not None else self.size
           A_						=      get_size_dict(UpperCamelCase__		,				default_to_square=UpperCamelCase__       )
           A_						=      crop_size if crop_size is not None else self.crop_size
           A_						=      get_size_dict(UpperCamelCase__		,				param_name="""crop_size"""       )

           A_						=      make_list_of_images(UpperCamelCase__       )

           if not valid_images(UpperCamelCase__       ):
                  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 or resample is None:
                  raise ValueError("""Size and resample must be specified if do_resize is True."""       )

           if do_center_crop and crop_pct is None:
                  raise ValueError("""Crop_pct must be specified if do_center_crop is True."""       )

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

           if do_normalize and (image_mean is None or image_std is None):
                  raise ValueError("""Image mean and std must be specified if do_normalize is True."""       )

           # All transformations expect numpy arrays.
           A_						=      [to_numpy_array(UpperCamelCase__       ) for image in images]

           if do_resize:
                  A_						=      [self.resize(image=UpperCamelCase__		,				size=UpperCamelCase__		,				crop_pct=UpperCamelCase__		,				resample=UpperCamelCase__       ) for image in images]

           if do_center_crop:
                  A_						=      [self.center_crop(image=UpperCamelCase__		,				size=UpperCamelCase__       ) for image in images]

           if do_rescale:
                  A_						=      [self.rescale(image=UpperCamelCase__		,				scale=UpperCamelCase__       ) for image in images]

           if do_normalize:
                  A_						=      [self.normalize(image=UpperCamelCase__		,				mean=UpperCamelCase__		,				std=UpperCamelCase__       ) for image in images]

           A_						=      [to_channel_dimension_format(UpperCamelCase__		,				UpperCamelCase__       ) for image in images]

           A_						=      {"""pixel_values""": images}
           return BatchFeature(data=UpperCamelCase__		,				tensor_type=UpperCamelCase__       )



'''simple docstring'''







import unittest

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

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class 					A__      (      _snake_case						,						unittest.TestCase			):
    lowercase						=		KandinskyVaaPriorPipeline
    lowercase						=		["prompt"]
    lowercase						=		["prompt", "negative_prompt"]
    lowercase						=		[
        "num_images_per_prompt",
        "generator",
        "num_inference_steps",
        "latents",
        "negative_prompt",
        "guidance_scale",
        "output_type",
        "return_dict",
    ]
    lowercase						=		False



    @property
    def 					snake_case_			(					self       )						->      Any:





           '''simple docstring'''
           return 32



    @property
    def 					snake_case_			(					self       )						->      Optional[int]:





           '''simple docstring'''
           return 32



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





           '''simple docstring'''
           return self.time_input_dim



    @property
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           return self.time_input_dim * 4



    @property
    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           return 100



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





           '''simple docstring'''
           A_						=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )
           return tokenizer



    @property
    def 					snake_case_			(					self       )						->      Union[str, Any]:





           '''simple docstring'''
           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-0_5		,				num_attention_heads=4		,				num_hidden_layers=5		,				pad_token_id=1		,				vocab_size=1000		,				)
           return CLIPTextModelWithProjection(UpperCamelCase__       )



    @property
    def 					snake_case_			(					self       )						->      Tuple:





           '''simple docstring'''
           torch.manual_seed(0       )

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

           A_						=      PriorTransformer(**UpperCamelCase__       )
           # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
           A_						=      nn.Parameter(torch.ones(model.clip_std.shape       )       )
           return model



    @property
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           torch.manual_seed(0       )
           A_						=      CLIPVisionConfig(
               hidden_size=self.text_embedder_hidden_size		,				image_size=224		,				projection_dim=self.text_embedder_hidden_size		,				intermediate_size=37		,				num_attention_heads=4		,				num_channels=3		,				num_hidden_layers=5		,				patch_size=14		,				)

           A_						=      CLIPVisionModelWithProjection(UpperCamelCase__       )
           return model



    @property
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           A_						=      CLIPImageProcessor(
               crop_size=224		,				do_center_crop=UpperCamelCase__		,				do_normalize=UpperCamelCase__		,				do_resize=UpperCamelCase__		,				image_mean=[0.48145466, 0.4578275, 0.40821073]		,				image_std=[0.26862954, 0.26130258, 0.27577711]		,				resample=3		,				size=224		,				)

           return image_processor



    def 					snake_case_			(					self       )						->      Union[str, Any]:





           '''simple docstring'''
           A_						=      self.dummy_prior
           A_						=      self.dummy_image_encoder
           A_						=      self.dummy_text_encoder
           A_						=      self.dummy_tokenizer
           A_						=      self.dummy_image_processor

           A_						=      UnCLIPScheduler(
               variance_type="""fixed_small_log"""		,				prediction_type="""sample"""		,				num_train_timesteps=1000		,				clip_sample=UpperCamelCase__		,				clip_sample_range=10.0		,				)

           A_						=      {
               """prior""": prior,
               """image_encoder""": image_encoder,
               """text_encoder""": text_encoder,
               """tokenizer""": tokenizer,
               """scheduler""": scheduler,
               """image_processor""": image_processor,
           }

           return components



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__=0       )						->      Optional[int]:





           '''simple docstring'''
           if str(UpperCamelCase__       ).startswith("""mps"""       ):
                  A_						=      torch.manual_seed(UpperCamelCase__       )
           else:
                  A_						=      torch.Generator(device=UpperCamelCase__       ).manual_seed(UpperCamelCase__       )
           A_						=      {
               """prompt""": """horse""",
               """generator""": generator,
               """guidance_scale""": 4.0,
               """num_inference_steps""": 2,
               """output_type""": """np""",
           }
           return inputs



    def 					snake_case_			(					self       )						->      List[str]:





           '''simple docstring'''
           A_						=      """cpu"""

           A_						=      self.get_dummy_components()

           A_						=      self.pipeline_class(**UpperCamelCase__       )
           A_						=      pipe.to(UpperCamelCase__       )

           pipe.set_progress_bar_config(disable=UpperCamelCase__       )

           A_						=      pipe(**self.get_dummy_inputs(UpperCamelCase__       )       )
           A_						=      output.image_embeds

           A_						=      pipe(
               **self.get_dummy_inputs(UpperCamelCase__       )		,				return_dict=UpperCamelCase__		,				)[0]

           A_						=      image[0, -10:]
           A_						=      image_from_tuple[0, -10:]

           assert image.shape == (1, 32)

           A_						=      np.array(
               [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]       )

           assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2
           assert np.abs(image_from_tuple_slice.flatten() - expected_slice       ).max() < 1e-2



    @skip_mps
    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           A_						=      torch_device == """cpu"""
           A_						=      True
           A_						=      False

           self._test_inference_batch_single_identical(
               test_max_difference=UpperCamelCase__		,				relax_max_difference=UpperCamelCase__		,				test_mean_pixel_difference=UpperCamelCase__		,				)







    @skip_mps
    def 					snake_case_			(					self       )						->      Union[str, Any]:





           '''simple docstring'''
           A_						=      torch_device == """cpu"""
           A_						=      False

           self._test_attention_slicing_forward_pass(
               test_max_difference=UpperCamelCase__		,				test_mean_pixel_difference=UpperCamelCase__		,				)
