


import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







class        snake_case_   (lowerCamelCase_	,						lowerCamelCase_	,						unittest.TestCase      ):
					UpperCAmelCase__					:    Optional[Any]          =			StableDiffusionXLImgaImgPipeline
					UpperCAmelCase__					:    List[str]          =			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
					UpperCAmelCase__					:    Optional[int]          =			PipelineTesterMixin.required_optional_params - {'''latents'''}
					UpperCAmelCase__					:    List[str]          =			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
					UpperCAmelCase__					:    int          =			IMAGE_TO_IMAGE_IMAGE_PARAMS
					UpperCAmelCase__					:    str          =			IMAGE_TO_IMAGE_IMAGE_PARAMS
					def  lowerCamelCase__(      self   :int      )					->    List[str]:
												torch.manual_seed(0      )
												a__             =     UNetaDConditionModel(
												    block_out_channels=(32, 64)			,layers_per_block=2			,sample_size=32			,in_channels=4			,out_channels=4			,down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')			,up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')			,attention_head_dim=(2, 4)			,use_linear_projection=__snake_case			,addition_embed_type='text_time'			,addition_time_embed_dim=8			,transformer_layers_per_block=(1, 2)			,projection_class_embeddings_input_dim=80			,cross_attention_dim=64			,)
												a__             =     EulerDiscreteScheduler(
												    beta_start=0.0_00_85			,beta_end=0.0_12			,steps_offset=1			,beta_schedule='scaled_linear'			,timestep_spacing='leading'			,)
												torch.manual_seed(0      )
												a__             =     AutoencoderKL(
												    block_out_channels=[32, 64]			,in_channels=3			,out_channels=3			,down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']			,up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']			,latent_channels=4			,sample_size=1_28			,)
												torch.manual_seed(0      )
												a__             =     CLIPTextConfig(
												    bos_token_id=0			,eos_token_id=2			,hidden_size=32			,intermediate_size=37			,layer_norm_eps=1E-05			,num_attention_heads=4			,num_hidden_layers=5			,pad_token_id=1			,vocab_size=10_00			,hidden_act='gelu'			,projection_dim=32			,)
												a__             =     CLIPTextModel(__snake_case      )
												a__             =     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'			,local_files_only=__snake_case      )

												a__             =     CLIPTextModelWithProjection(__snake_case      )
												a__             =     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'			,local_files_only=__snake_case      )

												a__             =     {
												    'unet': unet,
												    'scheduler': scheduler,
												    'vae': vae,
												    'text_encoder': text_encoder,
												    'tokenizer': tokenizer,
												    'text_encoder_2': text_encoder_a,
												    'tokenizer_2': tokenizer_a,
												    # "safety_checker": None,
												    # "feature_extractor": None,
												}
												return components
					def  lowerCamelCase__(      self   :Dict			,__snake_case   :Optional[Any]			,__snake_case   :Any=0      )					->    Tuple:
												a__             =     floats_tensor((1, 3, 32, 32)			,rng=random.Random(__snake_case      )      ).to(__snake_case      )
												a__             =     image / 2 + 0.5
												if str(__snake_case      ).startswith('mps'      ):
																			a__             =     torch.manual_seed(__snake_case      )
												else:
																			a__             =     torch.Generator(device=__snake_case      ).manual_seed(__snake_case      )
												a__             =     {
												    'prompt': 'A painting of a squirrel eating a burger',
												    'image': image,
												    'generator': generator,
												    'num_inference_steps': 2,
												    'guidance_scale': 5.0,
												    'output_type': 'numpy',
												    'strength': 0.75,
												}
												return inputs
					def  lowerCamelCase__(      self   :Tuple      )					->    int:
												a__             =     'cpu'  # ensure determinism for the device-dependent torch.Generator
												a__             =     self.get_dummy_components()
												a__             =     StableDiffusionXLImgaImgPipeline(**__snake_case      )
												a__             =     sd_pipe.to(__snake_case      )
												sd_pipe.set_progress_bar_config(disable=__snake_case      )

												a__             =     self.get_dummy_inputs(__snake_case      )
												a__             =     sd_pipe(**__snake_case      ).images
												a__             =     image[0, -3:, -3:, -1]

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

												a__             =     np.array([0.46_56, 0.48_40, 0.44_39, 0.66_98, 0.55_74, 0.45_24, 0.57_99, 0.59_43, 0.51_65]      )

												assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-2
					def  lowerCamelCase__(      self   :List[str]      )					->    Dict:
												super().test_attention_slicing_forward_pass(expected_max_diff=3E-3      )
					def  lowerCamelCase__(      self   :str      )					->    Optional[int]:
												super().test_inference_batch_single_identical(expected_max_diff=3E-3      )
					def  lowerCamelCase__(      self   :List[str]      )					->    Union[str, Any]:
												pass



					def  lowerCamelCase__(      self   :Optional[Any]      )					->    Tuple:
												a__             =     self.get_dummy_components()
												a__             =     StableDiffusionXLImgaImgPipeline(**__snake_case      )
												a__             =     sd_pipe.to(__snake_case      )
												a__             =     sd_pipe.to(__snake_case      )
												sd_pipe.set_progress_bar_config(disable=__snake_case      )

												# forward without prompt embeds
												a__             =     self.get_dummy_inputs(__snake_case      )
												a__             =     3 * ['this is a negative prompt']
												a__             =     negative_prompt
												a__             =     3 * [inputs['prompt']]

												a__             =     sd_pipe(**__snake_case      )
												a__             =     output.images[0, -3:, -3:, -1]

												# forward with prompt embeds
												a__             =     self.get_dummy_inputs(__snake_case      )
												a__             =     3 * ['this is a negative prompt']
												a__             =     3 * [inputs.pop('prompt'      )]

												(
												    (
												    a__
												)      ,					(
												    a__
												)      ,					(
												    a__
												)      ,					(
												    a__
												)      ,					
												)             =     sd_pipe.encode_prompt(__snake_case			,negative_prompt=__snake_case      )

												a__             =     sd_pipe(
												    **__snake_case			,prompt_embeds=__snake_case			,negative_prompt_embeds=__snake_case			,pooled_prompt_embeds=__snake_case			,negative_pooled_prompt_embeds=__snake_case			,)
												a__             =     output.images[0, -3:, -3:, -1]

												# make sure that it's equal
												assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()      ).max() < 1E-4







@slow
@require_torch_gpu
class        snake_case_   (unittest.TestCase      ):
					def  lowerCamelCase__(      self   :Optional[Any]      )					->    Union[str, Any]:
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()
					def  lowerCamelCase__(      self   :Optional[int]			,__snake_case   :Tuple			,__snake_case   :Any="cpu"			,__snake_case   :Optional[int]=torch.floataa			,__snake_case   :Optional[int]=0      )					->    Optional[int]:
												a__             =     torch.Generator(device=__snake_case      ).manual_seed(__snake_case      )
												a__             =     np.random.RandomState(__snake_case      ).standard_normal((1, 4, 64, 64)      )
												a__             =     torch.from_numpy(__snake_case      ).to(device=__snake_case			,dtype=__snake_case      )
												a__             =     {
												    'prompt': 'a photograph of an astronaut riding a horse',
												    'latents': latents,
												    'generator': generator,
												    'num_inference_steps': 3,
												    'guidance_scale': 7.5,
												    'output_type': 'numpy',
												}
												return inputs



					def  lowerCamelCase__(      self   :Any      )					->    Any:
												a__             =     DiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-base'      )
												pipe.to(__snake_case      )
												pipe.set_progress_bar_config(disable=__snake_case      )

												a__             =     self.get_inputs(__snake_case      )
												a__             =     pipe(**__snake_case      ).images
												a__             =     image[0, -3:, -3:, -1].flatten()

												assert image.shape == (1, 5_12, 5_12, 3)
												a__             =     np.array([0.4_94_93, 0.4_78_96, 0.4_07_98, 0.5_42_14, 0.5_32_12, 0.4_82_02, 0.4_76_56, 0.4_63_29, 0.4_85_06]      )
												assert np.abs(image_slice - expected_slice      ).max() < 7E-3








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,
    is_batched,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


snake_case							:       Union[str, Any]     			=      logging.get_logger(__name__)







class        snake_case_   (lowerCamelCase_      ):
					UpperCAmelCase__					:    List[Any]          =			['''pixel_values''']
					def __init__(      self   :Optional[int]			,__snake_case   :bool = True			,__snake_case   :Optional[Dict[str, int]] = None			,__snake_case   :PILImageResampling = PILImageResampling.BICUBIC			,__snake_case   :bool = True			,__snake_case   :bool = True			,__snake_case   :Union[int, float] = 1 / 2_55			,__snake_case   :Dict[str, int] = None			,__snake_case   :bool = True			,__snake_case   :Optional[Union[float, List[float]]] = None			,__snake_case   :Optional[Union[float, List[float]]] = None			,**__snake_case   :Dict			,)					->    None:
												super().__init__(**__snake_case      )
												a__             =     size if size is not None else {'height': 2_24, 'width': 2_24}
												a__             =     get_size_dict(__snake_case      )
												a__             =     crop_size if crop_size is not None else {'height': 2_24, 'width': 2_24}
												a__             =     get_size_dict(__snake_case			,default_to_square=__snake_case			,param_name='crop_size'      )

												a__             =     do_resize
												a__             =     do_rescale
												a__             =     do_normalize
												a__             =     do_center_crop
												a__             =     crop_size
												a__             =     size
												a__             =     resample
												a__             =     rescale_factor
												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  lowerCamelCase__(      self   :Dict			,__snake_case   :np.ndarray			,__snake_case   :Dict[str, int]			,__snake_case   :PILImageResampling = PILImageResampling.BILINEAR			,__snake_case   :Optional[Union[str, ChannelDimension]] = None			,**__snake_case   :List[Any]			,)					->    np.ndarray:
												a__             =     get_size_dict(__snake_case      )

												if "shortest_edge" in size:
																			a__             =     get_resize_output_image_size(__snake_case			,size=size['shortest_edge']			,default_to_square=__snake_case      )
																			# size = get_resize_output_image_size(image, size["shortest_edge"], size["longest_edge"])
												elif "height" in size and "width" in size:
																			a__             =     (size['height'], size['width'])
												else:
																			raise ValueError(F'Size must contain \'height\' and \'width\' keys or \'shortest_edge\' key. Got {size.keys()}'      )
												return resize(__snake_case			,size=__snake_case			,resample=__snake_case			,data_format=__snake_case			,**__snake_case      )
					def  lowerCamelCase__(      self   :Dict			,__snake_case   :np.ndarray			,__snake_case   :Dict[str, int]			,__snake_case   :Optional[Union[str, ChannelDimension]] = None			,**__snake_case   :Any			,)					->    np.ndarray:
												a__             =     get_size_dict(__snake_case      )
												if "height" not in size or "width" not in size:
																			raise ValueError(F'The `size` parameter must contain the keys (height, width). Got {size.keys()}'      )
												return center_crop(__snake_case			,size=(size['height'], size['width'])			,data_format=__snake_case			,**__snake_case      )
					def  lowerCamelCase__(      self   :List[Any]			,__snake_case   :np.ndarray			,__snake_case   :float			,__snake_case   :Optional[Union[str, ChannelDimension]] = None			,**__snake_case   :int      )					->    np.ndarray:
												return rescale(__snake_case			,scale=__snake_case			,data_format=__snake_case			,**__snake_case      )
					def  lowerCamelCase__(      self   :Tuple			,__snake_case   :np.ndarray			,__snake_case   :Union[float, List[float]]			,__snake_case   :Union[float, List[float]]			,__snake_case   :Optional[Union[str, ChannelDimension]] = None			,**__snake_case   :Any			,)					->    np.ndarray:
												return normalize(__snake_case			,mean=__snake_case			,std=__snake_case			,data_format=__snake_case			,**__snake_case      )



					def  lowerCamelCase__(      self   :Any			,__snake_case   :ImageInput			,__snake_case   :Optional[bool] = None			,__snake_case   :Dict[str, int] = None			,__snake_case   :PILImageResampling = None			,__snake_case   :bool = None			,__snake_case   :int = None			,__snake_case   :Optional[bool] = None			,__snake_case   :Optional[float] = None			,__snake_case   :Optional[bool] = None			,__snake_case   :Optional[Union[float, List[float]]] = None			,__snake_case   :Optional[Union[float, List[float]]] = None			,__snake_case   :Optional[Union[str, TensorType]] = None			,__snake_case   :Union[str, ChannelDimension] = ChannelDimension.FIRST			,**__snake_case   :Optional[int]			,)					->    BatchFeature:
												a__             =     do_resize if do_resize is not None else self.do_resize
												a__             =     do_rescale if do_rescale is not None else self.do_rescale
												a__             =     do_normalize if do_normalize is not None else self.do_normalize
												a__             =     do_center_crop if do_center_crop is not None else self.do_center_crop
												a__             =     crop_size if crop_size is not None else self.crop_size
												a__             =     get_size_dict(__snake_case			,param_name='crop_size'			,default_to_square=__snake_case      )
												a__             =     resample if resample is not None else self.resample
												a__             =     rescale_factor if rescale_factor is not None else self.rescale_factor
												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(__snake_case      )

												if not is_batched(__snake_case      ):
																			a__             =     [images]

												if not valid_images(__snake_case      ):
																			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.'      )

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

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

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

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

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

												a__             =     [to_channel_dimension_format(__snake_case			,__snake_case      ) for image in images]

												a__             =     {'pixel_values': images}
												return BatchFeature(data=__snake_case			,tensor_type=__snake_case      )




