


'''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 _SCREAMING_SNAKE_CASE      ()       ->       Union[str, Any]:






       """simple docstring"""







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

       return image

def _SCREAMING_SNAKE_CASE      (A       )       ->       Tuple:






       """simple docstring"""







       lowercase__      			= []
       # 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      (A  , A  , A       )       ->       Union[str, Any]:






       """simple docstring"""







       lowercase__      			= dct.pop(A       )
       lowercase__      			= val

def _SCREAMING_SNAKE_CASE      (A  , A       )       ->       int:






       """simple docstring"""







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

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

def _SCREAMING_SNAKE_CASE      (A  , A       )       ->       List[str]:






       """simple docstring"""







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

       lowercase__      			= BlipaConfig(vision_config=A  , text_config=A       )

       return config, image_size

@torch.no_grad()
def _SCREAMING_SNAKE_CASE      (A  , A=None  , A=False       )       ->       int:






       """simple docstring"""







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

       lowercase__      			= BlipaForConditionalGeneration(A       ).eval()

       lowercase__      			= {
           '''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'''),
       }

       lowercase__			,lowercase__      			= model_name_to_original[model_name]

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

       # update state dict keys
       lowercase__      			= original_model.state_dict()
       lowercase__      			= 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():
              lowercase__      			= state_dict.pop(A       )
              if key.startswith('''Qformer.bert'''       ):
                     lowercase__      			= key.replace('''Qformer.bert'''  , '''qformer'''       )
              if "attention.self" in key:
                     lowercase__      			= key.replace('''self'''  , '''attention'''       )
              if "opt_proj" in key:
                     lowercase__      			= key.replace('''opt_proj'''  , '''language_projection'''       )
              if "t5_proj" in key:
                     lowercase__      			= key.replace('''t5_proj'''  , '''language_projection'''       )
              if key.startswith('''opt'''       ):
                     lowercase__      			= key.replace('''opt'''  , '''language'''       )
              if key.startswith('''t5'''       ):
                     lowercase__      			= key.replace('''t5'''  , '''language'''       )
              lowercase__      			= val

       # read in qv biases
       read_in_q_v_bias(A  , A       )

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

       lowercase__      			= load_demo_image()
       lowercase__      			= vis_processors['''eval'''](A       ).unsqueeze(0       ).to(A       )
       lowercase__      			= tokenizer(['''\n''']  , return_tensors='''pt'''       ).input_ids.to(A       )

       # create processor
       lowercase__      			= BlipImageProcessor(
           size={'''height''': image_size, '''width''': image_size}  , image_mean=A  , image_std=A       )
       lowercase__      			= BlipaProcessor(image_processor=A  , tokenizer=A       )
       lowercase__      			= 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:
                     lowercase__      			= original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}       ).logits
                     lowercase__      			= hf_model(A  , A       ).logits
              else:
                     lowercase__      			= original_model(
                         {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}       ).logits
                     lowercase__      			= input_ids.masked_fill(input_ids == tokenizer.pad_token_id  , -100       )
                     lowercase__      			= 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":
              lowercase__      			= torch.tensor(
                  [[-41.5_850, -4.4_440, -8.9_922], [-47.4_322, -5.9_143, -1.7_340]]  , device=A       )
              assert torch.allclose(logits[0, :3, :3]  , A  , atol=1E-4       )
       elif model_name == "blip2-flan-t5-xl-coco":
              lowercase__      			= torch.tensor(
                  [[-57.0_109, -9.8_967, -12.6_280], [-68.6_578, -12.7_191, -10.5_065]]  , device=A       )
       else:
              # cast to same type
              lowercase__      			= logits.dtype
              assert torch.allclose(original_logits.to(A       )  , A  , atol=1E-2       )
       print('''Looks ok!'''       )

       print('''Generating a caption...'''       )
       lowercase__      			= ''''''
       lowercase__      			= tokenizer(A  , return_tensors='''pt'''       ).input_ids.to(A       )

       lowercase__      			= original_model.generate({'''image''': original_pixel_values}       )
       lowercase__      			= 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       )
       lowercase__      			= input_ids.shape[1]
       lowercase__      			= processor.batch_decode(outputs[:, prompt_length:]  , skip_special_tokens=A       )
       lowercase__      			= [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							:   List[Any]  						=	argparse.ArgumentParser()
 lowerCamelCase							:   Optional[int]  						=	[
     '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							:   List[Any]  						=	parser.parse_args()

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








'''simple docstring'''

import argparse

import torch
from torch import nn

from transformers import MBartConfig, MBartForConditionalGeneration

def _SCREAMING_SNAKE_CASE      (A       )       ->       Optional[Any]:






       """simple docstring"""







       lowercase__      			= [
           '''encoder.version''',
           '''decoder.version''',
           '''model.encoder.version''',
           '''model.decoder.version''',
           '''_float_tensor''',
           '''decoder.output_projection.weight''',
       ]
       for k in ignore_keys:
              state_dict.pop(A  , A       )

def _SCREAMING_SNAKE_CASE      (A       )       ->       List[str]:






       """simple docstring"""







       lowercase__			,lowercase__      			= emb.weight.shape
       lowercase__      			= nn.Linear(A  , A  , bias=A       )
       lowercase__      			= emb.weight.data
       return lin_layer

def _SCREAMING_SNAKE_CASE      (A  , A="facebook/mbart-large-en-ro"  , A=False  , A=False       )       ->       Union[str, Any]:






       """simple docstring"""







       lowercase__      			= torch.load(A  , map_location='''cpu'''       )['''model''']
       remove_ignore_keys_(A       )
       lowercase__      			= state_dict['''encoder.embed_tokens.weight'''].shape[0]

       lowercase__      			= MBartConfig.from_pretrained(A  , vocab_size=A       )
       if mbart_aa and finetuned:
              lowercase__      			= '''relu'''

       lowercase__      			= state_dict['''decoder.embed_tokens.weight''']
       lowercase__      			= MBartForConditionalGeneration(A       )
       model.model.load_state_dict(A       )

       if finetuned:
              lowercase__      			= make_linear_from_emb(model.model.shared       )

       return model


if __name__ == "__main__":
 lowerCamelCase							:   Any  						=	argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     'fairseq_path', type=str, help='bart.large, bart.large.cnn or a path to a model.pt on local filesystem.'
 )
 parser.add_argument('pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
 parser.add_argument(
     '--hf_config',
     default='facebook/mbart-large-cc25',
     type=str,
     help='Which huggingface architecture to use: mbart-large',
 )
 parser.add_argument('--mbart_50', action='store_true', help='whether the model is mMART-50 checkpoint')
 parser.add_argument('--finetuned', action='store_true', help='whether the model is a fine-tuned checkpoint')
 lowerCamelCase							:   Any  						=	parser.parse_args()
 lowerCamelCase							:   List[str]  						=	convert_fairseq_mbart_checkpoint_from_disk(
     args.fairseq_path, hf_config_path=args.hf_config, finetuned=args.finetuned, mbart_aa=args.mbart_aa
 )
 model.save_pretrained(args.pytorch_dump_folder_path)




