





'''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 						__magic_name__       (							)						->					List[str]:

   '''simple docstring'''







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

   return image

def 						__magic_name__       (							__UpperCAmelCase      )						->					Any:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE      							=       []
   # 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 						__magic_name__       (							__UpperCAmelCase							,	__UpperCAmelCase							,	__UpperCAmelCase      )						->					List[Any]:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE      							=       dct.pop(__UpperCAmelCase      )
   __SCREAMING_SNAKE_CASE      							=       val

def 						__magic_name__       (							__UpperCAmelCase							,	__UpperCAmelCase      )						->					Union[str, Any]:

   '''simple docstring'''







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

      # next, set bias in the state dict
      __SCREAMING_SNAKE_CASE      							=       torch.cat((q_bias, torch.zeros_like(__UpperCAmelCase							,	requires_grad=__UpperCAmelCase      ), v_bias)      )
      __SCREAMING_SNAKE_CASE      							=       qkv_bias

def 						__magic_name__       (							__UpperCAmelCase							,	__UpperCAmelCase      )						->					Any:

   '''simple docstring'''







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

   __SCREAMING_SNAKE_CASE      							=       BlipaConfig(vision_config=__UpperCAmelCase							,	text_config=__UpperCAmelCase      )

   return config, image_size

@torch.no_grad()
def 						__magic_name__       (							__UpperCAmelCase							,	__UpperCAmelCase=None							,	__UpperCAmelCase=False      )						->					str:

   '''simple docstring'''







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

   __SCREAMING_SNAKE_CASE      							=       BlipaForConditionalGeneration(__UpperCAmelCase      ).eval()

   __SCREAMING_SNAKE_CASE      							=       {
       """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      							=       model_name_to_original[model_name]

   # load original model
   print("""Loading original model..."""      )
   __SCREAMING_SNAKE_CASE      							=       """cuda""" if torch.cuda.is_available() else """cpu"""
   __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE      							=       load_model_and_preprocess(
       name=__UpperCAmelCase							,	model_type=__UpperCAmelCase							,	is_eval=__UpperCAmelCase							,	device=__UpperCAmelCase      )
   original_model.eval()
   print("""Done!"""      )

   # update state dict keys
   __SCREAMING_SNAKE_CASE      							=       original_model.state_dict()
   __SCREAMING_SNAKE_CASE      							=       create_rename_keys(__UpperCAmelCase      )
   for src, dest in rename_keys:
      rename_key(__UpperCAmelCase							,	__UpperCAmelCase							,	__UpperCAmelCase      )

   # some keys can be renamed efficiently
   for key, val in state_dict.copy().items():
      __SCREAMING_SNAKE_CASE      							=       state_dict.pop(__UpperCAmelCase      )
      if key.startswith("""Qformer.bert"""      ):
         __SCREAMING_SNAKE_CASE      							=       key.replace("""Qformer.bert"""							,	"""qformer"""      )
      if "attention.self" in key:
         __SCREAMING_SNAKE_CASE      							=       key.replace("""self"""							,	"""attention"""      )
      if "opt_proj" in key:
         __SCREAMING_SNAKE_CASE      							=       key.replace("""opt_proj"""							,	"""language_projection"""      )
      if "t5_proj" in key:
         __SCREAMING_SNAKE_CASE      							=       key.replace("""t5_proj"""							,	"""language_projection"""      )
      if key.startswith("""opt"""      ):
         __SCREAMING_SNAKE_CASE      							=       key.replace("""opt"""							,	"""language"""      )
      if key.startswith("""t5"""      ):
         __SCREAMING_SNAKE_CASE      							=       key.replace("""t5"""							,	"""language"""      )
      __SCREAMING_SNAKE_CASE      							=       val

   # read in qv biases
   read_in_q_v_bias(__UpperCAmelCase							,	__UpperCAmelCase      )

   __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE      							=       hf_model.load_state_dict(__UpperCAmelCase							,	strict=__UpperCAmelCase      )
   assert len(__UpperCAmelCase      ) == 0
   assert unexpected_keys == ["qformer.embeddings.position_ids"]

   __SCREAMING_SNAKE_CASE      							=       load_demo_image()
   __SCREAMING_SNAKE_CASE      							=       vis_processors["""eval"""](__UpperCAmelCase      ).unsqueeze(0      ).to(__UpperCAmelCase      )
   __SCREAMING_SNAKE_CASE      							=       tokenizer(["""\n"""]							,	return_tensors="""pt"""      ).input_ids.to(__UpperCAmelCase      )

   # create processor
   __SCREAMING_SNAKE_CASE      							=       BlipImageProcessor(
       size={"""height""": image_size, """width""": image_size}							,	image_mean=__UpperCAmelCase							,	image_std=__UpperCAmelCase      )
   __SCREAMING_SNAKE_CASE      							=       BlipaProcessor(image_processor=__UpperCAmelCase							,	tokenizer=__UpperCAmelCase      )
   __SCREAMING_SNAKE_CASE      							=       processor(images=__UpperCAmelCase							,	return_tensors="""pt"""      ).pixel_values.to(__UpperCAmelCase      )

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

   original_model.to(__UpperCAmelCase      )
   hf_model.to(__UpperCAmelCase      )
   with torch.no_grad():
      if "opt" in model_name:
         __SCREAMING_SNAKE_CASE      							=       original_model({"""image""": original_pixel_values, """text_input""": [""""""]}      ).logits
         __SCREAMING_SNAKE_CASE      							=       hf_model(__UpperCAmelCase							,	__UpperCAmelCase      ).logits
      else:
         __SCREAMING_SNAKE_CASE      							=       original_model(
             {"""image""": original_pixel_values, """text_input""": ["""\n"""], """text_output""": ["""\n"""]}      ).logits
         __SCREAMING_SNAKE_CASE      							=       input_ids.masked_fill(input_ids == tokenizer.pad_token_id							,	-100      )
         __SCREAMING_SNAKE_CASE      							=       hf_model(__UpperCAmelCase							,	__UpperCAmelCase							,	labels=__UpperCAmelCase      ).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      							=       torch.tensor(
          [[-41.5850, -4.4_4_4_0, -8.9_9_2_2], [-47.4322, -5.9_1_4_3, -1.7_3_4_0]]							,	device=__UpperCAmelCase      )
      assert torch.allclose(logits[0, :3, :3]							,	__UpperCAmelCase							,	atol=1e-4      )
   elif model_name == "blip2-flan-t5-xl-coco":
      __SCREAMING_SNAKE_CASE      							=       torch.tensor(
          [[-57.0109, -9.8_9_6_7, -12.6280], [-68.6578, -12.7191, -10.5065]]							,	device=__UpperCAmelCase      )
   else:
      # cast to same type
      __SCREAMING_SNAKE_CASE      							=       logits.dtype
      assert torch.allclose(original_logits.to(__UpperCAmelCase      )							,	__UpperCAmelCase							,	atol=1e-2      )
   print("""Looks ok!"""      )

   print("""Generating a caption..."""      )
   __SCREAMING_SNAKE_CASE      							=       """"""
   __SCREAMING_SNAKE_CASE      							=       tokenizer(__UpperCAmelCase							,	return_tensors="""pt"""      ).input_ids.to(__UpperCAmelCase      )

   __SCREAMING_SNAKE_CASE      							=       original_model.generate({"""image""": original_pixel_values}      )
   __SCREAMING_SNAKE_CASE      							=       hf_model.generate(
       __UpperCAmelCase							,	__UpperCAmelCase							,	do_sample=__UpperCAmelCase							,	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:"""							,	__UpperCAmelCase      )
   __SCREAMING_SNAKE_CASE      							=       input_ids.shape[1]
   __SCREAMING_SNAKE_CASE      							=       processor.batch_decode(outputs[:, prompt_length:]							,	skip_special_tokens=__UpperCAmelCase      )
   __SCREAMING_SNAKE_CASE      							=       [text.strip() for text in output_text]
   print("""HF generation:"""							,	__UpperCAmelCase      )

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

   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        =						argparse.ArgumentParser()
    a        =						[
        "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        =						parser.parse_args()

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












'''simple docstring'''

import os
import string
import sys


a        =						1 << 8

a        =						{
    "tab": ord("\t"),
    "newline": ord("\r"),
    "esc": 27,
    "up": 65 + ARROW_KEY_FLAG,
    "down": 66 + ARROW_KEY_FLAG,
    "right": 67 + ARROW_KEY_FLAG,
    "left": 68 + ARROW_KEY_FLAG,
    "mod_int": 91,
    "undefined": sys.maxsize,
    "interrupt": 3,
    "insert": 50,
    "delete": 51,
    "pg_up": 53,
    "pg_down": 54,
}

a        =						KEYMAP["up"]
a        =						KEYMAP["left"]

if sys.platform == "win32":
    a        =						[]
    a        =						{
        b"\xe0H": KEYMAP["up"] - ARROW_KEY_FLAG,
        b"\x00H": KEYMAP["up"] - ARROW_KEY_FLAG,
        b"\xe0P": KEYMAP["down"] - ARROW_KEY_FLAG,
        b"\x00P": KEYMAP["down"] - ARROW_KEY_FLAG,
        b"\xe0M": KEYMAP["right"] - ARROW_KEY_FLAG,
        b"\x00M": KEYMAP["right"] - ARROW_KEY_FLAG,
        b"\xe0K": KEYMAP["left"] - ARROW_KEY_FLAG,
        b"\x00K": KEYMAP["left"] - ARROW_KEY_FLAG,
    }

for i in range(10):
    a        =						ord(str(i))

def 						__magic_name__       (							)						->					Union[str, Any]:

   '''simple docstring'''







   if os.name == "nt":
      import msvcrt

      __SCREAMING_SNAKE_CASE      							=       """mbcs"""
      # Flush the keyboard buffer
      while msvcrt.kbhit():
         msvcrt.getch()
      if len(__UpperCAmelCase      ) == 0:
         # Read the keystroke
         __SCREAMING_SNAKE_CASE      							=       msvcrt.getch()

         # If it is a prefix char, get second part
         if ch in (b"\x00", b"\xe0"):
            __SCREAMING_SNAKE_CASE      							=       ch + msvcrt.getch()
            # Translate actual Win chars to bullet char types
            try:
               __SCREAMING_SNAKE_CASE      							=       chr(WIN_KEYMAP[cha]      )
               WIN_CH_BUFFER.append(chr(KEYMAP["""mod_int"""]      )      )
               WIN_CH_BUFFER.append(__UpperCAmelCase      )
               if ord(__UpperCAmelCase      ) in (
                   KEYMAP["insert"] - 1 << 9,
                   KEYMAP["delete"] - 1 << 9,
                   KEYMAP["pg_up"] - 1 << 9,
                   KEYMAP["pg_down"] - 1 << 9,
               ):
                  WIN_CH_BUFFER.append(chr(126      )      )
               __SCREAMING_SNAKE_CASE      							=       chr(KEYMAP["""esc"""]      )
            except KeyError:
               __SCREAMING_SNAKE_CASE      							=       cha[1]
         else:
            __SCREAMING_SNAKE_CASE      							=       ch.decode(__UpperCAmelCase      )
      else:
         __SCREAMING_SNAKE_CASE      							=       WIN_CH_BUFFER.pop(0      )
   elif os.name == "posix":
      import termios
      import tty

      __SCREAMING_SNAKE_CASE      							=       sys.stdin.fileno()
      __SCREAMING_SNAKE_CASE      							=       termios.tcgetattr(__UpperCAmelCase      )
      try:
         tty.setraw(__UpperCAmelCase      )
         __SCREAMING_SNAKE_CASE      							=       sys.stdin.read(1      )
      finally:
         termios.tcsetattr(__UpperCAmelCase							,	termios.TCSADRAIN							,	__UpperCAmelCase      )
   return ch







def 						__magic_name__       (							)						->					List[str]:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE      							=       get_raw_chars()
   if ord(__UpperCAmelCase      ) in [KEYMAP["interrupt"], KEYMAP["newline"]]:
      return char

   elif ord(__UpperCAmelCase      ) == KEYMAP["esc"]:
      __SCREAMING_SNAKE_CASE      							=       get_raw_chars()
      if ord(__UpperCAmelCase      ) == KEYMAP["mod_int"]:
         __SCREAMING_SNAKE_CASE      							=       get_raw_chars()
         if ord(__UpperCAmelCase      ) >= KEYMAP["arrow_begin"] - ARROW_KEY_FLAG and ord(__UpperCAmelCase      ) <= KEYMAP["arrow_end"] - ARROW_KEY_FLAG:
            return chr(ord(__UpperCAmelCase      ) + ARROW_KEY_FLAG      )
         else:
            return KEYMAP["undefined"]
      else:
         return get_raw_chars()

   else:
      if char in string.printable:
         return char
      else:
         return KEYMAP["undefined"]





