




"""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 	lowercase_				(			)						->    Tuple:
      lowerCAmelCase__      :			Optional[Any]   =	"""https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"""
      lowerCAmelCase__      :			Union[str, Any]   =	Image.open(requests.get(__UpperCAmelCase    ,		stream=__UpperCAmelCase				).raw				).convert("""RGB"""				)

      return image



def 	lowercase_				(			__UpperCAmelCase				)						->    Optional[Any]:
      lowerCAmelCase__      :			int   =	[]
      # 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 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase				)						->    List[str]:
      lowerCAmelCase__      :			Any   =	dct.pop(__UpperCAmelCase				)
      lowerCAmelCase__      :			str   =	val



def 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase				)						->    Union[str, Any]:
      for i in range(config.vision_config.num_hidden_layers				):
            # read in original q and v biases
            lowerCAmelCase__      :			str   =	state_dict.pop(f"""visual_encoder.blocks.{i}.attn.q_bias"""				)
            lowerCAmelCase__      :			Union[str, Any]   =	state_dict.pop(f"""visual_encoder.blocks.{i}.attn.v_bias"""				)

            # next, set bias in the state dict
            lowerCAmelCase__      :			List[Any]   =	torch.cat((q_bias, torch.zeros_like(__UpperCAmelCase    ,		requires_grad=__UpperCAmelCase				), v_bias)				)
            lowerCAmelCase__      :			Optional[int]   =	qkv_bias



def 	lowercase_				(			__UpperCAmelCase				)						->    Dict:
      lowerCAmelCase__      :			Optional[int]   =	364 if """coco""" in model_name else 224
      lowerCAmelCase__      :			Any   =	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:
            lowerCAmelCase__      :			Tuple   =	TaConfig.from_pretrained("""google/flan-t5-xl"""    ,		dense_act_fn="""gelu"""    ,		bos_token_id=1				).to_dict()
      elif "t5-xxl" in model_name:
            lowerCAmelCase__      :			List[str]   =	TaConfig.from_pretrained("""google/flan-t5-xxl"""    ,		dense_act_fn="""gelu"""    ,		bos_token_id=1				).to_dict()
      elif "vicuna-7b" in model_name:
            lowerCAmelCase__      :			Dict   =	LlamaConfig.from_pretrained("""decapoda-research/llama-7b-hf"""    ,		vocab_size=3_2001				).to_dict()
      elif "vicuna-13b" in model_name:
            lowerCAmelCase__      :			Any   =	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
      lowerCAmelCase__      :			Optional[Any]   =	InstructBlipQFormerConfig(vocab_size=3_0523				).to_dict()
      lowerCAmelCase__      :			Optional[Any]   =	InstructBlipConfig(vision_config=__UpperCAmelCase    ,		text_config=__UpperCAmelCase    ,		qformer_config=__UpperCAmelCase				)

      return config, image_size



@torch.no_grad()
def 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase=None    ,		__UpperCAmelCase=False				)						->    str:
      lowerCAmelCase__      :			str   =	AutoTokenizer.from_pretrained("""bert-base-uncased"""    ,		truncation_side="""left"""				)
      qformer_tokenizer.add_special_tokens({"""bos_token""": """[DEC]"""}				)

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

      lowerCAmelCase__ ,   lowerCAmelCase__      :			Optional[Any]   =	get_blipa_config(__UpperCAmelCase				)
      lowerCAmelCase__      :			Tuple   =	InstructBlipForConditionalGeneration(__UpperCAmelCase				).eval()

      lowerCAmelCase__      :			Tuple   =	{
          """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__      :			Union[str, Any]   =	model_name_to_original[model_name]

      # load original model
      print("""Loading original model..."""				)
      lowerCAmelCase__      :			Dict   =	"""cuda:1""" if torch.cuda.is_available() else """cpu"""
      lowerCAmelCase__      :			Dict   =	"""cuda:2""" if torch.cuda.is_available() else """cpu"""
      lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__      :			Union[str, Any]   =	load_model_and_preprocess(
          name=__UpperCAmelCase    ,		model_type=__UpperCAmelCase    ,		is_eval=__UpperCAmelCase    ,		device=__UpperCAmelCase				)
      original_model.eval()
      print("""Done!"""				)

      # update state dict keys
      lowerCAmelCase__      :			Optional[Any]   =	original_model.state_dict()
      lowerCAmelCase__      :			Tuple   =	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():
            lowerCAmelCase__      :			Optional[Any]   =	state_dict.pop(__UpperCAmelCase				)
            if key.startswith("""Qformer.bert"""				):
                  lowerCAmelCase__      :			str   =	key.replace("""Qformer.bert"""    ,		"""qformer"""				)
            if "attention.self" in key:
                  lowerCAmelCase__      :			Dict   =	key.replace("""self"""    ,		"""attention"""				)
            if "llm_proj" in key:
                  lowerCAmelCase__      :			int   =	key.replace("""llm_proj"""    ,		"""language_projection"""				)
            if "t5_proj" in key:
                  lowerCAmelCase__      :			List[Any]   =	key.replace("""t5_proj"""    ,		"""language_projection"""				)
            if key.startswith("""llm_model"""				):
                  lowerCAmelCase__      :			Optional[Any]   =	key.replace("""llm_model"""    ,		"""language_model"""				)
            if key.startswith("""t5"""				):
                  lowerCAmelCase__      :			Tuple   =	key.replace("""t5"""    ,		"""language"""				)
            lowerCAmelCase__      :			Optional[Any]   =	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				)

      lowerCAmelCase__      :			Union[str, Any]   =	load_demo_image()
      lowerCAmelCase__      :			Dict   =	"""What is unusual about this image?"""

      # create processor
      lowerCAmelCase__      :			Any   =	BlipImageProcessor(
          size={"""height""": image_size, """width""": image_size}    ,		image_mean=__UpperCAmelCase    ,		image_std=__UpperCAmelCase				)
      lowerCAmelCase__      :			List[Any]   =	InstructBlipProcessor(
          image_processor=__UpperCAmelCase    ,		tokenizer=__UpperCAmelCase    ,		qformer_tokenizer=__UpperCAmelCase    ,		)
      lowerCAmelCase__      :			Dict   =	processor(images=__UpperCAmelCase    ,		text=__UpperCAmelCase    ,		return_tensors="""pt"""				).to(__UpperCAmelCase				)

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

  _A													=      parser.parse_args()

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











"""simple docstring"""


def 	lowercase_				(			__UpperCAmelCase				)						->    str:
      if not all(char in """01""" for char in bin_string				):
            raise ValueError("""Non-binary value was passed to the function"""				)
      if not bin_string:
            raise ValueError("""Empty string was passed to the function"""				)
      lowerCAmelCase__      :			List[str]   =	""""""
      while len(__UpperCAmelCase				) % 3 != 0:
            lowerCAmelCase__      :			Optional[Any]   =	"""0""" + bin_string
      lowerCAmelCase__      :			List[Any]   =	[
          bin_string[index : index + 3]
          for index in range(len(__UpperCAmelCase				)				)
          if index % 3 == 0
      ]
      for bin_group in bin_string_in_3_list:
            lowerCAmelCase__      :			Optional[int]   =	0
            for index, val in enumerate(__UpperCAmelCase				):
                  oct_val += int(2 ** (2 - index) * int(__UpperCAmelCase				)				)
            oct_string += str(__UpperCAmelCase				)
      return oct_string


if __name__ == "__main__":
  from doctest import testmod

  testmod()





