



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    snake_case_				(    ):
	__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    snake_case_				(    _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    snake_case_				(    _SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   ):
	__lowercase		       =			dct.pop(_SCREAMING_SNAKE_CASE   )
	__lowercase		       =			val



def    snake_case_				(    _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    snake_case_				(    _SCREAMING_SNAKE_CASE   ):
	__lowercase		       =			3_6_4 if "coco" in model_name else 2_2_4
	__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_2_0_0_1   ).to_dict()
	elif "vicuna-13b" in model_name:
		__lowercase		       =			LlamaConfig.from_pretrained("decapoda-research/llama-13b-hf"  ,	vocab_size=3_2_0_0_1   ).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_0_5_2_3   ).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    snake_case_				(    _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  ,	-1_0_0   )
			__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=2_5_6  ,	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__":
	snake_case__					:     Optional[Any]    		=   argparse.ArgumentParser()
	snake_case__					:     Tuple    		=   [
	    """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""",
	)

	snake_case__					:     Dict    		=   parser.parse_args()

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








# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


snake_case__					:     List[str]    		=   {
    """configuration_mgp_str""": ["""MGP_STR_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MgpstrConfig"""],
    """processing_mgp_str""": ["""MgpstrProcessor"""],
    """tokenization_mgp_str""": ["""MgpstrTokenizer"""],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	snake_case__					:     Dict    		=   [
	    """MGP_STR_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """MgpstrModel""",
	    """MgpstrPreTrainedModel""",
	    """MgpstrForSceneTextRecognition""",
	]

if TYPE_CHECKING:
	from .configuration_mgp_str import MGP_STR_PRETRAINED_CONFIG_ARCHIVE_MAP, MgpstrConfig
	from .processing_mgp_str import MgpstrProcessor
	from .tokenization_mgp_str import MgpstrTokenizer

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_mgp_str import (
		    MGP_STR_PRETRAINED_MODEL_ARCHIVE_LIST,
		    MgpstrForSceneTextRecognition,
		    MgpstrModel,
		    MgpstrPreTrainedModel,
		)
else:
	import sys

	snake_case__					:     Union[str, Any]    		=   _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



