'''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
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD


def  _lowerCamelCase   (				):



					"""simple docstring"""





					UpperCAmelCase_							:       List[str]       =				'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
					UpperCAmelCase_							:       Dict       =				Image.open(requests.get(lowerCamelCase_     ,					stream=lowerCamelCase_       ).raw       ).convert('RGB'       )

					return image


def  _lowerCamelCase   (				lowerCamelCase_ :							Any       ):



					"""simple docstring"""





					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.layernorm.weight')       )
					rename_keys.append(('Qformer.bert.embeddings.LayerNorm.bias', 'qformer.layernorm.bias')       )

					# fmt: on
					return rename_keys


def  _lowerCamelCase   (				lowerCamelCase_ :							int     ,					lowerCamelCase_ :							Tuple     ,					lowerCamelCase_ :							Optional[int]       ):



					"""simple docstring"""





					UpperCAmelCase_							:       int       =				dct.pop(lowerCamelCase_       )
					UpperCAmelCase_							:       str       =				val


def  _lowerCamelCase   (				lowerCamelCase_ :							Optional[Any]     ,					lowerCamelCase_ :							Union[str, Any]       ):



					"""simple docstring"""





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

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


def  _lowerCamelCase   (				lowerCamelCase_ :							List[Any]     ,					lowerCamelCase_ :							Any       ):



					"""simple docstring"""





					UpperCAmelCase_							:       List[Any]       =				364 if 'coco' in model_name else 224
					UpperCAmelCase_							:       Tuple       =				BlipaVisionConfig(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 "opt-2.7b" in model_name:
										UpperCAmelCase_							:       Union[str, Any]       =				OPTConfig.from_pretrained('facebook/opt-2.7b'     ,					eos_token_id=lowerCamelCase_       ).to_dict()
					elif "opt-6.7b" in model_name:
										UpperCAmelCase_							:       List[str]       =				OPTConfig.from_pretrained('facebook/opt-6.7b'     ,					eos_token_id=lowerCamelCase_       ).to_dict()
					elif "t5-xl" in model_name:
										UpperCAmelCase_							:       str       =				TaConfig.from_pretrained('google/flan-t5-xl'     ,					dense_act_fn='gelu'     ,					bos_token_id=1       ).to_dict()
					elif "t5-xxl" in model_name:
										UpperCAmelCase_							:       List[Any]       =				TaConfig.from_pretrained('google/flan-t5-xxl'     ,					dense_act_fn='gelu'     ,					bos_token_id=1       ).to_dict()

					UpperCAmelCase_							:       Union[str, Any]       =				BlipaConfig(vision_config=lowerCamelCase_     ,					text_config=lowerCamelCase_       )

					return config, image_size


@torch.no_grad()
def  _lowerCamelCase   (				lowerCamelCase_ :							Union[str, Any]     ,					lowerCamelCase_ :							Optional[int]=None     ,					lowerCamelCase_ :							Tuple=False       ):



					"""simple docstring"""





					UpperCAmelCase_							:       Any       =				(
					    AutoTokenizer.from_pretrained('facebook/opt-2.7b'       )
					    if 'opt' in model_name
					    else AutoTokenizer.from_pretrained('google/flan-t5-xl'       )
					)
					UpperCAmelCase_							:       List[Any]       =				tokenizer('\n'     ,					add_special_tokens=lowerCamelCase_       ).input_ids[0]
					UpperCAmelCase_       ,      UpperCAmelCase_							:       Dict       =				get_blipa_config(lowerCamelCase_     ,					eos_token_id=lowerCamelCase_       )

					UpperCAmelCase_							:       str       =				BlipaForConditionalGeneration(lowerCamelCase_       ).eval()

					UpperCAmelCase_							:       Optional[int]       =				{
					    'blip2-opt-2.7b': ('blip2_opt', 'pretrain_opt2.7b'),
					    'blip2-opt-6.7b': ('blip2_opt', 'pretrain_opt6.7b'),
					    'blip2-opt-2.7b-coco': ('blip2_opt', 'caption_coco_opt2.7b'),
					    'blip2-opt-6.7b-coco': ('blip2_opt', 'caption_coco_opt6.7b'),
					    'blip2-flan-t5-xl': ('blip2_t5', 'pretrain_flant5xl'),
					    'blip2-flan-t5-xl-coco': ('blip2_t5', 'caption_coco_flant5xl'),
					    'blip2-flan-t5-xxl': ('blip2_t5', 'pretrain_flant5xxl'),
					}

					UpperCAmelCase_       ,      UpperCAmelCase_							:       int       =				model_name_to_original[model_name]

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

					# update state dict keys
					UpperCAmelCase_							:       Optional[int]       =				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_							:       str       =				state_dict.pop(lowerCamelCase_       )
										if key.startswith('Qformer.bert'       ):
															UpperCAmelCase_							:       Optional[int]       =				key.replace('Qformer.bert'     ,					'qformer'       )
										if "attention.self" in key:
															UpperCAmelCase_							:       List[Any]       =				key.replace('self'     ,					'attention'       )
										if "opt_proj" in key:
															UpperCAmelCase_							:       List[str]       =				key.replace('opt_proj'     ,					'language_projection'       )
										if "t5_proj" in key:
															UpperCAmelCase_							:       List[str]       =				key.replace('t5_proj'     ,					'language_projection'       )
										if key.startswith('opt'       ):
															UpperCAmelCase_							:       Any       =				key.replace('opt'     ,					'language'       )
										if key.startswith('t5'       ):
															UpperCAmelCase_							:       Any       =				key.replace('t5'     ,					'language'       )
										UpperCAmelCase_							:       Optional[Any]       =				val

					# read in qv biases
					read_in_q_v_bias(lowerCamelCase_     ,					lowerCamelCase_       )

					UpperCAmelCase_       ,      UpperCAmelCase_							:       Tuple       =				hf_model.load_state_dict(lowerCamelCase_     ,					strict=lowerCamelCase_       )
					assert len(lowerCamelCase_       ) == 0
					assert unexpected_keys == ["qformer.embeddings.position_ids"]

					UpperCAmelCase_							:       Tuple       =				load_demo_image()
					UpperCAmelCase_							:       Tuple       =				vis_processors['eval'](lowerCamelCase_       ).unsqueeze(0       ).to(lowerCamelCase_       )
					UpperCAmelCase_							:       Tuple       =				tokenizer(['\n']     ,					return_tensors='pt'       ).input_ids.to(lowerCamelCase_       )

					# create processor
					UpperCAmelCase_							:       List[Any]       =				BlipImageProcessor(
					    size={'height': image_size, 'width': image_size}     ,					image_mean=lowerCamelCase_     ,					image_std=lowerCamelCase_       )
					UpperCAmelCase_							:       Union[str, Any]       =				BlipaProcessor(image_processor=lowerCamelCase_     ,					tokenizer=lowerCamelCase_       )
					UpperCAmelCase_							:       Tuple       =				processor(images=lowerCamelCase_     ,					return_tensors='pt'       ).pixel_values.to(lowerCamelCase_       )

					# make sure processor creates exact same pixel values
					assert torch.allclose(lowerCamelCase_     ,					lowerCamelCase_       )

					original_model.to(lowerCamelCase_       )
					hf_model.to(lowerCamelCase_       )
					with torch.no_grad():
										if "opt" in model_name:
															UpperCAmelCase_							:       Tuple       =				original_model({'image': original_pixel_values, 'text_input': ['']}       ).logits
															UpperCAmelCase_							:       Any       =				hf_model(lowerCamelCase_     ,					lowerCamelCase_       ).logits
										else:
															UpperCAmelCase_							:       List[str]       =				original_model(
															    {'image': original_pixel_values, 'text_input': ['\n'], 'text_output': ['\n']}       ).logits
															UpperCAmelCase_							:       Tuple       =				input_ids.masked_fill(input_ids == tokenizer.pad_token_id     ,					-100       )
															UpperCAmelCase_							:       List[str]       =				hf_model(lowerCamelCase_     ,					lowerCamelCase_     ,					labels=lowerCamelCase_       ).logits

					assert original_logits.shape == logits.shape
					print('First values of original logits:'     ,					original_logits[0, :3, :3]       )
					print('First values of HF logits:'     ,					logits[0, :3, :3]       )

					# assert values
					if model_name == "blip2-flan-t5-xl":
										UpperCAmelCase_							:       int       =				torch.tensor(
										    [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]     ,					device=lowerCamelCase_       )
										assert torch.allclose(logits[0, :3, :3]     ,					lowerCamelCase_     ,					atol=1e-4       )
					elif model_name == "blip2-flan-t5-xl-coco":
										UpperCAmelCase_							:       Any       =				torch.tensor(
										    [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]     ,					device=lowerCamelCase_       )
					else:
										# cast to same type
										UpperCAmelCase_							:       Dict       =				logits.dtype
										assert torch.allclose(original_logits.to(lowerCamelCase_       )     ,					lowerCamelCase_     ,					atol=1e-2       )
					print('Looks ok!'       )

					print('Generating a caption...'       )
					UpperCAmelCase_							:       Optional[Any]       =				''
					UpperCAmelCase_							:       str       =				tokenizer(lowerCamelCase_     ,					return_tensors='pt'       ).input_ids.to(lowerCamelCase_       )

					UpperCAmelCase_							:       Union[str, Any]       =				original_model.generate({'image': original_pixel_values}       )
					UpperCAmelCase_							:       Optional[int]       =				hf_model.generate(
					    lowerCamelCase_     ,					lowerCamelCase_     ,					do_sample=lowerCamelCase_     ,					num_beams=5     ,					max_length=30     ,					min_length=1     ,					top_p=0.9     ,					repetition_penalty=1.0     ,					length_penalty=1.0     ,					temperature=1     ,					)
					print('Original generation:'     ,					lowerCamelCase_       )
					UpperCAmelCase_							:       Tuple       =				input_ids.shape[1]
					UpperCAmelCase_							:       Optional[Any]       =				processor.batch_decode(outputs[:, prompt_length:]     ,					skip_special_tokens=lowerCamelCase_       )
					UpperCAmelCase_							:       Optional[int]       =				[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'''nielsr/{model_name}'''       )
										hf_model.push_to_hub(F'''nielsr/{model_name}'''       )


if __name__ == "__main__":
	snake_case__					:      int				      =  argparse.ArgumentParser()
	snake_case__					:      Any				      =  [
	    '''blip2-opt-2.7b''',
	    '''blip2-opt-6.7b''',
	    '''blip2-opt-2.7b-coco''',
	    '''blip2-opt-6.7b-coco''',
	    '''blip2-flan-t5-xl''',
	    '''blip2-flan-t5-xl-coco''',
	    '''blip2-flan-t5-xxl''',
	]
	parser.add_argument(
	    '''--model_name''',
	    default='''blip2-opt-2.7b''',
	    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__					:      Any				      =  parser.parse_args()

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







'''simple docstring'''


def  _lowerCamelCase   (				lowerCamelCase_ :							str     ,					lowerCamelCase_ :							str = " "       ):



     """simple docstring"""





     UpperCAmelCase_							:       str       =				[]

     UpperCAmelCase_							:       List[Any]       =				0
     for index, char in enumerate(lowerCamelCase_       ):
          if char == separator:
               split_words.append(string[last_index:index]       )
               UpperCAmelCase_							:       Optional[Any]       =				index + 1
          elif index + 1 == len(lowerCamelCase_       ):
               split_words.append(string[last_index : index + 1]       )
     return split_words


if __name__ == "__main__":
 from doctest import testmod

 testmod()






