

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

	return image

def 						UpperCamelCase					(lowercase_: int					)     ->      List[Any]:
	A__							:    Tuple        =		[]
	# 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					(lowercase_: Optional[int]    ,      lowercase_: Optional[Any]    ,      lowercase_: Dict					)     ->      List[Any]:
	A__							:    int        =		dct.pop(SCREAMING_SNAKE_CASE__					)
	A__							:    List[Any]        =		val

def 						UpperCamelCase					(lowercase_: Optional[int]    ,      lowercase_: Any					)     ->      Tuple:





	for i in range(config.vision_config.num_hidden_layers					):
		# read in original q and v biases
		A__							:    str        =		state_dict.pop(f"""visual_encoder.blocks.{i}.attn.q_bias"""					)
		A__							:    Tuple        =		state_dict.pop(f"""visual_encoder.blocks.{i}.attn.v_bias"""					)

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

def 						UpperCamelCase					(lowercase_: str    ,      lowercase_: str					)     ->      Dict:
	A__							:    List[Any]        =		364 if """coco""" in model_name else 224
	A__							:    int        =		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:
		A__							:    List[Any]        =		OPTConfig.from_pretrained("""facebook/opt-2.7b"""    ,      eos_token_id=SCREAMING_SNAKE_CASE__					).to_dict()
	elif "opt-6.7b" in model_name:
		A__							:    List[Any]        =		OPTConfig.from_pretrained("""facebook/opt-6.7b"""    ,      eos_token_id=SCREAMING_SNAKE_CASE__					).to_dict()
	elif "t5-xl" in model_name:
		A__							:    Any        =		TaConfig.from_pretrained("""google/flan-t5-xl"""    ,      dense_act_fn="""gelu"""    ,      bos_token_id=1					).to_dict()
	elif "t5-xxl" in model_name:
		A__							:    List[Any]        =		TaConfig.from_pretrained("""google/flan-t5-xxl"""    ,      dense_act_fn="""gelu"""    ,      bos_token_id=1					).to_dict()

	A__							:    Any        =		BlipaConfig(vision_config=SCREAMING_SNAKE_CASE__    ,      text_config=SCREAMING_SNAKE_CASE__					)

	return config, image_size

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

	A__							:    Any        =		BlipaForConditionalGeneration(SCREAMING_SNAKE_CASE__					).eval()

	A__							:    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"""),
	}

	A__							:    Optional[int]        =		model_name_to_original[model_name]

	# load original model
	print("""Loading original model..."""					)
	A__							:    Dict        =		"""cuda""" if torch.cuda.is_available() else """cpu"""
	A__							:    List[str]        =		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
	A__							:    List[Any]        =		original_model.state_dict()
	A__							:    Any        =		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():
		A__							:    Union[str, Any]        =		state_dict.pop(SCREAMING_SNAKE_CASE__					)
		if key.startswith("""Qformer.bert"""					):
			A__							:    str        =		key.replace("""Qformer.bert"""    ,      """qformer"""					)
		if "attention.self" in key:
			A__							:    Any        =		key.replace("""self"""    ,      """attention"""					)
		if "opt_proj" in key:
			A__							:    List[Any]        =		key.replace("""opt_proj"""    ,      """language_projection"""					)
		if "t5_proj" in key:
			A__							:    List[Any]        =		key.replace("""t5_proj"""    ,      """language_projection"""					)
		if key.startswith("""opt"""					):
			A__							:    List[str]        =		key.replace("""opt"""    ,      """language"""					)
		if key.startswith("""t5"""					):
			A__							:    Any        =		key.replace("""t5"""    ,      """language"""					)
		A__							:    Optional[int]        =		val

	# read in qv biases
	read_in_q_v_bias(SCREAMING_SNAKE_CASE__    ,      SCREAMING_SNAKE_CASE__					)

	A__							:    Union[str, 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"]

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

	# create processor
	A__							:    Any        =		BlipImageProcessor(
	    size={"""height""": image_size, """width""": image_size}    ,      image_mean=SCREAMING_SNAKE_CASE__    ,      image_std=SCREAMING_SNAKE_CASE__					)
	A__							:    Optional[Any]        =		BlipaProcessor(image_processor=SCREAMING_SNAKE_CASE__    ,      tokenizer=SCREAMING_SNAKE_CASE__					)
	A__							:    Optional[int]        =		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:
			A__							:    List[Any]        =		original_model({"""image""": original_pixel_values, """text_input""": [""""""]}					).logits
			A__							:    List[str]        =		hf_model(SCREAMING_SNAKE_CASE__    ,      SCREAMING_SNAKE_CASE__					).logits
		else:
			A__							:    Dict        =		original_model(
			    {"""image""": original_pixel_values, """text_input""": ["""\n"""], """text_output""": ["""\n"""]}					).logits
			A__							:    Optional[Any]        =		input_ids.masked_fill(input_ids == tokenizer.pad_token_id    ,      -100					)
			A__							:    int        =		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":
		A__							:    List[Any]        =		torch.tensor(
		    [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]    ,      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":
		A__							:    int        =		torch.tensor(
		    [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]    ,      device=SCREAMING_SNAKE_CASE__					)
	else:
		# cast to same type
		A__							:    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..."""					)
	A__							:    Optional[Any]        =		""""""
	A__							:    Any        =		tokenizer(SCREAMING_SNAKE_CASE__    ,      return_tensors="""pt"""					).input_ids.to(SCREAMING_SNAKE_CASE__					)

	A__							:    int        =		original_model.generate({"""image""": original_pixel_values}					)
	A__							:    Any        =		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__					)
	A__							:    Optional[Any]        =		input_ids.shape[1]
	A__							:    Optional[int]        =		processor.batch_decode(outputs[:, prompt_length:]    ,      skip_special_tokens=SCREAMING_SNAKE_CASE__					)
	A__							:    Dict        =		[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__":
	A_       : Dict         =			argparse.ArgumentParser()
	A_       : 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',
	)

	A_       : int         =			parser.parse_args()

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







'''simple docstring'''


UpperCAmelCase_		:   Union[str, Any]							   =		'\n# Transformers installation\n! pip install transformers datasets\n# To install from source instead of the last release, comment the command above and uncomment the following one.\n# ! pip install git+https://github.com/huggingface/transformers.git\n'

UpperCAmelCase_		:   Dict							   =		[{'type': 'code', 'content': INSTALL_CONTENT}]
UpperCAmelCase_		:   Any							   =		{
    '{processor_class}': 'FakeProcessorClass',
    '{model_class}': 'FakeModelClass',
    '{object_class}': 'FakeObjectClass',
}




