





"""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  (   UpperCAmelCase__			,   UpperCAmelCase__			,   UpperCAmelCase__			,   unittest.TestCase		):



      '''simple docstring'''
      lowercase_      : Tuple	    =   StableDiffusionInstructPixaPixPipeline
      lowercase_      : Any	    =   TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width""", """cross_attention_kwargs"""}
      lowercase_      : Dict	    =   TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
      lowercase_      : int	    =   IMAGE_TO_IMAGE_IMAGE_PARAMS
      lowercase_      : Dict	    =   IMAGE_TO_IMAGE_IMAGE_PARAMS

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        torch.manual_seed(0  )
        A_  :       Dict										=  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					,      )
        A_  :       Union[str, Any]										=  PNDMScheduler(skip_prk_steps=snake_case_  )
        torch.manual_seed(0  )
        A_  :       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  )
        A_  :       str										=  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					,      )
        A_  :       int										=  CLIPTextModel(snake_case_  )
        A_  :       Dict										=  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'  )

        A_  :       Tuple										=  {
            'unet': unet,
            'scheduler': scheduler,
            'vae': vae,
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'safety_checker': None,
            'feature_extractor': None,
        }
        return components

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_=0  ):

        """simple docstring"""

        A_  :       Any										=  floats_tensor((1, 3, 3_2, 3_2)					,      rng=random.Random(snake_case_  )  ).to(snake_case_  )
        A_  :       Dict										=  image.cpu().permute(0					,      2					,      3					,      1  )[0]
        A_  :       Optional[Any]										=  Image.fromarray(np.uinta(snake_case_  )  ).convert('RGB'  )
        if str(snake_case_  ).startswith('mps'  ):
          A_  :       Optional[Any]										=  torch.manual_seed(snake_case_  )
        else:
          A_  :       int										=  torch.Generator(device=snake_case_  ).manual_seed(snake_case_  )
        A_  :       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 					lowerCamelCase_			(     self  ):

        """simple docstring"""

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

        A_  :       Tuple										=  self.get_dummy_inputs(snake_case_  )
        A_  :       Tuple										=  sd_pipe(**snake_case_  ).images
        A_  :       Dict										=  image[0, -3:, -3:, -1]
        assert image.shape == (1, 3_2, 3_2, 3)
        A_  :       int										=  np.array([0.75_26, 0.37_50, 0.45_47, 0.61_17, 0.58_66, 0.50_16, 0.43_27, 0.56_42, 0.48_15]  )

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Tuple										=  'cpu'  # ensure determinism for the device-dependent torch.Generator
        A_  :       Any										=  self.get_dummy_components()
        A_  :       Tuple										=  StableDiffusionInstructPixaPixPipeline(**snake_case_  )
        A_  :       Tuple										=  sd_pipe.to(snake_case_  )
        sd_pipe.set_progress_bar_config(disable=snake_case_  )

        A_  :       Dict										=  self.get_dummy_inputs(snake_case_  )
        A_  :       Tuple										=  'french fries'
        A_  :       List[str]										=  sd_pipe(**snake_case_					,      negative_prompt=snake_case_  )
        A_  :       str										=  output.images
        A_  :       Optional[int]										=  image[0, -3:, -3:, -1]

        assert image.shape == (1, 3_2, 3_2, 3)
        A_  :       Union[str, Any]										=  np.array([0.75_11, 0.36_42, 0.45_53, 0.62_36, 0.57_97, 0.50_13, 0.43_43, 0.56_11, 0.48_31]  )

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

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

        A_  :       Optional[Any]										=  self.get_dummy_inputs(snake_case_  )
        A_  :       Union[str, Any]										=  [inputs['prompt']] * 2

        A_  :       Tuple										=  np.array(inputs['image']  ).astype(np.floataa  ) / 2_55.0
        A_  :       int										=  torch.from_numpy(snake_case_  ).unsqueeze(0  ).to(snake_case_  )
        A_  :       Dict										=  image / 2 + 0.5
        A_  :       int										=  image.permute(0					,      3					,      1					,      2  )
        A_  :       Union[str, Any]										=  image.repeat(2					,      1					,      1					,      1  )

        A_  :       Tuple										=  sd_pipe(**snake_case_  ).images
        A_  :       Optional[Any]										=  image[-1, -3:, -3:, -1]

        assert image.shape == (2, 3_2, 3_2, 3)
        A_  :       Any										=  np.array([0.58_12, 0.57_48, 0.52_22, 0.59_08, 0.56_95, 0.71_74, 0.68_04, 0.55_23, 0.55_79]  )

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       str										=  'cpu'  # ensure determinism for the device-dependent torch.Generator
        A_  :       List[str]										=  self.get_dummy_components()
        A_  :       Dict										=  EulerAncestralDiscreteScheduler(
            beta_start=0.0_00_85					,      beta_end=0.0_12					,      beta_schedule='scaled_linear'  )
        A_  :       Dict										=  StableDiffusionInstructPixaPixPipeline(**snake_case_  )
        A_  :       Dict										=  sd_pipe.to(snake_case_  )
        sd_pipe.set_progress_bar_config(disable=snake_case_  )

        A_  :       Optional[Any]										=  self.get_dummy_inputs(snake_case_  )
        A_  :       Optional[int]										=  sd_pipe(**snake_case_  ).images
        A_  :       Tuple										=  image[0, -3:, -3:, -1]

        A_  :       List[Any]										=  [round(snake_case_					,      4  ) for x in image_slice.flatten().tolist()]
        print(','.join([str(snake_case_  ) for x in slice]  )  )

        assert image.shape == (1, 3_2, 3_2, 3)
        A_  :       List[str]										=  np.array([0.74_17, 0.38_42, 0.47_32, 0.57_76, 0.58_91, 0.51_39, 0.40_52, 0.56_73, 0.49_86]  )

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Optional[Any]										=  self.get_dummy_components()
        A_  :       Any										=  StableDiffusionInstructPixaPixPipeline(**snake_case_  )
        A_  :       Optional[int]										=  VaeImageProcessor(do_resize=snake_case_					,      do_normalize=snake_case_  )
        A_  :       List[Any]										=  pipe.to(snake_case_  )
        pipe.set_progress_bar_config(disable=snake_case_  )

        A_  :       Dict										=  pipe(**self.get_dummy_inputs_by_type(snake_case_					,      input_image_type='pt'  )  )[0]

        A_  :       int										=  components['vae']
        A_  :       int										=  self.get_dummy_inputs_by_type(snake_case_					,      input_image_type='pt'  )

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

        A_  :       Union[str, Any]										=  pipe(**snake_case_  )[0]

        A_  :       int										=  np.abs(out - out_latents_inputs  ).max()
        self.assertLess(snake_case_					,      1E-4					,      'passing latents as image input generate different result from passing image'  )


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



      '''simple docstring'''

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

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

      def 					lowerCamelCase_			(     self					,      snake_case_=0  ):

        """simple docstring"""

        A_  :       Optional[Any]										=  torch.manual_seed(snake_case_  )
        A_  :       List[Any]										=  load_image(
            'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg'  )
        A_  :       Any										=  {
            '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 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Optional[Any]										=  StableDiffusionInstructPixaPixPipeline.from_pretrained(
            'timbrooks/instruct-pix2pix'					,      safety_checker=snake_case_  )
        pipe.to(snake_case_  )
        pipe.set_progress_bar_config(disable=snake_case_  )
        pipe.enable_attention_slicing()

        A_  :       str										=  self.get_inputs()
        A_  :       List[str]										=  pipe(**snake_case_  ).images
        A_  :       List[Any]										=  image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        A_  :       Union[str, Any]										=  np.array([0.59_02, 0.60_15, 0.60_27, 0.59_83, 0.60_92, 0.60_61, 0.57_65, 0.57_85, 0.55_55]  )

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  StableDiffusionInstructPixaPixPipeline.from_pretrained(
            'timbrooks/instruct-pix2pix'					,      safety_checker=snake_case_  )
        A_  :       List[str]										=  LMSDiscreteScheduler.from_config(pipe.scheduler.config  )
        pipe.to(snake_case_  )
        pipe.set_progress_bar_config(disable=snake_case_  )
        pipe.enable_attention_slicing()

        A_  :       int										=  self.get_inputs()
        A_  :       str										=  pipe(**snake_case_  ).images
        A_  :       Dict										=  image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        A_  :       Union[str, Any]										=  np.array([0.65_78, 0.68_17, 0.69_72, 0.67_61, 0.68_56, 0.69_16, 0.64_28, 0.65_16, 0.63_01]  )

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       int										=  StableDiffusionInstructPixaPixPipeline.from_pretrained(
            'timbrooks/instruct-pix2pix'					,      safety_checker=snake_case_  )
        A_  :       Dict										=  DDIMScheduler.from_config(pipe.scheduler.config  )
        pipe.to(snake_case_  )
        pipe.set_progress_bar_config(disable=snake_case_  )
        pipe.enable_attention_slicing()

        A_  :       Optional[int]										=  self.get_inputs()
        A_  :       Dict										=  pipe(**snake_case_  ).images
        A_  :       str										=  image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        A_  :       Any										=  np.array([0.38_28, 0.38_34, 0.38_18, 0.37_92, 0.38_65, 0.37_52, 0.37_92, 0.38_47, 0.37_53]  )

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Tuple										=  0

        def callback_fn(snake_case_					,      snake_case_					,      snake_case_  ) -> None:
          A_  :       Union[str, Any]										=  True
          nonlocal number_of_steps
          number_of_steps += 1
          if step == 1:
            A_  :       str										=  latents.detach().cpu().numpy()
            assert latents.shape == (1, 4, 6_4, 6_4)
            A_  :       Dict										=  latents[0, -3:, -3:, -1]
            A_  :       str										=  np.array([-0.24_63, -0.46_44, -0.97_56, 1.51_76, 1.44_14, 0.78_66, 0.98_97, 0.85_21, 0.79_83]  )

            assert np.abs(latents_slice.flatten() - expected_slice  ).max() < 5E-2
          elif step == 2:
            A_  :       Dict										=  latents.detach().cpu().numpy()
            assert latents.shape == (1, 4, 6_4, 6_4)
            A_  :       Union[str, Any]										=  latents[0, -3:, -3:, -1]
            A_  :       Tuple										=  np.array([-0.26_44, -0.46_26, -0.96_53, 1.51_76, 1.45_51, 0.76_86, 0.98_05, 0.84_52, 0.81_15]  )

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

        A_  :       Dict										=  False

        A_  :       List[Any]										=  StableDiffusionInstructPixaPixPipeline.from_pretrained(
            'timbrooks/instruct-pix2pix'					,      safety_checker=snake_case_					,      torch_dtype=torch.floataa  )
        A_  :       List[Any]										=  pipe.to(snake_case_  )
        pipe.set_progress_bar_config(disable=snake_case_  )
        pipe.enable_attention_slicing()

        A_  :       List[str]										=  self.get_inputs()
        pipe(**snake_case_					,      callback=snake_case_					,      callback_steps=1  )
        assert callback_fn.has_been_called
        assert number_of_steps == 3

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        A_  :       Union[str, Any]										=  StableDiffusionInstructPixaPixPipeline.from_pretrained(
            'timbrooks/instruct-pix2pix'					,      safety_checker=snake_case_					,      torch_dtype=torch.floataa  )
        A_  :       Union[str, Any]										=  pipe.to(snake_case_  )
        pipe.set_progress_bar_config(disable=snake_case_  )
        pipe.enable_attention_slicing(1  )
        pipe.enable_sequential_cpu_offload()

        A_  :       Optional[int]										=  self.get_inputs()
        A_  :       Optional[Any]										=  pipe(**snake_case_  )

        A_  :       Any										=  torch.cuda.max_memory_allocated()
        # make sure that less than 2.2 GB is allocated
        assert mem_bytes < 2.2 * 1_0**9

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Dict										=  self.get_inputs()
        # resize to resolution that is divisible by 8 but not 16 or 32
        A_  :       Dict										=  inputs['image'].resize((5_0_4, 5_0_4)  )

        A_  :       List[Any]										=  'timbrooks/instruct-pix2pix'
        A_  :       Union[str, Any]										=  StableDiffusionInstructPixaPixPipeline.from_pretrained(
            snake_case_					,      safety_checker=snake_case_					,      )
        pipe.to(snake_case_  )
        pipe.set_progress_bar_config(disable=snake_case_  )
        pipe.enable_attention_slicing()

        A_  :       Any										=  pipe(**snake_case_  )
        A_  :       int										=  output.images[0]

        A_  :       List[str]										=  image[2_5_5:2_5_8, 3_8_3:3_8_6, -1]

        assert image.shape == (5_0_4, 5_0_4, 3)
        A_  :       Optional[Any]										=  np.array([0.27_26, 0.25_29, 0.26_64, 0.26_55, 0.26_41, 0.26_42, 0.25_91, 0.26_49, 0.25_90]  )

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






"""simple docstring"""




from typing import Optional, Union

import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models.modeling_utils import ModelMixin

class  _UpperCAmelCase  (   UpperCAmelCase__			,   UpperCAmelCase__		):



      '''simple docstring'''

      @register_to_config
      def __init__(     self					,      snake_case_ = 7_6_8					,      ):

        """simple docstring"""

        super().__init__()

        A_  :       Optional[int]										=  nn.Parameter(torch.zeros(1					,      snake_case_  )  )
        A_  :       Optional[int]										=  nn.Parameter(torch.ones(1					,      snake_case_  )  )

      def 					lowerCamelCase_			(     self					,      snake_case_ = None					,      snake_case_ = None					,      ):

        """simple docstring"""

        A_  :       str										=  nn.Parameter(self.mean.to(snake_case_  ).to(snake_case_  )  )
        A_  :       Optional[int]										=  nn.Parameter(self.std.to(snake_case_  ).to(snake_case_  )  )
        return self

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       Tuple										=  (embeds - self.mean) * 1.0 / self.std
        return embeds

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       List[str]										=  (embeds * self.std) + self.mean
        return embeds