





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"""


  a_      =			StableDiffusionInstructPixaPixPipeline
  a_      =			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width""", """cross_attention_kwargs"""}
  a_      =			TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
  a_      =			IMAGE_TO_IMAGE_IMAGE_PARAMS
  a_      =			IMAGE_TO_IMAGE_IMAGE_PARAMS






  def     lowercase		(    self			:		Union[str, Any]					)							->						Union[str, Any]:
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     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 ,  )
     __lowerCAmelCase							     =     PNDMScheduler(skip_prk_steps=lowerCAmelCase_					)
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     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					)
     __lowerCAmelCase							     =     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 ,  )
     __lowerCAmelCase							     =     CLIPTextModel(lowerCAmelCase_					)
     __lowerCAmelCase							     =     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)

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






  def     lowercase		(    self			:		Tuple ,  lowerCAmelCase_			:		List[str] ,  lowerCAmelCase_			:		Union[str, Any]=0					)							->						Union[str, Any]:
     __lowerCAmelCase							     =     floats_tensor((1, 3, 3_2, 3_2) ,  rng=random.Random(lowerCAmelCase_					)					).to(lowerCAmelCase_					)
     __lowerCAmelCase							     =     image.cpu().permute(0 ,  2 ,  3 ,  1					)[0]
     __lowerCAmelCase							     =     Image.fromarray(np.uinta(lowerCAmelCase_					)					).convert('RGB'					)
     if str(lowerCAmelCase_					).startswith('mps'					):
        __lowerCAmelCase							     =     torch.manual_seed(lowerCAmelCase_					)
     else:
        __lowerCAmelCase							     =     torch.Generator(device=lowerCAmelCase_					).manual_seed(lowerCAmelCase_					)
     __lowerCAmelCase							     =     {
         '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     lowercase		(    self			:		Dict					)							->						Any:
     __lowerCAmelCase							     =     'cpu'  # ensure determinism for the device-dependent torch.Generator
     __lowerCAmelCase							     =     self.get_dummy_components()
     __lowerCAmelCase							     =     StableDiffusionInstructPixaPixPipeline(**lowerCAmelCase_					)
     __lowerCAmelCase							     =     sd_pipe.to(lowerCAmelCase_					)
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

     __lowerCAmelCase							     =     self.get_dummy_inputs(lowerCAmelCase_					)
     __lowerCAmelCase							     =     sd_pipe(**lowerCAmelCase_					).images
     __lowerCAmelCase							     =     image[0, -3:, -3:, -1]
     assert image.shape == (1, 3_2, 3_2, 3)
     __lowerCAmelCase							     =     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     lowercase		(    self			:		Any					)							->						str:
     __lowerCAmelCase							     =     'cpu'  # ensure determinism for the device-dependent torch.Generator
     __lowerCAmelCase							     =     self.get_dummy_components()
     __lowerCAmelCase							     =     StableDiffusionInstructPixaPixPipeline(**lowerCAmelCase_					)
     __lowerCAmelCase							     =     sd_pipe.to(lowerCAmelCase_					)
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

     __lowerCAmelCase							     =     self.get_dummy_inputs(lowerCAmelCase_					)
     __lowerCAmelCase							     =     'french fries'
     __lowerCAmelCase							     =     sd_pipe(**lowerCAmelCase_ ,  negative_prompt=lowerCAmelCase_					)
     __lowerCAmelCase							     =     output.images
     __lowerCAmelCase							     =     image[0, -3:, -3:, -1]

     assert image.shape == (1, 3_2, 3_2, 3)
     __lowerCAmelCase							     =     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     lowercase		(    self			:		List[Any]					)							->						Dict:
     __lowerCAmelCase							     =     'cpu'  # ensure determinism for the device-dependent torch.Generator
     __lowerCAmelCase							     =     self.get_dummy_components()
     __lowerCAmelCase							     =     StableDiffusionInstructPixaPixPipeline(**lowerCAmelCase_					)
     __lowerCAmelCase							     =     sd_pipe.to(lowerCAmelCase_					)
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

     __lowerCAmelCase							     =     self.get_dummy_inputs(lowerCAmelCase_					)
     __lowerCAmelCase							     =     [inputs['prompt']] * 2

     __lowerCAmelCase							     =     np.array(inputs['image']					).astype(np.floataa					) / 2_55.0
     __lowerCAmelCase							     =     torch.from_numpy(lowerCAmelCase_					).unsqueeze(0					).to(lowerCAmelCase_					)
     __lowerCAmelCase							     =     image / 2 + 0.5
     __lowerCAmelCase							     =     image.permute(0 ,  3 ,  1 ,  2					)
     __lowerCAmelCase							     =     image.repeat(2 ,  1 ,  1 ,  1					)

     __lowerCAmelCase							     =     sd_pipe(**lowerCAmelCase_					).images
     __lowerCAmelCase							     =     image[-1, -3:, -3:, -1]

     assert image.shape == (2, 3_2, 3_2, 3)
     __lowerCAmelCase							     =     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     lowercase		(    self			:		Tuple					)							->						int:
     __lowerCAmelCase							     =     'cpu'  # ensure determinism for the device-dependent torch.Generator
     __lowerCAmelCase							     =     self.get_dummy_components()
     __lowerCAmelCase							     =     EulerAncestralDiscreteScheduler(
         beta_start=0.0_00_85 ,  beta_end=0.0_12 ,  beta_schedule='scaled_linear'					)
     __lowerCAmelCase							     =     StableDiffusionInstructPixaPixPipeline(**lowerCAmelCase_					)
     __lowerCAmelCase							     =     sd_pipe.to(lowerCAmelCase_					)
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

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

     __lowerCAmelCase							     =     [round(lowerCAmelCase_ ,  4					) for x in image_slice.flatten().tolist()]
     print(','.join([str(lowerCAmelCase_					) for x in slice]					)					)

     assert image.shape == (1, 3_2, 3_2, 3)
     __lowerCAmelCase							     =     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     lowercase		(    self			:		str					)							->						Tuple:
     super().test_inference_batch_single_identical(expected_max_diff=3e-3					)






  def     lowercase		(    self			:		str					)							->						List[Any]:
     __lowerCAmelCase							     =     self.get_dummy_components()
     __lowerCAmelCase							     =     StableDiffusionInstructPixaPixPipeline(**lowerCAmelCase_					)
     __lowerCAmelCase							     =     VaeImageProcessor(do_resize=lowerCAmelCase_ ,  do_normalize=lowerCAmelCase_					)
     __lowerCAmelCase							     =     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

     __lowerCAmelCase							     =     pipe(**self.get_dummy_inputs_by_type(lowerCAmelCase_ ,  input_image_type='pt'					)					)[0]

     __lowerCAmelCase							     =     components['vae']
     __lowerCAmelCase							     =     self.get_dummy_inputs_by_type(lowerCAmelCase_ ,  input_image_type='pt'					)

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

     __lowerCAmelCase							     =     pipe(**lowerCAmelCase_					)[0]

     __lowerCAmelCase							     =     np.abs(out - out_latents_inputs					).max()
     self.assertLess(lowerCAmelCase_ ,  1e-4 ,  'passing latents as image input generate different result from passing image'					)





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

  """simple docstring"""


  def     lowercase		(    self			:		Tuple					)							->						int:
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()






  def     lowercase		(    self			:		Optional[Any] ,  lowerCAmelCase_			:		Tuple=0					)							->						List[str]:
     __lowerCAmelCase							     =     torch.manual_seed(lowerCAmelCase_					)
     __lowerCAmelCase							     =     load_image(
         'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg'					)
     __lowerCAmelCase							     =     {
         '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     lowercase		(    self			:		Any					)							->						Tuple:
     __lowerCAmelCase							     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
         'timbrooks/instruct-pix2pix' ,  safety_checker=lowerCAmelCase_					)
     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)
     pipe.enable_attention_slicing()

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

     assert image.shape == (1, 5_1_2, 5_1_2, 3)
     __lowerCAmelCase							     =     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     lowercase		(    self			:		Optional[int]					)							->						Tuple:
     __lowerCAmelCase							     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
         'timbrooks/instruct-pix2pix' ,  safety_checker=lowerCAmelCase_					)
     __lowerCAmelCase							     =     LMSDiscreteScheduler.from_config(pipe.scheduler.config					)
     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)
     pipe.enable_attention_slicing()

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

     assert image.shape == (1, 5_1_2, 5_1_2, 3)
     __lowerCAmelCase							     =     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     lowercase		(    self			:		List[Any]					)							->						Tuple:
     __lowerCAmelCase							     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
         'timbrooks/instruct-pix2pix' ,  safety_checker=lowerCAmelCase_					)
     __lowerCAmelCase							     =     DDIMScheduler.from_config(pipe.scheduler.config					)
     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)
     pipe.enable_attention_slicing()

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

     assert image.shape == (1, 5_1_2, 5_1_2, 3)
     __lowerCAmelCase							     =     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     lowercase		(    self			:		Dict					)							->						Dict:
     __lowerCAmelCase							     =     0

     def callback_fn(lowerCAmelCase_			:		int ,  lowerCAmelCase_			:		int ,  lowerCAmelCase_			:		torch.FloatTensor					) -> None:
        __lowerCAmelCase							     =     True
        nonlocal number_of_steps
        number_of_steps += 1
        if step == 1:
           __lowerCAmelCase							     =     latents.detach().cpu().numpy()
           assert latents.shape == (1, 4, 6_4, 6_4)
           __lowerCAmelCase							     =     latents[0, -3:, -3:, -1]
           __lowerCAmelCase							     =     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:
           __lowerCAmelCase							     =     latents.detach().cpu().numpy()
           assert latents.shape == (1, 4, 6_4, 6_4)
           __lowerCAmelCase							     =     latents[0, -3:, -3:, -1]
           __lowerCAmelCase							     =     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

     __lowerCAmelCase							     =     False

     __lowerCAmelCase							     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
         'timbrooks/instruct-pix2pix' ,  safety_checker=lowerCAmelCase_ ,  torch_dtype=torch.floataa					)
     __lowerCAmelCase							     =     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)
     pipe.enable_attention_slicing()

     __lowerCAmelCase							     =     self.get_inputs()
     pipe(**lowerCAmelCase_ ,  callback=lowerCAmelCase_ ,  callback_steps=1					)
     assert callback_fn.has_been_called
     assert number_of_steps == 3






  def     lowercase		(    self			:		str					)							->						Any:
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()

     __lowerCAmelCase							     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
         'timbrooks/instruct-pix2pix' ,  safety_checker=lowerCAmelCase_ ,  torch_dtype=torch.floataa					)
     __lowerCAmelCase							     =     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)
     pipe.enable_attention_slicing(1					)
     pipe.enable_sequential_cpu_offload()

     __lowerCAmelCase							     =     self.get_inputs()
     __lowerCAmelCase							     =     pipe(**lowerCAmelCase_					)

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






  def     lowercase		(    self			:		Tuple					)							->						Union[str, Any]:
     __lowerCAmelCase							     =     self.get_inputs()
     # resize to resolution that is divisible by 8 but not 16 or 32
     __lowerCAmelCase							     =     inputs['image'].resize((5_0_4, 5_0_4)					)

     __lowerCAmelCase							     =     'timbrooks/instruct-pix2pix'
     __lowerCAmelCase							     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
         lowerCAmelCase_ ,  safety_checker=lowerCAmelCase_ ,  )
     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)
     pipe.enable_attention_slicing()

     __lowerCAmelCase							     =     pipe(**lowerCAmelCase_					)
     __lowerCAmelCase							     =     output.images[0]

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

     assert image.shape == (5_0_4, 5_0_4, 3)
     __lowerCAmelCase							     =     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








from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


_snake_case  :   Optional[Any]						   =  {'configuration_fnet': ['FNET_PRETRAINED_CONFIG_ARCHIVE_MAP', 'FNetConfig']}

try:
     if not is_sentencepiece_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _snake_case  :   int						   =  ['FNetTokenizer']

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _snake_case  :   Tuple						   =  ['FNetTokenizerFast']

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _snake_case  :   List[str]						   =  [
         'FNET_PRETRAINED_MODEL_ARCHIVE_LIST',
         'FNetForMaskedLM',
         'FNetForMultipleChoice',
         'FNetForNextSentencePrediction',
         'FNetForPreTraining',
         'FNetForQuestionAnswering',
         'FNetForSequenceClassification',
         'FNetForTokenClassification',
         'FNetLayer',
         'FNetModel',
         'FNetPreTrainedModel',
     ]


if TYPE_CHECKING:
     from .configuration_fnet import FNET_PRETRAINED_CONFIG_ARCHIVE_MAP, FNetConfig

     try:
          if not is_sentencepiece_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_fnet import FNetTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_fnet_fast import FNetTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_fnet import (
              FNET_PRETRAINED_MODEL_ARCHIVE_LIST,
              FNetForMaskedLM,
              FNetForMultipleChoice,
              FNetForNextSentencePrediction,
              FNetForPreTraining,
              FNetForQuestionAnswering,
              FNetForSequenceClassification,
              FNetForTokenClassification,
              FNetLayer,
              FNetModel,
              FNetPreTrainedModel,
          )


else:
     import sys

     _snake_case  :   str						   =  _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)

