



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 (     ):
			lowercase__					  =				"https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"
			lowercase__					  =				Image.open(requests.get(SCREAMING_SNAKE_CASE_						,  stream=SCREAMING_SNAKE_CASE_  ).raw  ).convert("RGB"  )

			return image



def __lowerCAmelCase (     SCREAMING_SNAKE_CASE_  ):
			lowercase__					  =				[]
			# 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_  ):
			lowercase__					  =				dct.pop(SCREAMING_SNAKE_CASE_  )
			lowercase__					  =				val



def __lowerCAmelCase (     SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  ):
			for i in range(config.vision_config.num_hidden_layers  ):
						# read in original q and v biases
						lowercase__					  =				state_dict.pop(f'''visual_encoder.blocks.{i}.attn.q_bias'''  )
						lowercase__					  =				state_dict.pop(f'''visual_encoder.blocks.{i}.attn.v_bias'''  )

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



def __lowerCAmelCase (     SCREAMING_SNAKE_CASE_  ):
			lowercase__					  =				364 if "coco" in model_name else 224
			lowercase__					  =				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:
						lowercase__					  =				TaConfig.from_pretrained("google/flan-t5-xl"						,  dense_act_fn="gelu"						,  bos_token_id=1  ).to_dict()
			elif "t5-xxl" in model_name:
						lowercase__					  =				TaConfig.from_pretrained("google/flan-t5-xxl"						,  dense_act_fn="gelu"						,  bos_token_id=1  ).to_dict()
			elif "vicuna-7b" in model_name:
						lowercase__					  =				LlamaConfig.from_pretrained("decapoda-research/llama-7b-hf"						,  vocab_size=3_2001  ).to_dict()
			elif "vicuna-13b" in model_name:
						lowercase__					  =				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
			lowercase__					  =				InstructBlipQFormerConfig(vocab_size=3_0523  ).to_dict()
			lowercase__					  =				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  ):
			lowercase__					  =				AutoTokenizer.from_pretrained("bert-base-uncased"						,  truncation_side="left"  )
			qformer_tokenizer.add_special_tokens({"bos_token": "[DEC]"}  )

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

			lowercase__  ,       lowercase__					  =				get_blipa_config(SCREAMING_SNAKE_CASE_  )
			lowercase__					  =				InstructBlipForConditionalGeneration(SCREAMING_SNAKE_CASE_  ).eval()

			lowercase__					  =				{
			    "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"),
			}

			lowercase__  ,       lowercase__					  =				model_name_to_original[model_name]

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

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

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

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









from transformers import HfArgumentParser, TensorFlowBenchmark, TensorFlowBenchmarkArguments



def __lowerCAmelCase (     ):
			lowercase__					  =				HfArgumentParser(SCREAMING_SNAKE_CASE_  )
			lowercase__					  =				parser.parse_args_into_dataclasses()[0]
			lowercase__					  =				TensorFlowBenchmark(args=SCREAMING_SNAKE_CASE_  )
			try:
						lowercase__					  =				parser.parse_args_into_dataclasses()[0]
			except ValueError as e:
						lowercase__					  =				"Arg --no_{0} is no longer used, please use --no-{0} instead."
						lowercase__					  =				" ".join(str(SCREAMING_SNAKE_CASE_  ).split(" "  )[:-1]  )
						lowercase__					  =				""
						lowercase__					  =				eval(str(SCREAMING_SNAKE_CASE_  ).split(" "  )[-1]  )
						lowercase__					  =				[]
						for arg in depreciated_args:
									# arg[2:] removes '--'
									if arg[2:] in TensorFlowBenchmark.deprecated_args:
												# arg[5:] removes '--no_'
												full_error_msg += arg_error_msg.format(arg[5:]  )
									else:
												wrong_args.append(SCREAMING_SNAKE_CASE_  )
						if len(SCREAMING_SNAKE_CASE_  ) > 0:
									lowercase__					  =				full_error_msg + begin_error_msg + str(SCREAMING_SNAKE_CASE_  )
						raise ValueError(SCREAMING_SNAKE_CASE_  )
			benchmark.run()


if __name__ == "__main__":
							main()




