




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_float32 (there's also the fix_lavis branch)
# also note: to convert Vicuna checkpoints, we had to include /home/niels/python_projects/checkpoints/FastChat/vicuna-7b in lavis/configs/models/blip2/blip2_instruct_vicuna7b.yaml
# same for Vicuna-13b
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipImageProcessor,
    InstructBlipConfig,
    InstructBlipForConditionalGeneration,
    InstructBlipProcessor,
    InstructBlipQFormerConfig,
    InstructBlipVisionConfig,
    LlamaConfig,
    LlamaTokenizerFast,
    TaConfig,
    TaTokenizerFast,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
def 		a_							( ):
  lowerCAmelCase__										=      '''https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'''
  lowerCAmelCase__										=      Image.open(requests.get(__lowerCAmelCase   , stream=__lowerCAmelCase				).raw				).convert('''RGB'''				)

  return image
def 		a_							( __lowerCAmelCase				):
  lowerCAmelCase__										=      []
  # 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.embeddings.layernorm.weight''')				)
  rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.bias''', '''qformer.embeddings.layernorm.bias''')				)

  # fmt: on
  return rename_keys
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase   , __lowerCAmelCase				):
  lowerCAmelCase__										=      dct.pop(__lowerCAmelCase				)
  lowerCAmelCase__										=      val
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase				):




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

    # next, set bias in the state dict
    lowerCAmelCase__										=      torch.cat((q_bias, torch.zeros_like(__lowerCAmelCase   , requires_grad=__lowerCAmelCase				), v_bias)				)
    lowerCAmelCase__										=      qkv_bias
def 		a_							( __lowerCAmelCase				):
  lowerCAmelCase__										=      3_64 if '''coco''' in model_name else 2_24
  lowerCAmelCase__										=      InstructBlipVisionConfig(image_size=__lowerCAmelCase				).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 "t5-xl" in model_name:
    lowerCAmelCase__										=      TaConfig.from_pretrained('''google/flan-t5-xl'''   , dense_act_fn='''gelu'''   , bos_token_id=1				).to_dict()
  elif "t5-xxl" in model_name:
    lowerCAmelCase__										=      TaConfig.from_pretrained('''google/flan-t5-xxl'''   , dense_act_fn='''gelu'''   , bos_token_id=1				).to_dict()
  elif "vicuna-7b" in model_name:
    lowerCAmelCase__										=      LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''   , vocab_size=3_20_01				).to_dict()
  elif "vicuna-13b" in model_name:
    lowerCAmelCase__										=      LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf'''   , vocab_size=3_20_01				).to_dict()
  else:
    raise ValueError('''Model name not supported'''				)

  # the authors add one special "[DEC]" token to the vocab of Q-Former, hence vocab size = 30522 + 1
  lowerCAmelCase__										=      InstructBlipQFormerConfig(vocab_size=3_05_23				).to_dict()
  lowerCAmelCase__										=      InstructBlipConfig(vision_config=__lowerCAmelCase   , text_config=__lowerCAmelCase   , qformer_config=__lowerCAmelCase				)

  return config, image_size
@torch.no_grad()
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase=None   , __lowerCAmelCase=False				):
  lowerCAmelCase__										=      AutoTokenizer.from_pretrained('''bert-base-uncased'''   , truncation_side='''left'''				)
  qformer_tokenizer.add_special_tokens({'''bos_token''': '''[DEC]'''}				)

  if "t5" in model_name:
    lowerCAmelCase__										=      TaTokenizerFast.from_pretrained('''google/flan-t5-xl'''   , truncation_side='''left'''				)
  elif "vicuna" in model_name:
    # the following was used in the original implementation:
    # tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b", use_fast=False, truncation_side="left")
    # tokenizer.add_special_tokens({"pad_token": "[PAD]"})
    # tokenizer.add_special_tokens({"bos_token": "</s>"})
    # tokenizer.add_special_tokens({"eos_token": "</s>"})
    # tokenizer.add_special_tokens({"unk_token": "</s>"})
    lowerCAmelCase__										=      LlamaTokenizerFast.from_pretrained(
        '''huggyllama/llama-7b'''   , truncation_side='''left'''   , bos_token='''</s>'''   , unk_token='''</s>'''				)
    tokenizer.add_special_tokens({'''pad_token''': '''[PAD]'''}				)

  lowerCAmelCase__				,      lowerCAmelCase__										=      get_blipa_config(__lowerCAmelCase				)
  lowerCAmelCase__										=      InstructBlipForConditionalGeneration(__lowerCAmelCase				).eval()

  lowerCAmelCase__										=      {
      '''instructblip-vicuna-7b''': ('''blip2_vicuna_instruct''', '''vicuna7b'''),
      '''instructblip-vicuna-13b''': ('''blip2_vicuna_instruct''', '''vicuna13b'''),
      '''instructblip-flan-t5-xl''': ('''blip2_t5_instruct''', '''flant5xl'''),
      '''instructblip-flan-t5-xxl''': ('''blip2_t5_instruct''', '''flant5xxl'''),
  }

  lowerCAmelCase__				,      lowerCAmelCase__										=      model_name_to_original[model_name]

  # load original model
  print('''Loading original model...'''				)
  lowerCAmelCase__										=      '''cuda:1''' if torch.cuda.is_available() else '''cpu'''
  lowerCAmelCase__										=      '''cuda:2''' if torch.cuda.is_available() else '''cpu'''
  lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__										=      load_model_and_preprocess(
      name=__lowerCAmelCase   , model_type=__lowerCAmelCase   , is_eval=__lowerCAmelCase   , device=__lowerCAmelCase				)
  original_model.eval()
  print('''Done!'''				)

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

  # some keys can be renamed efficiently
  for key, val in state_dict.copy().items():
    lowerCAmelCase__										=      state_dict.pop(__lowerCAmelCase				)
    if key.startswith('''Qformer.bert'''				):
      lowerCAmelCase__										=      key.replace('''Qformer.bert'''   , '''qformer'''				)
    if "attention.self" in key:
      lowerCAmelCase__										=      key.replace('''self'''   , '''attention'''				)
    if "llm_proj" in key:
      lowerCAmelCase__										=      key.replace('''llm_proj'''   , '''language_projection'''				)
    if "t5_proj" in key:
      lowerCAmelCase__										=      key.replace('''t5_proj'''   , '''language_projection'''				)
    if key.startswith('''llm_model'''				):
      lowerCAmelCase__										=      key.replace('''llm_model'''   , '''language_model'''				)
    if key.startswith('''t5'''				):
      lowerCAmelCase__										=      key.replace('''t5'''   , '''language'''				)
    lowerCAmelCase__										=      val

  # read in qv biases
  read_in_q_v_bias(__lowerCAmelCase   , __lowerCAmelCase				)

  # note: weights get loaded in torch.float32 by default
  hf_model.load_state_dict(__lowerCAmelCase   , strict=__lowerCAmelCase				)

  lowerCAmelCase__										=      load_demo_image()
  lowerCAmelCase__										=      '''What is unusual about this image?'''

  # create processor
  lowerCAmelCase__										=      BlipImageProcessor(
      size={'''height''': image_size, '''width''': image_size}   , image_mean=__lowerCAmelCase   , image_std=__lowerCAmelCase				)
  lowerCAmelCase__										=      InstructBlipProcessor(
      image_processor=__lowerCAmelCase   , tokenizer=__lowerCAmelCase   , qformer_tokenizer=__lowerCAmelCase   , )
  lowerCAmelCase__										=      processor(images=__lowerCAmelCase   , text=__lowerCAmelCase   , return_tensors='''pt'''				).to(__lowerCAmelCase				)

  # make sure processor creates exact same pixel values
  lowerCAmelCase__										=      vis_processors['''eval'''](__lowerCAmelCase				).unsqueeze(0				).to(__lowerCAmelCase				)
  lowerCAmelCase__										=      inputs.pixel_values
  assert torch.allclose(original_pixel_values.to(pixel_values.device				)   , __lowerCAmelCase				)

  original_model.to(__lowerCAmelCase				)
  hf_model.to(__lowerCAmelCase				)
  with torch.no_grad():
    if "vicuna" in model_name:
      lowerCAmelCase__										=      original_model({'''image''': original_pixel_values, '''text_input''': [prompt]}				).logits
      lowerCAmelCase__										=      hf_model(**__lowerCAmelCase				).logits
    else:
      lowerCAmelCase__										=      original_model(
          {'''image''': original_pixel_values, '''text_input''': [prompt], '''text_output''': ['''\n''']}				).logits
      lowerCAmelCase__										=      tokenizer('''\n'''   , return_tensors='''pt'''				).input_ids.to(__lowerCAmelCase				)
      lowerCAmelCase__										=      label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id   , -1_00				)
      lowerCAmelCase__										=      hf_model(**__lowerCAmelCase   , labels=__lowerCAmelCase				).logits

  print('''First values of original logits:'''   , original_logits[0, :3, :3]				)
  print('''First values of HF logits:'''   , logits[0, :3, :3]				)

  # assert values
  assert original_logits.shape == logits.shape
  lowerCAmelCase__										=      1E-4 if '''vicuna''' in model_name else 1E-5
  assert torch.allclose(original_logits.to(logits.device				)   , __lowerCAmelCase   , atol=__lowerCAmelCase				)
  print('''Looks ok!'''				)

  print('''Generating with original model...'''				)
  lowerCAmelCase__										=      original_model.generate({'''image''': original_pixel_values, '''prompt''': prompt}   , num_beams=5				)

  # important: we need to cast the weights of the HF model to the appropriate type
  print('''Generating with HF model...'''				)
  lowerCAmelCase__										=      hf_model.generate(
      **__lowerCAmelCase   , do_sample=__lowerCAmelCase   , num_beams=5   , max_length=2_56   , min_length=1   , top_p=0.9   , repetition_penalty=1.5   , length_penalty=1.0   , temperature=1   , )
  if "vicuna" in model_name:
    # convert output id 0 to 2 (eos_token_id)
    # TODO add this in the generate method?
    lowerCAmelCase__										=      2
  print('''Original generation:'''   , __lowerCAmelCase				)
  lowerCAmelCase__										=      processor.batch_decode(__lowerCAmelCase   , skip_special_tokens=__lowerCAmelCase				)
  lowerCAmelCase__										=      [text.strip() for text in output_text]
  print('''HF generation:'''   , __lowerCAmelCase				)

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

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


if __name__ == "__main__":
      __magic_name__  :  List[str]					=       argparse.ArgumentParser()
      __magic_name__  :  Any					=       [
          """instructblip-vicuna-7b""",
          """instructblip-vicuna-13b""",
          """instructblip-flan-t5-xl""",
          """instructblip-flan-t5-xxl""",
      ]
      parser.add_argument(
          """--model_name""",
          default="""instructblip-flan-t5-xl""",
          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""",
      )

      __magic_name__  :  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 SpeechaTextConfig, SpeechaTextForConditionalGeneration


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		List[Any]     ):







      '''simple docstring'''




      lowerCAmelCase    		=			[
          """encoder.version""",
          """decoder.version""",
          """model.encoder.version""",
          """model.decoder.version""",
          """decoder.output_projection.weight""",
          """_float_tensor""",
          """encoder.embed_positions._float_tensor""",
          """decoder.embed_positions._float_tensor""",
      ]
      for k in ignore_keys:
            state_dict.pop(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     )


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		Dict     ):







      '''simple docstring'''




      lowerCAmelCase    		=			list(s_dict.keys()     )
      for key in keys:
            if "transformer_layers" in key:
                  lowerCAmelCase    		=			s_dict.pop(SCREAMING_SNAKE_CASE     )
            elif "subsample" in key:
                  lowerCAmelCase    		=			s_dict.pop(SCREAMING_SNAKE_CASE     )


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		Any     ):







      '''simple docstring'''




      lowerCAmelCase		,						lowerCAmelCase    		=			emb.weight.shape
      lowerCAmelCase    		=			nn.Linear(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE				,		bias=SCREAMING_SNAKE_CASE     )
      lowerCAmelCase    		=			emb.weight.data
      return lin_layer


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		Optional[Any]				,		SCREAMING_SNAKE_CASE	:		Tuple     ):







      '''simple docstring'''




      lowerCAmelCase    		=			torch.load(SCREAMING_SNAKE_CASE				,		map_location="""cpu"""     )
      lowerCAmelCase    		=			mam_aaa["""args"""]
      lowerCAmelCase    		=			mam_aaa["""model"""]
      lowerCAmelCase    		=			state_dict["""decoder.output_projection.weight"""]

      remove_ignore_keys_(SCREAMING_SNAKE_CASE     )
      rename_keys(SCREAMING_SNAKE_CASE     )

      lowerCAmelCase    		=			state_dict["""decoder.embed_tokens.weight"""].shape[0]

      lowerCAmelCase    		=			args.share_decoder_input_output_embed

      lowerCAmelCase    		=			[int(SCREAMING_SNAKE_CASE     ) for i in args.conv_kernel_sizes.split(""","""     )]
      lowerCAmelCase    		=			SpeechaTextConfig(
          vocab_size=SCREAMING_SNAKE_CASE				,		max_source_positions=args.max_source_positions				,		max_target_positions=args.max_target_positions				,		encoder_layers=args.encoder_layers				,		decoder_layers=args.decoder_layers				,		encoder_attention_heads=args.encoder_attention_heads				,		decoder_attention_heads=args.decoder_attention_heads				,		encoder_ffn_dim=args.encoder_ffn_embed_dim				,		decoder_ffn_dim=args.decoder_ffn_embed_dim				,		d_model=args.encoder_embed_dim				,		dropout=args.dropout				,		attention_dropout=args.attention_dropout				,		activation_dropout=args.activation_dropout				,		activation_function="""relu"""				,		num_conv_layers=len(SCREAMING_SNAKE_CASE     )				,		conv_channels=args.conv_channels				,		conv_kernel_sizes=SCREAMING_SNAKE_CASE				,		input_feat_per_channel=args.input_feat_per_channel				,		input_channels=args.input_channels				,		tie_word_embeddings=SCREAMING_SNAKE_CASE				,		num_beams=5				,		max_length=2_00				,		use_cache=SCREAMING_SNAKE_CASE				,		decoder_start_token_id=2				,		early_stopping=SCREAMING_SNAKE_CASE				,		)

      lowerCAmelCase    		=			SpeechaTextForConditionalGeneration(SCREAMING_SNAKE_CASE     )
      lowerCAmelCase		,						lowerCAmelCase    		=			model.model.load_state_dict(SCREAMING_SNAKE_CASE				,		strict=SCREAMING_SNAKE_CASE     )
      if len(SCREAMING_SNAKE_CASE     ) > 0 and not set(SCREAMING_SNAKE_CASE     ) <= {
          "encoder.embed_positions.weights",
          "decoder.embed_positions.weights",
      }:
            raise ValueError(
                """Only `encoder.embed_positions.weights` and `decoder.embed_positions.weights`  are allowed to be missing,"""
                F' but all the following weights are missing {missing}'     )

      if tie_embeds:
            lowerCAmelCase    		=			make_linear_from_emb(model.model.decoder.embed_tokens     )
      else:
            lowerCAmelCase    		=			lm_head_weights

      model.save_pretrained(SCREAMING_SNAKE_CASE     )


if __name__ == "__main__":
       SCREAMING_SNAKE_CASE__   =  argparse.ArgumentParser()
       # Required parameters
       parser.add_argument("--fairseq_path", type=str, help="Path to the fairseq model (.pt) file.")
       parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
       SCREAMING_SNAKE_CASE__   =  parser.parse_args()
       convert_fairseq_sat_checkpoint_to_tfms(args.fairseq_path, args.pytorch_dump_folder_path)






