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






  '''simple docstring'''





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

  return image

def a__       (			A_   ):






  '''simple docstring'''





  __magic_name__								=			[]
  # 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 a__       (			A_,	A_,	A_   ):






  '''simple docstring'''





  __magic_name__								=			dct.pop(A_   )
  __magic_name__								=			val

def a__       (			A_,	A_   ):






  '''simple docstring'''





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

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

def a__       (			A_,	A_   ):






  '''simple docstring'''





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

  __magic_name__								=			BlipaConfig(vision_config=A_,	text_config=A_   )

  return config, image_size

@torch.no_grad()
def a__       (			A_,	A_=None,	A_=False   ):






  '''simple docstring'''





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

  __magic_name__								=			BlipaForConditionalGeneration(A_   ).eval()

  __magic_name__								=			{
      """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"""),
  }

  __magic_name__	,       __magic_name__								=			model_name_to_original[model_name]

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

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

  # read in qv biases
  read_in_q_v_bias(A_,	A_   )

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

  __magic_name__								=			load_demo_image()
  __magic_name__								=			vis_processors["""eval"""](A_   ).unsqueeze(0   ).to(A_   )
  __magic_name__								=			tokenizer(["""\n"""],	return_tensors="""pt"""   ).input_ids.to(A_   )

  # create processor
  __magic_name__								=			BlipImageProcessor(
      size={"""height""": image_size, """width""": image_size},	image_mean=A_,	image_std=A_   )
  __magic_name__								=			BlipaProcessor(image_processor=A_,	tokenizer=A_   )
  __magic_name__								=			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:
      __magic_name__								=			original_model({"""image""": original_pixel_values, """text_input""": [""""""]}   ).logits
      __magic_name__								=			hf_model(A_,	A_   ).logits
    else:
      __magic_name__								=			original_model(
          {"""image""": original_pixel_values, """text_input""": ["""\n"""], """text_output""": ["""\n"""]}   ).logits
      __magic_name__								=			input_ids.masked_fill(input_ids == tokenizer.pad_token_id,	-100   )
      __magic_name__								=			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":
    __magic_name__								=			torch.tensor(
        [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]],	device=A_   )
    assert torch.allclose(logits[0, :3, :3],	A_,	atol=1e-4   )
  elif model_name == "blip2-flan-t5-xl-coco":
    __magic_name__								=			torch.tensor(
        [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]],	device=A_   )
  else:
    # cast to same type
    __magic_name__								=			logits.dtype
    assert torch.allclose(original_logits.to(A_   ),	A_,	atol=1e-2   )
  print("""Looks ok!"""   )

  print("""Generating a caption..."""   )
  __magic_name__								=			""""""
  __magic_name__								=			tokenizer(A_,	return_tensors="""pt"""   ).input_ids.to(A_   )

  __magic_name__								=			original_model.generate({"""image""": original_pixel_values}   )
  __magic_name__								=			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_   )
  __magic_name__								=			input_ids.shape[1]
  __magic_name__								=			processor.batch_decode(outputs[:, prompt_length:],	skip_special_tokens=A_   )
  __magic_name__								=			[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     :	str       			=  argparse.ArgumentParser()
    __lowerCAmelCase     :	Tuple       			=  [
        '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     :	Tuple       			=  parser.parse_args()

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



import argparse

import torch

from transformers import GPTaConfig, GPTaModel, load_tf_weights_in_gpta
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

def a__       (			A_,	A_,	A_   ):






  '''simple docstring'''





  if gpta_config_file == "":
    __magic_name__								=			GPTaConfig()
  else:
    __magic_name__								=			GPTaConfig.from_json_file(A_   )
  __magic_name__								=			GPTaModel(A_   )

  # Load weights from numpy
  load_tf_weights_in_gpta(A_,	A_,	A_   )

  # Save pytorch-model
  __magic_name__								=			pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
  __magic_name__								=			pytorch_dump_folder_path + """/""" + CONFIG_NAME
  print(f'''Save PyTorch model to {pytorch_weights_dump_path}'''   )
  torch.save(model.state_dict(),	A_   )
  print(f'''Save configuration file to {pytorch_config_dump_path}'''   )
  with open(A_,	"""w""",	encoding="""utf-8"""   ) as f:
    f.write(config.to_json_string()   )


if __name__ == "__main__":
    __lowerCAmelCase     :	List[str]       			=  argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '--gpt2_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
    )
    parser.add_argument(
        '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
    )
    parser.add_argument(
        '--gpt2_config_file',
        default='',
        type=str,
        help=(
            'An optional config json file corresponding to the pre-trained OpenAI model. \n'
            'This specifies the model architecture.'
        ),
    )
    __lowerCAmelCase     :	Tuple       			=  parser.parse_args()
    convert_gpta_checkpoint_to_pytorch(args.gpta_checkpoint_path, args.gpta_config_file, args.pytorch_dump_folder_path)

