





"""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      UpperCAmelCase__				(    ):







      '''simple docstring'''




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

      return image


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		List[str]     ):







      '''simple docstring'''




      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      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		List[Any]				,		SCREAMING_SNAKE_CASE	:		int				,		SCREAMING_SNAKE_CASE	:		int     ):







      '''simple docstring'''




      lowerCAmelCase    		=			dct.pop(SCREAMING_SNAKE_CASE     )
      lowerCAmelCase    		=			val


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		int				,		SCREAMING_SNAKE_CASE	:		int     ):







      '''simple docstring'''




      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(SCREAMING_SNAKE_CASE				,		requires_grad=SCREAMING_SNAKE_CASE     ), v_bias)     )
            lowerCAmelCase    		=			qkv_bias


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		Optional[Any]     ):







      '''simple docstring'''




      lowerCAmelCase    		=			3_64 if """coco""" in model_name else 2_24
      lowerCAmelCase    		=			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:
            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=SCREAMING_SNAKE_CASE				,		text_config=SCREAMING_SNAKE_CASE				,		qformer_config=SCREAMING_SNAKE_CASE     )

      return config, image_size


@torch.no_grad()
def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		Union[str, Any]				,		SCREAMING_SNAKE_CASE	:		Union[str, Any]=None				,		SCREAMING_SNAKE_CASE	:		Any=False     ):







      '''simple docstring'''




      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(SCREAMING_SNAKE_CASE     )
      lowerCAmelCase    		=			InstructBlipForConditionalGeneration(SCREAMING_SNAKE_CASE     ).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=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
      lowerCAmelCase    		=			original_model.state_dict()
      lowerCAmelCase    		=			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():
            lowerCAmelCase    		=			state_dict.pop(SCREAMING_SNAKE_CASE     )
            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(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     )

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

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

      # make sure processor creates exact same pixel values
      lowerCAmelCase    		=			vis_processors["""eval"""](SCREAMING_SNAKE_CASE     ).unsqueeze(0     ).to(SCREAMING_SNAKE_CASE     )
      lowerCAmelCase    		=			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:
                  lowerCAmelCase    		=			original_model({"""image""": original_pixel_values, """text_input""": [prompt]}     ).logits
                  lowerCAmelCase    		=			hf_model(**SCREAMING_SNAKE_CASE     ).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(SCREAMING_SNAKE_CASE     )
                  lowerCAmelCase    		=			label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id				,		-1_00     )
                  lowerCAmelCase    		=			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
      lowerCAmelCase    		=			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..."""     )
      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(
          **SCREAMING_SNAKE_CASE				,		do_sample=SCREAMING_SNAKE_CASE				,		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:"""				,		SCREAMING_SNAKE_CASE     )
      lowerCAmelCase    		=			processor.batch_decode(SCREAMING_SNAKE_CASE				,		skip_special_tokens=SCREAMING_SNAKE_CASE     )
      lowerCAmelCase    		=			[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__":
       SCREAMING_SNAKE_CASE__   =  argparse.ArgumentParser()
       SCREAMING_SNAKE_CASE__   =  [
           "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__   =  parser.parse_args()

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













"""simple docstring"""





def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		str     ):







      '''simple docstring'''




      lowerCAmelCase    		=			""""""
      for ch in key:
            if ch == " " or ch not in key_no_dups and ch.isalpha():
                  key_no_dups += ch
      return key_no_dups


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		str     ):







      '''simple docstring'''




      lowerCAmelCase    		=			[chr(i + 65     ) for i in range(26     )]
      # Remove duplicate characters from key
      lowerCAmelCase    		=			remove_duplicates(key.upper()     )
      lowerCAmelCase    		=			len(SCREAMING_SNAKE_CASE     )
      # First fill cipher with key characters
      lowerCAmelCase    		=			{alphabet[i]: char for i, char in enumerate(SCREAMING_SNAKE_CASE     )}
      # Then map remaining characters in alphabet to
      # the alphabet from the beginning
      for i in range(len(SCREAMING_SNAKE_CASE     )				,		26     ):
            lowerCAmelCase    		=			alphabet[i - offset]
            # Ensure we are not mapping letters to letters previously mapped
            while char in key:
                  offset -= 1
                  lowerCAmelCase    		=			alphabet[i - offset]
            lowerCAmelCase    		=			char
      return cipher_alphabet


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		str				,		SCREAMING_SNAKE_CASE	:		dict[str, str]     ):







      '''simple docstring'''




      return "".join(cipher_map.get(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     ) for ch in message.upper()     )


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		str				,		SCREAMING_SNAKE_CASE	:		dict[str, str]     ):







      '''simple docstring'''




      lowerCAmelCase    		=			{v: k for k, v in cipher_map.items()}
      return "".join(rev_cipher_map.get(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     ) for ch in message.upper()     )


def      UpperCAmelCase__				(    ):







      '''simple docstring'''




      lowerCAmelCase    		=			input("""Enter message to encode or decode: """     ).strip()
      lowerCAmelCase    		=			input("""Enter keyword: """     ).strip()
      lowerCAmelCase    		=			input("""Encipher or decipher? E/D:"""     ).strip()[0].lower()
      try:
            lowerCAmelCase    		=			{"""e""": encipher, """d""": decipher}[option]
      except KeyError:
            raise KeyError("""invalid input option"""     )
      lowerCAmelCase    		=			create_cipher_map(SCREAMING_SNAKE_CASE     )
      print(func(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     )     )


if __name__ == "__main__":
       import doctest

       doctest.testmod()
       main()






