






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 					snake_case__				(				):
    lowerCAmelCase__ :int   				=					"https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"
    lowerCAmelCase__ :Optional[int]   				=					Image.open(requests.get(_UpperCamelCase     ,      stream=_UpperCamelCase							).raw							).convert("RGB"							)

    return image


def 					snake_case__				(				UpperCAmelCase				:	Dict							):
    lowerCAmelCase__ :Tuple   				=					[]
    # 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 					snake_case__				(				UpperCAmelCase				:	List[str]     ,      UpperCAmelCase				:	Any     ,      UpperCAmelCase				:	List[str]							):
    lowerCAmelCase__ :Optional[Any]   				=					dct.pop(_UpperCamelCase							)
    lowerCAmelCase__ :List[str]   				=					val


def 					snake_case__				(				UpperCAmelCase				:	List[Any]     ,      UpperCAmelCase				:	List[str]							):







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

        # next, set bias in the state dict
        lowerCAmelCase__ :Dict   				=					torch.cat((q_bias, torch.zeros_like(_UpperCamelCase     ,      requires_grad=_UpperCamelCase							), v_bias)							)
        lowerCAmelCase__ :Union[str, Any]   				=					qkv_bias


def 					snake_case__				(				UpperCAmelCase				:	Optional[Any]     ,      UpperCAmelCase				:	int							):
    lowerCAmelCase__ :Union[str, Any]   				=					3_6_4 if "coco" in model_name else 2_2_4
    lowerCAmelCase__ :Optional[Any]   				=					BlipaVisionConfig(image_size=_UpperCamelCase							).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:
        lowerCAmelCase__ :Optional[int]   				=					OPTConfig.from_pretrained("facebook/opt-2.7b"     ,      eos_token_id=_UpperCamelCase							).to_dict()
    elif "opt-6.7b" in model_name:
        lowerCAmelCase__ :Union[str, Any]   				=					OPTConfig.from_pretrained("facebook/opt-6.7b"     ,      eos_token_id=_UpperCamelCase							).to_dict()
    elif "t5-xl" in model_name:
        lowerCAmelCase__ :List[str]   				=					TaConfig.from_pretrained("google/flan-t5-xl"     ,      dense_act_fn="gelu"     ,      bos_token_id=1							).to_dict()
    elif "t5-xxl" in model_name:
        lowerCAmelCase__ :str   				=					TaConfig.from_pretrained("google/flan-t5-xxl"     ,      dense_act_fn="gelu"     ,      bos_token_id=1							).to_dict()

    lowerCAmelCase__ :Union[str, Any]   				=					BlipaConfig(vision_config=_UpperCamelCase     ,      text_config=_UpperCamelCase							)

    return config, image_size


@torch.no_grad()
def 					snake_case__				(				UpperCAmelCase				:	Any     ,      UpperCAmelCase				:	Tuple=None     ,      UpperCAmelCase				:	Tuple=False							):
    lowerCAmelCase__ :Optional[int]   				=					(
        AutoTokenizer.from_pretrained("facebook/opt-2.7b"							)
        if "opt" in model_name
        else AutoTokenizer.from_pretrained("google/flan-t5-xl"							)
    )
    lowerCAmelCase__ :str   				=					tokenizer("\n"     ,      add_special_tokens=_UpperCamelCase							).input_ids[0]
    lowerCAmelCase__				,lowerCAmelCase__ :Any   				=					get_blipa_config(_UpperCamelCase     ,      eos_token_id=_UpperCamelCase							)

    lowerCAmelCase__ :Any   				=					BlipaForConditionalGeneration(_UpperCamelCase							).eval()

    lowerCAmelCase__ :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"),
    }

    lowerCAmelCase__				,lowerCAmelCase__ :Tuple   				=					model_name_to_original[model_name]

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

    # update state dict keys
    lowerCAmelCase__ :List[str]   				=					original_model.state_dict()
    lowerCAmelCase__ :Tuple   				=					create_rename_keys(_UpperCamelCase							)
    for src, dest in rename_keys:
        rename_key(_UpperCamelCase     ,      _UpperCamelCase     ,      _UpperCamelCase							)

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

    # read in qv biases
    read_in_q_v_bias(_UpperCamelCase     ,      _UpperCamelCase							)

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

    lowerCAmelCase__ :Dict   				=					load_demo_image()
    lowerCAmelCase__ :str   				=					vis_processors["eval"](_UpperCamelCase							).unsqueeze(0							).to(_UpperCamelCase							)
    lowerCAmelCase__ :Optional[Any]   				=					tokenizer(["\n"]     ,      return_tensors="pt"							).input_ids.to(_UpperCamelCase							)

    # create processor
    lowerCAmelCase__ :Any   				=					BlipImageProcessor(
        size={"height": image_size, "width": image_size}     ,      image_mean=_UpperCamelCase     ,      image_std=_UpperCamelCase							)
    lowerCAmelCase__ :List[str]   				=					BlipaProcessor(image_processor=_UpperCamelCase     ,      tokenizer=_UpperCamelCase							)
    lowerCAmelCase__ :Optional[Any]   				=					processor(images=_UpperCamelCase     ,      return_tensors="pt"							).pixel_values.to(_UpperCamelCase							)

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

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

    print("Generating a caption..."							)
    lowerCAmelCase__ :Optional[int]   				=					""
    lowerCAmelCase__ :Tuple   				=					tokenizer(_UpperCamelCase     ,      return_tensors="pt"							).input_ids.to(_UpperCamelCase							)

    lowerCAmelCase__ :Optional[int]   				=					original_model.generate({"image": original_pixel_values}							)
    lowerCAmelCase__ :str   				=					hf_model.generate(
        _UpperCamelCase     ,      _UpperCamelCase     ,      do_sample=_UpperCamelCase     ,      num_beams=5     ,      max_length=3_0     ,      min_length=1     ,      top_p=0.9     ,      repetition_penalty=1.0     ,      length_penalty=1.0     ,      temperature=1     ,      )
    print("Original generation:"     ,      _UpperCamelCase							)
    lowerCAmelCase__ :List[str]   				=					input_ids.shape[1]
    lowerCAmelCase__ :int   				=					processor.batch_decode(outputs[:, prompt_length:]     ,      skip_special_tokens=_UpperCamelCase							)
    lowerCAmelCase__ :Optional[int]   				=					[text.strip() for text in output_text]
    print("HF generation:"     ,      _UpperCamelCase							)

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

    if push_to_hub:
        processor.push_to_hub(F'''nielsr/{model_name}'''							)
        hf_model.push_to_hub(F'''nielsr/{model_name}'''							)


if __name__ == "__main__":
 _a     : int   	=      argparse.ArgumentParser()
 _a     : 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""",
 )

 _a     : str   	=      parser.parse_args()

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







import argparse
import logging
import os
import sys

import numpy as np
import onnxruntime
import torch
from bart_onnx.generation_onnx import BARTBeamSearchGenerator
from bart_onnx.reduce_onnx_size import remove_dup_initializers

import transformers
from transformers import BartForConditionalGeneration, BartTokenizer


logging.basicConfig(
    format='''%(asctime)s | %(levelname)s | %(name)s |  [%(filename)s:%(lineno)d] %(message)s''',
    datefmt='''%Y-%m-%d %H:%M:%S''',
    level=os.environ.get('''LOGLEVEL''', '''INFO''').upper(),
    stream=sys.stdout,
)

a			:			Union[str, Any]												=					logging.getLogger(__name__)

a			:			Any												=					{'''facebook/bart-base''': BartForConditionalGeneration}
a			:			int												=					{'''facebook/bart-base''': BartTokenizer}


def 		lowercase_	(  ):


			'''simple docstring'''







			__lowercase			       = argparse.ArgumentParser(description='''Export Bart model + Beam Search to ONNX graph.'''	)
			parser.add_argument(
			    '''--validation_file'''							,      type=_UpperCamelCase							,      default=_UpperCamelCase							,      help='''A csv or a json file containing the validation data.'''	)
			parser.add_argument(
			    '''--max_length'''							,      type=_UpperCamelCase							,      default=5							,      help='''The maximum total input sequence length after tokenization.'''							,      )
			parser.add_argument(
			    '''--num_beams'''							,      type=_UpperCamelCase							,      default=_UpperCamelCase							,      help=(
			        '''Number of beams to use for evaluation. This argument will be '''
			        '''passed to ``model.generate``, which is used during ``evaluate`` and ``predict``.'''
			    )							,      )
			parser.add_argument(
			    '''--model_name_or_path'''							,      type=_UpperCamelCase							,      help='''Path to pretrained model or model identifier from huggingface.co/models.'''							,      required=_UpperCamelCase							,      )
			parser.add_argument(
			    '''--config_name'''							,      type=_UpperCamelCase							,      default=_UpperCamelCase							,      help='''Pretrained config name or path if not the same as model_name'''							,      )
			parser.add_argument(
			    '''--device'''							,      type=_UpperCamelCase							,      default='''cpu'''							,      help='''Device where the model will be run'''							,      )
			parser.add_argument('''--output_file_path'''							,      type=_UpperCamelCase							,      default=_UpperCamelCase							,      help='''Where to store the final ONNX file.'''	)

			__lowercase			       = parser.parse_args()

			return args


def 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase="cpu"	):


			'''simple docstring'''







			__lowercase			       = model_dict[model_name].from_pretrained(_UpperCamelCase	).to(_UpperCamelCase	)
			__lowercase			       = tokenizer_dict[model_name].from_pretrained(_UpperCamelCase	)

			if model_name in ["facebook/bart-base"]:
						__lowercase			       = 0
						__lowercase			       = None
						__lowercase			       = 0

			return huggingface_model, tokenizer


def 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			model.eval()

			__lowercase			       = None
			__lowercase			       = torch.jit.script(BARTBeamSearchGenerator(_UpperCamelCase	)	)

			with torch.no_grad():
						__lowercase			       = '''My friends are cool but they eat too many carbs.'''
						__lowercase			       = tokenizer([ARTICLE_TO_SUMMARIZE]							,      max_length=10_24							,      return_tensors='''pt'''	).to(model.device	)

						__lowercase			       = model.generate(
						    inputs['''input_ids''']							,      attention_mask=inputs['''attention_mask''']							,      num_beams=_UpperCamelCase							,      max_length=_UpperCamelCase							,      early_stopping=_UpperCamelCase							,      decoder_start_token_id=model.config.decoder_start_token_id							,      )

						torch.onnx.export(
						    _UpperCamelCase							,      (
						        inputs['''input_ids'''],
						        inputs['''attention_mask'''],
						        num_beams,
						        max_length,
						        model.config.decoder_start_token_id,
						    )							,      _UpperCamelCase							,      opset_version=14							,      input_names=['''input_ids''', '''attention_mask''', '''num_beams''', '''max_length''', '''decoder_start_token_id''']							,      output_names=['''output_ids''']							,      dynamic_axes={
						        '''input_ids''': {0: '''batch''', 1: '''seq'''},
						        '''output_ids''': {0: '''batch''', 1: '''seq_out'''},
						    }							,      example_outputs=_UpperCamelCase							,      )

						logger.info('''Model exported to {}'''.format(_UpperCamelCase	)	)

						__lowercase			       = remove_dup_initializers(os.path.abspath(_UpperCamelCase	)	)

						logger.info('''Deduplicated and optimized model written to {}'''.format(_UpperCamelCase	)	)

						__lowercase			       = onnxruntime.InferenceSession(_UpperCamelCase	)
						__lowercase			       = ort_sess.run(
						    _UpperCamelCase							,      {
						        '''input_ids''': inputs['''input_ids'''].cpu().numpy(),
						        '''attention_mask''': inputs['''attention_mask'''].cpu().numpy(),
						        '''num_beams''': np.array(_UpperCamelCase	),
						        '''max_length''': np.array(_UpperCamelCase	),
						        '''decoder_start_token_id''': np.array(model.config.decoder_start_token_id	),
						    }							,      )

						np.testing.assert_allclose(summary_ids.cpu().numpy()							,      ort_out[0]							,      rtol=1E-3							,      atol=1E-3	)

						logger.info('''Model outputs from torch and ONNX Runtime are similar.'''	)
						logger.info('''Success.'''	)


def 		lowercase_	(  ):


			'''simple docstring'''







			__lowercase			       = parse_args()
			__lowercase			       = 5
			__lowercase			       = 4

			# Make one log on every process with the configuration for debugging.
			logging.basicConfig(
			    format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s'''							,      datefmt='''%m/%d/%Y %H:%M:%S'''							,      level=logging.INFO							,      )

			logger.setLevel(logging.INFO	)
			transformers.utils.logging.set_verbosity_error()

			__lowercase			       = torch.device(args.device	)

			__lowercase       ,			__lowercase			       = load_model_tokenizer(args.model_name_or_path							,      _UpperCamelCase	)

			if model.config.decoder_start_token_id is None:
						raise ValueError('''Make sure that `config.decoder_start_token_id` is correctly defined'''	)

			model.to(_UpperCamelCase	)

			if args.max_length:
						__lowercase			       = args.max_length

			if args.num_beams:
						__lowercase			       = args.num_beams

			if args.output_file_path:
						__lowercase			       = args.output_file_path
			else:
						__lowercase			       = '''BART.onnx'''

			logger.info('''Exporting model to ONNX'''	)
			export_and_validate_model(_UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	)


if __name__ == "__main__":
	main()


