



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





 """simple docstring"""



 lowerCamelCase_        =      StableDiffusionInstructPixaPixPipeline
 lowerCamelCase_        =      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width', 'cross_attention_kwargs'}
 lowerCamelCase_        =      TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
 lowerCamelCase_        =      IMAGE_TO_IMAGE_IMAGE_PARAMS
 lowerCamelCase_        =      IMAGE_TO_IMAGE_IMAGE_PARAMS
 def lowercase_		(    self       )						->				str:




   """simple docstring"""







   torch.manual_seed(0       )
   _A							     =     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							     =     PNDMScheduler(skip_prk_steps=a       )
   torch.manual_seed(0       )
   _A							     =     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							     =     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							     =     CLIPTextModel(a       )
   _A							     =     CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )

   _A							     =     {
       '''unet''': unet,
       '''scheduler''': scheduler,
       '''vae''': vae,
       '''text_encoder''': text_encoder,
       '''tokenizer''': tokenizer,
       '''safety_checker''': None,
       '''feature_extractor''': None,
   }
   return components
 def lowercase_		(    self						,     a						,     a=0       )						->				Dict:




   """simple docstring"""







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




   """simple docstring"""







   _A							     =     '''cpu'''  # ensure determinism for the device-dependent torch.Generator
   _A							     =     self.get_dummy_components()
   _A							     =     StableDiffusionInstructPixaPixPipeline(**a       )
   _A							     =     sd_pipe.to(a       )
   sd_pipe.set_progress_bar_config(disable=a       )

   _A							     =     self.get_dummy_inputs(a       )
   _A							     =     sd_pipe(**a       ).images
   _A							     =     image[0, -3:, -3:, -1]
   assert image.shape == (1, 3_2, 3_2, 3)
   _A							     =     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 lowercase_		(    self       )						->				List[str]:




   """simple docstring"""







   _A							     =     '''cpu'''  # ensure determinism for the device-dependent torch.Generator
   _A							     =     self.get_dummy_components()
   _A							     =     StableDiffusionInstructPixaPixPipeline(**a       )
   _A							     =     sd_pipe.to(a       )
   sd_pipe.set_progress_bar_config(disable=a       )

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

   assert image.shape == (1, 3_2, 3_2, 3)
   _A							     =     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 lowercase_		(    self       )						->				Union[str, Any]:




   """simple docstring"""







   _A							     =     '''cpu'''  # ensure determinism for the device-dependent torch.Generator
   _A							     =     self.get_dummy_components()
   _A							     =     StableDiffusionInstructPixaPixPipeline(**a       )
   _A							     =     sd_pipe.to(a       )
   sd_pipe.set_progress_bar_config(disable=a       )

   _A							     =     self.get_dummy_inputs(a       )
   _A							     =     [inputs['''prompt''']] * 2

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

   _A							     =     sd_pipe(**a       ).images
   _A							     =     image[-1, -3:, -3:, -1]

   assert image.shape == (2, 3_2, 3_2, 3)
   _A							     =     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 lowercase_		(    self       )						->				List[str]:




   """simple docstring"""







   _A							     =     '''cpu'''  # ensure determinism for the device-dependent torch.Generator
   _A							     =     self.get_dummy_components()
   _A							     =     EulerAncestralDiscreteScheduler(
       beta_start=0.00_085						,     beta_end=0.012						,     beta_schedule='''scaled_linear'''       )
   _A							     =     StableDiffusionInstructPixaPixPipeline(**a       )
   _A							     =     sd_pipe.to(a       )
   sd_pipe.set_progress_bar_config(disable=a       )

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

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

   assert image.shape == (1, 3_2, 3_2, 3)
   _A							     =     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 lowercase_		(    self       )						->				List[Any]:




   """simple docstring"""







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



 def lowercase_		(    self       )						->				Dict:




   """simple docstring"""







   _A							     =     self.get_dummy_components()
   _A							     =     StableDiffusionInstructPixaPixPipeline(**a       )
   _A							     =     VaeImageProcessor(do_resize=a						,     do_normalize=a       )
   _A							     =     pipe.to(a       )
   pipe.set_progress_bar_config(disable=a       )

   _A							     =     pipe(**self.get_dummy_inputs_by_type(a						,     input_image_type='''pt'''       )       )[0]

   _A							     =     components['''vae''']
   _A							     =     self.get_dummy_inputs_by_type(a						,     input_image_type='''pt'''       )

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

   _A							     =     pipe(**a       )[0]

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







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





 """simple docstring"""



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




   """simple docstring"""







   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()
 def lowercase_		(    self						,     a=0       )						->				Union[str, Any]:




   """simple docstring"""







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




   """simple docstring"""







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

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

   assert image.shape == (1, 5_1_2, 5_1_2, 3)
   _A							     =     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 lowercase_		(    self       )						->				int:




   """simple docstring"""







   _A							     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
       '''timbrooks/instruct-pix2pix'''						,     safety_checker=a       )
   _A							     =     LMSDiscreteScheduler.from_config(pipe.scheduler.config       )
   pipe.to(a       )
   pipe.set_progress_bar_config(disable=a       )
   pipe.enable_attention_slicing()

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

   assert image.shape == (1, 5_1_2, 5_1_2, 3)
   _A							     =     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 lowercase_		(    self       )						->				int:




   """simple docstring"""







   _A							     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
       '''timbrooks/instruct-pix2pix'''						,     safety_checker=a       )
   _A							     =     DDIMScheduler.from_config(pipe.scheduler.config       )
   pipe.to(a       )
   pipe.set_progress_bar_config(disable=a       )
   pipe.enable_attention_slicing()

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

   assert image.shape == (1, 5_1_2, 5_1_2, 3)
   _A							     =     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 lowercase_		(    self       )						->				Union[str, Any]:




   """simple docstring"""







   _A							     =     0

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

   _A							     =     False

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

   _A							     =     self.get_inputs()
   pipe(**a						,     callback=a						,     callback_steps=1       )
   assert callback_fn.has_been_called
   assert number_of_steps == 3
 def lowercase_		(    self       )						->				Any:




   """simple docstring"""







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

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

   _A							     =     self.get_inputs()
   _A							     =     pipe(**a       )

   _A							     =     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:




   """simple docstring"""







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

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

   _A							     =     pipe(**a       )
   _A							     =     output.images[0]

   _A							     =     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							     =     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'''

import dataclasses
import json
import warnings
from dataclasses import dataclass, field
from time import time
from typing import List

from ..utils import logging


a   	=    logging.get_logger(__name__)

def     __magic_name__							(				__UpperCAmelCase=None				,						__UpperCAmelCase=None   )    ->       Dict:

   '''simple docstring'''







   return field(default_factory=lambda: default				,						metadata=__UpperCAmelCase   )







@dataclass
class __a			:
   __UpperCamelCase :			List[str]						   =						list_field(
       default=[], metadata={
           'help': (
               'Model checkpoints to be provided to the AutoModel classes. Leave blank to benchmark the base version'
               ' of all available models'
           )
       }, )

   __UpperCamelCase :			List[int]						   =						list_field(
       default=[8], metadata={'help': 'List of batch sizes for which memory and time performance will be evaluated'}      )

   __UpperCamelCase :			List[int]						   =						list_field(
       default=[8, 32, 128, 512], metadata={'help': 'List of sequence lengths for which memory and time performance will be evaluated'}, )

   __UpperCamelCase :			bool						   =						field(
       default=_snake_case, metadata={'help': 'Whether to benchmark inference of model. Inference can be disabled via --no-inference.'}, )
   __UpperCamelCase :			bool						   =						field(
       default=_snake_case, metadata={'help': 'Whether to run on available cuda devices. Cuda can be disabled via --no-cuda.'}, )
   __UpperCamelCase :			bool						   =						field(
       default=_snake_case, metadata={'help': 'Whether to run on available tpu devices. TPU can be disabled via --no-tpu.'}      )
   __UpperCamelCase :			bool						   =						field(default=_snake_case, metadata={'help': 'Use FP16 to accelerate inference.'}      )
   __UpperCamelCase :			bool						   =						field(default=_snake_case, metadata={'help': 'Benchmark training of model'}      )
   __UpperCamelCase :			bool						   =						field(default=_snake_case, metadata={'help': 'Verbose memory tracing'}      )
   __UpperCamelCase :			bool						   =						field(
       default=_snake_case, metadata={'help': 'Whether to perform speed measurements. Speed measurements can be disabled via --no-speed.'}, )
   __UpperCamelCase :			bool						   =						field(
       default=_snake_case, metadata={
           'help': 'Whether to perform memory measurements. Memory measurements can be disabled via --no-memory'
       }, )
   __UpperCamelCase :			bool						   =						field(default=_snake_case, metadata={'help': 'Trace memory line by line'}      )
   __UpperCamelCase :			bool						   =						field(default=_snake_case, metadata={'help': 'Save result to a CSV file'}      )
   __UpperCamelCase :			bool						   =						field(default=_snake_case, metadata={'help': 'Save all print statements in a log file'}      )
   __UpperCamelCase :			bool						   =						field(default=_snake_case, metadata={'help': 'Whether to print environment information'}      )
   __UpperCamelCase :			bool						   =						field(
       default=_snake_case, metadata={
           'help': (
               'Whether to use multiprocessing for memory and speed measurement. It is highly recommended to use'
               ' multiprocessing for accurate CPU and GPU memory measurements. This option should only be disabled'
               ' for debugging / testing and on TPU.'
           )
       }, )
   __UpperCamelCase :			str						   =						field(
       default=F'inference_time_{round(time()      )}.csv', metadata={'help': 'CSV filename used if saving time results to csv.'}, )
   __UpperCamelCase :			str						   =						field(
       default=F'inference_memory_{round(time()      )}.csv', metadata={'help': 'CSV filename used if saving memory results to csv.'}, )
   __UpperCamelCase :			str						   =						field(
       default=F'train_time_{round(time()      )}.csv', metadata={'help': 'CSV filename used if saving time results to csv for training.'}, )
   __UpperCamelCase :			str						   =						field(
       default=F'train_memory_{round(time()      )}.csv', metadata={'help': 'CSV filename used if saving memory results to csv for training.'}, )
   __UpperCamelCase :			str						   =						field(
       default=F'env_info_{round(time()      )}.csv', metadata={'help': 'CSV filename used if saving environment information.'}, )
   __UpperCamelCase :			str						   =						field(
       default=F'log_{round(time()      )}.csv', metadata={'help': 'Log filename used if print statements are saved in log.'}, )
   __UpperCamelCase :			int						   =						field(default=3, metadata={'help': 'Times an experiment will be run.'}      )
   __UpperCamelCase :			bool						   =						field(
       default=_snake_case, metadata={
           'help': (
               'Instead of loading the model as defined in `config.architectures` if exists, just load the pretrain'
               ' model weights.'
           )
       }, )





   def   UpperCAmelCase__  (       self					:      List[Any]     ):

      '''simple docstring'''

      warnings.warn(
          f"""The class {self.__class__} is deprecated. Hugging Face Benchmarking utils"""
          """ are deprecated in general and it is advised to use external Benchmarking libraries """
          """ to benchmark Transformer models."""					,lowerCamelCase					,)





   def   UpperCAmelCase__  (       self					:      int     ):

      '''simple docstring'''

      return json.dumps(dataclasses.asdict(self     )					,indent=2     )





   @property
   def   UpperCAmelCase__  (       self					:      str     ):

      '''simple docstring'''

      if len(self.models     ) <= 0:
         raise ValueError(
             """Please make sure you provide at least one model name / model identifier, *e.g.* `--models"""
             """ bert-base-cased` or `args.models = ['bert-base-cased']."""     )
      return self.models





   @property
   def   UpperCAmelCase__  (       self					:      Tuple     ):

      '''simple docstring'''

      if not self.multi_process:
         return False
      elif self.is_tpu:
         logger.info("""Multiprocessing is currently not possible on TPU."""     )
         return False
      else:
         return True





