


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 							_UpperCamelCase (						A,A,unittest.TestCase						):

   '''simple docstring'''







   a_				:							Optional[int]    = StableDiffusionXLImgaImgPipeline
   a_				:							Optional[Any]    = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
   a_				:							Tuple    = PipelineTesterMixin.required_optional_params - {"latents"}
   a_				:							str    = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   a_				:							Dict    = IMAGE_TO_IMAGE_IMAGE_PARAMS
   a_				:							str    = IMAGE_TO_IMAGE_IMAGE_PARAMS

   def   _snake_case    (  self  : Optional[Any]							):






         '''simple docstring'''



         torch.manual_seed(0							)
         __lowerCamelCase     :      Optional[int]       						=						UNetaDConditionModel(
             block_out_channels=(3_2, 6_4)       ,  layers_per_block=2       ,  sample_size=3_2       ,  in_channels=4       ,  out_channels=4       ,  down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")       ,  up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")       ,  attention_head_dim=(2, 4)       ,  use_linear_projection=_lowerCamelCase       ,  addition_embed_type="""text_time"""       ,  addition_time_embed_dim=8       ,  transformer_layers_per_block=(1, 2)       ,  projection_class_embeddings_input_dim=8_0       ,  cross_attention_dim=6_4       ,  )
         __lowerCamelCase     :      List[str]       						=						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     :      List[Any]       						=						AutoencoderKL(
             block_out_channels=[3_2, 6_4]       ,  in_channels=3       ,  out_channels=3       ,  down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]       ,  up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]       ,  latent_channels=4       ,  sample_size=1_2_8       ,  )
         torch.manual_seed(0							)
         __lowerCamelCase     :      Tuple       						=						CLIPTextConfig(
             bos_token_id=0       ,  eos_token_id=2       ,  hidden_size=3_2       ,  intermediate_size=3_7       ,  layer_norm_eps=1E-05       ,  num_attention_heads=4       ,  num_hidden_layers=5       ,  pad_token_id=1       ,  vocab_size=1_0_0_0       ,  hidden_act="""gelu"""       ,  projection_dim=3_2       ,  )
         __lowerCamelCase     :      Optional[Any]       						=						CLIPTextModel(_lowerCamelCase							)
         __lowerCamelCase     :      str       						=						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       ,  local_files_only=_lowerCamelCase							)

         __lowerCamelCase     :      Tuple       						=						CLIPTextModelWithProjection(_lowerCamelCase							)
         __lowerCamelCase     :      Any       						=						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       ,  local_files_only=_lowerCamelCase							)

         __lowerCamelCase     :      Tuple       						=						{
             """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   _snake_case    (  self  : str       ,  _lowerCamelCase  : Dict       ,  _lowerCamelCase  : Any=0							):






         '''simple docstring'''



         __lowerCamelCase     :      str       						=						floats_tensor((1, 3, 3_2, 3_2)       ,  rng=random.Random(_lowerCamelCase							)							).to(_lowerCamelCase							)
         __lowerCamelCase     :      Dict       						=						image / 2 + 0.5
         if str(_lowerCamelCase							).startswith("""mps"""							):
               __lowerCamelCase     :      List[str]       						=						torch.manual_seed(_lowerCamelCase							)
         else:
               __lowerCamelCase     :      str       						=						torch.Generator(device=_lowerCamelCase							).manual_seed(_lowerCamelCase							)
         __lowerCamelCase     :      Optional[int]       						=						{
             """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   _snake_case    (  self  : List[Any]							):






         '''simple docstring'''



         __lowerCamelCase     :      Tuple       						=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         __lowerCamelCase     :      Optional[Any]       						=						self.get_dummy_components()
         __lowerCamelCase     :      Optional[Any]       						=						StableDiffusionXLImgaImgPipeline(**_lowerCamelCase							)
         __lowerCamelCase     :      str       						=						sd_pipe.to(_lowerCamelCase							)
         sd_pipe.set_progress_bar_config(disable=_lowerCamelCase							)

         __lowerCamelCase     :      List[Any]       						=						self.get_dummy_inputs(_lowerCamelCase							)
         __lowerCamelCase     :      List[Any]       						=						sd_pipe(**_lowerCamelCase							).images
         __lowerCamelCase     :      Optional[int]       						=						image[0, -3:, -3:, -1]

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

         __lowerCamelCase     :      List[Any]       						=						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   _snake_case    (  self  : Tuple							):






         '''simple docstring'''



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

   def   _snake_case    (  self  : Optional[int]							):






         '''simple docstring'''



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

   def   _snake_case    (  self  : int							):






         '''simple docstring'''



         pass


   def   _snake_case    (  self  : Optional[Any]							):






         '''simple docstring'''



         __lowerCamelCase     :      Dict       						=						self.get_dummy_components()
         __lowerCamelCase     :      Dict       						=						StableDiffusionXLImgaImgPipeline(**_lowerCamelCase							)
         __lowerCamelCase     :      Optional[int]       						=						sd_pipe.to(_lowerCamelCase							)
         __lowerCamelCase     :      List[Any]       						=						sd_pipe.to(_lowerCamelCase							)
         sd_pipe.set_progress_bar_config(disable=_lowerCamelCase							)

         # forward without prompt embeds
         __lowerCamelCase     :      Tuple       						=						self.get_dummy_inputs(_lowerCamelCase							)
         __lowerCamelCase     :      str       						=						3 * ["""this is a negative prompt"""]
         __lowerCamelCase     :      Optional[Any]       						=						negative_prompt
         __lowerCamelCase     :      Tuple       						=						3 * [inputs["""prompt"""]]

         __lowerCamelCase     :      int       						=						sd_pipe(**_lowerCamelCase							)
         __lowerCamelCase     :      Dict       						=						output.images[0, -3:, -3:, -1]

         # forward with prompt embeds
         __lowerCamelCase     :      List[str]       						=						self.get_dummy_inputs(_lowerCamelCase							)
         __lowerCamelCase     :      Any       						=						3 * ["""this is a negative prompt"""]
         __lowerCamelCase     :      List[str]       						=						3 * [inputs.pop("""prompt"""							)]

         (
             __lowerCamelCase 
         )     :      int       						=						sd_pipe.encode_prompt(_lowerCamelCase       ,  negative_prompt=_lowerCamelCase							)

         __lowerCamelCase     :      Optional[Any]       						=						sd_pipe(
             **_lowerCamelCase       ,  prompt_embeds=_lowerCamelCase       ,  negative_prompt_embeds=_lowerCamelCase       ,  pooled_prompt_embeds=_lowerCamelCase       ,  negative_pooled_prompt_embeds=_lowerCamelCase       ,  )
         __lowerCamelCase     :      int       						=						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 							_UpperCamelCase (						unittest.TestCase						):

   '''simple docstring'''







   def   _snake_case    (  self  : Tuple							):






         '''simple docstring'''



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

   def   _snake_case    (  self  : int       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : int="cpu"       ,  _lowerCamelCase  : Optional[Any]=torch.floataa       ,  _lowerCamelCase  : Optional[Any]=0							):






         '''simple docstring'''



         __lowerCamelCase     :      str       						=						torch.Generator(device=_lowerCamelCase							).manual_seed(_lowerCamelCase							)
         __lowerCamelCase     :      str       						=						np.random.RandomState(_lowerCamelCase							).standard_normal((1, 4, 6_4, 6_4)							)
         __lowerCamelCase     :      Tuple       						=						torch.from_numpy(_lowerCamelCase							).to(device=_lowerCamelCase       ,  dtype=_lowerCamelCase							)
         __lowerCamelCase     :      Union[str, Any]       						=						{
             """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   _snake_case    (  self  : List[str]							):






         '''simple docstring'''



         __lowerCamelCase     :      Tuple       						=						DiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-base"""							)
         pipe.to(_lowerCamelCase							)
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

         __lowerCamelCase     :      Optional[Any]       						=						self.get_inputs(_lowerCamelCase							)
         __lowerCamelCase     :      List[Any]       						=						pipe(**_lowerCamelCase							).images
         __lowerCamelCase     :      Optional[int]       						=						image[0, -3:, -3:, -1].flatten()

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      Optional[Any]       						=						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





import re
import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin



class 							_UpperCamelCase (						A						):

   '''simple docstring'''







   a_				:							List[Any]    = ["image_processor", "tokenizer"]
   a_				:							Tuple    = "AutoImageProcessor"
   a_				:							Tuple    = "AutoTokenizer"

   def __init__(  self  : Union[str, Any]       ,  _lowerCamelCase  : List[Any]=None       ,  _lowerCamelCase  : Any=None       ,  **_lowerCamelCase  : List[str]							):






         '''simple docstring'''



         __lowerCamelCase     :      Tuple       						=						None
         if "feature_extractor" in kwargs:
               warnings.warn(
                   """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
                   """ instead."""       ,  _lowerCamelCase       ,  )
               __lowerCamelCase     :      List[Any]       						=						kwargs.pop("""feature_extractor"""							)

         __lowerCamelCase     :      Optional[Any]       						=						image_processor if image_processor is not None else feature_extractor
         if image_processor is None:
               raise ValueError("""You need to specify an `image_processor`."""							)
         if tokenizer is None:
               raise ValueError("""You need to specify a `tokenizer`."""							)

         super().__init__(_lowerCamelCase       ,  _lowerCamelCase							)
         __lowerCamelCase     :      List[Any]       						=						self.image_processor
         __lowerCamelCase     :      Any       						=						False

   def __call__(  self  : Union[str, Any]       ,  *_lowerCamelCase  : str       ,  **_lowerCamelCase  : Dict							):






         '''simple docstring'''



         if self._in_target_context_manager:
               return self.current_processor(*_lowerCamelCase       ,  **_lowerCamelCase							)

         __lowerCamelCase     :      List[str]       						=						kwargs.pop("""images"""       ,  _lowerCamelCase							)
         __lowerCamelCase     :      List[str]       						=						kwargs.pop("""text"""       ,  _lowerCamelCase							)
         if len(_lowerCamelCase							) > 0:
               __lowerCamelCase     :      Tuple       						=						args[0]
               __lowerCamelCase     :      List[Any]       						=						args[1:]

         if images is None and text is None:
               raise ValueError("""You need to specify either an `images` or `text` input to process."""							)

         if images is not None:
               __lowerCamelCase     :      Tuple       						=						self.image_processor(_lowerCamelCase       ,  *_lowerCamelCase       ,  **_lowerCamelCase							)
         if text is not None:
               __lowerCamelCase     :      int       						=						self.tokenizer(_lowerCamelCase       ,  **_lowerCamelCase							)

         if text is None:
               return inputs
         elif images is None:
               return encodings
         else:
               __lowerCamelCase     :      Optional[int]       						=						encodings["""input_ids"""]
               return inputs

   def   _snake_case    (  self  : List[Any]       ,  *_lowerCamelCase  : Tuple       ,  **_lowerCamelCase  : Union[str, Any]							):






         '''simple docstring'''



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

   def   _snake_case    (  self  : List[str]       ,  *_lowerCamelCase  : Optional[int]       ,  **_lowerCamelCase  : int							):






         '''simple docstring'''



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

   @contextmanager
   def   _snake_case    (  self  : Optional[Any]							):






         '''simple docstring'''



         warnings.warn(
             """`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your """
             """labels by using the argument `text` of the regular `__call__` method (either in the same call as """
             """your images inputs, or in a separate call."""							)
         __lowerCamelCase     :      Any       						=						True
         __lowerCamelCase     :      int       						=						self.tokenizer
         yield
         __lowerCamelCase     :      Dict       						=						self.image_processor
         __lowerCamelCase     :      List[Any]       						=						False

   def   _snake_case    (  self  : List[str]       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : List[str]=False       ,  _lowerCamelCase  : List[Any]=None							):






         '''simple docstring'''



         if added_vocab is None:
               __lowerCamelCase     :      List[str]       						=						self.tokenizer.get_added_vocab()

         __lowerCamelCase     :      List[Any]       						=						{}

         while tokens:
               __lowerCamelCase     :      Dict       						=						re.search(R"""<s_(.*?)>"""       ,  _lowerCamelCase       ,  re.IGNORECASE							)
               if start_token is None:
                     break
               __lowerCamelCase     :      Dict       						=						start_token.group(1							)
               __lowerCamelCase     :      List[str]       						=						re.search(RF"""</s_{key}>"""       ,  _lowerCamelCase       ,  re.IGNORECASE							)
               __lowerCamelCase     :      Optional[int]       						=						start_token.group()
               if end_token is None:
                     __lowerCamelCase     :      int       						=						tokens.replace(_lowerCamelCase       ,  """"""							)
               else:
                     __lowerCamelCase     :      Dict       						=						end_token.group()
                     __lowerCamelCase     :      Optional[Any]       						=						re.escape(_lowerCamelCase							)
                     __lowerCamelCase     :      List[str]       						=						re.escape(_lowerCamelCase							)
                     __lowerCamelCase     :      Any       						=						re.search(F"""{start_token_escaped}(.*?){end_token_escaped}"""       ,  _lowerCamelCase       ,  re.IGNORECASE							)
                     if content is not None:
                           __lowerCamelCase     :      Union[str, Any]       						=						content.group(1							).strip()
                           if r"<s_" in content and r"</s_" in content:  # non-leaf node
                                 __lowerCamelCase     :      Union[str, Any]       						=						self.tokenajson(_lowerCamelCase       ,  is_inner_value=_lowerCamelCase       ,  added_vocab=_lowerCamelCase							)
                                 if value:
                                       if len(_lowerCamelCase							) == 1:
                                             __lowerCamelCase     :      Optional[int]       						=						value[0]
                                       __lowerCamelCase     :      Union[str, Any]       						=						value
                           else:  # leaf nodes
                                 __lowerCamelCase     :      List[Any]       						=						[]
                                 for leaf in content.split(R"""<sep/>"""							):
                                       __lowerCamelCase     :      Optional[Any]       						=						leaf.strip()
                                       if leaf in added_vocab and leaf[0] == "<" and leaf[-2:] == "/>":
                                             __lowerCamelCase     :      str       						=						leaf[1:-2]  # for categorical special tokens
                                       output[key].append(_lowerCamelCase							)
                                 if len(output[key]							) == 1:
                                       __lowerCamelCase     :      Tuple       						=						output[key][0]

                     __lowerCamelCase     :      Tuple       						=						tokens[tokens.find(_lowerCamelCase							) + len(_lowerCamelCase							) :].strip()
                     if tokens[:6] == r"<sep/>":  # non-leaf nodes
                           return [output] + self.tokenajson(tokens[6:]       ,  is_inner_value=_lowerCamelCase       ,  added_vocab=_lowerCamelCase							)

         if len(_lowerCamelCase							):
               return [output] if is_inner_value else output
         else:
               return [] if is_inner_value else {"text_sequence": tokens}

   @property
   def   _snake_case    (  self  : Tuple							):






         '''simple docstring'''



         warnings.warn(
             """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""       ,  _lowerCamelCase       ,  )
         return self.image_processor_class


   @property
   def   _snake_case    (  self  : List[Any]							):






         '''simple docstring'''



         warnings.warn(
             """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""       ,  _lowerCamelCase       ,  )
         return self.image_processor



