
"""simple docstring"""




import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_url
from PIL import Image

from transformers import DPTConfig, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase__      :    List[Any]    							= logging.get_logger(__name__)

def lowercase_  (   _snake_case						):
							SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   DPTConfig()

							if "large" in checkpoint_url:
														SCREAMING_SNAKE_CASE__						:       Any           =   1_024
														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   4_096
														SCREAMING_SNAKE_CASE__						:       Tuple           =   24
														SCREAMING_SNAKE_CASE__						:       Any           =   16
														SCREAMING_SNAKE_CASE__						:       Any           =   [5, 11, 17, 23]
														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   [256, 512, 1_024, 1_024]
														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   (1, 384, 384)

							if "ade" in checkpoint_url:
														SCREAMING_SNAKE_CASE__						:       Optional[int]           =   True

														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   150
														SCREAMING_SNAKE_CASE__						:       int           =   """huggingface/label-files"""
														SCREAMING_SNAKE_CASE__						:       int           =   """ade20k-id2label.json"""
														SCREAMING_SNAKE_CASE__						:       Tuple           =   json.load(open(cached_download(hf_hub_url(_snake_case   ,_snake_case   ,repo_type="""dataset"""						)						)   ,"""r"""						)						)
														SCREAMING_SNAKE_CASE__						:       List[Any]           =   {int(_snake_case						): v for k, v in idalabel.items()}
														SCREAMING_SNAKE_CASE__						:       Optional[int]           =   idalabel
														SCREAMING_SNAKE_CASE__						:       List[Any]           =   {v: k for k, v in idalabel.items()}
														SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   [1, 150, 480, 480]

							return config, expected_shape

def lowercase_  (   _snake_case						):
							SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   ["""pretrained.model.head.weight""", """pretrained.model.head.bias"""]
							for k in ignore_keys:
														state_dict.pop(_snake_case   ,_snake_case						)

def lowercase_  (   _snake_case						):
							if (
							    "pretrained.model" in name
							    and "cls_token" not in name
							    and "pos_embed" not in name
							    and "patch_embed" not in name
							):
														SCREAMING_SNAKE_CASE__						:       List[Any]           =   name.replace("""pretrained.model"""   ,"""dpt.encoder"""						)
							if "pretrained.model" in name:
														SCREAMING_SNAKE_CASE__						:       List[Any]           =   name.replace("""pretrained.model"""   ,"""dpt.embeddings"""						)
							if "patch_embed" in name:
														SCREAMING_SNAKE_CASE__						:       Optional[int]           =   name.replace("""patch_embed"""   ,"""patch_embeddings"""						)
							if "pos_embed" in name:
														SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   name.replace("""pos_embed"""   ,"""position_embeddings"""						)
							if "attn.proj" in name:
														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   name.replace("""attn.proj"""   ,"""attention.output.dense"""						)
							if "proj" in name and "project" not in name:
														SCREAMING_SNAKE_CASE__						:       int           =   name.replace("""proj"""   ,"""projection"""						)
							if "blocks" in name:
														SCREAMING_SNAKE_CASE__						:       List[str]           =   name.replace("""blocks"""   ,"""layer"""						)
							if "mlp.fc1" in name:
														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   name.replace("""mlp.fc1"""   ,"""intermediate.dense"""						)
							if "mlp.fc2" in name:
														SCREAMING_SNAKE_CASE__						:       Any           =   name.replace("""mlp.fc2"""   ,"""output.dense"""						)
							if "norm1" in name:
														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   name.replace("""norm1"""   ,"""layernorm_before"""						)
							if "norm2" in name:
														SCREAMING_SNAKE_CASE__						:       str           =   name.replace("""norm2"""   ,"""layernorm_after"""						)
							if "scratch.output_conv" in name:
														SCREAMING_SNAKE_CASE__						:       List[str]           =   name.replace("""scratch.output_conv"""   ,"""head"""						)
							if "scratch" in name:
														SCREAMING_SNAKE_CASE__						:       List[Any]           =   name.replace("""scratch"""   ,"""neck"""						)
							if "layer1_rn" in name:
														SCREAMING_SNAKE_CASE__						:       List[Any]           =   name.replace("""layer1_rn"""   ,"""convs.0"""						)
							if "layer2_rn" in name:
														SCREAMING_SNAKE_CASE__						:       List[str]           =   name.replace("""layer2_rn"""   ,"""convs.1"""						)
							if "layer3_rn" in name:
														SCREAMING_SNAKE_CASE__						:       Any           =   name.replace("""layer3_rn"""   ,"""convs.2"""						)
							if "layer4_rn" in name:
														SCREAMING_SNAKE_CASE__						:       List[str]           =   name.replace("""layer4_rn"""   ,"""convs.3"""						)
							if "refinenet" in name:
														SCREAMING_SNAKE_CASE__						:       List[str]           =   int(name[len("""neck.refinenet"""						) : len("""neck.refinenet"""						) + 1]						)
														# tricky here: we need to map 4 to 0, 3 to 1, 2 to 2 and 1 to 3
														SCREAMING_SNAKE_CASE__						:       List[Any]           =   name.replace(f'''refinenet{layer_idx}'''   ,f'''fusion_stage.layers.{abs(layer_idx-4						)}'''						)
							if "out_conv" in name:
														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   name.replace("""out_conv"""   ,"""projection"""						)
							if "resConfUnit1" in name:
														SCREAMING_SNAKE_CASE__						:       List[str]           =   name.replace("""resConfUnit1"""   ,"""residual_layer1"""						)
							if "resConfUnit2" in name:
														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   name.replace("""resConfUnit2"""   ,"""residual_layer2"""						)
							if "conv1" in name:
														SCREAMING_SNAKE_CASE__						:       str           =   name.replace("""conv1"""   ,"""convolution1"""						)
							if "conv2" in name:
														SCREAMING_SNAKE_CASE__						:       str           =   name.replace("""conv2"""   ,"""convolution2"""						)
							# readout blocks
							if "pretrained.act_postprocess1.0.project.0" in name:
														SCREAMING_SNAKE_CASE__						:       str           =   name.replace("""pretrained.act_postprocess1.0.project.0"""   ,"""neck.reassemble_stage.readout_projects.0.0"""						)
							if "pretrained.act_postprocess2.0.project.0" in name:
														SCREAMING_SNAKE_CASE__						:       Dict           =   name.replace("""pretrained.act_postprocess2.0.project.0"""   ,"""neck.reassemble_stage.readout_projects.1.0"""						)
							if "pretrained.act_postprocess3.0.project.0" in name:
														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   name.replace("""pretrained.act_postprocess3.0.project.0"""   ,"""neck.reassemble_stage.readout_projects.2.0"""						)
							if "pretrained.act_postprocess4.0.project.0" in name:
														SCREAMING_SNAKE_CASE__						:       List[Any]           =   name.replace("""pretrained.act_postprocess4.0.project.0"""   ,"""neck.reassemble_stage.readout_projects.3.0"""						)
							# resize blocks
							if "pretrained.act_postprocess1.3" in name:
														SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   name.replace("""pretrained.act_postprocess1.3"""   ,"""neck.reassemble_stage.layers.0.projection"""						)
							if "pretrained.act_postprocess1.4" in name:
														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   name.replace("""pretrained.act_postprocess1.4"""   ,"""neck.reassemble_stage.layers.0.resize"""						)
							if "pretrained.act_postprocess2.3" in name:
														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   name.replace("""pretrained.act_postprocess2.3"""   ,"""neck.reassemble_stage.layers.1.projection"""						)
							if "pretrained.act_postprocess2.4" in name:
														SCREAMING_SNAKE_CASE__						:       List[str]           =   name.replace("""pretrained.act_postprocess2.4"""   ,"""neck.reassemble_stage.layers.1.resize"""						)
							if "pretrained.act_postprocess3.3" in name:
														SCREAMING_SNAKE_CASE__						:       str           =   name.replace("""pretrained.act_postprocess3.3"""   ,"""neck.reassemble_stage.layers.2.projection"""						)
							if "pretrained.act_postprocess4.3" in name:
														SCREAMING_SNAKE_CASE__						:       List[Any]           =   name.replace("""pretrained.act_postprocess4.3"""   ,"""neck.reassemble_stage.layers.3.projection"""						)
							if "pretrained.act_postprocess4.4" in name:
														SCREAMING_SNAKE_CASE__						:       Optional[int]           =   name.replace("""pretrained.act_postprocess4.4"""   ,"""neck.reassemble_stage.layers.3.resize"""						)
							if "pretrained" in name:
														SCREAMING_SNAKE_CASE__						:       Optional[int]           =   name.replace("""pretrained"""   ,"""dpt"""						)
							if "bn" in name:
														SCREAMING_SNAKE_CASE__						:       Tuple           =   name.replace("""bn"""   ,"""batch_norm"""						)
							if "head" in name:
														SCREAMING_SNAKE_CASE__						:       Optional[int]           =   name.replace("""head"""   ,"""head.head"""						)
							if "encoder.norm" in name:
														SCREAMING_SNAKE_CASE__						:       str           =   name.replace("""encoder.norm"""   ,"""layernorm"""						)
							if "auxlayer" in name:
														SCREAMING_SNAKE_CASE__						:       int           =   name.replace("""auxlayer"""   ,"""auxiliary_head.head"""						)

							return name

def lowercase_  (   _snake_case   ,_snake_case						):
							for i in range(config.num_hidden_layers						):
														# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
														SCREAMING_SNAKE_CASE__						:       List[str]           =   state_dict.pop(f'''dpt.encoder.layer.{i}.attn.qkv.weight'''						)
														SCREAMING_SNAKE_CASE__						:       str           =   state_dict.pop(f'''dpt.encoder.layer.{i}.attn.qkv.bias'''						)
														# next, add query, keys and values (in that order) to the state dict
														SCREAMING_SNAKE_CASE__						:       Optional[int]           =   in_proj_weight[: config.hidden_size, :]
														SCREAMING_SNAKE_CASE__						:       Any           =   in_proj_bias[: config.hidden_size]
														SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   in_proj_weight[
														    config.hidden_size : config.hidden_size * 2, :
														]
														SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   in_proj_bias[
														    config.hidden_size : config.hidden_size * 2
														]
														SCREAMING_SNAKE_CASE__						:       Dict           =   in_proj_weight[
														    -config.hidden_size :, :
														]
														SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   in_proj_bias[-config.hidden_size :]

def lowercase_  (   ):
							SCREAMING_SNAKE_CASE__						:       Optional[int]           =   """http://images.cocodataset.org/val2017/000000039769.jpg"""
							SCREAMING_SNAKE_CASE__						:       str           =   Image.open(requests.get(_snake_case   ,stream=_snake_case						).raw						)
							return im

@torch.no_grad()
def lowercase_  (   _snake_case   ,_snake_case   ,_snake_case   ,_snake_case						):
							SCREAMING_SNAKE_CASE__      ,      SCREAMING_SNAKE_CASE__						:       List[Any]           =   get_dpt_config(_snake_case						)
							# load original state_dict from URL
							SCREAMING_SNAKE_CASE__						:       Optional[int]           =   torch.hub.load_state_dict_from_url(_snake_case   ,map_location="""cpu"""						)
							# remove certain keys
							remove_ignore_keys_(_snake_case						)
							# rename keys
							for key in state_dict.copy().keys():
														SCREAMING_SNAKE_CASE__						:       Tuple           =   state_dict.pop(_snake_case						)
														SCREAMING_SNAKE_CASE__						:       List[Any]           =   val
							# read in qkv matrices
							read_in_q_k_v(_snake_case   ,_snake_case						)

							# load HuggingFace model
							SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   DPTForSemanticSegmentation(_snake_case						) if """ade""" in checkpoint_url else DPTForDepthEstimation(_snake_case						)
							model.load_state_dict(_snake_case						)
							model.eval()

							# Check outputs on an image
							SCREAMING_SNAKE_CASE__						:       Tuple           =   480 if """ade""" in checkpoint_url else 384
							SCREAMING_SNAKE_CASE__						:       str           =   DPTImageProcessor(size=_snake_case						)

							SCREAMING_SNAKE_CASE__						:       List[str]           =   prepare_img()
							SCREAMING_SNAKE_CASE__						:       Dict           =   image_processor(_snake_case   ,return_tensors="""pt"""						)

							# forward pass
							SCREAMING_SNAKE_CASE__						:       Optional[int]           =   model(**_snake_case						).logits if """ade""" in checkpoint_url else model(**_snake_case						).predicted_depth

							# Assert logits
							SCREAMING_SNAKE_CASE__						:       Any           =   torch.tensor([[6.3199, 6.3629, 6.4148], [6.3850, 6.3615, 6.4166], [6.3519, 6.3176, 6.3575]]						)
							if "ade" in checkpoint_url:
														SCREAMING_SNAKE_CASE__						:       int           =   torch.tensor([[4.0480, 4.2420, 4.4360], [4.3124, 4.5693, 4.8261], [4.5768, 4.8965, 5.2163]]						)
							assert outputs.shape == torch.Size(_snake_case						)
							assert (
							    torch.allclose(outputs[0, 0, :3, :3]   ,_snake_case   ,atol=1E-4						)
							    if "ade" in checkpoint_url
							    else torch.allclose(outputs[0, :3, :3]   ,_snake_case						)
							)

							Path(_snake_case						).mkdir(exist_ok=_snake_case						)
							print(f'''Saving model to {pytorch_dump_folder_path}'''						)
							model.save_pretrained(_snake_case						)
							print(f'''Saving image processor to {pytorch_dump_folder_path}'''						)
							image_processor.save_pretrained(_snake_case						)

							if push_to_hub:
														print("""Pushing model to hub..."""						)
														model.push_to_hub(
														    repo_path_or_name=Path(_snake_case   ,_snake_case						)   ,organization="""nielsr"""   ,commit_message="""Add model"""   ,use_temp_dir=_snake_case   ,)
														image_processor.push_to_hub(
														    repo_path_or_name=Path(_snake_case   ,_snake_case						)   ,organization="""nielsr"""   ,commit_message="""Add image processor"""   ,use_temp_dir=_snake_case   ,)


if __name__ == "__main__":
					UpperCAmelCase__      :    int    							= argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '--checkpoint_url',
					    default='https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt',
					    type=str,
					    help='URL of the original DPT checkpoint you\'d like to convert.',
					)
					parser.add_argument(
					    '--pytorch_dump_folder_path',
					    default=None,
					    type=str,
					    required=True,
					    help='Path to the output PyTorch model directory.',
					)
					parser.add_argument(
					    '--push_to_hub',
					    action='store_true',
					)
					parser.add_argument(
					    '--model_name',
					    default='dpt-large',
					    type=str,
					    help='Name of the model, in case you\'re pushing to the hub.',
					)

					UpperCAmelCase__      :    Any    							= parser.parse_args()
					convert_dpt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub, args.model_name)






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




