






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

			return image

def 				snake_case      (   A__						):
			UpperCAmelCase_     : Optional[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 				snake_case      (   A__	,A__	,A__						):
			UpperCAmelCase_     : Any			     =					dct.pop(_lowerCamelCase						)
			UpperCAmelCase_     : Optional[int]			     =					val

def 				snake_case      (   A__	,A__						):




			for i in range(config.vision_config.num_hidden_layers						):
						# read in original q and v biases
						UpperCAmelCase_     : Tuple			     =					state_dict.pop(F"""visual_encoder.blocks.{i}.attn.q_bias"""						)
						UpperCAmelCase_     : Optional[int]			     =					state_dict.pop(F"""visual_encoder.blocks.{i}.attn.v_bias"""						)

						# next, set bias in the state dict
						UpperCAmelCase_     : str			     =					torch.cat((q_bias, torch.zeros_like(_lowerCamelCase	,requires_grad=_lowerCamelCase						), v_bias)						)
						UpperCAmelCase_     : Any			     =					qkv_bias

def 				snake_case      (   A__						):
			UpperCAmelCase_     : List[str]			     =					3_64 if "coco" in model_name else 2_24
			UpperCAmelCase_     : Dict			     =					InstructBlipVisionConfig(image_size=_lowerCamelCase						).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_     : Any			     =					TaConfig.from_pretrained("google/flan-t5-xl"	,dense_act_fn="gelu"	,bos_token_id=1						).to_dict()
			elif "t5-xxl" in model_name:
						UpperCAmelCase_     : int			     =					TaConfig.from_pretrained("google/flan-t5-xxl"	,dense_act_fn="gelu"	,bos_token_id=1						).to_dict()
			elif "vicuna-7b" in model_name:
						UpperCAmelCase_     : Union[str, Any]			     =					LlamaConfig.from_pretrained("decapoda-research/llama-7b-hf"	,vocab_size=3_20_01						).to_dict()
			elif "vicuna-13b" in model_name:
						UpperCAmelCase_     : Any			     =					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
			UpperCAmelCase_     : int			     =					InstructBlipQFormerConfig(vocab_size=3_05_23						).to_dict()
			UpperCAmelCase_     : str			     =					InstructBlipConfig(vision_config=_lowerCamelCase	,text_config=_lowerCamelCase	,qformer_config=_lowerCamelCase						)

			return config, image_size

@torch.no_grad()
def 				snake_case      (   A__	,A__=None	,A__=False						):
			UpperCAmelCase_     : Optional[int]			     =					AutoTokenizer.from_pretrained("bert-base-uncased"	,truncation_side="left"						)
			qformer_tokenizer.add_special_tokens({"bos_token": "[DEC]"}						)

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

			UpperCAmelCase_     : Optional[int]			     =					get_blipa_config(_lowerCamelCase						)
			UpperCAmelCase_     : Optional[int]			     =					InstructBlipForConditionalGeneration(_lowerCamelCase						).eval()

			UpperCAmelCase_     : Dict			     =					{
			    "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_     : Optional[Any]			     =					model_name_to_original[model_name]

			# load original model
			print("Loading original model..."						)
			UpperCAmelCase_     : Optional[int]			     =					"cuda:1" if torch.cuda.is_available() else "cpu"
			UpperCAmelCase_     : Optional[int]			     =					"cuda:2" if torch.cuda.is_available() else "cpu"
			UpperCAmelCase_     : int			     =					load_model_and_preprocess(
			    name=_lowerCamelCase	,model_type=_lowerCamelCase	,is_eval=_lowerCamelCase	,device=_lowerCamelCase						)
			original_model.eval()
			print("Done!"						)

			# update state dict keys
			UpperCAmelCase_     : Union[str, Any]			     =					original_model.state_dict()
			UpperCAmelCase_     : Optional[Any]			     =					create_rename_keys(_lowerCamelCase						)
			for src, dest in rename_keys:
						rename_key(_lowerCamelCase	,_lowerCamelCase	,_lowerCamelCase						)

			# some keys can be renamed efficiently
			for key, val in state_dict.copy().items():
						UpperCAmelCase_     : Optional[int]			     =					state_dict.pop(_lowerCamelCase						)
						if key.startswith("Qformer.bert"						):
									UpperCAmelCase_     : List[Any]			     =					key.replace("Qformer.bert"	,"qformer"						)
						if "attention.self" in key:
									UpperCAmelCase_     : Optional[int]			     =					key.replace("self"	,"attention"						)
						if "llm_proj" in key:
									UpperCAmelCase_     : Union[str, Any]			     =					key.replace("llm_proj"	,"language_projection"						)
						if "t5_proj" in key:
									UpperCAmelCase_     : List[str]			     =					key.replace("t5_proj"	,"language_projection"						)
						if key.startswith("llm_model"						):
									UpperCAmelCase_     : Tuple			     =					key.replace("llm_model"	,"language_model"						)
						if key.startswith("t5"						):
									UpperCAmelCase_     : Optional[Any]			     =					key.replace("t5"	,"language"						)
						UpperCAmelCase_     : Optional[Any]			     =					val

			# read in qv biases
			read_in_q_v_bias(_lowerCamelCase	,_lowerCamelCase						)

			# note: weights get loaded in torch.float32 by default
			hf_model.load_state_dict(_lowerCamelCase	,strict=_lowerCamelCase						)

			UpperCAmelCase_     : int			     =					load_demo_image()
			UpperCAmelCase_     : Optional[int]			     =					"What is unusual about this image?"

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

			# make sure processor creates exact same pixel values
			UpperCAmelCase_     : Tuple			     =					vis_processors["eval"](_lowerCamelCase						).unsqueeze(0						).to(_lowerCamelCase						)
			UpperCAmelCase_     : Any			     =					inputs.pixel_values
			assert torch.allclose(original_pixel_values.to(pixel_values.device						)	,_lowerCamelCase						)

			original_model.to(_lowerCamelCase						)
			hf_model.to(_lowerCamelCase						)
			with torch.no_grad():
						if "vicuna" in model_name:
									UpperCAmelCase_     : List[Any]			     =					original_model({"image": original_pixel_values, "text_input": [prompt]}						).logits
									UpperCAmelCase_     : Union[str, Any]			     =					hf_model(**_lowerCamelCase						).logits
						else:
									UpperCAmelCase_     : Optional[int]			     =					original_model(
									    {"image": original_pixel_values, "text_input": [prompt], "text_output": ["\n"]}						).logits
									UpperCAmelCase_     : List[str]			     =					tokenizer("\n"	,return_tensors="pt"						).input_ids.to(_lowerCamelCase						)
									UpperCAmelCase_     : Optional[Any]			     =					label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id	,-1_00						)
									UpperCAmelCase_     : Tuple			     =					hf_model(**_lowerCamelCase	,labels=_lowerCamelCase						).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_     : Dict			     =					1e-4 if "vicuna" in model_name else 1e-5
			assert torch.allclose(original_logits.to(logits.device						)	,_lowerCamelCase	,atol=_lowerCamelCase						)
			print("Looks ok!"						)

			print("Generating with original model..."						)
			UpperCAmelCase_     : Dict			     =					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_     : Any			     =					hf_model.generate(
			    **_lowerCamelCase	,do_sample=_lowerCamelCase	,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?
						UpperCAmelCase_     : List[str]			     =					2
			print("Original generation:"	,_lowerCamelCase						)
			UpperCAmelCase_     : Optional[int]			     =					processor.batch_decode(_lowerCamelCase	,skip_special_tokens=_lowerCamelCase						)
			UpperCAmelCase_     : Optional[Any]			     =					[text.strip() for text in output_text]
			print("HF generation:"	,_lowerCamelCase						)

			if pytorch_dump_folder_path is not None:
						processor.save_pretrained(_lowerCamelCase						)
						hf_model.save_pretrained(_lowerCamelCase						)

			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_					=      argparse.ArgumentParser()
			lowerCamelCase_					=      [
			    '''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_					=      parser.parse_args()

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








import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_snake_case		=  logging.get_logger(__name__)

_snake_case		=  {
    "facebook/wav2vec2-base-960h": "https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/config.json",
    # See all Wav2Vec2 models at https://huggingface.co/models?filter=wav2vec2
}



class   UpperCAmelCase_    (       a):
							lowerCamelCase__   			=		'wav2vec2'







							def __init__(   self,			__a=32,			__a=768,			__a=12,			__a=12,			__a=3072,			__a="gelu",			__a=0.1,			__a=0.1,			__a=0.1,			__a=0.0,			__a=0.0,			__a=0.1,			__a=0.1,			__a=0.02,			__a=1E-5,			__a="group",			__a="gelu",			__a=(512, 512, 512, 512, 512, 512, 512),			__a=(5, 2, 2, 2, 2, 2, 2),			__a=(10, 3, 3, 3, 3, 2, 2),			__a=False,			__a=128,			__a=16,			__a=False,			__a=True,			__a=0.05,			__a=10,			__a=2,			__a=0.0,			__a=10,			__a=0,			__a=320,			__a=2,			__a=0.1,			__a=100,			__a=256,			__a=256,			__a=0.1,			__a="sum",			__a=False,			__a=False,			__a=256,			__a=(512, 512, 512, 512, 1500),			__a=(5, 3, 3, 1, 1),			__a=(1, 2, 3, 1, 1),			__a=512,			__a=0,			__a=1,			__a=2,			__a=False,			__a=3,			__a=2,			__a=3,			__a=None,			__a=None,			**__a,			):

														'''simple docstring'''

														super().__init__(**__a,			pad_token_id=__a,			bos_token_id=__a,			eos_token_id=__a)
														_lowerCAmelCase     :							str    =    hidden_size
														_lowerCAmelCase     :							Optional[int]    =    feat_extract_norm
														_lowerCAmelCase     :							Union[str, Any]    =    feat_extract_activation
														_lowerCAmelCase     :							Optional[Any]    =    list(__a)
														_lowerCAmelCase     :							List[str]    =    list(__a)
														_lowerCAmelCase     :							str    =    list(__a)
														_lowerCAmelCase     :							List[str]    =    conv_bias
														_lowerCAmelCase     :							str    =    num_conv_pos_embeddings
														_lowerCAmelCase     :							List[Any]    =    num_conv_pos_embedding_groups
														_lowerCAmelCase     :							str    =    len(self.conv_dim)
														_lowerCAmelCase     :							List[str]    =    num_hidden_layers
														_lowerCAmelCase     :							str    =    intermediate_size
														_lowerCAmelCase     :							Any    =    hidden_act
														_lowerCAmelCase     :							int    =    num_attention_heads
														_lowerCAmelCase     :							Optional[Any]    =    hidden_dropout
														_lowerCAmelCase     :							List[str]    =    attention_dropout
														_lowerCAmelCase     :							Tuple    =    activation_dropout
														_lowerCAmelCase     :							int    =    feat_proj_dropout
														_lowerCAmelCase     :							List[str]    =    final_dropout
														_lowerCAmelCase     :							int    =    layerdrop
														_lowerCAmelCase     :							int    =    layer_norm_eps
														_lowerCAmelCase     :							Union[str, Any]    =    initializer_range
														_lowerCAmelCase     :							str    =    vocab_size
														_lowerCAmelCase     :							Optional[Any]    =    do_stable_layer_norm
														_lowerCAmelCase     :							Any    =    use_weighted_layer_sum

														if (
														    (len(self.conv_stride) != self.num_feat_extract_layers)
														    or (len(self.conv_kernel) != self.num_feat_extract_layers)
														    or (len(self.conv_dim) != self.num_feat_extract_layers)
														):
																					raise ValueError(
																					    "Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =="
																					    " `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ="
																					    f" {len(self.conv_dim)}`, `len(config.conv_stride) = {len(self.conv_stride)}`,"
																					    f" `len(config.conv_kernel) = {len(self.conv_kernel)}`.")

														# fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
														_lowerCAmelCase     :							str    =    apply_spec_augment
														_lowerCAmelCase     :							Optional[Any]    =    mask_time_prob
														_lowerCAmelCase     :							Optional[int]    =    mask_time_length
														_lowerCAmelCase     :							List[str]    =    mask_time_min_masks
														_lowerCAmelCase     :							Optional[int]    =    mask_feature_prob
														_lowerCAmelCase     :							Optional[int]    =    mask_feature_length
														_lowerCAmelCase     :							List[str]    =    mask_feature_min_masks

														# parameters for pretraining with codevector quantized representations
														_lowerCAmelCase     :							Union[str, Any]    =    num_codevectors_per_group
														_lowerCAmelCase     :							str    =    num_codevector_groups
														_lowerCAmelCase     :							Optional[int]    =    contrastive_logits_temperature
														_lowerCAmelCase     :							Optional[int]    =    feat_quantizer_dropout
														_lowerCAmelCase     :							Optional[int]    =    num_negatives
														_lowerCAmelCase     :							Union[str, Any]    =    codevector_dim
														_lowerCAmelCase     :							Any    =    proj_codevector_dim
														_lowerCAmelCase     :							Optional[int]    =    diversity_loss_weight

														# ctc loss
														_lowerCAmelCase     :							Tuple    =    ctc_loss_reduction
														_lowerCAmelCase     :							Tuple    =    ctc_zero_infinity

														# adapter
														_lowerCAmelCase     :							List[Any]    =    add_adapter
														_lowerCAmelCase     :							List[str]    =    adapter_kernel_size
														_lowerCAmelCase     :							str    =    adapter_stride
														_lowerCAmelCase     :							List[str]    =    num_adapter_layers
														_lowerCAmelCase     :							str    =    output_hidden_size or hidden_size
														_lowerCAmelCase     :							Tuple    =    adapter_attn_dim

														# SequenceClassification-specific parameter. Feel free to ignore for other classes.
														_lowerCAmelCase     :							str    =    classifier_proj_size

														# XVector-specific parameters. Feel free to ignore for other classes.
														_lowerCAmelCase     :							str    =    list(__a)
														_lowerCAmelCase     :							Union[str, Any]    =    list(__a)
														_lowerCAmelCase     :							List[str]    =    list(__a)
														_lowerCAmelCase     :							Tuple    =    xvector_output_dim







							@property
							def    snake_case__      (   self):

														'''simple docstring'''

														return functools.reduce(operator.mul,			self.conv_stride,			1)
