






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 				lowerCamelCase__      (    ):
  __UpperCAmelCase							:	Optional[Any]           =  """https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"""
  __UpperCAmelCase							:	List[Any]           =  Image.open(requests.get(__lowerCamelCase  ,		stream=__lowerCamelCase		).raw		).convert("""RGB"""		)

  return image






def 				lowerCamelCase__      (    __lowerCamelCase			:     Tuple		):
  __UpperCAmelCase							:	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.embeddings.layernorm.weight""")		)
  rename_keys.append(("""Qformer.bert.embeddings.LayerNorm.bias""", """qformer.embeddings.layernorm.bias""")		)

  # fmt: on
  return rename_keys






def 				lowerCamelCase__      (    __lowerCamelCase			:     str  ,		__lowerCamelCase			:     List[str]  ,		__lowerCamelCase			:     List[str]		):
  __UpperCAmelCase							:	List[Any]           =  dct.pop(__lowerCamelCase		)
  __UpperCAmelCase							:	int           =  val






def 				lowerCamelCase__      (    __lowerCamelCase			:     List[str]  ,		__lowerCamelCase			:     str		):
  for i in range(config.vision_config.num_hidden_layers		):
    # read in original q and v biases
    __UpperCAmelCase							:	str           =  state_dict.pop(f"""visual_encoder.blocks.{i}.attn.q_bias"""		)
    __UpperCAmelCase							:	List[str]           =  state_dict.pop(f"""visual_encoder.blocks.{i}.attn.v_bias"""		)

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






def 				lowerCamelCase__      (    __lowerCamelCase			:     Union[str, Any]		):
  __UpperCAmelCase							:	Optional[int]           =  364 if """coco""" in model_name else 224
  __UpperCAmelCase							:	str           =  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:
    __UpperCAmelCase							:	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:
    __UpperCAmelCase							:	Optional[Any]           =  TaConfig.from_pretrained("""google/flan-t5-xxl"""  ,		dense_act_fn="""gelu"""  ,		bos_token_id=1		).to_dict()
  elif "vicuna-7b" in model_name:
    __UpperCAmelCase							:	Optional[Any]           =  LlamaConfig.from_pretrained("""decapoda-research/llama-7b-hf"""  ,		vocab_size=32001		).to_dict()
  elif "vicuna-13b" in model_name:
    __UpperCAmelCase							:	Optional[Any]           =  LlamaConfig.from_pretrained("""decapoda-research/llama-13b-hf"""  ,		vocab_size=32001		).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
  __UpperCAmelCase							:	Dict           =  InstructBlipQFormerConfig(vocab_size=30523		).to_dict()
  __UpperCAmelCase							:	Optional[int]           =  InstructBlipConfig(vision_config=__lowerCamelCase  ,		text_config=__lowerCamelCase  ,		qformer_config=__lowerCamelCase		)

  return config, image_size






@torch.no_grad()
def 				lowerCamelCase__      (    __lowerCamelCase			:     List[Any]  ,		__lowerCamelCase			:     Optional[int]=None  ,		__lowerCamelCase			:     List[Any]=False		):
  __UpperCAmelCase							:	Any           =  AutoTokenizer.from_pretrained("""bert-base-uncased"""  ,		truncation_side="""left"""		)
  qformer_tokenizer.add_special_tokens({"""bos_token""": """[DEC]"""}		)

  if "t5" in model_name:
    __UpperCAmelCase							:	Dict           =  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>"})
    __UpperCAmelCase							:	List[Any]           =  LlamaTokenizerFast.from_pretrained(
        """huggyllama/llama-7b"""  ,		truncation_side="""left"""  ,		bos_token="""</s>"""  ,		unk_token="""</s>"""		)
    tokenizer.add_special_tokens({"""pad_token""": """[PAD]"""}		)

  __UpperCAmelCase      , __UpperCAmelCase							:	Optional[Any]           =  get_blipa_config(__lowerCamelCase		)
  __UpperCAmelCase							:	Any           =  InstructBlipForConditionalGeneration(__lowerCamelCase		).eval()

  __UpperCAmelCase							:	List[str]           =  {
      """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"""),
  }

  __UpperCAmelCase      , __UpperCAmelCase							:	Dict           =  model_name_to_original[model_name]

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

  # update state dict keys
  __UpperCAmelCase							:	Optional[int]           =  original_model.state_dict()
  __UpperCAmelCase							:	List[str]           =  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():
    __UpperCAmelCase							:	Any           =  state_dict.pop(__lowerCamelCase		)
    if key.startswith("""Qformer.bert"""		):
      __UpperCAmelCase							:	Any           =  key.replace("""Qformer.bert"""  ,		"""qformer"""		)
    if "attention.self" in key:
      __UpperCAmelCase							:	str           =  key.replace("""self"""  ,		"""attention"""		)
    if "llm_proj" in key:
      __UpperCAmelCase							:	str           =  key.replace("""llm_proj"""  ,		"""language_projection"""		)
    if "t5_proj" in key:
      __UpperCAmelCase							:	int           =  key.replace("""t5_proj"""  ,		"""language_projection"""		)
    if key.startswith("""llm_model"""		):
      __UpperCAmelCase							:	Optional[int]           =  key.replace("""llm_model"""  ,		"""language_model"""		)
    if key.startswith("""t5"""		):
      __UpperCAmelCase							:	List[Any]           =  key.replace("""t5"""  ,		"""language"""		)
    __UpperCAmelCase							:	Optional[int]           =  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		)

  __UpperCAmelCase							:	List[Any]           =  load_demo_image()
  __UpperCAmelCase							:	Dict           =  """What is unusual about this image?"""

  # create processor
  __UpperCAmelCase							:	Optional[int]           =  BlipImageProcessor(
      size={"""height""": image_size, """width""": image_size}  ,		image_mean=__lowerCamelCase  ,		image_std=__lowerCamelCase		)
  __UpperCAmelCase							:	Optional[Any]           =  InstructBlipProcessor(
      image_processor=__lowerCamelCase  ,		tokenizer=__lowerCamelCase  ,		qformer_tokenizer=__lowerCamelCase  ,		)
  __UpperCAmelCase							:	str           =  processor(images=__lowerCamelCase  ,		text=__lowerCamelCase  ,		return_tensors="""pt"""		).to(__lowerCamelCase		)

  # make sure processor creates exact same pixel values
  __UpperCAmelCase							:	Optional[Any]           =  vis_processors["""eval"""](__lowerCamelCase		).unsqueeze(0		).to(__lowerCamelCase		)
  __UpperCAmelCase							:	int           =  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:
      __UpperCAmelCase							:	Union[str, Any]           =  original_model({"""image""": original_pixel_values, """text_input""": [prompt]}		).logits
      __UpperCAmelCase							:	int           =  hf_model(**__lowerCamelCase		).logits
    else:
      __UpperCAmelCase							:	str           =  original_model(
          {"""image""": original_pixel_values, """text_input""": [prompt], """text_output""": ["""\n"""]}		).logits
      __UpperCAmelCase							:	Optional[int]           =  tokenizer("""\n"""  ,		return_tensors="""pt"""		).input_ids.to(__lowerCamelCase		)
      __UpperCAmelCase							:	List[Any]           =  label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id  ,		-100		)
      __UpperCAmelCase							:	Union[str, Any]           =  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
  __UpperCAmelCase							:	Optional[int]           =  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..."""		)
  __UpperCAmelCase							:	List[Any]           =  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..."""		)
  __UpperCAmelCase							:	List[Any]           =  hf_model.generate(
      **__lowerCamelCase  ,		do_sample=__lowerCamelCase  ,		num_beams=5  ,		max_length=256  ,		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?
    __UpperCAmelCase							:	int           =  2
  print("""Original generation:"""  ,		__lowerCamelCase		)
  __UpperCAmelCase							:	int           =  processor.batch_decode(__lowerCamelCase  ,		skip_special_tokens=__lowerCamelCase		)
  __UpperCAmelCase							:	List[Any]           =  [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__":
       a		:       Any			     = argparse.ArgumentParser()
       a		:       Dict			     = [
           "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",
       )

       a		:       List[str]			     = parser.parse_args()

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










from __future__ import annotations

import typing
from collections.abc import Iterable

import numpy as np

a		:       Dict			     = typing.Union[Iterable[float], Iterable[int], np.ndarray]  # noqa: UP007
a		:       Any			     = typing.Union[np.floataa, int, float]  # noqa: UP007






def 				lowerCamelCase__      (    __lowerCamelCase			:     Vector  ,		__lowerCamelCase			:     Vector		):
  return np.sqrt(np.sum((np.asarray(__lowerCamelCase		) - np.asarray(__lowerCamelCase		)) ** 2		)		)






def 				lowerCamelCase__      (    __lowerCamelCase			:     Vector  ,		__lowerCamelCase			:     Vector		):
  return sum((va - va) ** 2 for va, va in zip(__lowerCamelCase  ,		__lowerCamelCase		)		) ** (1 / 2)


if __name__ == "__main__":






       def 				lowerCamelCase__      (    ):
         from timeit import timeit

         print("""Without Numpy"""		)
         print(
             timeit(
                 """euclidean_distance_no_np([1, 2, 3], [4, 5, 6])"""  ,		number=10000  ,		globals=globals()  ,		)		)
         print("""With Numpy"""		)
         print(
             timeit(
                 """euclidean_distance([1, 2, 3], [4, 5, 6])"""  ,		number=10000  ,		globals=globals()  ,		)		)







       benchmark()


