






import unittest

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

from diffusers import KandinskyPriorPipeline, 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 __lowerCamelCase				(   snake_case__  , unittest.TestCase):





	"""simple docstring"""






	UpperCamelCase__			=   KandinskyPriorPipeline
	UpperCamelCase__			=   ["prompt"]
	UpperCamelCase__			=   ["prompt", "negative_prompt"]
	UpperCamelCase__			=   [
	    "num_images_per_prompt",
	    "generator",
	    "num_inference_steps",
	    "latents",
	    "negative_prompt",
	    "guidance_scale",
	    "output_type",
	    "return_dict",
	]
	UpperCamelCase__			=   False

	@property
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							return 32

	@property
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							return 32

	@property
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							return self.time_input_dim

	@property
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							return self.time_input_dim * 4

	@property
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							return 100

	@property
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)
							return tokenizer

	@property
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							torch.manual_seed(0					)
							_UpperCAmelCase             =    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(UpperCAmelCase					)

	@property
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							torch.manual_seed(0					)

							_UpperCAmelCase             =    {
							    'num_attention_heads': 2,
							    'attention_head_dim': 12,
							    'embedding_dim': self.text_embedder_hidden_size,
							    'num_layers': 1,
							}

							_UpperCAmelCase             =    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
							_UpperCAmelCase             =    nn.Parameter(torch.ones(model.clip_std.shape					)					)
							return model

	@property
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							torch.manual_seed(0					)
							_UpperCAmelCase             =    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 ,					)

							_UpperCAmelCase             =    CLIPVisionModelWithProjection(UpperCAmelCase					)
							return model

	@property
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    CLIPImageProcessor(
							    crop_size=224 ,					do_center_crop=UpperCAmelCase ,					do_normalize=UpperCAmelCase ,					do_resize=UpperCAmelCase ,					image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73] ,					image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11] ,					resample=3 ,					size=224 ,					)

							return image_processor

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    self.dummy_prior
							_UpperCAmelCase             =    self.dummy_image_encoder
							_UpperCAmelCase             =    self.dummy_text_encoder
							_UpperCAmelCase             =    self.dummy_tokenizer
							_UpperCAmelCase             =    self.dummy_image_processor

							_UpperCAmelCase             =    UnCLIPScheduler(
							    variance_type='fixed_small_log' ,					prediction_type='sample' ,					num_train_timesteps=1000 ,					clip_sample=UpperCAmelCase ,					clip_sample_range=10.0 ,					)

							_UpperCAmelCase             =    {
							    'prior': prior,
							    'image_encoder': image_encoder,
							    'text_encoder': text_encoder,
							    'tokenizer': tokenizer,
							    'scheduler': scheduler,
							    'image_processor': image_processor,
							}

							return components

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase=0					):
							"""simple docstring"""







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

	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    'cpu'

							_UpperCAmelCase             =    self.get_dummy_components()

							_UpperCAmelCase             =    self.pipeline_class(**UpperCAmelCase					)
							_UpperCAmelCase             =    pipe.to(UpperCAmelCase					)

							pipe.set_progress_bar_config(disable=UpperCAmelCase					)

							_UpperCAmelCase             =    pipe(**self.get_dummy_inputs(UpperCAmelCase					)					)
							_UpperCAmelCase             =    output.image_embeds

							_UpperCAmelCase             =    pipe(
							    **self.get_dummy_inputs(UpperCAmelCase					) ,					return_dict=UpperCAmelCase ,					)[0]

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

							assert image.shape == (1, 32)

							_UpperCAmelCase             =    np.array(
							    [-0.05_32, 1.71_20, 0.36_56, -1.08_52, -0.89_46, -1.17_56, 0.43_48, 0.24_82, 0.51_46, -0.11_56]					)

							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  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    torch_device == 'cpu'
							_UpperCAmelCase             =    True
							_UpperCAmelCase             =    False

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

	@skip_mps
	def  UpperCamelCase					(							self					):
							"""simple docstring"""







							_UpperCAmelCase             =    torch_device == 'cpu'
							_UpperCAmelCase             =    False

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








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,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


_a   			=						logging.get_logger(__name__)


if is_vision_available():
		import PIL





class __lowerCamelCase				(   snake_case__):





	"""simple docstring"""






	UpperCamelCase__			=   ["pixel_values"]

	def __init__(							self ,					UpperCAmelCase = True ,					UpperCAmelCase = None ,					UpperCAmelCase = PILImageResampling.BICUBIC ,					UpperCAmelCase = True ,					UpperCAmelCase = None ,					UpperCAmelCase = True ,					UpperCAmelCase = 1 / 255 ,					UpperCAmelCase = True ,					UpperCAmelCase = None ,					UpperCAmelCase = None ,					UpperCAmelCase = True ,					**UpperCAmelCase ,					):
							"""simple docstring"""







							super().__init__(**UpperCAmelCase					)
							_UpperCAmelCase             =    size if size is not None else {'shortest_edge': 224}
							_UpperCAmelCase             =    get_size_dict(UpperCAmelCase ,					default_to_square=UpperCAmelCase					)
							_UpperCAmelCase             =    crop_size if crop_size is not None else {'height': 224, 'width': 224}
							_UpperCAmelCase             =    get_size_dict(UpperCAmelCase					)

							_UpperCAmelCase             =    do_resize
							_UpperCAmelCase             =    size
							_UpperCAmelCase             =    resample
							_UpperCAmelCase             =    do_center_crop
							_UpperCAmelCase             =    crop_size
							_UpperCAmelCase             =    do_rescale
							_UpperCAmelCase             =    rescale_factor
							_UpperCAmelCase             =    do_normalize
							_UpperCAmelCase             =    image_mean if image_mean is not None else OPENAI_CLIP_MEAN
							_UpperCAmelCase             =    image_std if image_std is not None else OPENAI_CLIP_STD
							_UpperCAmelCase             =    do_convert_rgb

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase = PILImageResampling.BICUBIC ,					UpperCAmelCase = None ,					**UpperCAmelCase ,					):
							"""simple docstring"""







							_UpperCAmelCase             =    get_size_dict(UpperCAmelCase ,					default_to_square=UpperCAmelCase					)
							_UpperCAmelCase             =    get_resize_output_image_size(
							    UpperCAmelCase ,					size=(size['height'], size['width']) ,					default_to_square=UpperCAmelCase					)
							return resize(UpperCAmelCase ,					size=UpperCAmelCase ,					resample=UpperCAmelCase ,					data_format=UpperCAmelCase ,					**UpperCAmelCase					)

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase = None ,					**UpperCAmelCase ,					):
							"""simple docstring"""







							_UpperCAmelCase             =    get_size_dict(UpperCAmelCase					)
							return center_crop(UpperCAmelCase ,					size=(size['height'], size['width']) ,					data_format=UpperCAmelCase ,					**UpperCAmelCase					)

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase = None ,					**UpperCAmelCase ,					):
							"""simple docstring"""







							return rescale(UpperCAmelCase ,					scale=UpperCAmelCase ,					data_format=UpperCAmelCase ,					**UpperCAmelCase					)

	def  UpperCamelCase					(							self ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase ,					UpperCAmelCase = None ,					**UpperCAmelCase ,					):
							"""simple docstring"""







							return normalize(UpperCAmelCase ,					mean=UpperCAmelCase ,					std=UpperCAmelCase ,					data_format=UpperCAmelCase ,					**UpperCAmelCase					)

	def  UpperCamelCase					(							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 ,					):
							"""simple docstring"""







							_UpperCAmelCase             =    do_resize if do_resize is not None else self.do_resize
							_UpperCAmelCase             =    size if size is not None else self.size
							_UpperCAmelCase             =    get_size_dict(UpperCAmelCase ,					default_to_square=UpperCAmelCase					)
							_UpperCAmelCase             =    resample if resample is not None else self.resample
							_UpperCAmelCase             =    do_center_crop if do_center_crop is not None else self.do_center_crop
							_UpperCAmelCase             =    crop_size if crop_size is not None else self.crop_size
							_UpperCAmelCase             =    get_size_dict(UpperCAmelCase					)
							_UpperCAmelCase             =    do_rescale if do_rescale is not None else self.do_rescale
							_UpperCAmelCase             =    rescale_factor if rescale_factor is not None else self.rescale_factor
							_UpperCAmelCase             =    do_normalize if do_normalize is not None else self.do_normalize
							_UpperCAmelCase             =    image_mean if image_mean is not None else self.image_mean
							_UpperCAmelCase             =    image_std if image_std is not None else self.image_std
							_UpperCAmelCase             =    do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

							_UpperCAmelCase             =    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:
													raise ValueError('Size must be specified if do_resize is True.'					)

							if do_center_crop and crop_size is None:
													raise ValueError('Crop size 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.'					)

							# PIL RGBA images are converted to RGB
							if do_convert_rgb:
													_UpperCAmelCase             =    [convert_to_rgb(UpperCAmelCase					) for image in images]

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

							if do_resize:
													_UpperCAmelCase             =    [self.resize(image=UpperCAmelCase ,					size=UpperCAmelCase ,					resample=UpperCAmelCase					) for image in images]

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

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

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

							_UpperCAmelCase             =    [to_channel_dimension_format(UpperCAmelCase ,					UpperCAmelCase					) for image in images]

							_UpperCAmelCase             =    {'pixel_values': images}
							return BatchFeature(data=UpperCAmelCase ,					tensor_type=UpperCAmelCase					)
