


'''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_   (  ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE :       int     		=      """https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"""
				_SCREAMING_SNAKE_CASE :       str     		=      Image.open(requests.get(SCREAMING_SNAKE_CASE__							,      stream=SCREAMING_SNAKE_CASE__     ).raw     ).convert("""RGB"""     )

				return image







def 			snake_case_   (  SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE :       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 			snake_case_   (  SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE :       int     		=      dct.pop(SCREAMING_SNAKE_CASE__     )
				_SCREAMING_SNAKE_CASE :       List[Any]     		=      val







def 			snake_case_   (  SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





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

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







def 			snake_case_   (  SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





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

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

				return config, image_size







@torch.no_grad()
def 			snake_case_   (  SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__=None							,      SCREAMING_SNAKE_CASE__=False     ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE :       Tuple     		=      (
				    AutoTokenizer.from_pretrained("""facebook/opt-2.7b"""     )
				    if """opt""" in model_name
				    else AutoTokenizer.from_pretrained("""google/flan-t5-xl"""     )
				)
				_SCREAMING_SNAKE_CASE :       Tuple     		=      tokenizer("""\n"""							,      add_special_tokens=SCREAMING_SNAKE_CASE__     ).input_ids[0]
				_SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE :       int     		=      get_blipa_config(SCREAMING_SNAKE_CASE__							,      eos_token_id=SCREAMING_SNAKE_CASE__     )

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

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

				_SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE :       Optional[int]     		=      model_name_to_original[model_name]

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

				# read in qv biases
				read_in_q_v_bias(SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__     )

				_SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE :       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"]

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

				# create processor
				_SCREAMING_SNAKE_CASE :       Any     		=      BlipImageProcessor(
				    size={"""height""": image_size, """width""": image_size}							,      image_mean=SCREAMING_SNAKE_CASE__							,      image_std=SCREAMING_SNAKE_CASE__     )
				_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      BlipaProcessor(image_processor=SCREAMING_SNAKE_CASE__							,      tokenizer=SCREAMING_SNAKE_CASE__     )
				_SCREAMING_SNAKE_CASE :       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:
												_SCREAMING_SNAKE_CASE :       List[Any]     		=      original_model({"""image""": original_pixel_values, """text_input""": [""""""]}     ).logits
												_SCREAMING_SNAKE_CASE :       List[str]     		=      hf_model(SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__     ).logits
								else:
												_SCREAMING_SNAKE_CASE :       Dict     		=      original_model(
												    {"""image""": original_pixel_values, """text_input""": ["""\n"""], """text_output""": ["""\n"""]}     ).logits
												_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      input_ids.masked_fill(input_ids == tokenizer.pad_token_id							,      -100     )
												_SCREAMING_SNAKE_CASE :       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":
								_SCREAMING_SNAKE_CASE :       List[Any]     		=      torch.tensor(
								    [[-4_1.5_8_5_0, -4.4_4_4_0, -8.9_9_2_2], [-4_7.4_3_2_2, -5.9_1_4_3, -1.7_3_4_0]]							,      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":
								_SCREAMING_SNAKE_CASE :       int     		=      torch.tensor(
								    [[-5_7.0_1_0_9, -9.8_9_6_7, -1_2.6_2_8_0], [-6_8.6_5_7_8, -1_2.7_1_9_1, -1_0.5_0_6_5]]							,      device=SCREAMING_SNAKE_CASE__     )
				else:
								# cast to same type
								_SCREAMING_SNAKE_CASE :       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..."""     )
				_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      """"""
				_SCREAMING_SNAKE_CASE :       Any     		=      tokenizer(SCREAMING_SNAKE_CASE__							,      return_tensors="""pt"""     ).input_ids.to(SCREAMING_SNAKE_CASE__     )

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


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',
}




