


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  (							)->Tuple:
   '''simple docstring'''






   snake_case_    			=    """https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"""
   snake_case_    			=    Image.open(requests.get(UpperCamelCase__						,  stream=UpperCamelCase__							).raw							).convert("RGB"							)

   return image



def 	_lowerCAmelCase  (							lowerCAmelCase_							:Tuple							)->str:
   '''simple docstring'''






   snake_case_    			=    []
   # 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_							:Tuple						,  lowerCAmelCase_							:List[Any]						,  lowerCAmelCase_							:List[str]							)->Dict:
   '''simple docstring'''






   snake_case_    			=    dct.pop(UpperCamelCase__							)
   snake_case_    			=    val



def 	_lowerCAmelCase  (							lowerCAmelCase_							:Tuple						,  lowerCAmelCase_							:Optional[int]							)->Any:
   '''simple docstring'''






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

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



def 	_lowerCAmelCase  (							lowerCAmelCase_							:Dict							)->int:
   '''simple docstring'''






   snake_case_    			=    364 if """coco""" in model_name else 224
   snake_case_    			=    InstructBlipVisionConfig(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 "t5-xl" in model_name:
      snake_case_    			=    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_    			=    TaConfig.from_pretrained("google/flan-t5-xxl"						,  dense_act_fn="gelu"						,  bos_token_id=1							).to_dict()
   elif "vicuna-7b" in model_name:
      snake_case_    			=    LlamaConfig.from_pretrained("decapoda-research/llama-7b-hf"						,  vocab_size=32_001							).to_dict()
   elif "vicuna-13b" in model_name:
      snake_case_    			=    LlamaConfig.from_pretrained("decapoda-research/llama-13b-hf"						,  vocab_size=32_001							).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
   snake_case_    			=    InstructBlipQFormerConfig(vocab_size=30_523							).to_dict()
   snake_case_    			=    InstructBlipConfig(vision_config=UpperCamelCase__						,  text_config=UpperCamelCase__						,  qformer_config=UpperCamelCase__							)

   return config, image_size



@torch.no_grad()
def 	_lowerCAmelCase  (							lowerCAmelCase_							:Union[str, Any]						,  lowerCAmelCase_							:Union[str, Any]=None						,  lowerCAmelCase_							:Dict=False							)->Union[str, Any]:
   '''simple docstring'''






   snake_case_    			=    AutoTokenizer.from_pretrained("bert-base-uncased"						,  truncation_side="left"							)
   qformer_tokenizer.add_special_tokens({"bos_token": "[DEC]"}							)

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

   snake_case_    			=    get_blipa_config(UpperCamelCase__							)
   snake_case_    			=    InstructBlipForConditionalGeneration(UpperCamelCase__							).eval()

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

   snake_case_    			=    model_name_to_original[model_name]

   # load original model
   print("Loading original model..."							)
   snake_case_    			=    """cuda:1""" if torch.cuda.is_available() else """cpu"""
   snake_case_    			=    """cuda:2""" if torch.cuda.is_available() else """cpu"""
   snake_case_    			=    load_model_and_preprocess(
       name=UpperCamelCase__						,  model_type=UpperCamelCase__						,  is_eval=UpperCamelCase__						,  device=UpperCamelCase__							)
   original_model.eval()
   print("Done!"							)

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

   # read in qv biases
   read_in_q_v_bias(UpperCamelCase__						,  UpperCamelCase__							)

   # note: weights get loaded in torch.float32 by default
   hf_model.load_state_dict(UpperCamelCase__						,  strict=UpperCamelCase__							)

   snake_case_    			=    load_demo_image()
   snake_case_    			=    """What is unusual about this image?"""

   # create processor
   snake_case_    			=    BlipImageProcessor(
       size={"height": image_size, "width": image_size}						,  image_mean=UpperCamelCase__						,  image_std=UpperCamelCase__							)
   snake_case_    			=    InstructBlipProcessor(
       image_processor=UpperCamelCase__						,  tokenizer=UpperCamelCase__						,  qformer_tokenizer=UpperCamelCase__						,  )
   snake_case_    			=    processor(images=UpperCamelCase__						,  text=UpperCamelCase__						,  return_tensors="pt"							).to(UpperCamelCase__							)

   # make sure processor creates exact same pixel values
   snake_case_    			=    vis_processors["""eval"""](UpperCamelCase__							).unsqueeze(0							).to(UpperCamelCase__							)
   snake_case_    			=    inputs.pixel_values
   assert torch.allclose(original_pixel_values.to(pixel_values.device							)						,  UpperCamelCase__							)

   original_model.to(UpperCamelCase__							)
   hf_model.to(UpperCamelCase__							)
   with torch.no_grad():
      if "vicuna" in model_name:
         snake_case_    			=    original_model({"image": original_pixel_values, "text_input": [prompt]}							).logits
         snake_case_    			=    hf_model(**UpperCamelCase__							).logits
      else:
         snake_case_    			=    original_model(
             {"image": original_pixel_values, "text_input": [prompt], "text_output": ["\n"]}							).logits
         snake_case_    			=    tokenizer("\n"						,  return_tensors="pt"							).input_ids.to(UpperCamelCase__							)
         snake_case_    			=    label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id						,  -100							)
         snake_case_    			=    hf_model(**UpperCamelCase__						,  labels=UpperCamelCase__							).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
   snake_case_    			=    1e-4 if """vicuna""" in model_name else 1e-5
   assert torch.allclose(original_logits.to(logits.device							)						,  UpperCamelCase__						,  atol=UpperCamelCase__							)
   print("Looks ok!"							)

   print("Generating with original model..."							)
   snake_case_    			=    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..."							)
   snake_case_    			=    hf_model.generate(
       **UpperCamelCase__						,  do_sample=UpperCamelCase__						,  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?
      snake_case_    			=    2
   print("Original generation:"						,  UpperCamelCase__							)
   snake_case_    			=    processor.batch_decode(UpperCamelCase__						,  skip_special_tokens=UpperCamelCase__							)
   snake_case_    			=    [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'''Salesforce/{model_name}'''							)
      hf_model.push_to_hub(F'''Salesforce/{model_name}'''							)


if __name__ == "__main__":
   SCREAMING_SNAKE_CASE		:List[Any]     						=   argparse.ArgumentParser()
   SCREAMING_SNAKE_CASE		: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''',
   )

   SCREAMING_SNAKE_CASE		:Optional[Any]     						=   parser.parse_args()

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












"""simple docstring"""




def 			lowerCAmelCase__     (     UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   ):



  '''simple docstring'''
  _a  :       Optional[Any]             =	[redshift, radiation_density, matter_density, dark_energy]
  if any(p < 0 for p in parameters				):
    raise ValueError("""All input parameters must be positive"""				)

  if any(p > 1 for p in parameters[1:4]				):
    raise ValueError("""Relative densities cannot be greater than one"""				)
  else:
    _a  :       Tuple             =	1 - (matter_density + radiation_density + dark_energy)

    _a  :       int             =	(
        radiation_density * (redshift + 1) ** 4
        + matter_density * (redshift + 1) ** 3
        + curvature * (redshift + 1) ** 2
        + dark_energy
    )

    _a  :       List[str]             =	hubble_constant * e_a ** (1 / 2)
    return hubble


if __name__ == "__main__":
 import doctest

 # run doctest
 doctest.testmod()

 # demo LCDM approximation
 _snake_case          =   0.3

 print(
     hubble_parameter(
         hubble_constant=68.3,
         radiation_density=1e-4,
         matter_density=matter_density,
         dark_energy=1 - matter_density,
         redshift=0,
     )
 )

