

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







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

      return image







def 							_lowerCamelCase   (			SCREAMING_SNAKE_CASE   ):
      '''simple docstring'''







      A_							       =		[]
      # 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   ):
      '''simple docstring'''







      A_							       =		dct.pop(SCREAMING_SNAKE_CASE   )
      A_							       =		val







def 							_lowerCamelCase   (			SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE   ):
      '''simple docstring'''







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

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







def 							_lowerCamelCase   (			SCREAMING_SNAKE_CASE   ):
      '''simple docstring'''







      A_							       =		364 if '''coco''' in model_name else 224
      A_							       =		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:
            A_							       =		TaConfig.from_pretrained('''google/flan-t5-xl'''		,				dense_act_fn='''gelu'''		,				bos_token_id=1   ).to_dict()
      elif "t5-xxl" in model_name:
            A_							       =		TaConfig.from_pretrained('''google/flan-t5-xxl'''		,				dense_act_fn='''gelu'''		,				bos_token_id=1   ).to_dict()
      elif "vicuna-7b" in model_name:
            A_							       =		LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''		,				vocab_size=32001   ).to_dict()
      elif "vicuna-13b" in model_name:
            A_							       =		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
      A_							       =		InstructBlipQFormerConfig(vocab_size=30523   ).to_dict()
      A_							       =		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   ):
      '''simple docstring'''







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

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

      A_       ,A_							       =		get_blipa_config(SCREAMING_SNAKE_CASE   )
      A_							       =		InstructBlipForConditionalGeneration(SCREAMING_SNAKE_CASE   ).eval()

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

      A_       ,A_							       =		model_name_to_original[model_name]

      # load original model
      print('''Loading original model...'''   )
      A_							       =		'''cuda:1''' if torch.cuda.is_available() else '''cpu'''
      A_							       =		'''cuda:2''' if torch.cuda.is_available() else '''cpu'''
      A_       ,A_       ,A_							       =		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
      A_							       =		original_model.state_dict()
      A_							       =		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():
            A_							       =		state_dict.pop(SCREAMING_SNAKE_CASE   )
            if key.startswith('''Qformer.bert'''   ):
                  A_							       =		key.replace('''Qformer.bert'''		,				'''qformer'''   )
            if "attention.self" in key:
                  A_							       =		key.replace('''self'''		,				'''attention'''   )
            if "llm_proj" in key:
                  A_							       =		key.replace('''llm_proj'''		,				'''language_projection'''   )
            if "t5_proj" in key:
                  A_							       =		key.replace('''t5_proj'''		,				'''language_projection'''   )
            if key.startswith('''llm_model'''   ):
                  A_							       =		key.replace('''llm_model'''		,				'''language_model'''   )
            if key.startswith('''t5'''   ):
                  A_							       =		key.replace('''t5'''		,				'''language'''   )
            A_							       =		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   )

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

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

      # make sure processor creates exact same pixel values
      A_							       =		vis_processors['''eval'''](SCREAMING_SNAKE_CASE   ).unsqueeze(0   ).to(SCREAMING_SNAKE_CASE   )
      A_							       =		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:
                  A_							       =		original_model({'''image''': original_pixel_values, '''text_input''': [prompt]}   ).logits
                  A_							       =		hf_model(**SCREAMING_SNAKE_CASE   ).logits
            else:
                  A_							       =		original_model(
                      {'''image''': original_pixel_values, '''text_input''': [prompt], '''text_output''': ['''\n''']}   ).logits
                  A_							       =		tokenizer('''\n'''		,				return_tensors='''pt'''   ).input_ids.to(SCREAMING_SNAKE_CASE   )
                  A_							       =		label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id		,				-100   )
                  A_							       =		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
      A_							       =		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...'''   )
      A_							       =		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...'''   )
      A_							       =		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?
            A_							       =		2
      print('''Original generation:'''		,				SCREAMING_SNAKE_CASE   )
      A_							       =		processor.batch_decode(SCREAMING_SNAKE_CASE		,				skip_special_tokens=SCREAMING_SNAKE_CASE   )
      A_							       =		[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__":
      __lowercase				      =						argparse.ArgumentParser()
      __lowercase				      =						[
          """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""",
      )

      __lowercase				      =						parser.parse_args()

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








__lowercase				      =						[4, 1, 7, 4, 2, 6, 4, 1, 5, 3, 7, 5]
__lowercase				      =						[3, 7, 7, 4, 2, 6, 4, 1, 5, 3, 7, 5]
__lowercase				      =						{
    0: """Sunday""",
    1: """Monday""",
    2: """Tuesday""",
    3: """Wednesday""",
    4: """Thursday""",
    5: """Friday""",
    6: """Saturday""",
}







def 							_lowerCamelCase   (			SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE   ):
      '''simple docstring'''







      assert len(str(SCREAMING_SNAKE_CASE   )   ) > 2, "year should be in YYYY format"
      assert 1 <= month <= 12, "month should be between 1 to 12"
      assert 1 <= day <= 31, "day should be between 1 to 31"

      # Doomsday algorithm:
      A_							       =		year // 100
      A_							       =		(5 * (century % 4) + 2) % 7
      A_							       =		year % 100
      A_							       =		centurian % 12
      A_							       =		(
          (centurian // 12) + centurian_m + (centurian_m // 4) + century_anchor
      ) % 7
      A_							       =		(
          DOOMSDAY_NOT_LEAP[month - 1]
          if (year % 4 != 0) or (centurian == 0 and (year % 400) == 0)
          else DOOMSDAY_LEAP[month - 1]
      )
      A_							       =		(dooms_day + day - day_anchor) % 7
      return WEEK_DAY_NAMES[week_day]


if __name__ == "__main__":
      import doctest

      doctest.testmod()





