






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    lowercase__   (	a__,    a__,    unittest.TestCase		):
	a_   			=StableDiffusionXLImgaImgPipeline
	a_   			=TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
	a_   			=PipelineTesterMixin.required_optional_params - {"""latents"""}
	a_   			=TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
	a_   			=IMAGE_TO_IMAGE_IMAGE_PARAMS
	a_   			=IMAGE_TO_IMAGE_IMAGE_PARAMS

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







							'''simple docstring'''





							torch.manual_seed(0				)
							lowerCAmelCase__       					=  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=_UpperCAmelCase						,			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						,			)
							lowerCAmelCase__       					=  EulerDiscreteScheduler(
							    beta_start=0.00_085						,			beta_end=0.012						,			steps_offset=1						,			beta_schedule="scaled_linear"						,			timestep_spacing="leading"						,			)
							torch.manual_seed(0				)
							lowerCAmelCase__       					=  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=128						,			)
							torch.manual_seed(0				)
							lowerCAmelCase__       					=  CLIPTextConfig(
							    bos_token_id=0						,			eos_token_id=2						,			hidden_size=32						,			intermediate_size=37						,			layer_norm_eps=1E-0_5						,			num_attention_heads=4						,			num_hidden_layers=5						,			pad_token_id=1						,			vocab_size=1000						,			hidden_act="gelu"						,			projection_dim=32						,			)
							lowerCAmelCase__       					=  CLIPTextModel(_UpperCAmelCase				)
							lowerCAmelCase__       					=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						,			local_files_only=_UpperCAmelCase				)

							lowerCAmelCase__       					=  CLIPTextModelWithProjection(_UpperCAmelCase				)
							lowerCAmelCase__       					=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						,			local_files_only=_UpperCAmelCase				)

							lowerCAmelCase__       					=  {
							    '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 						UpperCAmelCase     (							self						,			__UpperCAmelCase						,			__UpperCAmelCase=0				)->    Dict:







							'''simple docstring'''





							lowerCAmelCase__       					=  floats_tensor((1, 3, 32, 32)						,			rng=random.Random(_UpperCAmelCase				)				).to(_UpperCAmelCase				)
							lowerCAmelCase__       					=  image / 2 + 0.5
							if str(_UpperCAmelCase				).startswith("mps"				):
													lowerCAmelCase__       					=  torch.manual_seed(_UpperCAmelCase				)
							else:
													lowerCAmelCase__       					=  torch.Generator(device=_UpperCAmelCase				).manual_seed(_UpperCAmelCase				)
							lowerCAmelCase__       					=  {
							    '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 						UpperCAmelCase     (							self				)->    List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       					=  'cpu'  # ensure determinism for the device-dependent torch.Generator
							lowerCAmelCase__       					=  self.get_dummy_components()
							lowerCAmelCase__       					=  StableDiffusionXLImgaImgPipeline(**_UpperCAmelCase				)
							lowerCAmelCase__       					=  sd_pipe.to(_UpperCAmelCase				)
							sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase				)

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

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

							lowerCAmelCase__       					=  np.array([0.4_656, 0.4_840, 0.4_439, 0.6_698, 0.5_574, 0.4_524, 0.5_799, 0.5_943, 0.5_165]				)

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

	def 						UpperCAmelCase     (							self				)->    List[Any]:







							'''simple docstring'''





							super().test_attention_slicing_forward_pass(expected_max_diff=3E-3				)

	def 						UpperCAmelCase     (							self				)->    Optional[int]:







							'''simple docstring'''





							super().test_inference_batch_single_identical(expected_max_diff=3E-3				)

	def 						UpperCAmelCase     (							self				)->    Tuple:







							'''simple docstring'''





							pass


	def 						UpperCAmelCase     (							self				)->    Optional[Any]:







							'''simple docstring'''





							lowerCAmelCase__       					=  self.get_dummy_components()
							lowerCAmelCase__       					=  StableDiffusionXLImgaImgPipeline(**_UpperCAmelCase				)
							lowerCAmelCase__       					=  sd_pipe.to(_UpperCAmelCase				)
							lowerCAmelCase__       					=  sd_pipe.to(_UpperCAmelCase				)
							sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase				)

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

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

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

							(
							    lowerCAmelCase__
							)       					=  sd_pipe.encode_prompt(_UpperCAmelCase						,			negative_prompt=_UpperCAmelCase				)

							lowerCAmelCase__       					=  sd_pipe(
							    **_UpperCAmelCase						,			prompt_embeds=_UpperCAmelCase						,			negative_prompt_embeds=_UpperCAmelCase						,			pooled_prompt_embeds=_UpperCAmelCase						,			negative_pooled_prompt_embeds=_UpperCAmelCase						,			)
							lowerCAmelCase__       					=  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    lowercase__   (	unittest.TestCase		):

	def 						UpperCAmelCase     (							self				)->    Any:







							'''simple docstring'''





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

	def 						UpperCAmelCase     (							self						,			__UpperCAmelCase						,			__UpperCAmelCase="cpu"						,			__UpperCAmelCase=torch.floataa						,			__UpperCAmelCase=0				)->    List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       					=  torch.Generator(device=_UpperCAmelCase				).manual_seed(_UpperCAmelCase				)
							lowerCAmelCase__       					=  np.random.RandomState(_UpperCAmelCase				).standard_normal((1, 4, 64, 64)				)
							lowerCAmelCase__       					=  torch.from_numpy(_UpperCAmelCase				).to(device=_UpperCAmelCase						,			dtype=_UpperCAmelCase				)
							lowerCAmelCase__       					=  {
							    '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 						UpperCAmelCase     (							self				)->    List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       					=  DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-base"				)
							pipe.to(_UpperCAmelCase				)
							pipe.set_progress_bar_config(disable=_UpperCAmelCase				)

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

							assert image.shape == (1, 512, 512, 3)
							lowerCAmelCase__       					=  np.array([0.49_493, 0.47_896, 0.40_798, 0.54_214, 0.53_212, 0.48_202, 0.47_656, 0.46_329, 0.48_506]				)
							assert np.abs(image_slice - expected_slice				).max() < 7E-3













'''simple docstring'''




from typing import List, Optional, Union

from ...image_utils import ImageInput
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType


class  SCREAMING_SNAKE_CASE    (a__							):
     lowerCAmelCase					=       ['''image_processor''', '''tokenizer''']
     lowerCAmelCase					=       '''BlipImageProcessor'''
     lowerCAmelCase					=       ('''BertTokenizer''', '''BertTokenizerFast''')


     def __init__(						self						,	_UpperCAmelCase						,	_UpperCAmelCase):

          '''simple docstring'''






          __A						: int      =       False
          super().__init__(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: Optional[int]      =       self.image_processor


     def __call__(						self						,	_UpperCAmelCase = None						,	_UpperCAmelCase = None						,	_UpperCAmelCase = True						,	_UpperCAmelCase = False						,	_UpperCAmelCase = None						,	_UpperCAmelCase = None						,	_UpperCAmelCase = 0						,	_UpperCAmelCase = None						,	_UpperCAmelCase = None						,	_UpperCAmelCase = False						,	_UpperCAmelCase = False						,	_UpperCAmelCase = False						,	_UpperCAmelCase = False						,	_UpperCAmelCase = False						,	_UpperCAmelCase = True						,	_UpperCAmelCase = None						,	**_UpperCAmelCase						,	):

          '''simple docstring'''






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

          # Get only text
          if images is None:
               __A						: int      =       self.tokenizer
               __A						: Optional[Any]      =       self.tokenizer(
                   text=_UpperCAmelCase						,	add_special_tokens=_UpperCAmelCase						,	padding=_UpperCAmelCase						,	truncation=_UpperCAmelCase						,	max_length=_UpperCAmelCase						,	stride=_UpperCAmelCase						,	pad_to_multiple_of=_UpperCAmelCase						,	return_attention_mask=_UpperCAmelCase						,	return_overflowing_tokens=_UpperCAmelCase						,	return_special_tokens_mask=_UpperCAmelCase						,	return_offsets_mapping=_UpperCAmelCase						,	return_token_type_ids=_UpperCAmelCase						,	return_length=_UpperCAmelCase						,	verbose=_UpperCAmelCase						,	return_tensors=_UpperCAmelCase						,	**_UpperCAmelCase						,	)
               return text_encoding

          # add pixel_values
          __A						: List[Any]      =       self.image_processor(_UpperCAmelCase						,	return_tensors=_UpperCAmelCase)

          if text is not None:
               __A						: Optional[Any]      =       self.tokenizer(
                   text=_UpperCAmelCase						,	add_special_tokens=_UpperCAmelCase						,	padding=_UpperCAmelCase						,	truncation=_UpperCAmelCase						,	max_length=_UpperCAmelCase						,	stride=_UpperCAmelCase						,	pad_to_multiple_of=_UpperCAmelCase						,	return_attention_mask=_UpperCAmelCase						,	return_overflowing_tokens=_UpperCAmelCase						,	return_special_tokens_mask=_UpperCAmelCase						,	return_offsets_mapping=_UpperCAmelCase						,	return_token_type_ids=_UpperCAmelCase						,	return_length=_UpperCAmelCase						,	verbose=_UpperCAmelCase						,	return_tensors=_UpperCAmelCase						,	**_UpperCAmelCase						,	)
          else:
               __A						: int      =       None

          if text_encoding is not None:
               encoding_image_processor.update(_UpperCAmelCase)

          return encoding_image_processor


     def 							SCREAMING_SNAKE_CASE  (						self						,	*_UpperCAmelCase						,	**_UpperCAmelCase):

          '''simple docstring'''






          return self.tokenizer.batch_decode(*_UpperCAmelCase						,	**_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self						,	*_UpperCAmelCase						,	**_UpperCAmelCase):

          '''simple docstring'''






          return self.tokenizer.decode(*_UpperCAmelCase						,	**_UpperCAmelCase)


     @property
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[str]      =       self.tokenizer.model_input_names
          __A						: Tuple      =       self.image_processor.model_input_names
          return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))