


import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
  import tensorflow as tf
  from tensorflow.python.framework.errors_impl import ResourceExhaustedError

  from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
  import pyanvml.pyanvml as nvml

UpperCAmelCase__   	=							logging.get_logger(__name__)
def 				UpperCAmelCase_					(							__snake_case     ,						__snake_case						)     ->			str:







 """simple docstring"""


 def run_func(__snake_case						):
  @wraps(_snake_case						)
  def run_in_eager_mode(*__snake_case     ,						**__snake_case						):
   return func(*_snake_case     ,						**_snake_case						)

  @wraps(_snake_case						)
  @tf.function(experimental_compile=_snake_case						)
  def run_in_graph_mode(*__snake_case     ,						**__snake_case						):
   return func(*_snake_case     ,						**_snake_case						)

  if do_eager_mode is True:
   if use_xla is not False:
    raise ValueError(
        '''Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`.'''						)
   return run_in_eager_mode
  else:
   return run_in_graph_mode

 return run_func
def 				UpperCAmelCase_					(							__snake_case     ,						__snake_case     ,						__snake_case						)     ->			Dict:







 """simple docstring"""


 _lowercase							      =random.Random()
 _lowercase							      =[rng.randint(0     ,						vocab_size - 1						) for i in range(batch_size * sequence_length						)]
 return tf.constant(_snake_case     ,						shape=(batch_size, sequence_length)     ,						dtype=tf.intaa						)







class       lowerCamelCase__   (       a__):
     SCREAMING_SNAKE_CASE__      			=			42
     SCREAMING_SNAKE_CASE__      			=			42
     SCREAMING_SNAKE_CASE__      			=			"TensorFlow"





     @property
     def      __A					(self       )     ->   Union[str, Any]:
      return tf.__version__





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase       )     ->   float:
      _lowercase							      =self.args.strategy
      if strategy is None:
       raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''       )
      _lowercase							      =self._prepare_inference_func(SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__       )
      return self._measure_speed(_inference       )





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase       )     ->   float:
      _lowercase							      =self.args.strategy
      if strategy is None:
       raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''       )
      _lowercase							      =self._prepare_train_func(SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__       )
      return self._measure_speed(_train       )





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase       )     ->   [Memory, Optional[MemorySummary]]:




      if self.args.is_gpu:
       tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]		,     SCREAMING_SNAKE_CASE__       )
      _lowercase							      =self.args.strategy
      if strategy is None:
       raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''       )
      _lowercase							      =self._prepare_inference_func(SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__       )
      return self._measure_memory(_inference       )





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase       )     ->   [Memory, Optional[MemorySummary]]:




      if self.args.is_gpu:
       tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]		,     SCREAMING_SNAKE_CASE__       )
      _lowercase							      =self.args.strategy
      if strategy is None:
       raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''       )

      _lowercase							      =self._prepare_train_func(SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__       )
      return self._measure_memory(_train       )





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase       )     ->   Callable[[], None]:
      _lowercase							      =self.config_dict[model_name]

      if self.args.fpaa:
       raise NotImplementedError('''Mixed precision is currently not supported.'''       )

      _lowercase							      =(
          hasattr(SCREAMING_SNAKE_CASE__		,     '''architectures'''       )
          and isinstance(config.architectures		,     SCREAMING_SNAKE_CASE__       )
          and len(config.architectures       ) > 0
      )
      if not self.args.only_pretrain_model and has_model_class_in_config:
       try:
        _lowercase							      ="""TF""" + config.architectures[0]  # prepend 'TF' for tensorflow model
        _lowercase							      =__import__('''transformers'''		,     fromlist=[model_class]       )
        _lowercase							      =getattr(SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__       )
        _lowercase							      =model_cls(SCREAMING_SNAKE_CASE__       )
       except ImportError:
        raise ImportError(
            f"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
            ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''       )
      else:
       _lowercase							      =TF_MODEL_MAPPING[config.__class__](SCREAMING_SNAKE_CASE__       )

      # encoder-decoder has vocab size saved differently
      _lowercase							      =config.vocab_size if hasattr(SCREAMING_SNAKE_CASE__		,     '''vocab_size'''       ) else config.encoder.vocab_size
      _lowercase							      =random_input_ids(SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__       )

      @run_with_tf_optimizations(self.args.eager_mode		,     self.args.use_xla       )
      def encoder_decoder_forward():
       return model(SCREAMING_SNAKE_CASE__		,     decoder_input_ids=SCREAMING_SNAKE_CASE__		,     training=SCREAMING_SNAKE_CASE__       )

      @run_with_tf_optimizations(self.args.eager_mode		,     self.args.use_xla       )
      def encoder_forward():
       return model(SCREAMING_SNAKE_CASE__		,     training=SCREAMING_SNAKE_CASE__       )

      _lowercase							      =encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

      return _inference





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase       )     ->   Callable[[], None]:
      _lowercase							      =self.config_dict[model_name]

      if self.args.eager_mode is not False:
       raise ValueError('''Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`.'''       )

      if self.args.fpaa:
       raise NotImplementedError('''Mixed precision is currently not supported.'''       )

      _lowercase							      =(
          hasattr(SCREAMING_SNAKE_CASE__		,     '''architectures'''       )
          and isinstance(config.architectures		,     SCREAMING_SNAKE_CASE__       )
          and len(config.architectures       ) > 0
      )
      if not self.args.only_pretrain_model and has_model_class_in_config:
       try:
        _lowercase							      ="""TF""" + config.architectures[0]  # prepend 'TF' for tensorflow model
        _lowercase							      =__import__('''transformers'''		,     fromlist=[model_class]       )
        _lowercase							      =getattr(SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__       )
        _lowercase							      =model_cls(SCREAMING_SNAKE_CASE__       )
       except ImportError:
        raise ImportError(
            f"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
            ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''       )
      else:
       _lowercase							      =TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](SCREAMING_SNAKE_CASE__       )

      # encoder-decoder has vocab size saved differently
      _lowercase							      =config.vocab_size if hasattr(SCREAMING_SNAKE_CASE__		,     '''vocab_size'''       ) else config.encoder.vocab_size
      _lowercase							      =random_input_ids(SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__       )

      @run_with_tf_optimizations(self.args.eager_mode		,     self.args.use_xla       )
      def encoder_decoder_train():
       _lowercase							      =model(SCREAMING_SNAKE_CASE__		,     decoder_input_ids=SCREAMING_SNAKE_CASE__		,     labels=SCREAMING_SNAKE_CASE__		,     training=SCREAMING_SNAKE_CASE__       )[0]
       _lowercase							      =tf.gradients(SCREAMING_SNAKE_CASE__		,     model.trainable_variables       )
       return gradients

      @run_with_tf_optimizations(self.args.eager_mode		,     self.args.use_xla       )
      def encoder_train():
       _lowercase							      =model(SCREAMING_SNAKE_CASE__		,     labels=SCREAMING_SNAKE_CASE__		,     training=SCREAMING_SNAKE_CASE__       )[0]
       _lowercase							      =tf.gradients(SCREAMING_SNAKE_CASE__		,     model.trainable_variables       )
       return gradients

      _lowercase							      =encoder_decoder_train if config.is_encoder_decoder else encoder_train

      return _train





     def      __A					(self		,     UpperCAmelCase       )     ->   float:




      with self.args.strategy.scope():
       try:
        if self.args.is_tpu or self.args.use_xla:
         # run additional 10 times to stabilize compilation for tpu
         logger.info('''Do inference on TPU. Running model 5 times to stabilize compilation'''       )
         timeit.repeat(SCREAMING_SNAKE_CASE__		,     repeat=1		,     number=5       )

        # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
        _lowercase							      =timeit.repeat(
            SCREAMING_SNAKE_CASE__		,     repeat=self.args.repeat		,     number=1_0		,     )

        return min(SCREAMING_SNAKE_CASE__       ) / 10.0
       except ResourceExhaustedError as e:
        self.print_fn(f"Doesn\'t fit on GPU. {e}"       )





     def      __A					(self		,     UpperCAmelCase       )     ->   [Memory, MemorySummary]:
      logger.info(
          '''Note that TensorFlow allocates more memory than '''
          '''it might need to speed up computation. '''
          '''The memory reported here corresponds to the memory '''
          '''reported by `nvidia-smi`, which can vary depending '''
          '''on total available memory on the GPU that is used.'''       )
      with self.args.strategy.scope():
       try:
        if self.args.trace_memory_line_by_line:
         if not self.args.eager_mode:
          raise ValueError(
              '''`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory'''
              ''' consumption line by line.'''       )
         _lowercase							      =start_memory_tracing('''transformers'''       )

        if self.args.is_tpu:
         # tpu
         raise NotImplementedError(
             '''Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking'''
             ''' with `args.memory=False`'''       )
        elif self.args.is_gpu:
         # gpu
         if not is_pyanvml_available():
          logger.warning(
              '''py3nvml not installed, we won\'t log GPU memory usage. '''
              '''Install py3nvml (pip install py3nvml) to log information about GPU.'''       )
          _lowercase							      ="""N/A"""
         else:
          logger.info(
              '''Measuring total GPU usage on GPU device. Make sure to not have additional processes'''
              ''' running on the same GPU.'''       )
          # init nvml
          nvml.nvmlInit()
          func()
          _lowercase							      =nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx       )
          _lowercase							      =nvml.nvmlDeviceGetMemoryInfo(SCREAMING_SNAKE_CASE__       )
          _lowercase							      =meminfo.used
          _lowercase							      =Memory(SCREAMING_SNAKE_CASE__       )
          # shutdown nvml
          nvml.nvmlShutdown()
        else:
         # cpu
         if self.args.trace_memory_line_by_line:
          logger.info(
              '''When enabling line by line tracing, the max peak memory for CPU is inaccurate in'''
              ''' TensorFlow.'''       )
          _lowercase							      =None
         else:
          _lowercase							      =measure_peak_memory_cpu(SCREAMING_SNAKE_CASE__       )
          _lowercase							      =Memory(SCREAMING_SNAKE_CASE__       ) if isinstance(SCREAMING_SNAKE_CASE__		,     SCREAMING_SNAKE_CASE__       ) else memory_bytes
        if self.args.trace_memory_line_by_line:
         _lowercase							      =stop_memory_tracing(SCREAMING_SNAKE_CASE__       )
         if memory is None:
          _lowercase							      =summary.total
        else:
         _lowercase							      =None

        return memory, summary
       except ResourceExhaustedError as e:
        self.print_fn(f"Doesn\'t fit on GPU. {e}"       )
        return "N/A", None



"""simple docstring"""




import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin



@skip_mps
class 	lowerCAmelCase_						(a__ ,					a__ ,					unittest.TestCase  ):


				"""simple docstring"""



				__UpperCamelCase      :						Optional[Any]					       =   IFPipeline
				__UpperCamelCase      :						Dict					       =   TEXT_TO_IMAGE_PARAMS - {'''width''', '''height''', '''latents'''}
				__UpperCamelCase      :						Any					       =   TEXT_TO_IMAGE_BATCH_PARAMS
				__UpperCamelCase      :						Union[str, Any]					       =   PipelineTesterMixin.required_optional_params - {'''latents'''}



				def 	__magic_name__    (self    )					->		Union[str, Any]:



											"""simple docstring"""




											return self._get_dummy_components()



				def 	__magic_name__    (self   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__=0    )					->		List[Any]:



											"""simple docstring"""




											if str(SCREAMING_SNAKE_CASE__    ).startswith("""mps"""    ):
																		SCREAMING_SNAKE_CASE__						:       Dict           =   torch.manual_seed(SCREAMING_SNAKE_CASE__    )
											else:
																		SCREAMING_SNAKE_CASE__						:       Any           =   torch.Generator(device=SCREAMING_SNAKE_CASE__    ).manual_seed(SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       int           =   {
											    """prompt""": """A painting of a squirrel eating a burger""",
											    """generator""": generator,
											    """num_inference_steps""": 2,
											    """output_type""": """numpy""",
											}

											return inputs



				def 	__magic_name__    (self    )					->		Tuple:



											"""simple docstring"""




											self._test_save_load_optional_components()



				@unittest.skipIf(torch_device != """cuda"""   ,       reason="""float16 requires CUDA"""    )
				def 	__magic_name__    (self    )					->		List[str]:



											"""simple docstring"""




											super().test_save_load_floataa(expected_max_diff=1E-1    )



				def 	__magic_name__    (self    )					->		List[Any]:



											"""simple docstring"""




											self._test_attention_slicing_forward_pass(expected_max_diff=1E-2    )



				def 	__magic_name__    (self    )					->		Tuple:



											"""simple docstring"""




											self._test_save_load_local()



				def 	__magic_name__    (self    )					->		Dict:



											"""simple docstring"""




											self._test_inference_batch_single_identical(
											    expected_max_diff=1E-2   ,       )



				@unittest.skipIf(
				    torch_device != """cuda""" or not is_xformers_available()   ,       reason="""XFormers attention is only available with CUDA and `xformers` installed"""   ,       )
				def 	__magic_name__    (self    )					->		Optional[int]:



											"""simple docstring"""




											self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3    )






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


				"""simple docstring"""



				def 	__magic_name__    (self    )					->		Dict:



											"""simple docstring"""




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



				def 	__magic_name__    (self    )					->		Optional[int]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   IFPipeline.from_pretrained("""DeepFloyd/IF-I-XL-v1.0"""   ,       variant="""fp16"""   ,       torch_dtype=torch.floataa    )

											SCREAMING_SNAKE_CASE__						:       Dict           =   IFSuperResolutionPipeline.from_pretrained(
											    """DeepFloyd/IF-II-L-v1.0"""   ,       variant="""fp16"""   ,       torch_dtype=torch.floataa   ,       text_encoder=SCREAMING_SNAKE_CASE__   ,       tokenizer=SCREAMING_SNAKE_CASE__    )

											# pre compute text embeddings and remove T5 to save memory

											pipe_a.text_encoder.to("""cuda"""    )

											SCREAMING_SNAKE_CASE__      ,      SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   pipe_a.encode_prompt("""anime turtle"""   ,       device="""cuda"""    )

											del pipe_a.tokenizer
											del pipe_a.text_encoder
											gc.collect()

											SCREAMING_SNAKE_CASE__						:       List[str]           =   None
											SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   None

											pipe_a.enable_model_cpu_offload()
											pipe_a.enable_model_cpu_offload()

											pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()    )
											pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()    )

											self._test_if(SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    )

											pipe_a.remove_all_hooks()
											pipe_a.remove_all_hooks()

											# img2img

											SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   IFImgaImgPipeline(**pipe_a.components    )
											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   IFImgaImgSuperResolutionPipeline(**pipe_a.components    )

											pipe_a.enable_model_cpu_offload()
											pipe_a.enable_model_cpu_offload()

											pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()    )
											pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()    )

											self._test_if_imgaimg(SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    )

											pipe_a.remove_all_hooks()
											pipe_a.remove_all_hooks()

											# inpainting

											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   IFInpaintingPipeline(**pipe_a.components    )
											SCREAMING_SNAKE_CASE__						:       int           =   IFInpaintingSuperResolutionPipeline(**pipe_a.components    )

											pipe_a.enable_model_cpu_offload()
											pipe_a.enable_model_cpu_offload()

											pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()    )
											pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()    )

											self._test_if_inpainting(SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    )



				def 	__magic_name__    (self   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    )					->		Optional[int]:



											"""simple docstring"""




											_start_torch_memory_measurement()

											SCREAMING_SNAKE_CASE__						:       Tuple           =   torch.Generator(device="""cpu"""    ).manual_seed(0    )
											SCREAMING_SNAKE_CASE__						:       Dict           =   pipe_a(
											    prompt_embeds=SCREAMING_SNAKE_CASE__   ,       negative_prompt_embeds=SCREAMING_SNAKE_CASE__   ,       num_inference_steps=2   ,       generator=SCREAMING_SNAKE_CASE__   ,       output_type="""np"""   ,       )

											SCREAMING_SNAKE_CASE__						:       int           =   output.images[0]

											assert image.shape == (64, 64, 3)

											SCREAMING_SNAKE_CASE__						:       Any           =   torch.cuda.max_memory_allocated()
											assert mem_bytes < 13 * 10**9

											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   load_numpy(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy"""    )
											assert_mean_pixel_difference(SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    )

											# pipeline 2

											_start_torch_memory_measurement()

											SCREAMING_SNAKE_CASE__						:       int           =   torch.Generator(device="""cpu"""    ).manual_seed(0    )

											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   floats_tensor((1, 3, 64, 64)   ,       rng=random.Random(0    )    ).to(SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   pipe_a(
											    prompt_embeds=SCREAMING_SNAKE_CASE__   ,       negative_prompt_embeds=SCREAMING_SNAKE_CASE__   ,       image=SCREAMING_SNAKE_CASE__   ,       generator=SCREAMING_SNAKE_CASE__   ,       num_inference_steps=2   ,       output_type="""np"""   ,       )

											SCREAMING_SNAKE_CASE__						:       List[str]           =   output.images[0]

											assert image.shape == (2_56, 2_56, 3)

											SCREAMING_SNAKE_CASE__						:       List[Any]           =   torch.cuda.max_memory_allocated()
											assert mem_bytes < 4 * 10**9

											SCREAMING_SNAKE_CASE__						:       int           =   load_numpy(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy"""    )
											assert_mean_pixel_difference(SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    )



				def 	__magic_name__    (self   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    )					->		str:



											"""simple docstring"""




											_start_torch_memory_measurement()

											SCREAMING_SNAKE_CASE__						:       int           =   floats_tensor((1, 3, 64, 64)   ,       rng=random.Random(0    )    ).to(SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       Tuple           =   torch.Generator(device="""cpu"""    ).manual_seed(0    )

											SCREAMING_SNAKE_CASE__						:       Any           =   pipe_a(
											    prompt_embeds=SCREAMING_SNAKE_CASE__   ,       negative_prompt_embeds=SCREAMING_SNAKE_CASE__   ,       image=SCREAMING_SNAKE_CASE__   ,       num_inference_steps=2   ,       generator=SCREAMING_SNAKE_CASE__   ,       output_type="""np"""   ,       )

											SCREAMING_SNAKE_CASE__						:       Dict           =   output.images[0]

											assert image.shape == (64, 64, 3)

											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   torch.cuda.max_memory_allocated()
											assert mem_bytes < 10 * 10**9

											SCREAMING_SNAKE_CASE__						:       str           =   load_numpy(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy"""    )
											assert_mean_pixel_difference(SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    )

											# pipeline 2

											_start_torch_memory_measurement()

											SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   torch.Generator(device="""cpu"""    ).manual_seed(0    )

											SCREAMING_SNAKE_CASE__						:       List[Any]           =   floats_tensor((1, 3, 2_56, 2_56)   ,       rng=random.Random(0    )    ).to(SCREAMING_SNAKE_CASE__    )
											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   floats_tensor((1, 3, 64, 64)   ,       rng=random.Random(0    )    ).to(SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   pipe_a(
											    prompt_embeds=SCREAMING_SNAKE_CASE__   ,       negative_prompt_embeds=SCREAMING_SNAKE_CASE__   ,       image=SCREAMING_SNAKE_CASE__   ,       original_image=SCREAMING_SNAKE_CASE__   ,       generator=SCREAMING_SNAKE_CASE__   ,       num_inference_steps=2   ,       output_type="""np"""   ,       )

											SCREAMING_SNAKE_CASE__						:       List[Any]           =   output.images[0]

											assert image.shape == (2_56, 2_56, 3)

											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   torch.cuda.max_memory_allocated()
											assert mem_bytes < 4 * 10**9

											SCREAMING_SNAKE_CASE__						:       str           =   load_numpy(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy"""    )
											assert_mean_pixel_difference(SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    )



				def 	__magic_name__    (self   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    )					->		Union[str, Any]:



											"""simple docstring"""




											_start_torch_memory_measurement()

											SCREAMING_SNAKE_CASE__						:       Dict           =   floats_tensor((1, 3, 64, 64)   ,       rng=random.Random(0    )    ).to(SCREAMING_SNAKE_CASE__    )
											SCREAMING_SNAKE_CASE__						:       Dict           =   floats_tensor((1, 3, 64, 64)   ,       rng=random.Random(1    )    ).to(SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       Dict           =   torch.Generator(device="""cpu"""    ).manual_seed(0    )
											SCREAMING_SNAKE_CASE__						:       List[str]           =   pipe_a(
											    prompt_embeds=SCREAMING_SNAKE_CASE__   ,       negative_prompt_embeds=SCREAMING_SNAKE_CASE__   ,       image=SCREAMING_SNAKE_CASE__   ,       mask_image=SCREAMING_SNAKE_CASE__   ,       num_inference_steps=2   ,       generator=SCREAMING_SNAKE_CASE__   ,       output_type="""np"""   ,       )

											SCREAMING_SNAKE_CASE__						:       List[Any]           =   output.images[0]

											assert image.shape == (64, 64, 3)

											SCREAMING_SNAKE_CASE__						:       Tuple           =   torch.cuda.max_memory_allocated()
											assert mem_bytes < 10 * 10**9

											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   load_numpy(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy"""    )
											assert_mean_pixel_difference(SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    )

											# pipeline 2

											_start_torch_memory_measurement()

											SCREAMING_SNAKE_CASE__						:       int           =   torch.Generator(device="""cpu"""    ).manual_seed(0    )

											SCREAMING_SNAKE_CASE__						:       Dict           =   floats_tensor((1, 3, 64, 64)   ,       rng=random.Random(0    )    ).to(SCREAMING_SNAKE_CASE__    )
											SCREAMING_SNAKE_CASE__						:       int           =   floats_tensor((1, 3, 2_56, 2_56)   ,       rng=random.Random(0    )    ).to(SCREAMING_SNAKE_CASE__    )
											SCREAMING_SNAKE_CASE__						:       int           =   floats_tensor((1, 3, 2_56, 2_56)   ,       rng=random.Random(1    )    ).to(SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       int           =   pipe_a(
											    prompt_embeds=SCREAMING_SNAKE_CASE__   ,       negative_prompt_embeds=SCREAMING_SNAKE_CASE__   ,       image=SCREAMING_SNAKE_CASE__   ,       mask_image=SCREAMING_SNAKE_CASE__   ,       original_image=SCREAMING_SNAKE_CASE__   ,       generator=SCREAMING_SNAKE_CASE__   ,       num_inference_steps=2   ,       output_type="""np"""   ,       )

											SCREAMING_SNAKE_CASE__						:       Dict           =   output.images[0]

											assert image.shape == (2_56, 2_56, 3)

											SCREAMING_SNAKE_CASE__						:       List[str]           =   torch.cuda.max_memory_allocated()
											assert mem_bytes < 4 * 10**9

											SCREAMING_SNAKE_CASE__						:       Any           =   load_numpy(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy"""    )
											assert_mean_pixel_difference(SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    )






def lowercase_  (   ):
							torch.cuda.empty_cache()
							torch.cuda.reset_max_memory_allocated()
							torch.cuda.reset_peak_memory_stats()




