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

    return image






def __lowerCAmelCase (							UpperCamelCase__				)	->    Dict:
    __lowerCamelCase		=  []
    # 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 (							UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				)	->    str:
    __lowerCamelCase		=  dct.pop(_A				)
    __lowerCamelCase		=  val






def __lowerCAmelCase (							UpperCamelCase__				,  UpperCamelCase__				)	->    List[str]:







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

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






def __lowerCAmelCase (							UpperCamelCase__				)	->    int:
    __lowerCamelCase		=  3_64 if '''coco''' in model_name else 2_24
    __lowerCamelCase		=  InstructBlipVisionConfig(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 "t5-xl" in model_name:
        __lowerCamelCase		=  TaConfig.from_pretrained('''google/flan-t5-xl'''				,  dense_act_fn='''gelu'''				,  bos_token_id=1				).to_dict()
    elif "t5-xxl" in model_name:
        __lowerCamelCase		=  TaConfig.from_pretrained('''google/flan-t5-xxl'''				,  dense_act_fn='''gelu'''				,  bos_token_id=1				).to_dict()
    elif "vicuna-7b" in model_name:
        __lowerCamelCase		=  LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''				,  vocab_size=3_20_01				).to_dict()
    elif "vicuna-13b" in model_name:
        __lowerCamelCase		=  LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf'''				,  vocab_size=3_20_01				).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
    __lowerCamelCase		=  InstructBlipQFormerConfig(vocab_size=3_05_23				).to_dict()
    __lowerCamelCase		=  InstructBlipConfig(vision_config=_A				,  text_config=_A				,  qformer_config=_A				)

    return config, image_size






@torch.no_grad()
def __lowerCAmelCase (							UpperCamelCase__				,  UpperCamelCase__=None				,  UpperCamelCase__=False				)	->    List[Any]:
    __lowerCamelCase		=  AutoTokenizer.from_pretrained('''bert-base-uncased'''				,  truncation_side='''left'''				)
    qformer_tokenizer.add_special_tokens({'''bos_token''': '''[DEC]'''}				)

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

    __lowerCamelCase	,  __lowerCamelCase		=  get_blipa_config(_A				)
    __lowerCamelCase		=  InstructBlipForConditionalGeneration(_A				).eval()

    __lowerCamelCase		=  {
        '''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'''),
    }

    __lowerCamelCase	,  __lowerCamelCase		=  model_name_to_original[model_name]

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

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

    # read in qv biases
    read_in_q_v_bias(_A				,  _A				)

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

    __lowerCamelCase		=  load_demo_image()
    __lowerCamelCase		=  '''What is unusual about this image?'''

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

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

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

    print('''Generating with original model...'''				)
    __lowerCamelCase		=  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...'''				)
    __lowerCamelCase		=  hf_model.generate(
        **_A				,  do_sample=_A				,  num_beams=5				,  max_length=2_56				,  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?
        __lowerCamelCase		=  2
    print('''Original generation:'''				,  _A				)
    __lowerCamelCase		=  processor.batch_decode(_A				,  skip_special_tokens=_A				)
    __lowerCamelCase		=  [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"""Salesforce/{model_name}"""				)
        hf_model.push_to_hub(f"""Salesforce/{model_name}"""				)


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

  __UpperCAmelCase     							=parser.parse_args()

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



def 			UpperCAmelCase__				(			):





  '''simple docstring'''







  return [
      a * b * (10_00 - a - b)
      for a in range(1	,      9_99							)
      for b in range(_A	,      9_99							)
      if (a * a + b * b == (10_00 - a - b) ** 2)
  ][0]


if __name__ == "__main__":
       print(f"""{solution() = }""")
