


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





     '''simple docstring'''






     __UpperCAmelCase     :			int  							=      '''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
     __UpperCAmelCase     :			Optional[Any]  							=      Image.open(requests.get(_lowercase      , stream=_lowercase  ).raw  ).convert('''RGB'''  )

     return image




def      _a     (						_lowercase :     List[str]  ):





     '''simple docstring'''






     __UpperCAmelCase     :			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      _a     (						_lowercase :     Dict      , _lowercase :     Any      , _lowercase :     int  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Optional[Any]  							=      dct.pop(_lowercase  )
     __UpperCAmelCase     :			Optional[int]  							=      val




def      _a     (						_lowercase :     Optional[Any]      , _lowercase :     Dict  ):





     '''simple docstring'''






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

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




def      _a     (						_lowercase :     Dict      , _lowercase :     Any  ):





     '''simple docstring'''






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

     __UpperCAmelCase     :			Dict  							=      BlipaConfig(vision_config=_lowercase      , text_config=_lowercase  )

     return config, image_size




@torch.no_grad()
def      _a     (						_lowercase :     Optional[Any]      , _lowercase :     Union[str, Any]=None      , _lowercase :     Union[str, Any]=False  ):





     '''simple docstring'''






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

     __UpperCAmelCase     :			Any  							=      BlipaForConditionalGeneration(_lowercase  ).eval()

     __UpperCAmelCase     :			List[str]  							=      {
         '''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     :			Optional[int]  							=      model_name_to_original[model_name]

     # load original model
     print('''Loading original model...'''  )
     __UpperCAmelCase     :			Any  							=      '''cuda''' if torch.cuda.is_available() else '''cpu'''
     __UpperCAmelCase							,       __UpperCAmelCase							,       __UpperCAmelCase     :			int  							=      load_model_and_preprocess(
         name=_lowercase      , model_type=_lowercase      , is_eval=_lowercase      , device=_lowercase  )
     original_model.eval()
     print('''Done!'''  )

     # update state dict keys
     __UpperCAmelCase     :			Optional[int]  							=      original_model.state_dict()
     __UpperCAmelCase     :			Union[str, Any]  							=      create_rename_keys(_lowercase  )
     for src, dest in rename_keys:
          rename_key(_lowercase      , _lowercase      , _lowercase  )

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

     # read in qv biases
     read_in_q_v_bias(_lowercase      , _lowercase  )

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

     __UpperCAmelCase     :			Optional[int]  							=      load_demo_image()
     __UpperCAmelCase     :			Any  							=      vis_processors['''eval'''](_lowercase  ).unsqueeze(0  ).to(_lowercase  )
     __UpperCAmelCase     :			List[str]  							=      tokenizer(['''\n''']      , return_tensors='''pt'''  ).input_ids.to(_lowercase  )

     # create processor
     __UpperCAmelCase     :			Dict  							=      BlipImageProcessor(
         size={'''height''': image_size, '''width''': image_size}      , image_mean=_lowercase      , image_std=_lowercase  )
     __UpperCAmelCase     :			List[Any]  							=      BlipaProcessor(image_processor=_lowercase      , tokenizer=_lowercase  )
     __UpperCAmelCase     :			Union[str, Any]  							=      processor(images=_lowercase      , return_tensors='''pt'''  ).pixel_values.to(_lowercase  )

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

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

     print('''Generating a caption...'''  )
     __UpperCAmelCase     :			Any  							=      ''''''
     __UpperCAmelCase     :			List[Any]  							=      tokenizer(_lowercase      , return_tensors='''pt'''  ).input_ids.to(_lowercase  )

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

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

     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      :str			    =       [
          "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      :Union[str, Any]			    =       parser.parse_args()

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



'''simple docstring'''



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase      :List[Any]			    =       {
    "configuration_lxmert": ["LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "LxmertConfig"],
    "tokenization_lxmert": ["LxmertTokenizer"],
}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase      :List[Any]			    =       ["LxmertTokenizerFast"]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase      :List[Any]			    =       [
          "LxmertEncoder",
          "LxmertForPreTraining",
          "LxmertForQuestionAnswering",
          "LxmertModel",
          "LxmertPreTrainedModel",
          "LxmertVisualFeatureEncoder",
          "LxmertXLayer",
      ]

try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase      :List[str]			    =       [
          "TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST",
          "TFLxmertForPreTraining",
          "TFLxmertMainLayer",
          "TFLxmertModel",
          "TFLxmertPreTrainedModel",
          "TFLxmertVisualFeatureEncoder",
      ]


if TYPE_CHECKING:
      from .configuration_lxmert import LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP, LxmertConfig
      from .tokenization_lxmert import LxmertTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_lxmert_fast import LxmertTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_lxmert import (
                LxmertEncoder,
                LxmertForPreTraining,
                LxmertForQuestionAnswering,
                LxmertModel,
                LxmertPreTrainedModel,
                LxmertVisualFeatureEncoder,
                LxmertXLayer,
            )

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_lxmert import (
                TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFLxmertForPreTraining,
                TFLxmertMainLayer,
                TFLxmertModel,
                TFLxmertPreTrainedModel,
                TFLxmertVisualFeatureEncoder,
            )

else:
      import sys

      __UpperCAmelCase      :int			    =       _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)