





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

						return image





def        lowerCamelCase_      (  SCREAMING_SNAKE_CASE__      :			List[str]   ) ->	Optional[Any]:
						UpperCAmelCase_			:      Any     						=  []
						# 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        lowerCamelCase_      (  SCREAMING_SNAKE_CASE__      :			Union[str, Any],			SCREAMING_SNAKE_CASE__      :			Union[str, Any],			SCREAMING_SNAKE_CASE__      :			List[Any]   ) ->	Any:
						UpperCAmelCase_			:      int     						=  dct.pop(SCREAMING_SNAKE_CASE__   )
						UpperCAmelCase_			:      Union[str, Any]     						=  val





def        lowerCamelCase_      (  SCREAMING_SNAKE_CASE__      :			List[Any],			SCREAMING_SNAKE_CASE__      :			Tuple   ) ->	List[str]:
						for i in range(config.vision_config.num_hidden_layers   ):
												# read in original q and v biases
												UpperCAmelCase_			:      Optional[int]     						=  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_			:      Dict     						=  torch.cat((q_bias, torch.zeros_like(SCREAMING_SNAKE_CASE__,			requires_grad=SCREAMING_SNAKE_CASE__   ), v_bias)   )
												UpperCAmelCase_			:      Union[str, Any]     						=  qkv_bias





def        lowerCamelCase_      (  SCREAMING_SNAKE_CASE__      :			Dict,			SCREAMING_SNAKE_CASE__      :			Optional[int]   ) ->	str:
						UpperCAmelCase_			:      Optional[int]     						=  364 if '''coco''' in model_name else 224
						UpperCAmelCase_			:      List[str]     						=  BlipaVisionConfig(image_size=SCREAMING_SNAKE_CASE__   ).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_			:      Dict     						=  OPTConfig.from_pretrained('''facebook/opt-2.7b''',			eos_token_id=SCREAMING_SNAKE_CASE__   ).to_dict()
						elif "opt-6.7b" in model_name:
												UpperCAmelCase_			:      int     						=  OPTConfig.from_pretrained('''facebook/opt-6.7b''',			eos_token_id=SCREAMING_SNAKE_CASE__   ).to_dict()
						elif "t5-xl" in model_name:
												UpperCAmelCase_			:      Any     						=  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_			:      str     						=  BlipaConfig(vision_config=SCREAMING_SNAKE_CASE__,			text_config=SCREAMING_SNAKE_CASE__   )

						return config, image_size





@torch.no_grad()
def        lowerCamelCase_      (  SCREAMING_SNAKE_CASE__      :			str,			SCREAMING_SNAKE_CASE__      :			List[str]=None,			SCREAMING_SNAKE_CASE__      :			Tuple=False   ) ->	Optional[Any]:
						UpperCAmelCase_			:      Union[str, Any]     						=  (
						    AutoTokenizer.from_pretrained('''facebook/opt-2.7b'''   )
						    if '''opt''' in model_name
						    else AutoTokenizer.from_pretrained('''google/flan-t5-xl'''   )
						)
						UpperCAmelCase_			:      Union[str, Any]     						=  tokenizer('''\n''',			add_special_tokens=SCREAMING_SNAKE_CASE__   ).input_ids[0]
						UpperCAmelCase_						,       UpperCAmelCase_			:      List[str]     						=  get_blipa_config(SCREAMING_SNAKE_CASE__,			eos_token_id=SCREAMING_SNAKE_CASE__   )

						UpperCAmelCase_			:      Optional[int]     						=  BlipaForConditionalGeneration(SCREAMING_SNAKE_CASE__   ).eval()

						UpperCAmelCase_			:      Optional[int]     						=  {
						    '''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_			:      List[str]     						=  model_name_to_original[model_name]

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

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

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

						# read in qv biases
						read_in_q_v_bias(SCREAMING_SNAKE_CASE__,			SCREAMING_SNAKE_CASE__   )

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

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

						# create processor
						UpperCAmelCase_			:      Any     						=  BlipImageProcessor(
						    size={'''height''': image_size, '''width''': image_size},			image_mean=SCREAMING_SNAKE_CASE__,			image_std=SCREAMING_SNAKE_CASE__   )
						UpperCAmelCase_			:      Tuple     						=  BlipaProcessor(image_processor=SCREAMING_SNAKE_CASE__,			tokenizer=SCREAMING_SNAKE_CASE__   )
						UpperCAmelCase_			:      Dict     						=  processor(images=SCREAMING_SNAKE_CASE__,			return_tensors='''pt'''   ).pixel_values.to(SCREAMING_SNAKE_CASE__   )

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

						original_model.to(SCREAMING_SNAKE_CASE__   )
						hf_model.to(SCREAMING_SNAKE_CASE__   )
						with torch.no_grad():
												if "opt" in model_name:
																		UpperCAmelCase_			:      List[Any]     						=  original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}   ).logits
																		UpperCAmelCase_			:      Any     						=  hf_model(SCREAMING_SNAKE_CASE__,			SCREAMING_SNAKE_CASE__   ).logits
												else:
																		UpperCAmelCase_			:      List[Any]     						=  original_model(
																		    {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}   ).logits
																		UpperCAmelCase_			:      Dict     						=  input_ids.masked_fill(input_ids == tokenizer.pad_token_id,			-100   )
																		UpperCAmelCase_			:      List[str]     						=  hf_model(SCREAMING_SNAKE_CASE__,			SCREAMING_SNAKE_CASE__,			labels=SCREAMING_SNAKE_CASE__   ).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.58_50, -4.44_40, -8.99_22], [-47.43_22, -5.91_43, -1.73_40]],			device=SCREAMING_SNAKE_CASE__   )
												assert torch.allclose(logits[0, :3, :3],			SCREAMING_SNAKE_CASE__,			atol=1E-4   )
						elif model_name == "blip2-flan-t5-xl-coco":
												UpperCAmelCase_			:      Optional[Any]     						=  torch.tensor(
												    [[-57.01_09, -9.89_67, -12.62_80], [-68.65_78, -12.71_91, -10.50_65]],			device=SCREAMING_SNAKE_CASE__   )
						else:
												# cast to same type
												UpperCAmelCase_			:      List[str]     						=  logits.dtype
												assert torch.allclose(original_logits.to(SCREAMING_SNAKE_CASE__   ),			SCREAMING_SNAKE_CASE__,			atol=1E-2   )
						print('''Looks ok!'''   )

						print('''Generating a caption...'''   )
						UpperCAmelCase_			:      Tuple     						=  ''''''
						UpperCAmelCase_			:      Any     						=  tokenizer(SCREAMING_SNAKE_CASE__,			return_tensors='''pt'''   ).input_ids.to(SCREAMING_SNAKE_CASE__   )

						UpperCAmelCase_			:      int     						=  original_model.generate({'''image''': original_pixel_values}   )
						UpperCAmelCase_			:      str     						=  hf_model.generate(
						    SCREAMING_SNAKE_CASE__,			SCREAMING_SNAKE_CASE__,			do_sample=SCREAMING_SNAKE_CASE__,			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:''',			SCREAMING_SNAKE_CASE__   )
						UpperCAmelCase_			:      str     						=  input_ids.shape[1]
						UpperCAmelCase_			:      Tuple     						=  processor.batch_decode(outputs[:, prompt_length:],			skip_special_tokens=SCREAMING_SNAKE_CASE__   )
						UpperCAmelCase_			:      Union[str, Any]     						=  [text.strip() for text in output_text]
						print('''HF generation:''',			SCREAMING_SNAKE_CASE__   )

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

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


if __name__ == "__main__":
					snake_case_     :       int						=  argparse.ArgumentParser()
					snake_case_     :       int						=  [
					    "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",
					)

					snake_case_     :       List[str]						=  parser.parse_args()

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












'''simple docstring'''







import unittest

import numpy as np
import torch

from diffusers import KarrasVePipeline, KarrasVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()


class 					__a      (unittest.TestCase     ):

							@property
							def     UpperCAmelCase__	(							self							:     Dict   )					->       str:







													"""simple docstring"""
													torch.manual_seed(0   )
													UpperCAmelCase_			:      Optional[Any]     						=  UNetaDModel(
													    block_out_channels=(32, 64)       ,     layers_per_block=2       ,     sample_size=32       ,     in_channels=3       ,     out_channels=3       ,     down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')       ,     up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')       ,     )
													return model



							def     UpperCAmelCase__	(							self							:     Dict   )					->       Tuple:







													"""simple docstring"""
													UpperCAmelCase_			:      List[Any]     						=  self.dummy_uncond_unet
													UpperCAmelCase_			:      Dict     						=  KarrasVeScheduler()

													UpperCAmelCase_			:      Union[str, Any]     						=  KarrasVePipeline(unet=__magic_name__       ,     scheduler=__magic_name__   )
													pipe.to(__magic_name__   )
													pipe.set_progress_bar_config(disable=__magic_name__   )

													UpperCAmelCase_			:      Dict     						=  torch.manual_seed(0   )
													UpperCAmelCase_			:      Optional[int]     						=  pipe(num_inference_steps=2       ,     generator=__magic_name__       ,     output_type='''numpy'''   ).images

													UpperCAmelCase_			:      Tuple     						=  torch.manual_seed(0   )
													UpperCAmelCase_			:      str     						=  pipe(num_inference_steps=2       ,     generator=__magic_name__       ,     output_type='''numpy'''       ,     return_dict=__magic_name__   )[0]

													UpperCAmelCase_			:      Union[str, Any]     						=  image[0, -3:, -3:, -1]
													UpperCAmelCase_			:      Optional[Any]     						=  image_from_tuple[0, -3:, -3:, -1]

													assert image.shape == (1, 32, 32, 3)
													UpperCAmelCase_			:      Dict     						=  np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]   )

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





@slow
@require_torch
class 					__a      (unittest.TestCase     ):

							def     UpperCAmelCase__	(							self							:     int   )					->       Union[str, Any]:







													"""simple docstring"""
													UpperCAmelCase_			:      List[str]     						=  '''google/ncsnpp-celebahq-256'''
													UpperCAmelCase_			:      List[str]     						=  UNetaDModel.from_pretrained(__magic_name__   )
													UpperCAmelCase_			:      List[Any]     						=  KarrasVeScheduler()

													UpperCAmelCase_			:      Any     						=  KarrasVePipeline(unet=__magic_name__       ,     scheduler=__magic_name__   )
													pipe.to(__magic_name__   )
													pipe.set_progress_bar_config(disable=__magic_name__   )

													UpperCAmelCase_			:      Dict     						=  torch.manual_seed(0   )
													UpperCAmelCase_			:      Optional[Any]     						=  pipe(num_inference_steps=20       ,     generator=__magic_name__       ,     output_type='''numpy'''   ).images

													UpperCAmelCase_			:      Any     						=  image[0, -3:, -3:, -1]
													assert image.shape == (1, 2_56, 2_56, 3)
													UpperCAmelCase_			:      Optional[Any]     						=  np.array([0.5_7_8, 0.5_8_1_1, 0.5_9_2_4, 0.5_8_0_9, 0.5_8_7, 0.5_8_8_6, 0.5_8_6_1, 0.5_8_0_2, 0.5_8_6]   )

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





