






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

       return image




def       _SCREAMING_SNAKE_CASE (							__snake_case )				->						List[Any]:
       _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       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case   ,      __snake_case )				->						Dict:
       _UpperCAmelCase							=    dct.pop(__snake_case )
       _UpperCAmelCase							=    val




def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case )				->						List[Any]:
       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(__snake_case   ,      requires_grad=__snake_case ), v_bias) )
              _UpperCAmelCase							=    qkv_bias




def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case )				->						List[str]:
       _UpperCAmelCase							=    3_6_4 if """coco""" in model_name else 2_2_4
       _UpperCAmelCase							=    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:
              _UpperCAmelCase							=    OPTConfig.from_pretrained("""facebook/opt-2.7b"""   ,      eos_token_id=__snake_case ).to_dict()
       elif "opt-6.7b" in model_name:
              _UpperCAmelCase							=    OPTConfig.from_pretrained("""facebook/opt-6.7b"""   ,      eos_token_id=__snake_case ).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=__snake_case   ,      text_config=__snake_case )

       return config, image_size




@torch.no_grad()
def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case=None   ,      __snake_case=False )				->						Optional[Any]:
       _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=__snake_case ).input_ids[0]
       _UpperCAmelCase							=    get_blipa_config(__snake_case   ,      eos_token_id=__snake_case )

       _UpperCAmelCase							=    BlipaForConditionalGeneration(__snake_case ).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							=    model_name_to_original[model_name]

       # load original model
       print("""Loading original model...""" )
       _UpperCAmelCase							=    """cuda""" if torch.cuda.is_available() else """cpu"""
       _UpperCAmelCase							=    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
       _UpperCAmelCase							=    original_model.state_dict()
       _UpperCAmelCase							=    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():
              _UpperCAmelCase							=    state_dict.pop(__snake_case )
              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(__snake_case   ,      __snake_case )

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

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

       # create processor
       _UpperCAmelCase							=    BlipImageProcessor(
           size={"""height""": image_size, """width""": image_size}   ,      image_mean=__snake_case   ,      image_std=__snake_case )
       _UpperCAmelCase							=    BlipaProcessor(image_processor=__snake_case   ,      tokenizer=__snake_case )
       _UpperCAmelCase							=    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:
                     _UpperCAmelCase							=    original_model({"""image""": original_pixel_values, """text_input""": [""""""]} ).logits
                     _UpperCAmelCase							=    hf_model(__snake_case   ,      __snake_case ).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   ,      -1_0_0 )
                     _UpperCAmelCase							=    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":
              _UpperCAmelCase							=    torch.tensor(
                  [[-4_1.5_8_5_0, -4.4440, -8.9922], [-4_7.4_3_2_2, -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":
              _UpperCAmelCase							=    torch.tensor(
                  [[-5_7.0_1_0_9, -9.8967, -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=__snake_case )
       else:
              # cast to same type
              _UpperCAmelCase							=    logits.dtype
              assert torch.allclose(original_logits.to(__snake_case )   ,      __snake_case   ,      atol=1E-2 )
       print("""Looks ok!""" )

       print("""Generating a caption...""" )
       _UpperCAmelCase							=    """"""
       _UpperCAmelCase							=    tokenizer(__snake_case   ,      return_tensors="""pt""" ).input_ids.to(__snake_case )

       _UpperCAmelCase							=    original_model.generate({"""image""": original_pixel_values} )
       _UpperCAmelCase							=    hf_model.generate(
           __snake_case   ,      __snake_case   ,      do_sample=__snake_case   ,      num_beams=5   ,      max_length=3_0   ,      min_length=1   ,      top_p=0.9   ,      repetition_penalty=1.0   ,      length_penalty=1.0   ,      temperature=1   ,      )
       print("""Original generation:"""   ,      __snake_case )
       _UpperCAmelCase							=    input_ids.shape[1]
       _UpperCAmelCase							=    processor.batch_decode(outputs[:, prompt_length:]   ,      skip_special_tokens=__snake_case )
       _UpperCAmelCase							=    [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:     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:     Optional[int]   					=					parser.parse_args()

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





from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCAmelCase_		:					Tuple          =  logging.get_logger(__name__)

UpperCAmelCase_		:					List[str]          =  {
    "facebook/xlm-roberta-xl": "https://huggingface.co/facebook/xlm-roberta-xl/resolve/main/config.json",
    "facebook/xlm-roberta-xxl": "https://huggingface.co/facebook/xlm-roberta-xxl/resolve/main/config.json",
    # See all XLM-RoBERTa-XL models at https://huggingface.co/models?filter=xlm-roberta-xl
}




class 		__A   (      UpperCamelCase__      ):
						UpperCamelCase      =       """xlm-roberta-xl"""

						def __init__(		self :Dict     ,			__snake_case :Optional[Any]=25_08_80     ,			__snake_case :List[Any]=25_60     ,			__snake_case :Optional[Any]=36     ,			__snake_case :Any=32     ,			__snake_case :int=1_02_40     ,			__snake_case :List[Any]="gelu"     ,			__snake_case :Union[str, Any]=0.1     ,			__snake_case :Optional[Any]=0.1     ,			__snake_case :str=5_14     ,			__snake_case :Union[str, Any]=1     ,			__snake_case :Optional[int]=0.02     ,			__snake_case :str=1E-05     ,			__snake_case :str=1     ,			__snake_case :int=0     ,			__snake_case :Tuple=2     ,			__snake_case :Optional[int]="absolute"     ,			__snake_case :str=True     ,			__snake_case :Any=None     ,			**__snake_case :Dict     ,			):
										'''simple docstring'''

										super().__init__(pad_token_id=__snake_case     ,			bos_token_id=__snake_case     ,			eos_token_id=__snake_case     ,			**__snake_case			)
										__magic_name__						:						List[str]										=vocab_size
										__magic_name__						:						List[str]										=hidden_size
										__magic_name__						:						Union[str, Any]										=num_hidden_layers
										__magic_name__						:						Any										=num_attention_heads
										__magic_name__						:						Any										=hidden_act
										__magic_name__						:						List[str]										=intermediate_size
										__magic_name__						:						Any										=hidden_dropout_prob
										__magic_name__						:						Union[str, Any]										=attention_probs_dropout_prob
										__magic_name__						:						Any										=max_position_embeddings
										__magic_name__						:						Any										=type_vocab_size
										__magic_name__						:						List[str]										=initializer_range
										__magic_name__						:						Optional[int]										=layer_norm_eps
										__magic_name__						:						Dict										=position_embedding_type
										__magic_name__						:						Any										=use_cache
										__magic_name__						:						Dict										=classifier_dropout




class 		__A   (      UpperCamelCase__      ):

						@property
						def      A__    (		self :Dict			):
										'''simple docstring'''

										if self.task == "multiple-choice":
														__magic_name__						:						str										={0: """batch""", 1: """choice""", 2: """sequence"""}
										else:
														__magic_name__						:						Optional[Any]										={0: """batch""", 1: """sequence"""}
										return OrderedDict(
										    [
										        ("""input_ids""", dynamic_axis),
										        ("""attention_mask""", dynamic_axis),
										    ]			)




