






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 	lowerCamelCase__      (  ):

  """simple docstring"""






  SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"
  SCREAMING_SNAKE_CASE							:   Any						    =     Image.open(requests.get(lowercase    ,  stream=lowercase					).raw					).convert("RGB"					)

  return image


def 	lowerCamelCase__      (  lowercase					):

  """simple docstring"""






  SCREAMING_SNAKE_CASE							:   int						    =     []
  # 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 	lowerCamelCase__      (  lowercase    ,  lowercase    ,  lowercase					):

  """simple docstring"""






  SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     dct.pop(lowercase					)
  SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     val


def 	lowerCamelCase__      (  lowercase    ,  lowercase					):

  """simple docstring"""






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

    # next, set bias in the state dict
    SCREAMING_SNAKE_CASE							:   Optional[int]						    =     torch.cat((q_bias, torch.zeros_like(lowercase    ,  requires_grad=lowercase					), v_bias)					)
    SCREAMING_SNAKE_CASE							:   Dict						    =     qkv_bias


def 	lowerCamelCase__      (  lowercase    ,  lowercase					):

  """simple docstring"""






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

  SCREAMING_SNAKE_CASE							:   Tuple						    =     BlipaConfig(vision_config=lowercase    ,  text_config=lowercase					)

  return config, image_size


@torch.no_grad()
def 	lowerCamelCase__      (  lowercase    ,  lowercase=None    ,  lowercase=False					):

  """simple docstring"""






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

  SCREAMING_SNAKE_CASE							:   int						    =     BlipaForConditionalGeneration(lowercase					).eval()

  SCREAMING_SNAKE_CASE							:   Optional[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"),
  }

  SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE							:   Any						    =     model_name_to_original[model_name]

  # load original model
  print("Loading original model..."					)
  SCREAMING_SNAKE_CASE							:   Tuple						    =     "cuda" if torch.cuda.is_available() else "cpu"
  SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE							:   List[Any]						    =     load_model_and_preprocess(
      name=lowercase    ,  model_type=lowercase    ,  is_eval=lowercase    ,  device=lowercase					)
  original_model.eval()
  print("Done!"					)

  # update state dict keys
  SCREAMING_SNAKE_CASE							:   Tuple						    =     original_model.state_dict()
  SCREAMING_SNAKE_CASE							:   int						    =     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():
    SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     state_dict.pop(lowercase					)
    if key.startswith("Qformer.bert"					):
      SCREAMING_SNAKE_CASE							:   Any						    =     key.replace("Qformer.bert"    ,  "qformer"					)
    if "attention.self" in key:
      SCREAMING_SNAKE_CASE							:   Dict						    =     key.replace("self"    ,  "attention"					)
    if "opt_proj" in key:
      SCREAMING_SNAKE_CASE							:   Dict						    =     key.replace("opt_proj"    ,  "language_projection"					)
    if "t5_proj" in key:
      SCREAMING_SNAKE_CASE							:   int						    =     key.replace("t5_proj"    ,  "language_projection"					)
    if key.startswith("opt"					):
      SCREAMING_SNAKE_CASE							:   Optional[int]						    =     key.replace("opt"    ,  "language"					)
    if key.startswith("t5"					):
      SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     key.replace("t5"    ,  "language"					)
    SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     val

  # read in qv biases
  read_in_q_v_bias(lowercase    ,  lowercase					)

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

  SCREAMING_SNAKE_CASE							:   Dict						    =     load_demo_image()
  SCREAMING_SNAKE_CASE							:   str						    =     vis_processors["eval"](lowercase					).unsqueeze(0					).to(lowercase					)
  SCREAMING_SNAKE_CASE							:   int						    =     tokenizer(["\n"]    ,  return_tensors="pt"					).input_ids.to(lowercase					)

  # create processor
  SCREAMING_SNAKE_CASE							:   List[Any]						    =     BlipImageProcessor(
      size={"height": image_size, "width": image_size}    ,  image_mean=lowercase    ,  image_std=lowercase					)
  SCREAMING_SNAKE_CASE							:   Optional[int]						    =     BlipaProcessor(image_processor=lowercase    ,  tokenizer=lowercase					)
  SCREAMING_SNAKE_CASE							:   Dict						    =     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:
      SCREAMING_SNAKE_CASE							:   List[str]						    =     original_model({"image": original_pixel_values, "text_input": [""]}					).logits
      SCREAMING_SNAKE_CASE							:   int						    =     hf_model(lowercase    ,  lowercase					).logits
    else:
      SCREAMING_SNAKE_CASE							:   Dict						    =     original_model(
          {"image": original_pixel_values, "text_input": ["\n"], "text_output": ["\n"]}					).logits
      SCREAMING_SNAKE_CASE							:   List[Any]						    =     input_ids.masked_fill(input_ids == tokenizer.pad_token_id    ,  -100					)
      SCREAMING_SNAKE_CASE							:   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":
    SCREAMING_SNAKE_CASE							:   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":
    SCREAMING_SNAKE_CASE							:   Dict						    =     torch.tensor(
        [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]    ,  device=lowercase					)
  else:
    # cast to same type
    SCREAMING_SNAKE_CASE							:   Tuple						    =     logits.dtype
    assert torch.allclose(original_logits.to(lowercase					)    ,  lowercase    ,  atol=1E-2					)
  print("Looks ok!"					)

  print("Generating a caption..."					)
  SCREAMING_SNAKE_CASE							:   str						    =     ""
  SCREAMING_SNAKE_CASE							:   int						    =     tokenizer(lowercase    ,  return_tensors="pt"					).input_ids.to(lowercase					)

  SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     original_model.generate({"image": original_pixel_values}					)
  SCREAMING_SNAKE_CASE							:   List[str]						    =     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					)
  SCREAMING_SNAKE_CASE							:   List[str]						    =     input_ids.shape[1]
  SCREAMING_SNAKE_CASE							:   Dict						    =     processor.batch_decode(outputs[:, prompt_length:]    ,  skip_special_tokens=lowercase					)
  SCREAMING_SNAKE_CASE							:   Any						    =     [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__":
   snake_case     					=		argparse.ArgumentParser()
   snake_case     					=		[
       """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""",
   )

   snake_case     					=		parser.parse_args()

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








"""simple docstring"""


import os
from collections import namedtuple

import pytest

from datasets import ClassLabel, Features, Sequence, Value
from datasets.commands.test import TestCommand
from datasets.info import DatasetInfo, DatasetInfosDict


A_    							=					namedtuple(
    """_TestCommandArgs""",
    [
        """dataset""",
        """name""",
        """cache_dir""",
        """data_dir""",
        """all_configs""",
        """save_infos""",
        """ignore_verifications""",
        """force_redownload""",
        """clear_cache""",
    ],
    defaults=[None, None, None, False, False, False, False, False],
)






def  lowercase						(	lowerCAmelCase__	,lowerCAmelCase__						):
			return (abs(source - target						) / target) < 0.01






@pytest.mark.integration
def  lowercase						(	lowerCAmelCase__						):
			lowerCamelCase_					=  _TestCommandArgs(dataset=lowerCAmelCase__	,all_configs=lowerCAmelCase__	,save_infos=lowerCAmelCase__						)
			lowerCamelCase_					=  TestCommand(*lowerCAmelCase__						)
			test_command.run()
			lowerCamelCase_					=  os.path.join(lowerCAmelCase__	,'''README.md'''						)
			assert os.path.exists(lowerCAmelCase__						)
			lowerCamelCase_					=  DatasetInfosDict.from_directory(lowerCAmelCase__						)
			lowerCamelCase_					=  DatasetInfosDict(
			    {
			        '''default''': DatasetInfo(
			            features=Features(
			                {
			                    '''tokens''': Sequence(Value('''string'''						)						),
			                    '''ner_tags''': Sequence(
			                        ClassLabel(names=['''O''', '''B-PER''', '''I-PER''', '''B-ORG''', '''I-ORG''', '''B-LOC''', '''I-LOC''']						)						),
			                    '''langs''': Sequence(Value('''string'''						)						),
			                    '''spans''': Sequence(Value('''string'''						)						),
			                }						)	,splits=[
			                {
			                    '''name''': '''train''',
			                    '''num_bytes''': 2_351_563,
			                    '''num_examples''': 10_000,
			                },
			                {
			                    '''name''': '''validation''',
			                    '''num_bytes''': 238_418,
			                    '''num_examples''': 1_000,
			                },
			            ]	,download_size=3_940_680	,dataset_size=2_589_981	,)
			    }						)
			assert dataset_infos.keys() == expected_dataset_infos.keys()
			for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
						lowerCamelCase_		, lowerCamelCase_					=  getattr(dataset_infos['''default''']	,lowerCAmelCase__						), getattr(expected_dataset_infos['''default''']	,lowerCAmelCase__						)
						if key == "num_bytes":
									assert is_apercent_close(lowerCAmelCase__	,lowerCAmelCase__						)
						elif key == "splits":
									assert list(lowerCAmelCase__						) == list(lowerCAmelCase__						)
									for split in result:
												assert result[split].name == expected[split].name
												assert result[split].num_examples == expected[split].num_examples
												assert is_apercent_close(result[split].num_bytes	,expected[split].num_bytes						)
						else:
									result == expected



