






'''simple docstring'''





import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
def __UpperCAmelCase   (					)    ->   Optional[Any]:
			UpperCAmelCase_      :       Union[str, Any]          =       '''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
			UpperCAmelCase_      :       Tuple          =       Image.open(requests.get(A			,       stream=A		).raw		).convert('''RGB'''		)

			return image
def __UpperCAmelCase   (					A					:			Dict		)    ->   List[Any]:
			UpperCAmelCase_      :       str          =       []
			# fmt: off

			# vision encoder
			rename_keys.append(('''visual_encoder.cls_token''', '''vision_model.embeddings.class_embedding''')		)
			rename_keys.append(('''visual_encoder.pos_embed''', '''vision_model.embeddings.position_embedding''')		)
			rename_keys.append(('''visual_encoder.patch_embed.proj.weight''', '''vision_model.embeddings.patch_embedding.weight''')		)
			rename_keys.append(('''visual_encoder.patch_embed.proj.bias''', '''vision_model.embeddings.patch_embedding.bias''')		)
			rename_keys.append(('''ln_vision.weight''', '''vision_model.post_layernorm.weight''')		)
			rename_keys.append(('''ln_vision.bias''', '''vision_model.post_layernorm.bias''')		)

			for i in range(config.vision_config.num_hidden_layers		):
						rename_keys.append((F"visual_encoder.blocks.{i}.norm1.weight", F"vision_model.encoder.layers.{i}.layer_norm1.weight")		)
						rename_keys.append((F"visual_encoder.blocks.{i}.norm1.bias", F"vision_model.encoder.layers.{i}.layer_norm1.bias")		)
						rename_keys.append((F"visual_encoder.blocks.{i}.norm2.weight", F"vision_model.encoder.layers.{i}.layer_norm2.weight")		)
						rename_keys.append((F"visual_encoder.blocks.{i}.norm2.bias", F"vision_model.encoder.layers.{i}.layer_norm2.bias")		)
						rename_keys.append((F"visual_encoder.blocks.{i}.attn.qkv.weight", F"vision_model.encoder.layers.{i}.self_attn.qkv.weight")		)
						rename_keys.append((F"visual_encoder.blocks.{i}.attn.proj.weight", F"vision_model.encoder.layers.{i}.self_attn.projection.weight",)		)
						rename_keys.append((F"visual_encoder.blocks.{i}.attn.proj.bias", F"vision_model.encoder.layers.{i}.self_attn.projection.bias")		)
						rename_keys.append((F"visual_encoder.blocks.{i}.mlp.fc1.weight", F"vision_model.encoder.layers.{i}.mlp.fc1.weight")		)
						rename_keys.append((F"visual_encoder.blocks.{i}.mlp.fc1.bias", F"vision_model.encoder.layers.{i}.mlp.fc1.bias")		)
						rename_keys.append((F"visual_encoder.blocks.{i}.mlp.fc2.weight", F"vision_model.encoder.layers.{i}.mlp.fc2.weight")		)
						rename_keys.append((F"visual_encoder.blocks.{i}.mlp.fc2.bias", F"vision_model.encoder.layers.{i}.mlp.fc2.bias")		)

			# QFormer
			rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.weight''', '''qformer.layernorm.weight''')		)
			rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.bias''', '''qformer.layernorm.bias''')		)

			# fmt: on
			return rename_keys
def __UpperCAmelCase   (					A					:			Tuple			,       A					:			str			,       A					:			Union[str, Any]		)    ->   Union[str, Any]:
			UpperCAmelCase_      :       Union[str, Any]          =       dct.pop(A		)
			UpperCAmelCase_      :       int          =       val
def __UpperCAmelCase   (					A					:			int			,       A					:			List[str]		)    ->   Any:
			for i in range(config.vision_config.num_hidden_layers		):
						# read in original q and v biases
						UpperCAmelCase_      :       Optional[Any]          =       state_dict.pop(F"visual_encoder.blocks.{i}.attn.q_bias"		)
						UpperCAmelCase_      :       Any          =       state_dict.pop(F"visual_encoder.blocks.{i}.attn.v_bias"		)

						# next, set bias in the state dict
						UpperCAmelCase_      :       Any          =       torch.cat((q_bias, torch.zeros_like(A			,       requires_grad=A		), v_bias)		)
						UpperCAmelCase_      :       List[Any]          =       qkv_bias
def __UpperCAmelCase   (					A					:			Union[str, Any]			,       A					:			Tuple		)    ->   Union[str, Any]:
			UpperCAmelCase_      :       str          =       3_6_4 if '''coco''' in model_name else 2_2_4
			UpperCAmelCase_      :       Dict          =       BlipaVisionConfig(image_size=A		).to_dict()

			# make sure the models have proper bos_token_id and eos_token_id set (important for generation)
			# seems like flan-T5 models don't have bos_token_id properly set?
			if "opt-2.7b" in model_name:
						UpperCAmelCase_      :       Any          =       OPTConfig.from_pretrained('''facebook/opt-2.7b'''			,       eos_token_id=A		).to_dict()
			elif "opt-6.7b" in model_name:
						UpperCAmelCase_      :       Dict          =       OPTConfig.from_pretrained('''facebook/opt-6.7b'''			,       eos_token_id=A		).to_dict()
			elif "t5-xl" in model_name:
						UpperCAmelCase_      :       Dict          =       TaConfig.from_pretrained('''google/flan-t5-xl'''			,       dense_act_fn='''gelu'''			,       bos_token_id=1		).to_dict()
			elif "t5-xxl" in model_name:
						UpperCAmelCase_      :       str          =       TaConfig.from_pretrained('''google/flan-t5-xxl'''			,       dense_act_fn='''gelu'''			,       bos_token_id=1		).to_dict()

			UpperCAmelCase_      :       Any          =       BlipaConfig(vision_config=A			,       text_config=A		)

			return config, image_size
@torch.no_grad()
def __UpperCAmelCase   (					A					:			int			,       A					:			Optional[int]=None			,       A					:			Dict=False		)    ->   List[Any]:
			UpperCAmelCase_      :       Any          =       (
			    AutoTokenizer.from_pretrained('''facebook/opt-2.7b'''		)
			    if '''opt''' in model_name
			    else AutoTokenizer.from_pretrained('''google/flan-t5-xl'''		)
			)
			UpperCAmelCase_      :       Optional[int]          =       tokenizer('''\n'''			,       add_special_tokens=A		).input_ids[0]
			UpperCAmelCase_    ,      UpperCAmelCase_      :       int          =       get_blipa_config(A			,       eos_token_id=A		)

			UpperCAmelCase_      :       List[str]          =       BlipaForConditionalGeneration(A		).eval()

			UpperCAmelCase_      :       List[str]          =       {
			    '''blip2-opt-2.7b''': ('''blip2_opt''', '''pretrain_opt2.7b'''),
			    '''blip2-opt-6.7b''': ('''blip2_opt''', '''pretrain_opt6.7b'''),
			    '''blip2-opt-2.7b-coco''': ('''blip2_opt''', '''caption_coco_opt2.7b'''),
			    '''blip2-opt-6.7b-coco''': ('''blip2_opt''', '''caption_coco_opt6.7b'''),
			    '''blip2-flan-t5-xl''': ('''blip2_t5''', '''pretrain_flant5xl'''),
			    '''blip2-flan-t5-xl-coco''': ('''blip2_t5''', '''caption_coco_flant5xl'''),
			    '''blip2-flan-t5-xxl''': ('''blip2_t5''', '''pretrain_flant5xxl'''),
			}

			UpperCAmelCase_    ,      UpperCAmelCase_      :       int          =       model_name_to_original[model_name]

			# load original model
			print('''Loading original model...'''		)
			UpperCAmelCase_      :       Any          =       '''cuda''' if torch.cuda.is_available() else '''cpu'''
			UpperCAmelCase_    ,      UpperCAmelCase_    ,      UpperCAmelCase_      :       int          =       load_model_and_preprocess(
			    name=A			,       model_type=A			,       is_eval=A			,       device=A		)
			original_model.eval()
			print('''Done!'''		)

			# update state dict keys
			UpperCAmelCase_      :       List[str]          =       original_model.state_dict()
			UpperCAmelCase_      :       Any          =       create_rename_keys(A		)
			for src, dest in rename_keys:
						rename_key(A			,       A			,       A		)

			# some keys can be renamed efficiently
			for key, val in state_dict.copy().items():
						UpperCAmelCase_      :       List[Any]          =       state_dict.pop(A		)
						if key.startswith('''Qformer.bert'''		):
									UpperCAmelCase_      :       List[Any]          =       key.replace('''Qformer.bert'''			,       '''qformer'''		)
						if "attention.self" in key:
									UpperCAmelCase_      :       Tuple          =       key.replace('''self'''			,       '''attention'''		)
						if "opt_proj" in key:
									UpperCAmelCase_      :       List[str]          =       key.replace('''opt_proj'''			,       '''language_projection'''		)
						if "t5_proj" in key:
									UpperCAmelCase_      :       Dict          =       key.replace('''t5_proj'''			,       '''language_projection'''		)
						if key.startswith('''opt'''		):
									UpperCAmelCase_      :       Any          =       key.replace('''opt'''			,       '''language'''		)
						if key.startswith('''t5'''		):
									UpperCAmelCase_      :       Tuple          =       key.replace('''t5'''			,       '''language'''		)
						UpperCAmelCase_      :       Optional[int]          =       val

			# read in qv biases
			read_in_q_v_bias(A			,       A		)

			UpperCAmelCase_    ,      UpperCAmelCase_      :       int          =       hf_model.load_state_dict(A			,       strict=A		)
			assert len(A		) == 0
			assert unexpected_keys == ["qformer.embeddings.position_ids"]

			UpperCAmelCase_      :       Any          =       load_demo_image()
			UpperCAmelCase_      :       Dict          =       vis_processors['''eval'''](A		).unsqueeze(0		).to(A		)
			UpperCAmelCase_      :       List[Any]          =       tokenizer(['''\n''']			,       return_tensors='''pt'''		).input_ids.to(A		)

			# create processor
			UpperCAmelCase_      :       List[Any]          =       BlipImageProcessor(
			    size={'''height''': image_size, '''width''': image_size}			,       image_mean=A			,       image_std=A		)
			UpperCAmelCase_      :       Optional[Any]          =       BlipaProcessor(image_processor=A			,       tokenizer=A		)
			UpperCAmelCase_      :       Any          =       processor(images=A			,       return_tensors='''pt'''		).pixel_values.to(A		)

			# make sure processor creates exact same pixel values
			assert torch.allclose(A			,       A		)

			original_model.to(A		)
			hf_model.to(A		)
			with torch.no_grad():
						if "opt" in model_name:
									UpperCAmelCase_      :       List[Any]          =       original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}		).logits
									UpperCAmelCase_      :       Optional[int]          =       hf_model(A			,       A		).logits
						else:
									UpperCAmelCase_      :       Optional[Any]          =       original_model(
									    {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}		).logits
									UpperCAmelCase_      :       int          =       input_ids.masked_fill(input_ids == tokenizer.pad_token_id			,       -1_0_0		)
									UpperCAmelCase_      :       Dict          =       hf_model(A			,       A			,       labels=A		).logits

			assert original_logits.shape == logits.shape
			print('''First values of original logits:'''			,       original_logits[0, :3, :3]		)
			print('''First values of HF logits:'''			,       logits[0, :3, :3]		)

			# assert values
			if model_name == "blip2-flan-t5-xl":
						UpperCAmelCase_      :       Any          =       torch.tensor(
						    [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]			,       device=A		)
						assert torch.allclose(logits[0, :3, :3]			,       A			,       atol=1e-4		)
			elif model_name == "blip2-flan-t5-xl-coco":
						UpperCAmelCase_      :       Union[str, Any]          =       torch.tensor(
						    [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]			,       device=A		)
			else:
						# cast to same type
						UpperCAmelCase_      :       Dict          =       logits.dtype
						assert torch.allclose(original_logits.to(A		)			,       A			,       atol=1e-2		)
			print('''Looks ok!'''		)

			print('''Generating a caption...'''		)
			UpperCAmelCase_      :       Optional[int]          =       ''''''
			UpperCAmelCase_      :       Tuple          =       tokenizer(A			,       return_tensors='''pt'''		).input_ids.to(A		)

			UpperCAmelCase_      :       Optional[Any]          =       original_model.generate({'''image''': original_pixel_values}		)
			UpperCAmelCase_      :       str          =       hf_model.generate(
			    A			,       A			,       do_sample=A			,       num_beams=5			,       max_length=3_0			,       min_length=1			,       top_p=0.9			,       repetition_penalty=1.0			,       length_penalty=1.0			,       temperature=1			,       )
			print('''Original generation:'''			,       A		)
			UpperCAmelCase_      :       List[str]          =       input_ids.shape[1]
			UpperCAmelCase_      :       Dict          =       processor.batch_decode(outputs[:, prompt_length:]			,       skip_special_tokens=A		)
			UpperCAmelCase_      :       Union[str, Any]          =       [text.strip() for text in output_text]
			print('''HF generation:'''			,       A		)

			if pytorch_dump_folder_path is not None:
						processor.save_pretrained(A		)
						hf_model.save_pretrained(A		)

			if push_to_hub:
						processor.push_to_hub(F"nielsr/{model_name}"		)
						hf_model.push_to_hub(F"nielsr/{model_name}"		)


if __name__ == "__main__":
						_UpperCamelCase						:						List[Any]  =					argparse.ArgumentParser()
						_UpperCamelCase						:						Dict  =					[
						    'blip2-opt-2.7b',
						    'blip2-opt-6.7b',
						    'blip2-opt-2.7b-coco',
						    'blip2-opt-6.7b-coco',
						    'blip2-flan-t5-xl',
						    'blip2-flan-t5-xl-coco',
						    'blip2-flan-t5-xxl',
						]
						parser.add_argument(
						    '--model_name',
						    default='blip2-opt-2.7b',
						    choices=choices,
						    type=str,
						    help='Path to hf config.json of model to convert',
						)
						parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
						parser.add_argument(
						    '--push_to_hub',
						    action='store_true',
						    help='Whether to push the model and processor to the hub after converting',
						)

						_UpperCamelCase						:						int  =					parser.parse_args()

						convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)










'''simple docstring'''





import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class  snake_case__      (					UpperCamelCase     ,				UpperCamelCase     ,				UpperCamelCase     ,				unittest.TestCase):
					a_        =					StableDiffusionInstructPixaPixPipeline
					a_        =					TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width", "cross_attention_kwargs"}
					a_        =					TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
					a_        =					IMAGE_TO_IMAGE_IMAGE_PARAMS
					a_        =					IMAGE_TO_IMAGE_IMAGE_PARAMS




					def 		A      (  self  :		str      )    ->							Optional[Any]:
								torch.manual_seed(0      )
								UpperCAmelCase_      :       Tuple          =       UNetaDConditionModel(
								    block_out_channels=(32, 64)				, layers_per_block=2				, sample_size=32				, in_channels=8				, out_channels=4				, down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')				, up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')				, cross_attention_dim=32				, )
								UpperCAmelCase_      :       Any          =       PNDMScheduler(skip_prk_steps=_A      )
								torch.manual_seed(0      )
								UpperCAmelCase_      :       int          =       AutoencoderKL(
								    block_out_channels=[32, 64]				, in_channels=3				, out_channels=3				, down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']				, up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']				, latent_channels=4				, )
								torch.manual_seed(0      )
								UpperCAmelCase_      :       Tuple          =       CLIPTextConfig(
								    bos_token_id=0				, eos_token_id=2				, hidden_size=32				, intermediate_size=37				, layer_norm_eps=1e-05				, num_attention_heads=4				, num_hidden_layers=5				, pad_token_id=1				, vocab_size=10_00				, )
								UpperCAmelCase_      :       List[str]          =       CLIPTextModel(_A      )
								UpperCAmelCase_      :       List[str]          =       CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      )

								UpperCAmelCase_      :       Union[str, Any]          =       {
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								    '''vae''': vae,
								    '''text_encoder''': text_encoder,
								    '''tokenizer''': tokenizer,
								    '''safety_checker''': None,
								    '''feature_extractor''': None,
								}
								return components




					def 		A      (  self  :		List[str]				, _A  :		Tuple				, _A  :		str=0      )    ->							int:
								UpperCAmelCase_      :       Tuple          =       floats_tensor((1, 3, 32, 32)				, rng=random.Random(_A      )      ).to(_A      )
								UpperCAmelCase_      :       Dict          =       image.cpu().permute(0				, 2				, 3				, 1      )[0]
								UpperCAmelCase_      :       Optional[Any]          =       Image.fromarray(np.uinta(_A      )      ).convert('''RGB'''      )
								if str(_A      ).startswith('''mps'''      ):
											UpperCAmelCase_      :       str          =       torch.manual_seed(_A      )
								else:
											UpperCAmelCase_      :       Optional[Any]          =       torch.Generator(device=_A      ).manual_seed(_A      )
								UpperCAmelCase_      :       List[str]          =       {
								    '''prompt''': '''A painting of a squirrel eating a burger''',
								    '''image''': image,
								    '''generator''': generator,
								    '''num_inference_steps''': 2,
								    '''guidance_scale''': 6.0,
								    '''image_guidance_scale''': 1,
								    '''output_type''': '''numpy''',
								}
								return inputs




					def 		A      (  self  :		Tuple      )    ->							List[Any]:
								UpperCAmelCase_      :       List[str]          =       '''cpu'''  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase_      :       Any          =       self.get_dummy_components()
								UpperCAmelCase_      :       int          =       StableDiffusionInstructPixaPixPipeline(**_A      )
								UpperCAmelCase_      :       Tuple          =       sd_pipe.to(_A      )
								sd_pipe.set_progress_bar_config(disable=_A      )

								UpperCAmelCase_      :       str          =       self.get_dummy_inputs(_A      )
								UpperCAmelCase_      :       Union[str, Any]          =       sd_pipe(**_A      ).images
								UpperCAmelCase_      :       int          =       image[0, -3:, -3:, -1]
								assert image.shape == (1, 32, 32, 3)
								UpperCAmelCase_      :       Union[str, Any]          =       np.array([0.7_526, 0.3_750, 0.4_547, 0.6_117, 0.5_866, 0.5_016, 0.4_327, 0.5_642, 0.4_815]      )

								assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-3




					def 		A      (  self  :		int      )    ->							str:
								UpperCAmelCase_      :       List[Any]          =       '''cpu'''  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase_      :       List[str]          =       self.get_dummy_components()
								UpperCAmelCase_      :       Dict          =       StableDiffusionInstructPixaPixPipeline(**_A      )
								UpperCAmelCase_      :       str          =       sd_pipe.to(_A      )
								sd_pipe.set_progress_bar_config(disable=_A      )

								UpperCAmelCase_      :       Any          =       self.get_dummy_inputs(_A      )
								UpperCAmelCase_      :       List[Any]          =       '''french fries'''
								UpperCAmelCase_      :       Optional[Any]          =       sd_pipe(**_A				, negative_prompt=_A      )
								UpperCAmelCase_      :       Tuple          =       output.images
								UpperCAmelCase_      :       Union[str, Any]          =       image[0, -3:, -3:, -1]

								assert image.shape == (1, 32, 32, 3)
								UpperCAmelCase_      :       Dict          =       np.array([0.7_511, 0.3_642, 0.4_553, 0.6_236, 0.5_797, 0.5_013, 0.4_343, 0.5_611, 0.4_831]      )

								assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-3




					def 		A      (  self  :		Dict      )    ->							str:
								UpperCAmelCase_      :       List[str]          =       '''cpu'''  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase_      :       List[Any]          =       self.get_dummy_components()
								UpperCAmelCase_      :       Any          =       StableDiffusionInstructPixaPixPipeline(**_A      )
								UpperCAmelCase_      :       int          =       sd_pipe.to(_A      )
								sd_pipe.set_progress_bar_config(disable=_A      )

								UpperCAmelCase_      :       Any          =       self.get_dummy_inputs(_A      )
								UpperCAmelCase_      :       Any          =       [inputs['''prompt''']] * 2

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

								UpperCAmelCase_      :       List[Any]          =       sd_pipe(**_A      ).images
								UpperCAmelCase_      :       str          =       image[-1, -3:, -3:, -1]

								assert image.shape == (2, 32, 32, 3)
								UpperCAmelCase_      :       Optional[Any]          =       np.array([0.5_812, 0.5_748, 0.5_222, 0.5_908, 0.5_695, 0.7_174, 0.6_804, 0.5_523, 0.5_579]      )

								assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-3




					def 		A      (  self  :		List[str]      )    ->							List[Any]:
								UpperCAmelCase_      :       Union[str, Any]          =       '''cpu'''  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase_      :       Tuple          =       self.get_dummy_components()
								UpperCAmelCase_      :       List[str]          =       EulerAncestralDiscreteScheduler(
								    beta_start=0.00_085				, beta_end=0.012				, beta_schedule='''scaled_linear'''      )
								UpperCAmelCase_      :       List[Any]          =       StableDiffusionInstructPixaPixPipeline(**_A      )
								UpperCAmelCase_      :       Dict          =       sd_pipe.to(_A      )
								sd_pipe.set_progress_bar_config(disable=_A      )

								UpperCAmelCase_      :       Tuple          =       self.get_dummy_inputs(_A      )
								UpperCAmelCase_      :       str          =       sd_pipe(**_A      ).images
								UpperCAmelCase_      :       Dict          =       image[0, -3:, -3:, -1]

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

								assert image.shape == (1, 32, 32, 3)
								UpperCAmelCase_      :       Any          =       np.array([0.7_417, 0.3_842, 0.4_732, 0.5_776, 0.5_891, 0.5_139, 0.4_052, 0.5_673, 0.4_986]      )

								assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-3




					def 		A      (  self  :		str      )    ->							List[str]:
								super().test_inference_batch_single_identical(expected_max_diff=3e-3      )





					def 		A      (  self  :		List[Any]      )    ->							Optional[Any]:
								UpperCAmelCase_      :       Tuple          =       self.get_dummy_components()
								UpperCAmelCase_      :       Tuple          =       StableDiffusionInstructPixaPixPipeline(**_A      )
								UpperCAmelCase_      :       Any          =       VaeImageProcessor(do_resize=_A				, do_normalize=_A      )
								UpperCAmelCase_      :       Union[str, Any]          =       pipe.to(_A      )
								pipe.set_progress_bar_config(disable=_A      )

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

								UpperCAmelCase_      :       Union[str, Any]          =       components['''vae''']
								UpperCAmelCase_      :       Any          =       self.get_dummy_inputs_by_type(_A				, input_image_type='''pt'''      )

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

								UpperCAmelCase_      :       List[Any]          =       pipe(**_A      )[0]

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






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




					def 		A      (  self  :		Any      )    ->							Optional[Any]:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()




					def 		A      (  self  :		List[str]				, _A  :		Optional[Any]=0      )    ->							List[Any]:
								UpperCAmelCase_      :       Dict          =       torch.manual_seed(_A      )
								UpperCAmelCase_      :       Union[str, Any]          =       load_image(
								    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg'''      )
								UpperCAmelCase_      :       Optional[int]          =       {
								    '''prompt''': '''turn him into a cyborg''',
								    '''image''': image,
								    '''generator''': generator,
								    '''num_inference_steps''': 3,
								    '''guidance_scale''': 7.5,
								    '''image_guidance_scale''': 1.0,
								    '''output_type''': '''numpy''',
								}
								return inputs




					def 		A      (  self  :		Optional[int]      )    ->							Optional[Any]:
								UpperCAmelCase_      :       List[Any]          =       StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    '''timbrooks/instruct-pix2pix'''				, safety_checker=_A      )
								pipe.to(_A      )
								pipe.set_progress_bar_config(disable=_A      )
								pipe.enable_attention_slicing()

								UpperCAmelCase_      :       List[str]          =       self.get_inputs()
								UpperCAmelCase_      :       int          =       pipe(**_A      ).images
								UpperCAmelCase_      :       Dict          =       image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 5_12, 5_12, 3)
								UpperCAmelCase_      :       Any          =       np.array([0.5_902, 0.6_015, 0.6_027, 0.5_983, 0.6_092, 0.6_061, 0.5_765, 0.5_785, 0.5_555]      )

								assert np.abs(expected_slice - image_slice      ).max() < 1e-3




					def 		A      (  self  :		Optional[Any]      )    ->							Optional[Any]:
								UpperCAmelCase_      :       Optional[Any]          =       StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    '''timbrooks/instruct-pix2pix'''				, safety_checker=_A      )
								UpperCAmelCase_      :       Dict          =       LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
								pipe.to(_A      )
								pipe.set_progress_bar_config(disable=_A      )
								pipe.enable_attention_slicing()

								UpperCAmelCase_      :       List[Any]          =       self.get_inputs()
								UpperCAmelCase_      :       str          =       pipe(**_A      ).images
								UpperCAmelCase_      :       Optional[int]          =       image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 5_12, 5_12, 3)
								UpperCAmelCase_      :       List[str]          =       np.array([0.6_578, 0.6_817, 0.6_972, 0.6_761, 0.6_856, 0.6_916, 0.6_428, 0.6_516, 0.6_301]      )

								assert np.abs(expected_slice - image_slice      ).max() < 1e-3




					def 		A      (  self  :		List[str]      )    ->							Union[str, Any]:
								UpperCAmelCase_      :       str          =       StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    '''timbrooks/instruct-pix2pix'''				, safety_checker=_A      )
								UpperCAmelCase_      :       str          =       DDIMScheduler.from_config(pipe.scheduler.config      )
								pipe.to(_A      )
								pipe.set_progress_bar_config(disable=_A      )
								pipe.enable_attention_slicing()

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

								assert image.shape == (1, 5_12, 5_12, 3)
								UpperCAmelCase_      :       Dict          =       np.array([0.3_828, 0.3_834, 0.3_818, 0.3_792, 0.3_865, 0.3_752, 0.3_792, 0.3_847, 0.3_753]      )

								assert np.abs(expected_slice - image_slice      ).max() < 1e-3




					def 		A      (  self  :		Tuple      )    ->							List[Any]:
								UpperCAmelCase_      :       List[str]          =       0

								def callback_fn(_A  :		int				, _A  :		int				, _A  :		torch.FloatTensor      ) -> None:
											UpperCAmelCase_      :       Tuple          =       True
											nonlocal number_of_steps
											number_of_steps += 1
											if step == 1:
														UpperCAmelCase_      :       List[str]          =       latents.detach().cpu().numpy()
														assert latents.shape == (1, 4, 64, 64)
														UpperCAmelCase_      :       Optional[Any]          =       latents[0, -3:, -3:, -1]
														UpperCAmelCase_      :       List[Any]          =       np.array([-0.2_463, -0.4_644, -0.9_756, 1.5_176, 1.4_414, 0.7_866, 0.9_897, 0.8_521, 0.7_983]      )

														assert np.abs(latents_slice.flatten() - expected_slice      ).max() < 5e-2
											elif step == 2:
														UpperCAmelCase_      :       Union[str, Any]          =       latents.detach().cpu().numpy()
														assert latents.shape == (1, 4, 64, 64)
														UpperCAmelCase_      :       Dict          =       latents[0, -3:, -3:, -1]
														UpperCAmelCase_      :       Tuple          =       np.array([-0.2_644, -0.4_626, -0.9_653, 1.5_176, 1.4_551, 0.7_686, 0.9_805, 0.8_452, 0.8_115]      )

														assert np.abs(latents_slice.flatten() - expected_slice      ).max() < 5e-2

								UpperCAmelCase_      :       Dict          =       False

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

								UpperCAmelCase_      :       Dict          =       self.get_inputs()
								pipe(**_A				, callback=_A				, callback_steps=1      )
								assert callback_fn.has_been_called
								assert number_of_steps == 3




					def 		A      (  self  :		int      )    ->							int:
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								UpperCAmelCase_      :       Union[str, Any]          =       StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    '''timbrooks/instruct-pix2pix'''				, safety_checker=_A				, torch_dtype=torch.floataa      )
								UpperCAmelCase_      :       Tuple          =       pipe.to(_A      )
								pipe.set_progress_bar_config(disable=_A      )
								pipe.enable_attention_slicing(1      )
								pipe.enable_sequential_cpu_offload()

								UpperCAmelCase_      :       Dict          =       self.get_inputs()
								UpperCAmelCase_      :       List[str]          =       pipe(**_A      )

								UpperCAmelCase_      :       Dict          =       torch.cuda.max_memory_allocated()
								# make sure that less than 2.2 GB is allocated
								assert mem_bytes < 2.2 * 10**9





					def 		A      (  self  :		str      )    ->							int:
								UpperCAmelCase_      :       List[Any]          =       self.get_inputs()
								# resize to resolution that is divisible by 8 but not 16 or 32
								UpperCAmelCase_      :       List[str]          =       inputs['''image'''].resize((5_04, 5_04)      )

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

								UpperCAmelCase_      :       Optional[Any]          =       pipe(**_A      )
								UpperCAmelCase_      :       Union[str, Any]          =       output.images[0]

								UpperCAmelCase_      :       List[str]          =       image[2_55:2_58, 3_83:3_86, -1]

								assert image.shape == (5_04, 5_04, 3)
								UpperCAmelCase_      :       List[str]          =       np.array([0.2_726, 0.2_529, 0.2_664, 0.2_655, 0.2_641, 0.2_642, 0.2_591, 0.2_649, 0.2_590]      )

								assert np.abs(image_slice.flatten() - expected_slice      ).max() < 5e-3


