


'''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 						UpperCAmelCase__(				):






			"""simple docstring"""




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

			return image






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: List[Any] ):






			"""simple docstring"""




			__A=	[]
			# 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 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Union[str, Any],_SCREAMING_SNAKE_CASE					: Tuple,_SCREAMING_SNAKE_CASE					: Any ):






			"""simple docstring"""




			__A=	dct.pop(_SCREAMING_SNAKE_CASE )
			__A=	val






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Optional[int],_SCREAMING_SNAKE_CASE					: List[Any] ):






			"""simple docstring"""




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

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






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Tuple ):






			"""simple docstring"""




			__A=	364 if 'coco' in model_name else 224
			__A=	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:
						__A=	TaConfig.from_pretrained('google/flan-t5-xl',dense_act_fn='gelu',bos_token_id=1 ).to_dict()
			elif "t5-xxl" in model_name:
						__A=	TaConfig.from_pretrained('google/flan-t5-xxl',dense_act_fn='gelu',bos_token_id=1 ).to_dict()
			elif "vicuna-7b" in model_name:
						__A=	LlamaConfig.from_pretrained('decapoda-research/llama-7b-hf',vocab_size=3_2001 ).to_dict()
			elif "vicuna-13b" in model_name:
						__A=	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
			__A=	InstructBlipQFormerConfig(vocab_size=3_0523 ).to_dict()
			__A=	InstructBlipConfig(vision_config=_SCREAMING_SNAKE_CASE,text_config=_SCREAMING_SNAKE_CASE,qformer_config=_SCREAMING_SNAKE_CASE )

			return config, image_size






@torch.no_grad()
def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: List[str],_SCREAMING_SNAKE_CASE					: Optional[int]=None,_SCREAMING_SNAKE_CASE					: Tuple=False ):






			"""simple docstring"""




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

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

			__A,	__A=	get_blipa_config(_SCREAMING_SNAKE_CASE )
			__A=	InstructBlipForConditionalGeneration(_SCREAMING_SNAKE_CASE ).eval()

			__A=	{
			    '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'),
			}

			__A,	__A=	model_name_to_original[model_name]

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

			__A=	load_demo_image()
			__A=	'What is unusual about this image?'

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

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

					UpperCAmelCase__  			=				parser.parse_args()

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







'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_torch_available,
)


UpperCAmelCase__  			=				{
    '''configuration_trocr''': ['''TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''TrOCRConfig'''],
    '''processing_trocr''': ['''TrOCRProcessor'''],
}


try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCAmelCase__  			=				[
					    '''TROCR_PRETRAINED_MODEL_ARCHIVE_LIST''',
					    '''TrOCRForCausalLM''',
					    '''TrOCRPreTrainedModel''',
					]


if TYPE_CHECKING:
					from .configuration_trocr import TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP, TrOCRConfig
					from .processing_trocr import TrOCRProcessor

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_trocr import TROCR_PRETRAINED_MODEL_ARCHIVE_LIST, TrOCRForCausalLM, TrOCRPreTrainedModel

else:
					import sys

					UpperCAmelCase__  			=				_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



