

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

  return image
def      _lowerCAmelCase(	__A			):
  UpperCAmelCase					      =  []
  # 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(	__A   ,       __A   ,       __A			):
  UpperCAmelCase					      =  dct.pop(__A			)
  UpperCAmelCase					      =  val
def      _lowerCAmelCase(	__A   ,       __A			):
  for i in range(config.vision_config.num_hidden_layers			):
    # read in original q and v biases
    UpperCAmelCase					      =  state_dict.pop(F"visual_encoder.blocks.{i}.attn.q_bias"			)
    UpperCAmelCase					      =  state_dict.pop(F"visual_encoder.blocks.{i}.attn.v_bias"			)

    # next, set bias in the state dict
    UpperCAmelCase					      =  torch.cat((q_bias, torch.zeros_like(__A   ,       requires_grad=__A			), v_bias)			)
    UpperCAmelCase					      =  qkv_bias
def      _lowerCAmelCase(	__A   ,       __A			):
  UpperCAmelCase					      =  364 if "coco" in model_name else 224
  UpperCAmelCase					      =  BlipaVisionConfig(image_size=__A			).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					      =  OPTConfig.from_pretrained("facebook/opt-2.7b"   ,       eos_token_id=__A			).to_dict()
  elif "opt-6.7b" in model_name:
    UpperCAmelCase					      =  OPTConfig.from_pretrained("facebook/opt-6.7b"   ,       eos_token_id=__A			).to_dict()
  elif "t5-xl" in model_name:
    UpperCAmelCase					      =  TaConfig.from_pretrained("google/flan-t5-xl"   ,       dense_act_fn="gelu"   ,       bos_token_id=1			).to_dict()
  elif "t5-xxl" in model_name:
    UpperCAmelCase					      =  TaConfig.from_pretrained("google/flan-t5-xxl"   ,       dense_act_fn="gelu"   ,       bos_token_id=1			).to_dict()

  UpperCAmelCase					      =  BlipaConfig(vision_config=__A   ,       text_config=__A			)

  return config, image_size
@torch.no_grad()
def      _lowerCAmelCase(	__A   ,       __A=None   ,       __A=False			):
  UpperCAmelCase					      =  (
      AutoTokenizer.from_pretrained("facebook/opt-2.7b"			)
      if "opt" in model_name
      else AutoTokenizer.from_pretrained("google/flan-t5-xl"			)
  )
  UpperCAmelCase					      =  tokenizer("\n"   ,       add_special_tokens=__A			).input_ids[0]
  UpperCAmelCase      ,	UpperCAmelCase					      =  get_blipa_config(__A   ,       eos_token_id=__A			)

  UpperCAmelCase					      =  BlipaForConditionalGeneration(__A			).eval()

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

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

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

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

  # read in qv biases
  read_in_q_v_bias(__A   ,       __A			)

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

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

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

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

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

  print("Generating a caption..."			)
  UpperCAmelCase					      =  ""
  UpperCAmelCase					      =  tokenizer(__A   ,       return_tensors="pt"			).input_ids.to(__A			)

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

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

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


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

      lowerCAmelCase__       = parser.parse_args()

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






from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase__       = logging.get_logger(__name__)

lowerCAmelCase__       = {
    "unc-nlp/lxmert-base-uncased": "https://huggingface.co/unc-nlp/lxmert-base-uncased/resolve/main/config.json",
}






class 						__magic_name__    (       _snake_case    ):
 UpperCAmelCase							    =						"""lxmert"""
 UpperCAmelCase							    =						{}


 def __init__(						self		: int			,						lowerCAmelCase__		: Any=3_0_5_2_2			,						lowerCAmelCase__		: List[str]=7_6_8			,						lowerCAmelCase__		: Union[str, Any]=1_2			,						lowerCAmelCase__		: List[Any]=9_5_0_0			,						lowerCAmelCase__		: Any=1_6_0_0			,						lowerCAmelCase__		: Union[str, Any]=4_0_0			,						lowerCAmelCase__		: Tuple=3_0_7_2			,						lowerCAmelCase__		: Dict="gelu"			,						lowerCAmelCase__		: Tuple=0.1			,						lowerCAmelCase__		: Tuple=0.1			,						lowerCAmelCase__		: int=5_1_2			,						lowerCAmelCase__		: List[str]=2			,						lowerCAmelCase__		: List[str]=0.02			,						lowerCAmelCase__		: str=1e-1_2			,						lowerCAmelCase__		: str=9			,						lowerCAmelCase__		: int=5			,						lowerCAmelCase__		: Optional[int]=5			,						lowerCAmelCase__		: List[Any]=2_0_4_8			,						lowerCAmelCase__		: Any=4			,						lowerCAmelCase__		: Dict=6.67			,						lowerCAmelCase__		: Any=True			,						lowerCAmelCase__		: Union[str, Any]=True			,						lowerCAmelCase__		: Any=True			,						lowerCAmelCase__		: Tuple=True			,						lowerCAmelCase__		: Optional[Any]=True			,						lowerCAmelCase__		: Optional[int]=True			,						lowerCAmelCase__		: Tuple=True			,						**lowerCAmelCase__		: List[Any]			,						)    ->      Dict:
   UpperCAmelCase					      =  vocab_size
   UpperCAmelCase					      =  hidden_size
   UpperCAmelCase					      =  num_attention_heads
   UpperCAmelCase					      =  hidden_act
   UpperCAmelCase					      =  intermediate_size
   UpperCAmelCase					      =  hidden_dropout_prob
   UpperCAmelCase					      =  attention_probs_dropout_prob
   UpperCAmelCase					      =  max_position_embeddings
   UpperCAmelCase					      =  type_vocab_size
   UpperCAmelCase					      =  initializer_range
   UpperCAmelCase					      =  layer_norm_eps
   UpperCAmelCase					      =  num_qa_labels
   UpperCAmelCase					      =  num_object_labels
   UpperCAmelCase					      =  num_attr_labels
   UpperCAmelCase					      =  l_layers
   UpperCAmelCase					      =  x_layers
   UpperCAmelCase					      =  r_layers
   UpperCAmelCase					      =  visual_feat_dim
   UpperCAmelCase					      =  visual_pos_dim
   UpperCAmelCase					      =  visual_loss_normalizer
   UpperCAmelCase					      =  task_matched
   UpperCAmelCase					      =  task_mask_lm
   UpperCAmelCase					      =  task_obj_predict
   UpperCAmelCase					      =  task_qa
   UpperCAmelCase					      =  visual_obj_loss
   UpperCAmelCase					      =  visual_attr_loss
   UpperCAmelCase					      =  visual_feat_loss
   UpperCAmelCase					      =  {"vision": r_layers, "cross_encoder": x_layers, "language": l_layers}
   super().__init__(**lowerCAmelCase__    )



