


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




      """simple docstring"""

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

      return image






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      )->						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					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )->						Optional[int]:




      """simple docstring"""

      snake_case_ =   dct.pop(SCREAMING_SNAKE_CASE      )
      snake_case_ =   val






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )->						Union[str, 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(SCREAMING_SNAKE_CASE      ,      requires_grad=SCREAMING_SNAKE_CASE      ), v_bias)      )
            snake_case_ =   qkv_bias






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      )->						List[str]:




      """simple docstring"""

      snake_case_ =   364 if '''coco''' in model_name else 224
      snake_case_ =   InstructBlipVisionConfig(image_size=SCREAMING_SNAKE_CASE      ).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=3_2001      ).to_dict()
      elif "vicuna-13b" in model_name:
            snake_case_ =   LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf'''      ,      vocab_size=3_2001      ).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=3_0523      ).to_dict()
      snake_case_ =   InstructBlipConfig(vision_config=SCREAMING_SNAKE_CASE      ,      text_config=SCREAMING_SNAKE_CASE      ,      qformer_config=SCREAMING_SNAKE_CASE      )

      return config, image_size






@torch.no_grad()
def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE=None      ,      SCREAMING_SNAKE_CASE=False      )->						List[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_ ,					snake_case_ =   get_blipa_config(SCREAMING_SNAKE_CASE      )
      snake_case_ =   InstructBlipForConditionalGeneration(SCREAMING_SNAKE_CASE      ).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_ ,					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_ ,					snake_case_ ,					snake_case_ =   load_model_and_preprocess(
          name=SCREAMING_SNAKE_CASE      ,      model_type=SCREAMING_SNAKE_CASE      ,      is_eval=SCREAMING_SNAKE_CASE      ,      device=SCREAMING_SNAKE_CASE      )
      original_model.eval()
      print('''Done!'''      )

      # update state dict keys
      snake_case_ =   original_model.state_dict()
      snake_case_ =   create_rename_keys(SCREAMING_SNAKE_CASE      )
      for src, dest in rename_keys:
            rename_key(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )

      # some keys can be renamed efficiently
      for key, val in state_dict.copy().items():
            snake_case_ =   state_dict.pop(SCREAMING_SNAKE_CASE      )
            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(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )

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

      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=SCREAMING_SNAKE_CASE      ,      image_std=SCREAMING_SNAKE_CASE      )
      snake_case_ =   InstructBlipProcessor(
          image_processor=SCREAMING_SNAKE_CASE      ,      tokenizer=SCREAMING_SNAKE_CASE      ,      qformer_tokenizer=SCREAMING_SNAKE_CASE      ,      )
      snake_case_ =   processor(images=SCREAMING_SNAKE_CASE      ,      text=SCREAMING_SNAKE_CASE      ,      return_tensors='''pt'''      ).to(SCREAMING_SNAKE_CASE      )

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

      original_model.to(SCREAMING_SNAKE_CASE      )
      hf_model.to(SCREAMING_SNAKE_CASE      )
      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(**SCREAMING_SNAKE_CASE      ).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(SCREAMING_SNAKE_CASE      )
                  snake_case_ =   label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id      ,      -100      )
                  snake_case_ =   hf_model(**SCREAMING_SNAKE_CASE      ,      labels=SCREAMING_SNAKE_CASE      ).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      )      ,      SCREAMING_SNAKE_CASE      ,      atol=SCREAMING_SNAKE_CASE      )
      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(
          **SCREAMING_SNAKE_CASE      ,      do_sample=SCREAMING_SNAKE_CASE      ,      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:'''      ,      SCREAMING_SNAKE_CASE      )
      snake_case_ =   processor.batch_decode(SCREAMING_SNAKE_CASE      ,      skip_special_tokens=SCREAMING_SNAKE_CASE      )
      snake_case_ =   [text.strip() for text in output_text]
      print('''HF generation:'''      ,      SCREAMING_SNAKE_CASE      )

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

      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  __lowerCAmelCase					(SCREAMING_SNAKE_CASE = 3      ,      SCREAMING_SNAKE_CASE = 7      ,      SCREAMING_SNAKE_CASE = 100_0000      )->						int:




      """simple docstring"""

      snake_case_ =   0
      snake_case_ =   1

      for current_denominator in range(1      ,      limit + 1      ):
            snake_case_ =   current_denominator * numerator // denominator
            if current_denominator % denominator == 0:
                  current_numerator -= 1
            if current_numerator * max_denominator > current_denominator * max_numerator:
                  snake_case_ =   current_numerator
                  snake_case_ =   current_denominator
      return max_numerator


if __name__ == "__main__":
   print(solution(numerator=3, denominator=7, limit=100_0000))