


'''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     UpperCamelCase_      (   ):
  '''simple docstring'''





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

  return image







def     UpperCamelCase_      (   A__ :       Any ):
  '''simple docstring'''





  lowerCAmelCase_      :				List[str]      =       []
  # 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_      (   A__ :       Optional[int]							,   A__ :       int							,   A__ :       int ):
  '''simple docstring'''





  lowerCAmelCase_      :				Optional[int]      =       dct.pop(snake_case_ )
  lowerCAmelCase_      :				Optional[int]      =       val







def     UpperCamelCase_      (   A__ :       Optional[Any]							,   A__ :       Optional[Any] ):
  '''simple docstring'''





  for i in range(config.vision_config.num_hidden_layers ):
    # read in original q and v biases
    lowerCAmelCase_      :				int      =       state_dict.pop(f'visual_encoder.blocks.{i}.attn.q_bias' )
    lowerCAmelCase_      :				Optional[int]      =       state_dict.pop(f'visual_encoder.blocks.{i}.attn.v_bias' )

    # next, set bias in the state dict
    lowerCAmelCase_      :				Dict      =       torch.cat((q_bias, torch.zeros_like(snake_case_							,   requires_grad=snake_case_ ), v_bias) )
    lowerCAmelCase_      :				str      =       qkv_bias







def     UpperCamelCase_      (   A__ :       List[str]							,   A__ :       Optional[int] ):
  '''simple docstring'''





  lowerCAmelCase_      :				str      =       3_64 if """coco""" in model_name else 2_24
  lowerCAmelCase_      :				int      =       BlipaVisionConfig(image_size=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:
    lowerCAmelCase_      :				List[str]      =       OPTConfig.from_pretrained("""facebook/opt-2.7b"""							,   eos_token_id=snake_case_ ).to_dict()
  elif "opt-6.7b" in model_name:
    lowerCAmelCase_      :				str      =       OPTConfig.from_pretrained("""facebook/opt-6.7b"""							,   eos_token_id=snake_case_ ).to_dict()
  elif "t5-xl" in model_name:
    lowerCAmelCase_      :				List[Any]      =       TaConfig.from_pretrained("""google/flan-t5-xl"""							,   dense_act_fn="""gelu"""							,   bos_token_id=1 ).to_dict()
  elif "t5-xxl" in model_name:
    lowerCAmelCase_      :				Dict      =       TaConfig.from_pretrained("""google/flan-t5-xxl"""							,   dense_act_fn="""gelu"""							,   bos_token_id=1 ).to_dict()

  lowerCAmelCase_      :				Tuple      =       BlipaConfig(vision_config=snake_case_							,   text_config=snake_case_ )

  return config, image_size







@torch.no_grad()
def     UpperCamelCase_      (   A__ :       int							,   A__ :       Optional[int]=None							,   A__ :       Optional[Any]=False ):
  '''simple docstring'''





  lowerCAmelCase_      :				List[Any]      =       (
      AutoTokenizer.from_pretrained("""facebook/opt-2.7b""" )
      if """opt""" in model_name
      else AutoTokenizer.from_pretrained("""google/flan-t5-xl""" )
  )
  lowerCAmelCase_      :				Tuple      =       tokenizer("""\n"""							,   add_special_tokens=snake_case_ ).input_ids[0]
  lowerCAmelCase_,							lowerCAmelCase_      :				Union[str, Any]      =       get_blipa_config(snake_case_							,   eos_token_id=snake_case_ )

  lowerCAmelCase_      :				List[Any]      =       BlipaForConditionalGeneration(snake_case_ ).eval()

  lowerCAmelCase_      :				Union[str, Any]      =       {
      """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"""),
  }

  lowerCAmelCase_,							lowerCAmelCase_      :				Union[str, Any]      =       model_name_to_original[model_name]

  # load original model
  print("""Loading original model...""" )
  lowerCAmelCase_      :				Union[str, Any]      =       """cuda""" if torch.cuda.is_available() else """cpu"""
  lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_      :				int      =       load_model_and_preprocess(
      name=snake_case_							,   model_type=snake_case_							,   is_eval=snake_case_							,   device=snake_case_ )
  original_model.eval()
  print("""Done!""" )

  # update state dict keys
  lowerCAmelCase_      :				Any      =       original_model.state_dict()
  lowerCAmelCase_      :				int      =       create_rename_keys(snake_case_ )
  for src, dest in rename_keys:
    rename_key(snake_case_							,   snake_case_							,   snake_case_ )

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

  # read in qv biases
  read_in_q_v_bias(snake_case_							,   snake_case_ )

  lowerCAmelCase_,							lowerCAmelCase_      :				List[Any]      =       hf_model.load_state_dict(snake_case_							,   strict=snake_case_ )
  assert len(snake_case_ ) == 0
  assert unexpected_keys == ["qformer.embeddings.position_ids"]

  lowerCAmelCase_      :				Dict      =       load_demo_image()
  lowerCAmelCase_      :				Dict      =       vis_processors["""eval"""](snake_case_ ).unsqueeze(0 ).to(snake_case_ )
  lowerCAmelCase_      :				Optional[int]      =       tokenizer(["""\n"""]							,   return_tensors="""pt""" ).input_ids.to(snake_case_ )

  # create processor
  lowerCAmelCase_      :				List[str]      =       BlipImageProcessor(
      size={"""height""": image_size, """width""": image_size}							,   image_mean=snake_case_							,   image_std=snake_case_ )
  lowerCAmelCase_      :				Optional[int]      =       BlipaProcessor(image_processor=snake_case_							,   tokenizer=snake_case_ )
  lowerCAmelCase_      :				int      =       processor(images=snake_case_							,   return_tensors="""pt""" ).pixel_values.to(snake_case_ )

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

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

  print("""Generating a caption...""" )
  lowerCAmelCase_      :				Dict      =       """"""
  lowerCAmelCase_      :				Optional[int]      =       tokenizer(snake_case_							,   return_tensors="""pt""" ).input_ids.to(snake_case_ )

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

  if pytorch_dump_folder_path is not None:
    processor.save_pretrained(snake_case_ )
    hf_model.save_pretrained(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 :  Any    							=				argparse.ArgumentParser()
  __A :  Optional[Any]    							=				[
      '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 :  List[str]    							=				parser.parse_args()

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








'''simple docstring'''






import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from .tokenization_lxmert import LxmertTokenizer


__A :  Any    							=				{"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}

__A :  int    							=				{
    "vocab_file": {
        "unc-nlp/lxmert-base-uncased": "https://huggingface.co/unc-nlp/lxmert-base-uncased/resolve/main/vocab.txt",
    },
    "tokenizer_file": {
        "unc-nlp/lxmert-base-uncased": (
            "https://huggingface.co/unc-nlp/lxmert-base-uncased/resolve/main/tokenizer.json"
        ),
    },
}

__A :  List[Any]    							=				{
    "unc-nlp/lxmert-base-uncased": 512,
}

__A :  Optional[Any]    							=				{
    "unc-nlp/lxmert-base-uncased": {"do_lower_case": True},
}






class        __snake_case	(       _SCREAMING_SNAKE_CASE):







    """simple docstring"""



    lowercase       			= VOCAB_FILES_NAMES
    lowercase       			= PRETRAINED_VOCAB_FILES_MAP
    lowercase       			= PRETRAINED_INIT_CONFIGURATION
    lowercase       			= PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    lowercase       			= LxmertTokenizer

    def __init__(       self   :  Dict   ,						lowerCamelCase   :  Tuple=None   ,						lowerCamelCase   :  List[str]=None   ,						lowerCamelCase   :  Dict=True   ,						lowerCamelCase   :  List[Any]="[UNK]"   ,						lowerCamelCase   :  Tuple="[SEP]"   ,						lowerCamelCase   :  str="[PAD]"   ,						lowerCamelCase   :  Any="[CLS]"   ,						lowerCamelCase   :  str="[MASK]"   ,						lowerCamelCase   :  Dict=True   ,						lowerCamelCase   :  Union[str, Any]=None   ,						**lowerCamelCase   :  Optional[int]   ,						)       ->  Optional[Any]:
      super().__init__(
          lowerCamelCase   ,						tokenizer_file=lowerCamelCase   ,						do_lower_case=lowerCamelCase   ,						unk_token=lowerCamelCase   ,						sep_token=lowerCamelCase   ,						pad_token=lowerCamelCase   ,						cls_token=lowerCamelCase   ,						mask_token=lowerCamelCase   ,						tokenize_chinese_chars=lowerCamelCase   ,						strip_accents=lowerCamelCase   ,						**lowerCamelCase   ,						)

      lowerCAmelCase_      :				Optional[Any]      =       json.loads(self.backend_tokenizer.normalizer.__getstate__()		)
      if (
          normalizer_state.get("""lowercase"""   ,						lowerCamelCase		) != do_lower_case
          or normalizer_state.get("""strip_accents"""   ,						lowerCamelCase		) != strip_accents
          or normalizer_state.get("""handle_chinese_chars"""   ,						lowerCamelCase		) != tokenize_chinese_chars
      ):
        lowerCAmelCase_      :				Tuple      =       getattr(lowerCamelCase   ,						normalizer_state.pop("""type"""		)		)
        lowerCAmelCase_      :				Optional[Any]      =       do_lower_case
        lowerCAmelCase_      :				Dict      =       strip_accents
        lowerCAmelCase_      :				List[Any]      =       tokenize_chinese_chars
        lowerCAmelCase_      :				List[Any]      =       normalizer_class(**lowerCamelCase		)

      lowerCAmelCase_      :				Tuple      =       do_lower_case

    def __lowercase       (       self   :  Tuple   ,						lowerCamelCase   :  str   ,						lowerCamelCase   :  Optional[Any]=None		)       ->  int:
      lowerCAmelCase_      :				Any      =       [self.cls_token_id] + token_ids_a + [self.sep_token_id]

      if token_ids_a:
        output += token_ids_a + [self.sep_token_id]

      return output

    def __lowercase       (       self   :  str   ,						lowerCamelCase   :  List[int]   ,						lowerCamelCase   :  Optional[List[int]] = None		)       ->  List[int]:
      lowerCAmelCase_      :				str      =       [self.sep_token_id]
      lowerCAmelCase_      :				Optional[int]      =       [self.cls_token_id]
      if token_ids_a is None:
        return len(cls + token_ids_a + sep		) * [0]
      return len(cls + token_ids_a + sep		) * [0] + len(token_ids_a + sep		) * [1]





    def __lowercase       (       self   :  Tuple   ,						lowerCamelCase   :  str   ,						lowerCamelCase   :  Optional[str] = None		)       ->  Tuple[str]:
      lowerCAmelCase_      :				int      =       self._tokenizer.model.save(lowerCamelCase   ,						name=lowerCamelCase		)
      return tuple(lowerCamelCase		)



