




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"""







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

	return image



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Dict	):




	"""simple docstring"""







	UpperCAmelCase__							 =						[]
	# 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		(		_lowerCAmelCase							:   Tuple						,							_lowerCAmelCase							:   Tuple						,							_lowerCAmelCase							:   Optional[int]	):




	"""simple docstring"""







	UpperCAmelCase__							 =						dct.pop(SCREAMING_SNAKE_CASE_	)
	UpperCAmelCase__							 =						val



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Dict						,							_lowerCAmelCase							:   List[Any]	):




	"""simple docstring"""







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

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



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Union[str, Any]	):




	"""simple docstring"""







	UpperCAmelCase__							 =						364 if "coco" in model_name else 224
	UpperCAmelCase__							 =						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:
		UpperCAmelCase__							 =						TaConfig.from_pretrained("google/flan-t5-xl"						,							dense_act_fn="gelu"						,							bos_token_id=1	).to_dict()
	elif "t5-xxl" in model_name:
		UpperCAmelCase__							 =						TaConfig.from_pretrained("google/flan-t5-xxl"						,							dense_act_fn="gelu"						,							bos_token_id=1	).to_dict()
	elif "vicuna-7b" in model_name:
		UpperCAmelCase__							 =						LlamaConfig.from_pretrained("decapoda-research/llama-7b-hf"						,							vocab_size=3_2001	).to_dict()
	elif "vicuna-13b" in model_name:
		UpperCAmelCase__							 =						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
	UpperCAmelCase__							 =						InstructBlipQFormerConfig(vocab_size=3_0523	).to_dict()
	UpperCAmelCase__							 =						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		(		_lowerCAmelCase							:   Tuple						,							_lowerCAmelCase							:   str=None						,							_lowerCAmelCase							:   List[str]=False	):




	"""simple docstring"""







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

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

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

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

	UpperCAmelCase__ , UpperCAmelCase__							 =						model_name_to_original[model_name]

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

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

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

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

		_lowerCAmelCase						:   Optional[int]		=   parser.parse_args()

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











import argparse

from transformers import BigBirdConfig, BigBirdForPreTraining, BigBirdForQuestionAnswering, load_tf_weights_in_big_bird
from transformers.utils import logging


logging.set_verbosity_info()



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Any						,							_lowerCAmelCase							:   List[str]						,							_lowerCAmelCase							:   Optional[int]						,							_lowerCAmelCase							:   Union[str, Any]	):




 """simple docstring"""







 UpperCAmelCase__							 =						BigBirdConfig.from_json_file(_lowerCAmelCase	)
 print(F'''Building PyTorch model from configuration: {config}'''	)

 if is_trivia_qa:
  UpperCAmelCase__							 =						BigBirdForQuestionAnswering(_lowerCAmelCase	)
 else:
  UpperCAmelCase__							 =						BigBirdForPreTraining(_lowerCAmelCase	)

 # Load weights from tf checkpoint
 load_tf_weights_in_big_bird(_lowerCAmelCase						,							_lowerCAmelCase						,							is_trivia_qa=_lowerCAmelCase	)

 # Save pytorch-model
 print(F'''Save PyTorch model to {pytorch_dump_path}'''	)
 model.save_pretrained(_lowerCAmelCase	)


if __name__ == "__main__":
  _lowerCAmelCase						:   List[Any]		=   argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
  )
  parser.add_argument(
      "--big_bird_config_file",
      default=None,
      type=str,
      required=True,
      help=(
          "The config json file corresponding to the pre-trained BERT model. \n"
          "This specifies the model architecture."
      ),
  )
  parser.add_argument(
      "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
  )
  parser.add_argument(
      "--is_trivia_qa", action="store_true", help="Whether to convert a model with a trivia_qa head."
  )
  _lowerCAmelCase						:   Dict		=   parser.parse_args()
  convert_tf_checkpoint_to_pytorch(
      args.tf_checkpoint_path, args.big_bird_config_file, args.pytorch_dump_path, args.is_trivia_qa
  )





