


"""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   __lowerCAmelCase						(						)	->						List[Any]:


    """simple docstring"""





    snake_case    :      Dict      					=    "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"
    snake_case    :      int      					=    Image.open(requests.get(lowercase						,							stream=lowercase ).raw ).convert("RGB" )

    return image







def   __lowerCAmelCase						(						lowercase  :  Optional[int] )	->						List[str]:


    """simple docstring"""





    snake_case    :      Dict      					=    []
    # 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  :  str						,							lowercase  :  Dict						,							lowercase  :  Tuple )	->						Optional[Any]:


    """simple docstring"""





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







def   __lowerCAmelCase						(						lowercase  :  int						,							lowercase  :  Dict )	->						Optional[int]:


    """simple docstring"""





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

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







def   __lowerCAmelCase						(						lowercase  :  Tuple						,							lowercase  :  int )	->						List[str]:


    """simple docstring"""





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

    snake_case    :      int      					=    BlipaConfig(vision_config=lowercase						,							text_config=lowercase )

    return config, image_size







@torch.no_grad()
def   __lowerCAmelCase						(						lowercase  :  Dict						,							lowercase  :  List[str]=None						,							lowercase  :  str=False )	->						Any:


    """simple docstring"""





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

    snake_case    :      Any      					=    BlipaForConditionalGeneration(lowercase ).eval()

    snake_case    :      Tuple      					=    {
        "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"),
    }

    snake_case					,snake_case    :      Union[str, Any]      					=    model_name_to_original[model_name]

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

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

    # read in qv biases
    read_in_q_v_bias(lowercase						,							lowercase )

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

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

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

    print("Generating a caption..." )
    snake_case    :      Dict      					=    ""
    snake_case    :      List[str]      					=    tokenizer(lowercase						,							return_tensors="pt" ).input_ids.to(lowercase )

    snake_case    :      Optional[int]      					=    original_model.generate({"image": original_pixel_values} )
    snake_case    :      Optional[int]      					=    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 )
    snake_case    :      Optional[Any]      					=    input_ids.shape[1]
    snake_case    :      Optional[int]      					=    processor.batch_decode(outputs[:, prompt_length:]						,							skip_special_tokens=lowercase )
    snake_case    :      int      					=    [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"""



from __future__ import annotations

import unittest

import numpy as np

from transformers import BlipTextConfig
from transformers.testing_utils import require_tf, slow
from transformers.utils import is_tf_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask


if is_tf_available():
  import tensorflow as tf

  from transformers import TFBlipTextModel
  from transformers.models.blip.modeling_tf_blip import TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST







class _lowerCAmelCase				:







   def __init__(	self						,	UpperCamelCase__						,	UpperCamelCase__=12						,	UpperCamelCase__=7						,	UpperCamelCase__=True						,	UpperCamelCase__=True						,	UpperCamelCase__=True						,	UpperCamelCase__=99						,	UpperCamelCase__=32						,	UpperCamelCase__=32						,	UpperCamelCase__=2						,	UpperCamelCase__=4						,	UpperCamelCase__=37						,	UpperCamelCase__=0.1						,	UpperCamelCase__=0.1						,	UpperCamelCase__=512						,	UpperCamelCase__=0.02						,	UpperCamelCase__=0						,	UpperCamelCase__=None						,	)	->		Dict:

       '''simple docstring'''




       snake_case    :      Optional[Any]      					=    parent
       snake_case    :      Dict      					=    batch_size
       snake_case    :      List[str]      					=    seq_length
       snake_case    :      Dict      					=    is_training
       snake_case    :      Optional[Any]      					=    use_input_mask
       snake_case    :      Optional[int]      					=    use_labels
       snake_case    :      Tuple      					=    vocab_size
       snake_case    :      Optional[Any]      					=    hidden_size
       snake_case    :      Optional[Any]      					=    projection_dim
       snake_case    :      List[Any]      					=    num_hidden_layers
       snake_case    :      List[Any]      					=    num_attention_heads
       snake_case    :      int      					=    intermediate_size
       snake_case    :      str      					=    dropout
       snake_case    :      List[Any]      					=    attention_dropout
       snake_case    :      Any      					=    max_position_embeddings
       snake_case    :      List[Any]      					=    initializer_range
       snake_case    :      Any      					=    scope
       snake_case    :      Union[str, Any]      					=    bos_token_id







   def  lowerCamelCase		(	self						)	->		str:

       '''simple docstring'''




       snake_case    :      str      					=    ids_tensor([self.batch_size, self.seq_length]						,	self.vocab_size						)

       snake_case    :      int      					=    None
       if self.use_input_mask:
           snake_case    :      Optional[int]      					=    random_attention_mask([self.batch_size, self.seq_length]						)

       if input_mask is not None:
           snake_case    :      Tuple      					=    input_mask.numpy()
           snake_case					,snake_case    :      str      					=    input_mask.shape
           snake_case    :      Tuple      					=    np.random.randint(1						,	seq_length - 1						,	size=(batch_size,)						)
           for batch_idx, start_index in enumerate(UpperCamelCase__						):
               snake_case    :      int      					=    1
               snake_case    :      Tuple      					=    0

       snake_case    :      Union[str, Any]      					=    self.get_config()

       return config, input_ids, tf.convert_to_tensor(UpperCamelCase__						)







   def  lowerCamelCase		(	self						)	->		Union[str, Any]:

       '''simple docstring'''




       return BlipTextConfig(
           vocab_size=self.vocab_size						,	hidden_size=self.hidden_size						,	projection_dim=self.projection_dim						,	num_hidden_layers=self.num_hidden_layers						,	num_attention_heads=self.num_attention_heads						,	intermediate_size=self.intermediate_size						,	dropout=self.dropout						,	attention_dropout=self.attention_dropout						,	max_position_embeddings=self.max_position_embeddings						,	initializer_range=self.initializer_range						,	bos_token_id=self.bos_token_id						,	)







   def  lowerCamelCase		(	self						,	UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						)	->		Tuple:

       '''simple docstring'''




       snake_case    :      str      					=    TFBlipTextModel(config=UpperCamelCase__						)
       snake_case    :      Any      					=    model(UpperCamelCase__						,	attention_mask=UpperCamelCase__						,	training=UpperCamelCase__						)
       snake_case    :      Optional[int]      					=    model(UpperCamelCase__						,	training=UpperCamelCase__						)
       self.parent.assertEqual(result.last_hidden_state.shape						,	(self.batch_size, self.seq_length, self.hidden_size)						)
       self.parent.assertEqual(result.pooler_output.shape						,	(self.batch_size, self.hidden_size)						)







   def  lowerCamelCase		(	self						)	->		List[str]:

       '''simple docstring'''




       snake_case    :      Tuple      					=    self.prepare_config_and_inputs()
       snake_case					,snake_case					,snake_case    :      Tuple      					=    config_and_inputs
       snake_case    :      str      					=    {"input_ids": input_ids, "attention_mask": input_mask}
       return config, inputs_dict







@require_tf
class _lowerCAmelCase				(      snake_case_    ,							unittest.TestCase	):
   __UpperCAmelCase :       Any							  =  (TFBlipTextModel,) if is_tf_available() else ()
   __UpperCAmelCase :       Any							  =  False
   __UpperCAmelCase :       Dict							  =  False
   __UpperCAmelCase :       List[Any]							  =  False







   def  lowerCamelCase		(	self						)	->		List[str]:

       '''simple docstring'''




       snake_case    :      List[Any]      					=    BlipTextModelTester(self						)
       snake_case    :      Optional[int]      					=    ConfigTester(self						,	config_class=UpperCamelCase__						,	hidden_size=37						)







   def  lowerCamelCase		(	self						)	->		int:

       '''simple docstring'''




       self.config_tester.run_common_tests()







   def  lowerCamelCase		(	self						)	->		str:

       '''simple docstring'''




       snake_case    :      List[str]      					=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*UpperCamelCase__						)







   def  lowerCamelCase		(	self						)	->		Union[str, Any]:

       '''simple docstring'''




       pass







   def  lowerCamelCase		(	self						)	->		List[str]:

       '''simple docstring'''




       pass







   @unittest.skip(reason="Blip does not use inputs_embeds"						)
   def  lowerCamelCase		(	self						)	->		Tuple:

       '''simple docstring'''




       pass







   @unittest.skip(reason="BlipTextModel has no base class and is not available in MODEL_MAPPING"						)
   def  lowerCamelCase		(	self						)	->		Optional[Any]:

       '''simple docstring'''




       pass







   @unittest.skip(reason="BlipTextModel has no base class and is not available in MODEL_MAPPING"						)
   def  lowerCamelCase		(	self						)	->		Any:

       '''simple docstring'''




       pass







   @slow
   def  lowerCamelCase		(	self						)	->		int:

       '''simple docstring'''




       for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           snake_case    :      List[str]      					=    TFBlipTextModel.from_pretrained(UpperCamelCase__						)
           self.assertIsNotNone(UpperCamelCase__						)







   def  lowerCamelCase		(	self						,	UpperCamelCase__=True						)	->		Optional[int]:

       '''simple docstring'''




       super().test_pt_tf_model_equivalence(allow_missing_keys=UpperCamelCase__						)





