


'''simple docstring'''



import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileViTConfig,
    MobileViTForImageClassification,
    MobileViTForSemanticSegmentation,
    MobileViTImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCAmelCase      :Any			    =       logging.get_logger(__name__)




def      _a     (						_lowercase :     int  ):





     '''simple docstring'''






     __UpperCAmelCase     :			List[Any]  							=      MobileViTConfig()

     # size of the architecture
     if "mobilevit_s" in mobilevit_name:
          __UpperCAmelCase     :			Union[str, Any]  							=      [144, 192, 240]
          __UpperCAmelCase     :			Tuple  							=      [16, 32, 64, 96, 128, 160, 640]
     elif "mobilevit_xs" in mobilevit_name:
          __UpperCAmelCase     :			Any  							=      [96, 120, 144]
          __UpperCAmelCase     :			str  							=      [16, 32, 48, 64, 80, 96, 384]
     elif "mobilevit_xxs" in mobilevit_name:
          __UpperCAmelCase     :			str  							=      [64, 80, 96]
          __UpperCAmelCase     :			int  							=      [16, 16, 24, 48, 64, 80, 320]
          __UpperCAmelCase     :			int  							=      0.05
          __UpperCAmelCase     :			List[Any]  							=      2.0

     if mobilevit_name.startswith('''deeplabv3_'''  ):
          __UpperCAmelCase     :			Optional[int]  							=      512
          __UpperCAmelCase     :			Any  							=      16
          __UpperCAmelCase     :			str  							=      21
          __UpperCAmelCase     :			Dict  							=      '''pascal-voc-id2label.json'''
     else:
          __UpperCAmelCase     :			str  							=      1000
          __UpperCAmelCase     :			List[Any]  							=      '''imagenet-1k-id2label.json'''

     __UpperCAmelCase     :			Optional[int]  							=      '''huggingface/label-files'''
     __UpperCAmelCase     :			Any  							=      json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE_      , SCREAMING_SNAKE_CASE_      , repo_type='''dataset'''  )      , '''r'''  )  )
     __UpperCAmelCase     :			List[Any]  							=      {int(SCREAMING_SNAKE_CASE_  ): v for k, v in idalabel.items()}
     __UpperCAmelCase     :			Tuple  							=      idalabel
     __UpperCAmelCase     :			Union[str, Any]  							=      {v: k for k, v in idalabel.items()}

     return config




def      _a     (						_lowercase :     List[Any]      , _lowercase :     List[str]=False  ):





     '''simple docstring'''






     for i in range(1      , 6  ):
          if F'layer_{i}.' in name:
               __UpperCAmelCase     :			Dict  							=      name.replace(F'layer_{i}.'      , F'encoder.layer.{i - 1}.'  )

     if "conv_1." in name:
          __UpperCAmelCase     :			Tuple  							=      name.replace('''conv_1.'''      , '''conv_stem.'''  )
     if ".block." in name:
          __UpperCAmelCase     :			Optional[Any]  							=      name.replace('''.block.'''      , '''.'''  )

     if "exp_1x1" in name:
          __UpperCAmelCase     :			Optional[Any]  							=      name.replace('''exp_1x1'''      , '''expand_1x1'''  )
     if "red_1x1" in name:
          __UpperCAmelCase     :			Tuple  							=      name.replace('''red_1x1'''      , '''reduce_1x1'''  )
     if ".local_rep.conv_3x3." in name:
          __UpperCAmelCase     :			List[str]  							=      name.replace('''.local_rep.conv_3x3.'''      , '''.conv_kxk.'''  )
     if ".local_rep.conv_1x1." in name:
          __UpperCAmelCase     :			Union[str, Any]  							=      name.replace('''.local_rep.conv_1x1.'''      , '''.conv_1x1.'''  )
     if ".norm." in name:
          __UpperCAmelCase     :			List[Any]  							=      name.replace('''.norm.'''      , '''.normalization.'''  )
     if ".conv." in name:
          __UpperCAmelCase     :			int  							=      name.replace('''.conv.'''      , '''.convolution.'''  )
     if ".conv_proj." in name:
          __UpperCAmelCase     :			List[str]  							=      name.replace('''.conv_proj.'''      , '''.conv_projection.'''  )

     for i in range(0      , 2  ):
          for j in range(0      , 4  ):
               if F'.{i}.{j}.' in name:
                    __UpperCAmelCase     :			List[str]  							=      name.replace(F'.{i}.{j}.'      , F'.{i}.layer.{j}.'  )

     for i in range(2      , 6  ):
          for j in range(0      , 4  ):
               if F'.{i}.{j}.' in name:
                    __UpperCAmelCase     :			str  							=      name.replace(F'.{i}.{j}.'      , F'.{i}.'  )
                    if "expand_1x1" in name:
                         __UpperCAmelCase     :			List[str]  							=      name.replace('''expand_1x1'''      , '''downsampling_layer.expand_1x1'''  )
                    if "conv_3x3" in name:
                         __UpperCAmelCase     :			Optional[Any]  							=      name.replace('''conv_3x3'''      , '''downsampling_layer.conv_3x3'''  )
                    if "reduce_1x1" in name:
                         __UpperCAmelCase     :			Dict  							=      name.replace('''reduce_1x1'''      , '''downsampling_layer.reduce_1x1'''  )

     for i in range(2      , 5  ):
          if F'.global_rep.{i}.weight' in name:
               __UpperCAmelCase     :			Any  							=      name.replace(F'.global_rep.{i}.weight'      , '''.layernorm.weight'''  )
          if F'.global_rep.{i}.bias' in name:
               __UpperCAmelCase     :			str  							=      name.replace(F'.global_rep.{i}.bias'      , '''.layernorm.bias'''  )

     if ".global_rep." in name:
          __UpperCAmelCase     :			List[Any]  							=      name.replace('''.global_rep.'''      , '''.transformer.'''  )
     if ".pre_norm_mha.0." in name:
          __UpperCAmelCase     :			Tuple  							=      name.replace('''.pre_norm_mha.0.'''      , '''.layernorm_before.'''  )
     if ".pre_norm_mha.1.out_proj." in name:
          __UpperCAmelCase     :			List[Any]  							=      name.replace('''.pre_norm_mha.1.out_proj.'''      , '''.attention.output.dense.'''  )
     if ".pre_norm_ffn.0." in name:
          __UpperCAmelCase     :			str  							=      name.replace('''.pre_norm_ffn.0.'''      , '''.layernorm_after.'''  )
     if ".pre_norm_ffn.1." in name:
          __UpperCAmelCase     :			List[Any]  							=      name.replace('''.pre_norm_ffn.1.'''      , '''.intermediate.dense.'''  )
     if ".pre_norm_ffn.4." in name:
          __UpperCAmelCase     :			Dict  							=      name.replace('''.pre_norm_ffn.4.'''      , '''.output.dense.'''  )
     if ".transformer." in name:
          __UpperCAmelCase     :			List[str]  							=      name.replace('''.transformer.'''      , '''.transformer.layer.'''  )

     if ".aspp_layer." in name:
          __UpperCAmelCase     :			Tuple  							=      name.replace('''.aspp_layer.'''      , '''.'''  )
     if ".aspp_pool." in name:
          __UpperCAmelCase     :			Any  							=      name.replace('''.aspp_pool.'''      , '''.'''  )
     if "seg_head." in name:
          __UpperCAmelCase     :			Optional[int]  							=      name.replace('''seg_head.'''      , '''segmentation_head.'''  )
     if "segmentation_head.classifier.classifier." in name:
          __UpperCAmelCase     :			Any  							=      name.replace('''segmentation_head.classifier.classifier.'''      , '''segmentation_head.classifier.'''  )

     if "classifier.fc." in name:
          __UpperCAmelCase     :			List[Any]  							=      name.replace('''classifier.fc.'''      , '''classifier.'''  )
     elif (not base_model) and ("segmentation_head." not in name):
          __UpperCAmelCase     :			int  							=      '''mobilevit.''' + name

     return name




def      _a     (						_lowercase :     List[str]      , _lowercase :     Optional[Any]      , _lowercase :     List[str]=False  ):





     '''simple docstring'''






     if base_model:
          __UpperCAmelCase     :			Optional[Any]  							=      ''''''
     else:
          __UpperCAmelCase     :			Optional[Any]  							=      '''mobilevit.'''

     for key in orig_state_dict.copy().keys():
          __UpperCAmelCase     :			Dict  							=      orig_state_dict.pop(SCREAMING_SNAKE_CASE_  )

          if key[:8] == "encoder.":
               __UpperCAmelCase     :			Tuple  							=      key[8:]

          if "qkv" in key:
               __UpperCAmelCase     :			List[Any]  							=      key.split('''.'''  )
               __UpperCAmelCase     :			str  							=      int(key_split[0][6:]  ) - 1
               __UpperCAmelCase     :			Union[str, Any]  							=      int(key_split[3]  )
               __UpperCAmelCase     :			Dict  							=      model.get_submodule(F'{model_prefix}encoder.layer.{layer_num}'  )
               __UpperCAmelCase     :			Tuple  							=      layer.transformer.layer[transformer_num].attention.attention.all_head_size
               __UpperCAmelCase     :			List[str]  							=      (
                   F'{model_prefix}encoder.layer.{layer_num}.transformer.layer.{transformer_num}.attention.attention.'
               )
               if "weight" in key:
                    __UpperCAmelCase     :			Tuple  							=      val[:dim, :]
                    __UpperCAmelCase     :			Union[str, Any]  							=      val[dim : dim * 2, :]
                    __UpperCAmelCase     :			List[Any]  							=      val[-dim:, :]
               else:
                    __UpperCAmelCase     :			Optional[Any]  							=      val[:dim]
                    __UpperCAmelCase     :			Any  							=      val[dim : dim * 2]
                    __UpperCAmelCase     :			List[str]  							=      val[-dim:]
          else:
               __UpperCAmelCase     :			Union[str, Any]  							=      val

     return orig_state_dict




def      _a     (						):





     '''simple docstring'''






     __UpperCAmelCase     :			Union[str, Any]  							=      '''http://images.cocodataset.org/val2017/000000039769.jpg'''
     __UpperCAmelCase     :			Dict  							=      Image.open(requests.get(SCREAMING_SNAKE_CASE_      , stream=SCREAMING_SNAKE_CASE_  ).raw  )
     return im




@torch.no_grad()
def      _a     (						_lowercase :     Optional[int]      , _lowercase :     Any      , _lowercase :     Optional[Any]      , _lowercase :     List[Any]=False  ):





     '''simple docstring'''






     __UpperCAmelCase     :			List[str]  							=      get_mobilevit_config(SCREAMING_SNAKE_CASE_  )

     # load original state_dict
     __UpperCAmelCase     :			Dict  							=      torch.load(SCREAMING_SNAKE_CASE_      , map_location='''cpu'''  )

     # load 🤗 model
     if mobilevit_name.startswith('''deeplabv3_'''  ):
          __UpperCAmelCase     :			Tuple  							=      MobileViTForSemanticSegmentation(SCREAMING_SNAKE_CASE_  ).eval()
     else:
          __UpperCAmelCase     :			Dict  							=      MobileViTForImageClassification(SCREAMING_SNAKE_CASE_  ).eval()

     __UpperCAmelCase     :			int  							=      convert_state_dict(SCREAMING_SNAKE_CASE_      , SCREAMING_SNAKE_CASE_  )
     model.load_state_dict(SCREAMING_SNAKE_CASE_  )

     # Check outputs on an image, prepared by MobileViTImageProcessor
     __UpperCAmelCase     :			List[Any]  							=      MobileViTImageProcessor(crop_size=config.image_size      , size=config.image_size + 32  )
     __UpperCAmelCase     :			Dict  							=      image_processor(images=prepare_img()      , return_tensors='''pt'''  )
     __UpperCAmelCase     :			Optional[int]  							=      model(**SCREAMING_SNAKE_CASE_  )
     __UpperCAmelCase     :			int  							=      outputs.logits

     if mobilevit_name.startswith('''deeplabv3_'''  ):
          assert logits.shape == (1, 21, 32, 32)

          if mobilevit_name == "deeplabv3_mobilevit_s":
               __UpperCAmelCase     :			int  							=      torch.tensor(
                   [
                       [[6.2065, 6.1292, 6.2070], [6.1079, 6.1254, 6.1747], [6.0042, 6.1071, 6.1034]],
                       [[-6.9253, -6.8653, -7.0398], [-7.3218, -7.3983, -7.3670], [-7.1961, -7.2482, -7.1569]],
                       [[-4.4723, -4.4348, -4.3769], [-5.3629, -5.4632, -5.4598], [-5.1587, -5.3402, -5.5059]],
                   ]  )
          elif mobilevit_name == "deeplabv3_mobilevit_xs":
               __UpperCAmelCase     :			Tuple  							=      torch.tensor(
                   [
                       [[5.4449, 5.5733, 5.6314], [5.1815, 5.3930, 5.5963], [5.1656, 5.4333, 5.4853]],
                       [[-9.4423, -9.7766, -9.6714], [-9.1581, -9.5720, -9.5519], [-9.1006, -9.6458, -9.5703]],
                       [[-7.7721, -7.3716, -7.1583], [-8.4599, -8.0624, -7.7944], [-8.4172, -7.8366, -7.5025]],
                   ]  )
          elif mobilevit_name == "deeplabv3_mobilevit_xxs":
               __UpperCAmelCase     :			Union[str, Any]  							=      torch.tensor(
                   [
                       [[6.9811, 6.9743, 7.3123], [7.1777, 7.1931, 7.3938], [7.5633, 7.8050, 7.8901]],
                       [[-10.5536, -10.2332, -10.2924], [-10.2336, -9.8624, -9.5964], [-10.8840, -10.8158, -10.6659]],
                       [[-3.4938, -3.0631, -2.8620], [-3.4205, -2.8135, -2.6875], [-3.4179, -2.7945, -2.8750]],
                   ]  )
          else:
               raise ValueError(F'Unknown mobilevit_name: {mobilevit_name}'  )

          assert torch.allclose(logits[0, :3, :3, :3]      , SCREAMING_SNAKE_CASE_      , atol=1E-4  )
     else:
          assert logits.shape == (1, 1000)

          if mobilevit_name == "mobilevit_s":
               __UpperCAmelCase     :			List[Any]  							=      torch.tensor([-0.9866, 0.2392, -1.1241]  )
          elif mobilevit_name == "mobilevit_xs":
               __UpperCAmelCase     :			Optional[Any]  							=      torch.tensor([-2.4761, -0.9399, -1.9587]  )
          elif mobilevit_name == "mobilevit_xxs":
               __UpperCAmelCase     :			str  							=      torch.tensor([-1.9364, -1.2327, -0.4653]  )
          else:
               raise ValueError(F'Unknown mobilevit_name: {mobilevit_name}'  )

          assert torch.allclose(logits[0, :3]      , SCREAMING_SNAKE_CASE_      , atol=1E-4  )

     Path(SCREAMING_SNAKE_CASE_  ).mkdir(exist_ok=SCREAMING_SNAKE_CASE_  )
     print(F'Saving model {mobilevit_name} to {pytorch_dump_folder_path}'  )
     model.save_pretrained(SCREAMING_SNAKE_CASE_  )
     print(F'Saving image processor to {pytorch_dump_folder_path}'  )
     image_processor.save_pretrained(SCREAMING_SNAKE_CASE_  )

     if push_to_hub:
          __UpperCAmelCase     :			int  							=      {
              '''mobilevit_s''': '''mobilevit-small''',
              '''mobilevit_xs''': '''mobilevit-x-small''',
              '''mobilevit_xxs''': '''mobilevit-xx-small''',
              '''deeplabv3_mobilevit_s''': '''deeplabv3-mobilevit-small''',
              '''deeplabv3_mobilevit_xs''': '''deeplabv3-mobilevit-x-small''',
              '''deeplabv3_mobilevit_xxs''': '''deeplabv3-mobilevit-xx-small''',
          }

          print('''Pushing to the hub...'''  )
          __UpperCAmelCase     :			Dict  							=      model_mapping[mobilevit_name]
          image_processor.push_to_hub(SCREAMING_SNAKE_CASE_      , organization='''apple'''  )
          model.push_to_hub(SCREAMING_SNAKE_CASE_      , organization='''apple'''  )


if __name__ == "__main__":
      __UpperCAmelCase      :Optional[Any]			    =       argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          "--mobilevit_name",
          default="mobilevit_s",
          type=str,
          help=(
              "Name of the MobileViT model you\'d like to convert. Should be one of \'mobilevit_s\', \'mobilevit_xs\',"
              " \'mobilevit_xxs\', \'deeplabv3_mobilevit_s\', \'deeplabv3_mobilevit_xs\', \'deeplabv3_mobilevit_xxs\'."
          ),
      )
      parser.add_argument(
          "--checkpoint_path", required=True, type=str, help="Path to the original state dict (.pt file)."
      )
      parser.add_argument(
          "--pytorch_dump_folder_path", required=True, type=str, help="Path to the output PyTorch model directory."
      )
      parser.add_argument(
          "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
      )

      __UpperCAmelCase      :Dict			    =       parser.parse_args()
      convert_movilevit_checkpoint(
          args.mobilevit_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
      )




'''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      a    (		_a     ,  _a     ,  unittest.TestCase	):


     """simple docstring"""






     SCREAMING_SNAKE_CASE	:			Dict							       =       IFPipeline
     SCREAMING_SNAKE_CASE	:			Optional[int]							       =       TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
     SCREAMING_SNAKE_CASE	:			Optional[int]							       =       TEXT_TO_IMAGE_BATCH_PARAMS
     SCREAMING_SNAKE_CASE	:			str							       =       PipelineTesterMixin.required_optional_params - {"latents"}

     def 			lowerCamelCase__			(    self					:		Tuple      )						->							Optional[Any]:
          return self._get_dummy_components()

     def 			lowerCamelCase__			(    self					:		Optional[Any]							,			snake_case					:		Optional[int]							,			snake_case					:		List[Any]=0      )						->							Optional[Any]:
          if str(snake_case      ).startswith('''mps'''      ):
               __UpperCAmelCase     :			Optional[Any]  							=      torch.manual_seed(snake_case      )
          else:
               __UpperCAmelCase     :			Dict  							=      torch.Generator(device=snake_case      ).manual_seed(snake_case      )

          __UpperCAmelCase     :			str  							=      {
              '''prompt''': '''A painting of a squirrel eating a burger''',
              '''generator''': generator,
              '''num_inference_steps''': 2,
              '''output_type''': '''numpy''',
          }

          return inputs

     def 			lowerCamelCase__			(    self					:		List[str]      )						->							Optional[int]:
          self._test_save_load_optional_components()

     @unittest.skipIf(torch_device != '''cuda'''							,			reason='''float16 requires CUDA'''      )
     def 			lowerCamelCase__			(    self					:		Dict      )						->							Union[str, Any]:
          # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
          super().test_save_load_floataa(expected_max_diff=1E-1      )

     def 			lowerCamelCase__			(    self					:		List[str]      )						->							Tuple:
          self._test_attention_slicing_forward_pass(expected_max_diff=1E-2      )

     def 			lowerCamelCase__			(    self					:		Any      )						->							List[str]:
          self._test_save_load_local()

     def 			lowerCamelCase__			(    self					:		Union[str, Any]      )						->							str:
          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 			lowerCamelCase__			(    self					:		str      )						->							Optional[Any]:
          self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3      )



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


     """simple docstring"""






     def 			lowerCamelCase__			(    self					:		Dict      )						->							Dict:
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()

     def 			lowerCamelCase__			(    self					:		Any      )						->							Tuple:
          # if

          __UpperCAmelCase     :			Union[str, Any]  							=      IFPipeline.from_pretrained('''DeepFloyd/IF-I-XL-v1.0'''							,			variant='''fp16'''							,			torch_dtype=torch.floataa      )

          __UpperCAmelCase     :			List[str]  							=      IFSuperResolutionPipeline.from_pretrained(
              '''DeepFloyd/IF-II-L-v1.0'''							,			variant='''fp16'''							,			torch_dtype=torch.floataa							,			text_encoder=snake_case							,			tokenizer=snake_case      )

          # pre compute text embeddings and remove T5 to save memory

          pipe_a.text_encoder.to('''cuda'''      )

          __UpperCAmelCase							,       __UpperCAmelCase     :			Tuple  							=      pipe_a.encode_prompt('''anime turtle'''							,			device='''cuda'''      )

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

          __UpperCAmelCase     :			Any  							=      None
          __UpperCAmelCase     :			Optional[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(snake_case							,			snake_case							,			snake_case							,			snake_case      )

          pipe_a.remove_all_hooks()
          pipe_a.remove_all_hooks()

          # img2img

          __UpperCAmelCase     :			Union[str, Any]  							=      IFImgaImgPipeline(**pipe_a.components      )
          __UpperCAmelCase     :			int  							=      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(snake_case							,			snake_case							,			snake_case							,			snake_case      )

          pipe_a.remove_all_hooks()
          pipe_a.remove_all_hooks()

          # inpainting

          __UpperCAmelCase     :			List[str]  							=      IFInpaintingPipeline(**pipe_a.components      )
          __UpperCAmelCase     :			List[str]  							=      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(snake_case							,			snake_case							,			snake_case							,			snake_case      )

     def 			lowerCamelCase__			(    self					:		List[str]							,			snake_case					:		Any							,			snake_case					:		Union[str, Any]							,			snake_case					:		Optional[int]							,			snake_case					:		str      )						->							Optional[int]:
          # pipeline 1

          _start_torch_memory_measurement()

          __UpperCAmelCase     :			Dict  							=      torch.Generator(device='''cpu'''      ).manual_seed(0      )
          __UpperCAmelCase     :			List[str]  							=      pipe_a(
              prompt_embeds=snake_case							,			negative_prompt_embeds=snake_case							,			num_inference_steps=2							,			generator=snake_case							,			output_type='''np'''							,			)

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

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

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

          __UpperCAmelCase     :			Optional[Any]  							=      load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'''      )
          assert_mean_pixel_difference(snake_case							,			snake_case      )

          # pipeline 2

          _start_torch_memory_measurement()

          __UpperCAmelCase     :			int  							=      torch.Generator(device='''cpu'''      ).manual_seed(0      )

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

          __UpperCAmelCase     :			List[Any]  							=      pipe_a(
              prompt_embeds=snake_case							,			negative_prompt_embeds=snake_case							,			image=snake_case							,			generator=snake_case							,			num_inference_steps=2							,			output_type='''np'''							,			)

          __UpperCAmelCase     :			Optional[Any]  							=      output.images[0]

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

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

          __UpperCAmelCase     :			Tuple  							=      load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'''      )
          assert_mean_pixel_difference(snake_case							,			snake_case      )

     def 			lowerCamelCase__			(    self					:		Optional[int]							,			snake_case					:		Optional[int]							,			snake_case					:		Union[str, Any]							,			snake_case					:		str							,			snake_case					:		Dict      )						->							str:
          # pipeline 1

          _start_torch_memory_measurement()

          __UpperCAmelCase     :			Optional[Any]  							=      floats_tensor((1, 3, 64, 64)							,			rng=random.Random(0      )      ).to(snake_case      )

          __UpperCAmelCase     :			Dict  							=      torch.Generator(device='''cpu'''      ).manual_seed(0      )

          __UpperCAmelCase     :			Dict  							=      pipe_a(
              prompt_embeds=snake_case							,			negative_prompt_embeds=snake_case							,			image=snake_case							,			num_inference_steps=2							,			generator=snake_case							,			output_type='''np'''							,			)

          __UpperCAmelCase     :			int  							=      output.images[0]

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

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

          __UpperCAmelCase     :			Optional[Any]  							=      load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'''      )
          assert_mean_pixel_difference(snake_case							,			snake_case      )

          # pipeline 2

          _start_torch_memory_measurement()

          __UpperCAmelCase     :			Optional[int]  							=      torch.Generator(device='''cpu'''      ).manual_seed(0      )

          __UpperCAmelCase     :			Optional[Any]  							=      floats_tensor((1, 3, 256, 256)							,			rng=random.Random(0      )      ).to(snake_case      )
          __UpperCAmelCase     :			List[str]  							=      floats_tensor((1, 3, 64, 64)							,			rng=random.Random(0      )      ).to(snake_case      )

          __UpperCAmelCase     :			List[str]  							=      pipe_a(
              prompt_embeds=snake_case							,			negative_prompt_embeds=snake_case							,			image=snake_case							,			original_image=snake_case							,			generator=snake_case							,			num_inference_steps=2							,			output_type='''np'''							,			)

          __UpperCAmelCase     :			int  							=      output.images[0]

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

          __UpperCAmelCase     :			int  							=      torch.cuda.max_memory_allocated()
          assert mem_bytes < 4 * 10**9

          __UpperCAmelCase     :			Any  							=      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(snake_case							,			snake_case      )






     def 			lowerCamelCase__			(    self					:		str							,			snake_case					:		Union[str, Any]							,			snake_case					:		Optional[int]							,			snake_case					:		Optional[Any]							,			snake_case					:		Optional[Any]      )						->							Optional[int]:
          # pipeline 1

          _start_torch_memory_measurement()

          __UpperCAmelCase     :			List[Any]  							=      floats_tensor((1, 3, 64, 64)							,			rng=random.Random(0      )      ).to(snake_case      )
          __UpperCAmelCase     :			List[str]  							=      floats_tensor((1, 3, 64, 64)							,			rng=random.Random(1      )      ).to(snake_case      )

          __UpperCAmelCase     :			Optional[Any]  							=      torch.Generator(device='''cpu'''      ).manual_seed(0      )
          __UpperCAmelCase     :			Dict  							=      pipe_a(
              prompt_embeds=snake_case							,			negative_prompt_embeds=snake_case							,			image=snake_case							,			mask_image=snake_case							,			num_inference_steps=2							,			generator=snake_case							,			output_type='''np'''							,			)

          __UpperCAmelCase     :			Union[str, Any]  							=      output.images[0]

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

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

          __UpperCAmelCase     :			Tuple  							=      load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'''      )
          assert_mean_pixel_difference(snake_case							,			snake_case      )

          # pipeline 2

          _start_torch_memory_measurement()

          __UpperCAmelCase     :			int  							=      torch.Generator(device='''cpu'''      ).manual_seed(0      )

          __UpperCAmelCase     :			Tuple  							=      floats_tensor((1, 3, 64, 64)							,			rng=random.Random(0      )      ).to(snake_case      )
          __UpperCAmelCase     :			Any  							=      floats_tensor((1, 3, 256, 256)							,			rng=random.Random(0      )      ).to(snake_case      )
          __UpperCAmelCase     :			Tuple  							=      floats_tensor((1, 3, 256, 256)							,			rng=random.Random(1      )      ).to(snake_case      )

          __UpperCAmelCase     :			Union[str, Any]  							=      pipe_a(
              prompt_embeds=snake_case							,			negative_prompt_embeds=snake_case							,			image=snake_case							,			mask_image=snake_case							,			original_image=snake_case							,			generator=snake_case							,			num_inference_steps=2							,			output_type='''np'''							,			)

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

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

          __UpperCAmelCase     :			Any  							=      torch.cuda.max_memory_allocated()
          assert mem_bytes < 4 * 10**9

          __UpperCAmelCase     :			str  							=      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(snake_case							,			snake_case      )




def      _a     (						):





     '''simple docstring'''






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