




"""simple docstring"""




import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, 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_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class      UpperCAmelCase_  (    a_	,      a_	,      a_	,      unittest.TestCase						):
      __SCREAMING_SNAKE_CASE						:       Dict								=			StableDiffusionInstructPixaPixPipeline
      __SCREAMING_SNAKE_CASE						:       List[Any]								=			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width', 'cross_attention_kwargs'}
      __SCREAMING_SNAKE_CASE						:       Dict								=			TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
      __SCREAMING_SNAKE_CASE						:       int								=			IMAGE_TO_IMAGE_IMAGE_PARAMS
      __SCREAMING_SNAKE_CASE						:       str								=			IMAGE_TO_IMAGE_IMAGE_PARAMS



      def 						snake_case_ (					self					:     Any						):
           torch.manual_seed(0						)
           _UpperCAmelCase      :     Union[str, Any]        =   UNetaDConditionModel(
               block_out_channels=(3_2, 6_4)					,     layers_per_block=2					,     sample_size=3_2					,     in_channels=8					,     out_channels=4					,     down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")					,     up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")					,     cross_attention_dim=3_2					,     )
           _UpperCAmelCase      :     Union[str, Any]        =   PNDMScheduler(skip_prk_steps=lowercase_						)
           torch.manual_seed(0						)
           _UpperCAmelCase      :     Tuple        =   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					,     )
           torch.manual_seed(0						)
           _UpperCAmelCase      :     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					,     )
           _UpperCAmelCase      :     str        =   CLIPTextModel(lowercase_						)
           _UpperCAmelCase      :     List[Any]        =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)

           _UpperCAmelCase      :     Optional[Any]        =   {
               '''unet''': unet,
               '''scheduler''': scheduler,
               '''vae''': vae,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''safety_checker''': None,
               '''feature_extractor''': None,
           }
           return components



      def 						snake_case_ (					self					:     List[Any]					,     A					:     Dict					,     A					:     Tuple=0						):
           _UpperCAmelCase      :     Dict        =   floats_tensor((1, 3, 3_2, 3_2)					,     rng=random.Random(lowercase_						)						).to(lowercase_						)
           _UpperCAmelCase      :     Tuple        =   image.cpu().permute(0					,     2					,     3					,     1						)[0]
           _UpperCAmelCase      :     Dict        =   Image.fromarray(np.uinta(lowercase_						)						).convert("RGB"						)
           if str(lowercase_						).startswith("mps"						):
                _UpperCAmelCase      :     Any        =   torch.manual_seed(lowercase_						)
           else:
                _UpperCAmelCase      :     List[Any]        =   torch.Generator(device=lowercase_						).manual_seed(lowercase_						)
           _UpperCAmelCase      :     Dict        =   {
               '''prompt''': '''A painting of a squirrel eating a burger''',
               '''image''': image,
               '''generator''': generator,
               '''num_inference_steps''': 2,
               '''guidance_scale''': 6.0,
               '''image_guidance_scale''': 1,
               '''output_type''': '''numpy''',
           }
           return inputs



      def 						snake_case_ (					self					:     str						):
           _UpperCAmelCase      :     int        =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           _UpperCAmelCase      :     Union[str, Any]        =   self.get_dummy_components()
           _UpperCAmelCase      :     List[str]        =   StableDiffusionInstructPixaPixPipeline(**lowercase_						)
           _UpperCAmelCase      :     Dict        =   sd_pipe.to(lowercase_						)
           sd_pipe.set_progress_bar_config(disable=lowercase_						)

           _UpperCAmelCase      :     Union[str, Any]        =   self.get_dummy_inputs(lowercase_						)
           _UpperCAmelCase      :     Union[str, Any]        =   sd_pipe(**lowercase_						).images
           _UpperCAmelCase      :     str        =   image[0, -3:, -3:, -1]
           assert image.shape == (1, 3_2, 3_2, 3)
           _UpperCAmelCase      :     Union[str, Any]        =   np.array([0.7_526, 0.3_750, 0.4_547, 0.6_117, 0.5_866, 0.5_016, 0.4_327, 0.5_642, 0.4_815]						)

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



      def 						snake_case_ (					self					:     List[str]						):
           _UpperCAmelCase      :     List[Any]        =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           _UpperCAmelCase      :     int        =   self.get_dummy_components()
           _UpperCAmelCase      :     List[Any]        =   StableDiffusionInstructPixaPixPipeline(**lowercase_						)
           _UpperCAmelCase      :     List[str]        =   sd_pipe.to(lowercase_						)
           sd_pipe.set_progress_bar_config(disable=lowercase_						)

           _UpperCAmelCase      :     Tuple        =   self.get_dummy_inputs(lowercase_						)
           _UpperCAmelCase      :     List[Any]        =   '''french fries'''
           _UpperCAmelCase      :     int        =   sd_pipe(**lowercase_					,     negative_prompt=lowercase_						)
           _UpperCAmelCase      :     List[Any]        =   output.images
           _UpperCAmelCase      :     int        =   image[0, -3:, -3:, -1]

           assert image.shape == (1, 3_2, 3_2, 3)
           _UpperCAmelCase      :     Union[str, Any]        =   np.array([0.7_511, 0.3_642, 0.4_553, 0.6_236, 0.5_797, 0.5_013, 0.4_343, 0.5_611, 0.4_831]						)

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



      def 						snake_case_ (					self					:     Union[str, Any]						):
           _UpperCAmelCase      :     Dict        =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           _UpperCAmelCase      :     List[str]        =   self.get_dummy_components()
           _UpperCAmelCase      :     str        =   StableDiffusionInstructPixaPixPipeline(**lowercase_						)
           _UpperCAmelCase      :     List[Any]        =   sd_pipe.to(lowercase_						)
           sd_pipe.set_progress_bar_config(disable=lowercase_						)

           _UpperCAmelCase      :     List[str]        =   self.get_dummy_inputs(lowercase_						)
           _UpperCAmelCase      :     str        =   [inputs['''prompt''']] * 2

           _UpperCAmelCase      :     Dict        =   np.array(inputs["image"]						).astype(np.floataa						) / 255.0
           _UpperCAmelCase      :     Union[str, Any]        =   torch.from_numpy(lowercase_						).unsqueeze(0						).to(lowercase_						)
           _UpperCAmelCase      :     Any        =   image / 2 + 0.5
           _UpperCAmelCase      :     Optional[int]        =   image.permute(0					,     3					,     1					,     2						)
           _UpperCAmelCase      :     Any        =   image.repeat(2					,     1					,     1					,     1						)

           _UpperCAmelCase      :     Optional[int]        =   sd_pipe(**lowercase_						).images
           _UpperCAmelCase      :     List[str]        =   image[-1, -3:, -3:, -1]

           assert image.shape == (2, 3_2, 3_2, 3)
           _UpperCAmelCase      :     List[Any]        =   np.array([0.5_812, 0.5_748, 0.5_222, 0.5_908, 0.5_695, 0.7_174, 0.6_804, 0.5_523, 0.5_579]						)

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



      def 						snake_case_ (					self					:     List[str]						):
           _UpperCAmelCase      :     Dict        =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           _UpperCAmelCase      :     Optional[int]        =   self.get_dummy_components()
           _UpperCAmelCase      :     int        =   EulerAncestralDiscreteScheduler(
               beta_start=0.00_085					,     beta_end=0.012					,     beta_schedule="scaled_linear"						)
           _UpperCAmelCase      :     Optional[int]        =   StableDiffusionInstructPixaPixPipeline(**lowercase_						)
           _UpperCAmelCase      :     int        =   sd_pipe.to(lowercase_						)
           sd_pipe.set_progress_bar_config(disable=lowercase_						)

           _UpperCAmelCase      :     Optional[int]        =   self.get_dummy_inputs(lowercase_						)
           _UpperCAmelCase      :     Dict        =   sd_pipe(**lowercase_						).images
           _UpperCAmelCase      :     Tuple        =   image[0, -3:, -3:, -1]

           _UpperCAmelCase      :     Dict        =   [round(lowercase_					,     4						) for x in image_slice.flatten().tolist()]
           print(",".join([str(lowercase_						) for x in slice]						)						)

           assert image.shape == (1, 3_2, 3_2, 3)
           _UpperCAmelCase      :     str        =   np.array([0.7_417, 0.3_842, 0.4_732, 0.5_776, 0.5_891, 0.5_139, 0.4_052, 0.5_673, 0.4_986]						)

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



      def 						snake_case_ (					self					:     int						):
           super().test_inference_batch_single_identical(expected_max_diff=3e-3						)



      def 						snake_case_ (					self					:     Dict						):
           _UpperCAmelCase      :     Dict        =   self.get_dummy_components()
           _UpperCAmelCase      :     List[str]        =   StableDiffusionInstructPixaPixPipeline(**lowercase_						)
           _UpperCAmelCase      :     Optional[int]        =   VaeImageProcessor(do_resize=lowercase_					,     do_normalize=lowercase_						)
           _UpperCAmelCase      :     Optional[int]        =   pipe.to(lowercase_						)
           pipe.set_progress_bar_config(disable=lowercase_						)

           _UpperCAmelCase      :     Dict        =   pipe(**self.get_dummy_inputs_by_type(lowercase_					,     input_image_type="pt"						)						)[0]

           _UpperCAmelCase      :     int        =   components['''vae''']
           _UpperCAmelCase      :     Optional[Any]        =   self.get_dummy_inputs_by_type(lowercase_					,     input_image_type="pt"						)

           for image_param in self.image_latents_params:
                if image_param in inputs.keys():
                     _UpperCAmelCase      :     Dict        =   vae.encode(inputs[image_param]						).latent_dist.mode()

           _UpperCAmelCase      :     Dict        =   pipe(**lowercase_						)[0]

           _UpperCAmelCase      :     Dict        =   np.abs(out - out_latents_inputs						).max()
           self.assertLess(lowercase_					,     1e-4					,     "passing latents as image input generate different result from passing image"						)







@slow
@require_torch_gpu
class      UpperCAmelCase_  (    unittest.TestCase						):



      def 						snake_case_ (					self					:     str						):
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()



      def 						snake_case_ (					self					:     Any					,     A					:     Optional[int]=0						):
           _UpperCAmelCase      :     Any        =   torch.manual_seed(lowercase_						)
           _UpperCAmelCase      :     Union[str, Any]        =   load_image(
               "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg"						)
           _UpperCAmelCase      :     int        =   {
               '''prompt''': '''turn him into a cyborg''',
               '''image''': image,
               '''generator''': generator,
               '''num_inference_steps''': 3,
               '''guidance_scale''': 7.5,
               '''image_guidance_scale''': 1.0,
               '''output_type''': '''numpy''',
           }
           return inputs



      def 						snake_case_ (					self					:     Optional[int]						):
           _UpperCAmelCase      :     Optional[int]        =   StableDiffusionInstructPixaPixPipeline.from_pretrained(
               "timbrooks/instruct-pix2pix"					,     safety_checker=lowercase_						)
           pipe.to(lowercase_						)
           pipe.set_progress_bar_config(disable=lowercase_						)
           pipe.enable_attention_slicing()

           _UpperCAmelCase      :     Union[str, Any]        =   self.get_inputs()
           _UpperCAmelCase      :     Tuple        =   pipe(**lowercase_						).images
           _UpperCAmelCase      :     Dict        =   image[0, -3:, -3:, -1].flatten()

           assert image.shape == (1, 5_1_2, 5_1_2, 3)
           _UpperCAmelCase      :     Any        =   np.array([0.5_902, 0.6_015, 0.6_027, 0.5_983, 0.6_092, 0.6_061, 0.5_765, 0.5_785, 0.5_555]						)

           assert np.abs(expected_slice - image_slice						).max() < 1e-3



      def 						snake_case_ (					self					:     Any						):
           _UpperCAmelCase      :     str        =   StableDiffusionInstructPixaPixPipeline.from_pretrained(
               "timbrooks/instruct-pix2pix"					,     safety_checker=lowercase_						)
           _UpperCAmelCase      :     Optional[Any]        =   LMSDiscreteScheduler.from_config(pipe.scheduler.config						)
           pipe.to(lowercase_						)
           pipe.set_progress_bar_config(disable=lowercase_						)
           pipe.enable_attention_slicing()

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

           assert image.shape == (1, 5_1_2, 5_1_2, 3)
           _UpperCAmelCase      :     Optional[Any]        =   np.array([0.6_578, 0.6_817, 0.6_972, 0.6_761, 0.6_856, 0.6_916, 0.6_428, 0.6_516, 0.6_301]						)

           assert np.abs(expected_slice - image_slice						).max() < 1e-3



      def 						snake_case_ (					self					:     Union[str, Any]						):
           _UpperCAmelCase      :     List[Any]        =   StableDiffusionInstructPixaPixPipeline.from_pretrained(
               "timbrooks/instruct-pix2pix"					,     safety_checker=lowercase_						)
           _UpperCAmelCase      :     Any        =   DDIMScheduler.from_config(pipe.scheduler.config						)
           pipe.to(lowercase_						)
           pipe.set_progress_bar_config(disable=lowercase_						)
           pipe.enable_attention_slicing()

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

           assert image.shape == (1, 5_1_2, 5_1_2, 3)
           _UpperCAmelCase      :     Tuple        =   np.array([0.3_828, 0.3_834, 0.3_818, 0.3_792, 0.3_865, 0.3_752, 0.3_792, 0.3_847, 0.3_753]						)

           assert np.abs(expected_slice - image_slice						).max() < 1e-3



      def 						snake_case_ (					self					:     Optional[Any]						):
           _UpperCAmelCase      :     Optional[Any]        =   0

           def callback_fn(A					:     int					,     A					:     int					,     A					:     torch.FloatTensor						) -> None:
                _UpperCAmelCase      :     Union[str, Any]        =   True
                nonlocal number_of_steps
                number_of_steps += 1
                if step == 1:
                     _UpperCAmelCase      :     Optional[Any]        =   latents.detach().cpu().numpy()
                     assert latents.shape == (1, 4, 6_4, 6_4)
                     _UpperCAmelCase      :     int        =   latents[0, -3:, -3:, -1]
                     _UpperCAmelCase      :     Tuple        =   np.array([-0.2_463, -0.4_644, -0.9_756, 1.5_176, 1.4_414, 0.7_866, 0.9_897, 0.8_521, 0.7_983]						)

                     assert np.abs(latents_slice.flatten() - expected_slice						).max() < 5e-2
                elif step == 2:
                     _UpperCAmelCase      :     Tuple        =   latents.detach().cpu().numpy()
                     assert latents.shape == (1, 4, 6_4, 6_4)
                     _UpperCAmelCase      :     Dict        =   latents[0, -3:, -3:, -1]
                     _UpperCAmelCase      :     str        =   np.array([-0.2_644, -0.4_626, -0.9_653, 1.5_176, 1.4_551, 0.7_686, 0.9_805, 0.8_452, 0.8_115]						)

                     assert np.abs(latents_slice.flatten() - expected_slice						).max() < 5e-2

           _UpperCAmelCase      :     Union[str, Any]        =   False

           _UpperCAmelCase      :     Dict        =   StableDiffusionInstructPixaPixPipeline.from_pretrained(
               "timbrooks/instruct-pix2pix"					,     safety_checker=lowercase_					,     torch_dtype=torch.floataa						)
           _UpperCAmelCase      :     Optional[int]        =   pipe.to(lowercase_						)
           pipe.set_progress_bar_config(disable=lowercase_						)
           pipe.enable_attention_slicing()

           _UpperCAmelCase      :     Union[str, Any]        =   self.get_inputs()
           pipe(**lowercase_					,     callback=lowercase_					,     callback_steps=1						)
           assert callback_fn.has_been_called
           assert number_of_steps == 3



      def 						snake_case_ (					self					:     int						):
           torch.cuda.empty_cache()
           torch.cuda.reset_max_memory_allocated()
           torch.cuda.reset_peak_memory_stats()

           _UpperCAmelCase      :     List[Any]        =   StableDiffusionInstructPixaPixPipeline.from_pretrained(
               "timbrooks/instruct-pix2pix"					,     safety_checker=lowercase_					,     torch_dtype=torch.floataa						)
           _UpperCAmelCase      :     str        =   pipe.to(lowercase_						)
           pipe.set_progress_bar_config(disable=lowercase_						)
           pipe.enable_attention_slicing(1						)
           pipe.enable_sequential_cpu_offload()

           _UpperCAmelCase      :     Optional[Any]        =   self.get_inputs()
           _UpperCAmelCase      :     Optional[int]        =   pipe(**lowercase_						)

           _UpperCAmelCase      :     Optional[int]        =   torch.cuda.max_memory_allocated()
           # make sure that less than 2.2 GB is allocated
           assert mem_bytes < 2.2 * 1_0**9



      def 						snake_case_ (					self					:     str						):
           _UpperCAmelCase      :     int        =   self.get_inputs()
           # resize to resolution that is divisible by 8 but not 16 or 32
           _UpperCAmelCase      :     Optional[int]        =   inputs['''image'''].resize((5_0_4, 5_0_4)						)

           _UpperCAmelCase      :     Optional[int]        =   '''timbrooks/instruct-pix2pix'''
           _UpperCAmelCase      :     int        =   StableDiffusionInstructPixaPixPipeline.from_pretrained(
               lowercase_					,     safety_checker=lowercase_					,     )
           pipe.to(lowercase_						)
           pipe.set_progress_bar_config(disable=lowercase_						)
           pipe.enable_attention_slicing()

           _UpperCAmelCase      :     Optional[Any]        =   pipe(**lowercase_						)
           _UpperCAmelCase      :     Optional[Any]        =   output.images[0]

           _UpperCAmelCase      :     Any        =   image[2_5_5:2_5_8, 3_8_3:3_8_6, -1]

           assert image.shape == (5_0_4, 5_0_4, 3)
           _UpperCAmelCase      :     List[Any]        =   np.array([0.2_726, 0.2_529, 0.2_664, 0.2_655, 0.2_641, 0.2_642, 0.2_591, 0.2_649, 0.2_590]						)

           assert np.abs(image_slice.flatten() - expected_slice						).max() < 5e-3





"""simple docstring"""




# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re

from ..utils import cached_file


# docstyle-ignore
_lowerCAmelCase    : Any							=							"\nHuman: <<task>>\n\nAssistant: "


_lowerCAmelCase    : str							=							"huggingface-tools/default-prompts"
_lowerCAmelCase    : Union[str, Any]							=							{"chat": "chat_prompt_template.txt", "run": "run_prompt_template.txt"}

def 						__snake_case						(				SCREAMING_SNAKE_CASE__     :						Union[str, Any]						,  SCREAMING_SNAKE_CASE__     :						Optional[int]						,  SCREAMING_SNAKE_CASE__     :						int="run"    )  ->       int:





     '''simple docstring'''

     if prompt_or_repo_id is None:
          _UpperCAmelCase      :     Optional[int]        =   DEFAULT_PROMPTS_REPO

     # prompt is considered a repo ID when it does not contain any kind of space
     if re.search("\\s"						,  SCREAMING_SNAKE_CASE__    ) is not None:
          return prompt_or_repo_id

     _UpperCAmelCase      :     Dict        =   cached_file(
         SCREAMING_SNAKE_CASE__						,  PROMPT_FILES[mode]						,  repo_type="dataset"						,  user_agent={"agent": agent_name}    )
     with open(SCREAMING_SNAKE_CASE__						,  "r"						,  encoding="utf-8"    ) as f:
          return f.read()
