






"""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 			_snake_case    (					)    ->     str:
							lowerCamelCase_	:							Optional[int]									="https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"
							lowerCamelCase_	:							str									=Image.open(requests.get(lowerCamelCase__    ,						stream=lowerCamelCase__  ).raw  ).convert("RGB"  )

							return image

def 			_snake_case    (					lowerCamelCase__     :   Optional[int]  )    ->     Any:
							lowerCamelCase_	:							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 			_snake_case    (					lowerCamelCase__     :   Dict    ,						lowerCamelCase__     :   Optional[Any]    ,						lowerCamelCase__     :   Tuple  )    ->     Optional[Any]:
							lowerCamelCase_	:							str									=dct.pop(lowerCamelCase__  )
							lowerCamelCase_	:							Optional[Any]									=val

def 			_snake_case    (					lowerCamelCase__     :   str    ,						lowerCamelCase__     :   Optional[Any]  )    ->     str:
							for i in range(config.vision_config.num_hidden_layers  ):
														# read in original q and v biases
														lowerCamelCase_	:							Optional[Any]									=state_dict.pop(F"""visual_encoder.blocks.{i}.attn.q_bias"""  )
														lowerCamelCase_	:							int									=state_dict.pop(F"""visual_encoder.blocks.{i}.attn.v_bias"""  )

														# next, set bias in the state dict
														lowerCamelCase_	:							Optional[int]									=torch.cat((q_bias, torch.zeros_like(lowerCamelCase__    ,						requires_grad=lowerCamelCase__  ), v_bias)  )
														lowerCamelCase_	:							Any									=qkv_bias

def 			_snake_case    (					lowerCamelCase__     :   Optional[Any]    ,						lowerCamelCase__     :   List[str]  )    ->     Any:
							lowerCamelCase_	:							List[Any]									=364 if "coco" in model_name else 224
							lowerCamelCase_	:							Optional[Any]									=BlipaVisionConfig(image_size=lowerCamelCase__  ).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:
														lowerCamelCase_	:							Tuple									=OPTConfig.from_pretrained("facebook/opt-2.7b"    ,						eos_token_id=lowerCamelCase__  ).to_dict()
							elif "opt-6.7b" in model_name:
														lowerCamelCase_	:							Any									=OPTConfig.from_pretrained("facebook/opt-6.7b"    ,						eos_token_id=lowerCamelCase__  ).to_dict()
							elif "t5-xl" in model_name:
														lowerCamelCase_	:							Union[str, Any]									=TaConfig.from_pretrained("google/flan-t5-xl"    ,						dense_act_fn="gelu"    ,						bos_token_id=1  ).to_dict()
							elif "t5-xxl" in model_name:
														lowerCamelCase_	:							Any									=TaConfig.from_pretrained("google/flan-t5-xxl"    ,						dense_act_fn="gelu"    ,						bos_token_id=1  ).to_dict()

							lowerCamelCase_	:							Optional[int]									=BlipaConfig(vision_config=lowerCamelCase__    ,						text_config=lowerCamelCase__  )

							return config, image_size

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

							lowerCamelCase_	:							Union[str, Any]									=BlipaForConditionalGeneration(lowerCamelCase__  ).eval()

							lowerCamelCase_	:							Union[str, Any]									={
							    "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"),
							}

							lowerCamelCase_	:							Optional[int]									=model_name_to_original[model_name]

							# load original model
							print("Loading original model..."  )
							lowerCamelCase_	:							int									="cuda" if torch.cuda.is_available() else "cpu"
							lowerCamelCase_	:							Optional[int]									=load_model_and_preprocess(
							    name=lowerCamelCase__    ,						model_type=lowerCamelCase__    ,						is_eval=lowerCamelCase__    ,						device=lowerCamelCase__  )
							original_model.eval()
							print("Done!"  )

							# update state dict keys
							lowerCamelCase_	:							int									=original_model.state_dict()
							lowerCamelCase_	:							Optional[int]									=create_rename_keys(lowerCamelCase__  )
							for src, dest in rename_keys:
														rename_key(lowerCamelCase__    ,						lowerCamelCase__    ,						lowerCamelCase__  )

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

							# read in qv biases
							read_in_q_v_bias(lowerCamelCase__    ,						lowerCamelCase__  )

							lowerCamelCase_	:							List[str]									=hf_model.load_state_dict(lowerCamelCase__    ,						strict=lowerCamelCase__  )
							assert len(lowerCamelCase__  ) == 0
							assert unexpected_keys == ["qformer.embeddings.position_ids"]

							lowerCamelCase_	:							List[str]									=load_demo_image()
							lowerCamelCase_	:							Union[str, Any]									=vis_processors["eval"](lowerCamelCase__  ).unsqueeze(0  ).to(lowerCamelCase__  )
							lowerCamelCase_	:							List[str]									=tokenizer(["\n"]    ,						return_tensors="pt"  ).input_ids.to(lowerCamelCase__  )

							# create processor
							lowerCamelCase_	:							int									=BlipImageProcessor(
							    size={"height": image_size, "width": image_size}    ,						image_mean=lowerCamelCase__    ,						image_std=lowerCamelCase__  )
							lowerCamelCase_	:							List[str]									=BlipaProcessor(image_processor=lowerCamelCase__    ,						tokenizer=lowerCamelCase__  )
							lowerCamelCase_	:							Any									=processor(images=lowerCamelCase__    ,						return_tensors="pt"  ).pixel_values.to(lowerCamelCase__  )

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

							original_model.to(lowerCamelCase__  )
							hf_model.to(lowerCamelCase__  )
							with torch.no_grad():
														if "opt" in model_name:
																					lowerCamelCase_	:							List[Any]									=original_model({"image": original_pixel_values, "text_input": [""]}  ).logits
																					lowerCamelCase_	:							Optional[int]									=hf_model(lowerCamelCase__    ,						lowerCamelCase__  ).logits
														else:
																					lowerCamelCase_	:							str									=original_model(
																					    {"image": original_pixel_values, "text_input": ["\n"], "text_output": ["\n"]}  ).logits
																					lowerCamelCase_	:							List[Any]									=input_ids.masked_fill(input_ids == tokenizer.pad_token_id    ,						-100  )
																					lowerCamelCase_	:							List[str]									=hf_model(lowerCamelCase__    ,						lowerCamelCase__    ,						labels=lowerCamelCase__  ).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":
														lowerCamelCase_	:							Tuple									=torch.tensor(
														    [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]    ,						device=lowerCamelCase__  )
														assert torch.allclose(logits[0, :3, :3]    ,						lowerCamelCase__    ,						atol=1e-4  )
							elif model_name == "blip2-flan-t5-xl-coco":
														lowerCamelCase_	:							List[str]									=torch.tensor(
														    [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]    ,						device=lowerCamelCase__  )
							else:
														# cast to same type
														lowerCamelCase_	:							List[Any]									=logits.dtype
														assert torch.allclose(original_logits.to(lowerCamelCase__  )    ,						lowerCamelCase__    ,						atol=1e-2  )
							print("Looks ok!"  )

							print("Generating a caption..."  )
							lowerCamelCase_	:							Optional[int]									=""
							lowerCamelCase_	:							List[str]									=tokenizer(lowerCamelCase__    ,						return_tensors="pt"  ).input_ids.to(lowerCamelCase__  )

							lowerCamelCase_	:							List[Any]									=original_model.generate({"image": original_pixel_values}  )
							lowerCamelCase_	:							List[Any]									=hf_model.generate(
							    lowerCamelCase__    ,						lowerCamelCase__    ,						do_sample=lowerCamelCase__    ,						num_beams=5    ,						max_length=30    ,						min_length=1    ,						top_p=0.9    ,						repetition_penalty=1.0    ,						length_penalty=1.0    ,						temperature=1    ,						)
							print("Original generation:"    ,						lowerCamelCase__  )
							lowerCamelCase_	:							Tuple									=input_ids.shape[1]
							lowerCamelCase_	:							int									=processor.batch_decode(outputs[:, prompt_length:]    ,						skip_special_tokens=lowerCamelCase__  )
							lowerCamelCase_	:							Optional[Any]									=[text.strip() for text in output_text]
							print("HF generation:"    ,						lowerCamelCase__  )

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

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


if __name__ == "__main__":
						A__	:    Optional[int]			 = argparse.ArgumentParser()
						A__	:    Tuple			 = [
						    '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',
						)

						A__	:    Union[str, Any]			 = 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 transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class        lowercase__      (							unittest.TestCase       ):


			def      UpperCAmelCase__						(				self					:  Any ):
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()


			@property
			def      UpperCAmelCase__						(				self					:  Dict ):
										lowerCamelCase_	:							Optional[Any]									=1
										lowerCamelCase_	:							Union[str, Any]									=3
										lowerCamelCase_	:							Dict									=(32, 32)

										lowerCamelCase_	:							List[Any]									=floats_tensor((batch_size, num_channels) + sizes       ,	rng=random.Random(0 ) ).to(snake_case__ )
										return image


			@property
			def      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										torch.manual_seed(0 )
										lowerCamelCase_	:							Dict									=UNetaDConditionModel(
										    block_out_channels=(32, 64)       ,	layers_per_block=2       ,	sample_size=32       ,	in_channels=4       ,	out_channels=4       ,	down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")       ,	up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")       ,	cross_attention_dim=32       ,	)
										return model


			@property
			def      UpperCAmelCase__						(				self					:  List[Any] ):
										torch.manual_seed(0 )
										lowerCamelCase_	:							Union[str, Any]									=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       ,	)
										return model


			@property
			def      UpperCAmelCase__						(				self					:  Optional[Any] ):
										torch.manual_seed(0 )
										lowerCamelCase_	:							Any									=RobertaSeriesConfig(
										    hidden_size=32       ,	project_dim=32       ,	intermediate_size=37       ,	layer_norm_eps=1E-05       ,	num_attention_heads=4       ,	num_hidden_layers=5       ,	pad_token_id=1       ,	vocab_size=5006       ,	)
										return RobertaSeriesModelWithTransformation(snake_case__ )


			@property
			def      UpperCAmelCase__						(				self					:  int ):
										def extract(*snake_case__					:  Dict       ,	**snake_case__					:  int ):







																	class        lowercase__      :


																				def __init__(				self					:  Optional[Any] ):
																											lowerCamelCase_	:							Union[str, Any]									=torch.ones([0] )


																				def      UpperCAmelCase__						(				self					:  List[Any]       ,	snake_case__					:  Any ):
																											self.pixel_values.to(snake_case__ )
																											return self


																	return Out()

										return extract


			def      UpperCAmelCase__						(				self					:  Optional[Any] ):
										lowerCamelCase_	:							Dict									="cpu"  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_	:							Dict									=self.dummy_cond_unet
										lowerCamelCase_	:							List[str]									=PNDMScheduler(skip_prk_steps=snake_case__ )
										lowerCamelCase_	:							List[Any]									=self.dummy_vae
										lowerCamelCase_	:							Any									=self.dummy_text_encoder
										lowerCamelCase_	:							List[Any]									=XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta" )
										lowerCamelCase_	:							List[str]									=77

										lowerCamelCase_	:							Optional[int]									=self.dummy_image.to(snake_case__ )
										lowerCamelCase_	:							Any									=init_image / 2 + 0.5

										# make sure here that pndm scheduler skips prk
										lowerCamelCase_	:							List[Any]									=AltDiffusionImgaImgPipeline(
										    unet=snake_case__       ,	scheduler=snake_case__       ,	vae=snake_case__       ,	text_encoder=snake_case__       ,	tokenizer=snake_case__       ,	safety_checker=snake_case__       ,	feature_extractor=self.dummy_extractor       ,	)
										lowerCamelCase_	:							int									=VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor       ,	do_normalize=snake_case__ )
										lowerCamelCase_	:							Optional[int]									=alt_pipe.to(snake_case__ )
										alt_pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							Any									="A painting of a squirrel eating a burger"
										lowerCamelCase_	:							Union[str, Any]									=torch.Generator(device=snake_case__ ).manual_seed(0 )
										lowerCamelCase_	:							str									=alt_pipe(
										    [prompt]       ,	generator=snake_case__       ,	guidance_scale=6.0       ,	num_inference_steps=2       ,	output_type="np"       ,	image=snake_case__       ,	)

										lowerCamelCase_	:							List[Any]									=output.images

										lowerCamelCase_	:							str									=torch.Generator(device=snake_case__ ).manual_seed(0 )
										lowerCamelCase_	:							Optional[Any]									=alt_pipe(
										    [prompt]       ,	generator=snake_case__       ,	guidance_scale=6.0       ,	num_inference_steps=2       ,	output_type="np"       ,	image=snake_case__       ,	return_dict=snake_case__       ,	)[0]

										lowerCamelCase_	:							Tuple									=image[0, -3:, -3:, -1]
										lowerCamelCase_	:							Union[str, Any]									=image_from_tuple[0, -3:, -3:, -1]

										assert image.shape == (1, 32, 32, 3)
										lowerCamelCase_	:							List[str]									=np.array([0.4_427, 0.3_731, 0.4_249, 0.4_941, 0.4_546, 0.4_148, 0.4_193, 0.4_666, 0.4_499] )

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


			@unittest.skipIf(torch_device != "cuda"       ,	"This test requires a GPU" )
			def      UpperCAmelCase__						(				self					:  Tuple ):
										lowerCamelCase_	:							Dict									=self.dummy_cond_unet
										lowerCamelCase_	:							Any									=PNDMScheduler(skip_prk_steps=snake_case__ )
										lowerCamelCase_	:							Tuple									=self.dummy_vae
										lowerCamelCase_	:							Union[str, Any]									=self.dummy_text_encoder
										lowerCamelCase_	:							Optional[Any]									=XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta" )
										lowerCamelCase_	:							Tuple									=77

										lowerCamelCase_	:							str									=self.dummy_image.to(snake_case__ )

										# put models in fp16
										lowerCamelCase_	:							Optional[Any]									=unet.half()
										lowerCamelCase_	:							Dict									=vae.half()
										lowerCamelCase_	:							str									=bert.half()

										# make sure here that pndm scheduler skips prk
										lowerCamelCase_	:							Optional[int]									=AltDiffusionImgaImgPipeline(
										    unet=snake_case__       ,	scheduler=snake_case__       ,	vae=snake_case__       ,	text_encoder=snake_case__       ,	tokenizer=snake_case__       ,	safety_checker=snake_case__       ,	feature_extractor=self.dummy_extractor       ,	)
										lowerCamelCase_	:							List[Any]									=VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor       ,	do_normalize=snake_case__ )
										lowerCamelCase_	:							int									=alt_pipe.to(snake_case__ )
										alt_pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							Tuple									="A painting of a squirrel eating a burger"
										lowerCamelCase_	:							Tuple									=torch.manual_seed(0 )
										lowerCamelCase_	:							List[Any]									=alt_pipe(
										    [prompt]       ,	generator=snake_case__       ,	num_inference_steps=2       ,	output_type="np"       ,	image=snake_case__       ,	).images

										assert image.shape == (1, 32, 32, 3)


			@unittest.skipIf(torch_device != "cuda"       ,	"This test requires a GPU" )
			def      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										lowerCamelCase_	:							int									=load_image(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
										    "/img2img/sketch-mountains-input.jpg" )
										# resize to resolution that is divisible by 8 but not 16 or 32
										lowerCamelCase_	:							Any									=init_image.resize((760, 504) )

										lowerCamelCase_	:							List[str]									="BAAI/AltDiffusion"
										lowerCamelCase_	:							List[str]									=AltDiffusionImgaImgPipeline.from_pretrained(
										    snake_case__       ,	safety_checker=snake_case__       ,	)
										pipe.to(snake_case__ )
										pipe.set_progress_bar_config(disable=snake_case__ )
										pipe.enable_attention_slicing()

										lowerCamelCase_	:							Optional[int]									="A fantasy landscape, trending on artstation"

										lowerCamelCase_	:							List[Any]									=torch.manual_seed(0 )
										lowerCamelCase_	:							Optional[int]									=pipe(
										    prompt=snake_case__       ,	image=snake_case__       ,	strength=0.75       ,	guidance_scale=7.5       ,	generator=snake_case__       ,	output_type="np"       ,	)
										lowerCamelCase_	:							Optional[int]									=output.images[0]

										lowerCamelCase_	:							Tuple									=image[255:258, 383:386, -1]

										assert image.shape == (504, 760, 3)
										lowerCamelCase_	:							Optional[int]									=np.array([0.9_358, 0.9_397, 0.9_599, 0.9_901, 1.0_000, 1.0_000, 0.9_882, 1.0_000, 1.0_000] )

										assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2







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


			def      UpperCAmelCase__						(				self					:  str ):
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()


			def      UpperCAmelCase__						(				self					:  List[Any] ):
										lowerCamelCase_	:							str									=load_image(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
										    "/img2img/sketch-mountains-input.jpg" )
										lowerCamelCase_	:							Any									=init_image.resize((768, 512) )
										lowerCamelCase_	:							Optional[int]									=load_numpy(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy" )

										lowerCamelCase_	:							Dict									="BAAI/AltDiffusion"
										lowerCamelCase_	:							List[str]									=AltDiffusionImgaImgPipeline.from_pretrained(
										    snake_case__       ,	safety_checker=snake_case__       ,	)
										pipe.to(snake_case__ )
										pipe.set_progress_bar_config(disable=snake_case__ )
										pipe.enable_attention_slicing()

										lowerCamelCase_	:							Optional[Any]									="A fantasy landscape, trending on artstation"

										lowerCamelCase_	:							Dict									=torch.manual_seed(0 )
										lowerCamelCase_	:							int									=pipe(
										    prompt=snake_case__       ,	image=snake_case__       ,	strength=0.75       ,	guidance_scale=7.5       ,	generator=snake_case__       ,	output_type="np"       ,	)
										lowerCamelCase_	:							int									=output.images[0]

										assert image.shape == (512, 768, 3)
										# img2img is flaky across GPUs even in fp32, so using MAE here
										assert np.abs(expected_image - image ).max() < 1E-2



