

"""simple docstring"""




import argparse
import json

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

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)



def 	UpperCAmelCase__   (				A__	,    A__=False   )      ->      Tuple:




   """simple docstring"""
   lowerCamelCase__						=			[]
   for i in range(config.num_hidden_layers   ):
      # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
      rename_keys.append((f'module.blocks.{i}.norm1.weight', f'vit.encoder.layer.{i}.layernorm_before.weight')   )
      rename_keys.append((f'module.blocks.{i}.norm1.bias', f'vit.encoder.layer.{i}.layernorm_before.bias')   )
      rename_keys.append(
          (f'module.blocks.{i}.attn.proj.weight', f'vit.encoder.layer.{i}.attention.output.dense.weight')   )
      rename_keys.append((f'module.blocks.{i}.attn.proj.bias', f'vit.encoder.layer.{i}.attention.output.dense.bias')   )
      rename_keys.append((f'module.blocks.{i}.norm2.weight', f'vit.encoder.layer.{i}.layernorm_after.weight')   )
      rename_keys.append((f'module.blocks.{i}.norm2.bias', f'vit.encoder.layer.{i}.layernorm_after.bias')   )
      rename_keys.append((f'module.blocks.{i}.mlp.fc1.weight', f'vit.encoder.layer.{i}.intermediate.dense.weight')   )
      rename_keys.append((f'module.blocks.{i}.mlp.fc1.bias', f'vit.encoder.layer.{i}.intermediate.dense.bias')   )
      rename_keys.append((f'module.blocks.{i}.mlp.fc2.weight', f'vit.encoder.layer.{i}.output.dense.weight')   )
      rename_keys.append((f'module.blocks.{i}.mlp.fc2.bias', f'vit.encoder.layer.{i}.output.dense.bias')   )

   # projection layer + position embeddings
   rename_keys.extend(
       [
           ("module.cls_token", "vit.embeddings.cls_token"),
           ("module.patch_embed.proj.weight", "vit.embeddings.patch_embeddings.projection.weight"),
           ("module.patch_embed.proj.bias", "vit.embeddings.patch_embeddings.projection.bias"),
           ("module.pos_embed", "vit.embeddings.position_embeddings"),
       ]   )

   if base_model:
      # layernorm + pooler
      rename_keys.extend(
          [
              ("module.norm.weight", "layernorm.weight"),
              ("module.norm.bias", "layernorm.bias"),
          ]   )

      # if just the base model, we should remove "vit" from all keys that start with "vit"
      lowerCamelCase__						=			[(pair[0], pair[1][4:]) if pair[1].startswith("vit"   ) else pair for pair in rename_keys]
   else:
      # layernorm + classification head
      rename_keys.extend(
          [
              ("norm.weight", "vit.layernorm.weight"),
              ("norm.bias", "vit.layernorm.bias"),
              ("head.weight", "classifier.weight"),
              ("head.bias", "classifier.bias"),
          ]   )

   return rename_keys



def 	UpperCAmelCase__   (				A__	,    A__	,    A__=False   )      ->      int:




   """simple docstring"""
   for i in range(config.num_hidden_layers   ):
      if base_model:
         lowerCamelCase__						=			""
      else:
         lowerCamelCase__						=			"vit."
      # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
      lowerCamelCase__						=			state_dict.pop(f'module.blocks.{i}.attn.qkv.weight'   )
      lowerCamelCase__						=			state_dict.pop(f'module.blocks.{i}.attn.qkv.bias'   )
      # next, add query, keys and values (in that order) to the state dict
      lowerCamelCase__						=			in_proj_weight[
          : config.hidden_size, :
      ]
      lowerCamelCase__						=			in_proj_bias[: config.hidden_size]
      lowerCamelCase__						=			in_proj_weight[
          config.hidden_size : config.hidden_size * 2, :
      ]
      lowerCamelCase__						=			in_proj_bias[
          config.hidden_size : config.hidden_size * 2
      ]
      lowerCamelCase__						=			in_proj_weight[
          -config.hidden_size :, :
      ]
      lowerCamelCase__						=			in_proj_bias[-config.hidden_size :]



def 	UpperCAmelCase__   (				A__   )      ->      List[str]:




   """simple docstring"""
   lowerCamelCase__						=			["head.weight", "head.bias"]
   for k in ignore_keys:
      state_dict.pop(_SCREAMING_SNAKE_CASE	,    _SCREAMING_SNAKE_CASE   )



def 	UpperCAmelCase__   (				A__   )      ->      List[Any]:




   """simple docstring"""
   # projection head is used in the self-supervised pre-training in MSN,
   # for downstream task it's not needed.
   lowerCamelCase__						=			[
       "module.fc.fc1.weight",
       "module.fc.fc1.bias",
       "module.fc.bn1.weight",
       "module.fc.bn1.bias",
       "module.fc.bn1.running_mean",
       "module.fc.bn1.running_var",
       "module.fc.bn1.num_batches_tracked",
       "module.fc.fc2.weight",
       "module.fc.fc2.bias",
       "module.fc.bn2.weight",
       "module.fc.bn2.bias",
       "module.fc.bn2.running_mean",
       "module.fc.bn2.running_var",
       "module.fc.bn2.num_batches_tracked",
       "module.fc.fc3.weight",
       "module.fc.fc3.bias",
   ]
   for k in ignore_keys:
      state_dict.pop(_SCREAMING_SNAKE_CASE	,    _SCREAMING_SNAKE_CASE   )



def 	UpperCAmelCase__   (				A__	,    A__	,    A__   )      ->      Optional[int]:




   """simple docstring"""
   lowerCamelCase__						=			dct.pop(_SCREAMING_SNAKE_CASE   )
   lowerCamelCase__						=			val



def 	UpperCAmelCase__   (				A__	,    A__   )      ->      List[str]:




   """simple docstring"""
   lowerCamelCase__						=			ViTMSNConfig()
   lowerCamelCase__						=			1000

   lowerCamelCase__						=			"datasets/huggingface/label-files"
   lowerCamelCase__						=			"imagenet-1k-id2label.json"
   lowerCamelCase__						=			json.load(open(hf_hub_download(_SCREAMING_SNAKE_CASE	,    _SCREAMING_SNAKE_CASE   )	,    "r"   )   )
   lowerCamelCase__						=			{int(_SCREAMING_SNAKE_CASE   ): v for k, v in idalabel.items()}
   lowerCamelCase__						=			idalabel
   lowerCamelCase__						=			{v: k for k, v in idalabel.items()}

   if "s16" in checkpoint_url:
      lowerCamelCase__						=			384
      lowerCamelCase__						=			1536
      lowerCamelCase__						=			6
   elif "l16" in checkpoint_url:
      lowerCamelCase__						=			1024
      lowerCamelCase__						=			4096
      lowerCamelCase__						=			24
      lowerCamelCase__						=			16
      lowerCamelCase__						=			0.1
   elif "b4" in checkpoint_url:
      lowerCamelCase__						=			4
   elif "l7" in checkpoint_url:
      lowerCamelCase__						=			7
      lowerCamelCase__						=			1024
      lowerCamelCase__						=			4096
      lowerCamelCase__						=			24
      lowerCamelCase__						=			16
      lowerCamelCase__						=			0.1

   lowerCamelCase__						=			ViTMSNModel(_SCREAMING_SNAKE_CASE   )

   lowerCamelCase__						=			torch.hub.load_state_dict_from_url(_SCREAMING_SNAKE_CASE	,    map_location="cpu"   )["target_encoder"]

   lowerCamelCase__						=			ViTImageProcessor(size=config.image_size   )

   remove_projection_head(_SCREAMING_SNAKE_CASE   )
   lowerCamelCase__						=			create_rename_keys(_SCREAMING_SNAKE_CASE	,    base_model=_SCREAMING_SNAKE_CASE   )

   for src, dest in rename_keys:
      rename_key(_SCREAMING_SNAKE_CASE	,    _SCREAMING_SNAKE_CASE	,    _SCREAMING_SNAKE_CASE   )
   read_in_q_k_v(_SCREAMING_SNAKE_CASE	,    _SCREAMING_SNAKE_CASE	,    base_model=_SCREAMING_SNAKE_CASE   )

   model.load_state_dict(_SCREAMING_SNAKE_CASE   )
   model.eval()

   lowerCamelCase__						=			"http://images.cocodataset.org/val2017/000000039769.jpg"

   lowerCamelCase__						=			Image.open(requests.get(_SCREAMING_SNAKE_CASE	,    stream=_SCREAMING_SNAKE_CASE   ).raw   )
   lowerCamelCase__						=			ViTImageProcessor(
       size=config.image_size	,    image_mean=_SCREAMING_SNAKE_CASE	,    image_std=_SCREAMING_SNAKE_CASE   )
   lowerCamelCase__						=			image_processor(images=_SCREAMING_SNAKE_CASE	,    return_tensors="pt"   )

   # forward pass
   torch.manual_seed(2   )
   lowerCamelCase__						=			model(**_SCREAMING_SNAKE_CASE   )
   lowerCamelCase__						=			outputs.last_hidden_state

   # The following Colab Notebook was used to generate these outputs:
   # https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
   if "s16" in checkpoint_url:
      lowerCamelCase__						=			torch.tensor([[-1.0915, -1.4876, -1.1809]]   )
   elif "b16" in checkpoint_url:
      lowerCamelCase__						=			torch.tensor([[14.2889, -18.9045, 11.7281]]   )
   elif "l16" in checkpoint_url:
      lowerCamelCase__						=			torch.tensor([[41.5028, -22.8681, 45.6475]]   )
   elif "b4" in checkpoint_url:
      lowerCamelCase__						=			torch.tensor([[-4.3868, 5.2932, -0.4137]]   )
   else:
      lowerCamelCase__						=			torch.tensor([[-0.1792, -0.6465, 2.4263]]   )

   # verify logits
   assert torch.allclose(last_hidden_state[:, 0, :3]	,    _SCREAMING_SNAKE_CASE	,    atol=1E-4   )

   print(f'Saving model 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 __name__ == "__main__":
    SCREAMING_SNAKE_CASE_     :			Any	  =			argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '''--checkpoint_url''',
        default='''https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar''',
        type=str,
        help='''URL of the checkpoint you\'d like to convert.''',
    )
    parser.add_argument(
        '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
    )

    SCREAMING_SNAKE_CASE_     :			Any	  =			parser.parse_args()
    convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)







"""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					):
      __a								=	IFPipeline
      __a								=	TEXT_TO_IMAGE_PARAMS - {'width', 'height', 'latents'}
      __a								=	TEXT_TO_IMAGE_BATCH_PARAMS
      __a								=	PipelineTesterMixin.required_optional_params - {'latents'}





      def 					_lowerCamelCase  (  self		)   ->					Tuple:
         return self._get_dummy_components()





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__=0		)   ->					Union[str, Any]:
         if str(SCREAMING_SNAKE_CASE__		).startswith("mps"		):
            lowerCamelCase__						=			torch.manual_seed(SCREAMING_SNAKE_CASE__		)
         else:
            lowerCamelCase__						=			torch.Generator(device=SCREAMING_SNAKE_CASE__		).manual_seed(SCREAMING_SNAKE_CASE__		)

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

         return inputs





      def 					_lowerCamelCase  (  self		)   ->					str:
         self._test_save_load_optional_components()





      @unittest.skipIf(torch_device != "cuda"    ,						reason="float16 requires CUDA"		)
      def 					_lowerCamelCase  (  self		)   ->					Optional[int]:
         # 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		)   ->					Union[str, Any]:
         self._test_attention_slicing_forward_pass(expected_max_diff=1e-2		)





      def 					_lowerCamelCase  (  self		)   ->					int:
         self._test_save_load_local()





      def 					_lowerCamelCase  (  self		)   ->					List[Any]:
         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		)   ->					List[str]:
         self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3		)






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





      def 					_lowerCamelCase  (  self		)   ->					Union[str, Any]:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()





      def 					_lowerCamelCase  (  self		)   ->					str:
         # if

         lowerCamelCase__						=			IFPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0"    ,						variant="fp16"    ,						torch_dtype=torch.floataa		)

         lowerCamelCase__						=			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"		)

         lowerCamelCase__					,       lowerCamelCase__						=			pipe_a.encode_prompt("anime turtle"    ,						device="cuda"		)

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

         lowerCamelCase__						=			None
         lowerCamelCase__						=			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

         lowerCamelCase__						=			IFImgaImgPipeline(**pipe_a.components		)
         lowerCamelCase__						=			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

         lowerCamelCase__						=			IFInpaintingPipeline(**pipe_a.components		)
         lowerCamelCase__						=			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 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					int:
         # pipeline 1

         _start_torch_memory_measurement()

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

         lowerCamelCase__						=			output.images[0]

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

         lowerCamelCase__						=			torch.cuda.max_memory_allocated()
         assert mem_bytes < 13 * 10**9

         lowerCamelCase__						=			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()

         lowerCamelCase__						=			torch.Generator(device="cpu"		).manual_seed(0		)

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

         lowerCamelCase__						=			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"    ,						)

         lowerCamelCase__						=			output.images[0]

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

         lowerCamelCase__						=			torch.cuda.max_memory_allocated()
         assert mem_bytes < 4 * 10**9

         lowerCamelCase__						=			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 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					Tuple:
         # pipeline 1

         _start_torch_memory_measurement()

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

         lowerCamelCase__						=			torch.Generator(device="cpu"		).manual_seed(0		)

         lowerCamelCase__						=			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"    ,						)

         lowerCamelCase__						=			output.images[0]

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

         lowerCamelCase__						=			torch.cuda.max_memory_allocated()
         assert mem_bytes < 10 * 10**9

         lowerCamelCase__						=			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()

         lowerCamelCase__						=			torch.Generator(device="cpu"		).manual_seed(0		)

         lowerCamelCase__						=			floats_tensor((1, 3, 256, 256)    ,						rng=random.Random(0		)		).to(SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			floats_tensor((1, 3, 64, 64)    ,						rng=random.Random(0		)		).to(SCREAMING_SNAKE_CASE__		)

         lowerCamelCase__						=			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"    ,						)

         lowerCamelCase__						=			output.images[0]

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

         lowerCamelCase__						=			torch.cuda.max_memory_allocated()
         assert mem_bytes < 4 * 10**9

         lowerCamelCase__						=			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 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					List[Any]:
         # pipeline 1

         _start_torch_memory_measurement()

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

         lowerCamelCase__						=			torch.Generator(device="cpu"		).manual_seed(0		)
         lowerCamelCase__						=			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"    ,						)

         lowerCamelCase__						=			output.images[0]

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

         lowerCamelCase__						=			torch.cuda.max_memory_allocated()
         assert mem_bytes < 10 * 10**9

         lowerCamelCase__						=			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()

         lowerCamelCase__						=			torch.Generator(device="cpu"		).manual_seed(0		)

         lowerCamelCase__						=			floats_tensor((1, 3, 64, 64)    ,						rng=random.Random(0		)		).to(SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			floats_tensor((1, 3, 256, 256)    ,						rng=random.Random(0		)		).to(SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			floats_tensor((1, 3, 256, 256)    ,						rng=random.Random(1		)		).to(SCREAMING_SNAKE_CASE__		)

         lowerCamelCase__						=			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"    ,						)

         lowerCamelCase__						=			output.images[0]

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

         lowerCamelCase__						=			torch.cuda.max_memory_allocated()
         assert mem_bytes < 4 * 10**9

         lowerCamelCase__						=			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 	UpperCAmelCase__   (				)      ->      Union[str, Any]:




   """simple docstring"""
   torch.cuda.empty_cache()
   torch.cuda.reset_max_memory_allocated()
   torch.cuda.reset_peak_memory_stats()




