


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( ):
				_SCREAMING_SNAKE_CASE      	='''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
				_SCREAMING_SNAKE_CASE      	=Image.open(requests.get(A_						,stream=A_).raw).convert('''RGB''')

				return image







def     lowerCamelCase( a__):
				_SCREAMING_SNAKE_CASE      	=[]
				# 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( a__						,a__						,a__):
				_SCREAMING_SNAKE_CASE      	=dct.pop(A_)
				_SCREAMING_SNAKE_CASE      	=val







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

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







def     lowerCamelCase( a__						,a__):
				_SCREAMING_SNAKE_CASE      	=364 if '''coco''' in model_name else 224
				_SCREAMING_SNAKE_CASE      	=BlipaVisionConfig(image_size=A_).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:
								_SCREAMING_SNAKE_CASE      	=OPTConfig.from_pretrained('''facebook/opt-2.7b'''						,eos_token_id=A_).to_dict()
				elif "opt-6.7b" in model_name:
								_SCREAMING_SNAKE_CASE      	=OPTConfig.from_pretrained('''facebook/opt-6.7b'''						,eos_token_id=A_).to_dict()
				elif "t5-xl" in model_name:
								_SCREAMING_SNAKE_CASE      	=TaConfig.from_pretrained('''google/flan-t5-xl'''						,dense_act_fn='''gelu'''						,bos_token_id=1).to_dict()
				elif "t5-xxl" in model_name:
								_SCREAMING_SNAKE_CASE      	=TaConfig.from_pretrained('''google/flan-t5-xxl'''						,dense_act_fn='''gelu'''						,bos_token_id=1).to_dict()

				_SCREAMING_SNAKE_CASE      	=BlipaConfig(vision_config=A_						,text_config=A_)

				return config, image_size







@torch.no_grad()
def     lowerCamelCase( a__						,a__=None						,a__=False):
				_SCREAMING_SNAKE_CASE      	=(
				    AutoTokenizer.from_pretrained('''facebook/opt-2.7b''')
				    if '''opt''' in model_name
				    else AutoTokenizer.from_pretrained('''google/flan-t5-xl''')
				)
				_SCREAMING_SNAKE_CASE      	=tokenizer('''\n'''						,add_special_tokens=A_).input_ids[0]
				_SCREAMING_SNAKE_CASE      	=get_blipa_config(A_						,eos_token_id=A_)

				_SCREAMING_SNAKE_CASE      	=BlipaForConditionalGeneration(A_).eval()

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

				_SCREAMING_SNAKE_CASE      	=model_name_to_original[model_name]

				# load original model
				print('''Loading original model...''')
				_SCREAMING_SNAKE_CASE      	='''cuda''' if torch.cuda.is_available() else '''cpu'''
				_SCREAMING_SNAKE_CASE      	=load_model_and_preprocess(
				    name=A_						,model_type=A_						,is_eval=A_						,device=A_)
				original_model.eval()
				print('''Done!''')

				# update state dict keys
				_SCREAMING_SNAKE_CASE      	=original_model.state_dict()
				_SCREAMING_SNAKE_CASE      	=create_rename_keys(A_)
				for src, dest in rename_keys:
								rename_key(A_						,A_						,A_)

				# some keys can be renamed efficiently
				for key, val in state_dict.copy().items():
								_SCREAMING_SNAKE_CASE      	=state_dict.pop(A_)
								if key.startswith('''Qformer.bert'''):
												_SCREAMING_SNAKE_CASE      	=key.replace('''Qformer.bert'''						,'''qformer''')
								if "attention.self" in key:
												_SCREAMING_SNAKE_CASE      	=key.replace('''self'''						,'''attention''')
								if "opt_proj" in key:
												_SCREAMING_SNAKE_CASE      	=key.replace('''opt_proj'''						,'''language_projection''')
								if "t5_proj" in key:
												_SCREAMING_SNAKE_CASE      	=key.replace('''t5_proj'''						,'''language_projection''')
								if key.startswith('''opt'''):
												_SCREAMING_SNAKE_CASE      	=key.replace('''opt'''						,'''language''')
								if key.startswith('''t5'''):
												_SCREAMING_SNAKE_CASE      	=key.replace('''t5'''						,'''language''')
								_SCREAMING_SNAKE_CASE      	=val

				# read in qv biases
				read_in_q_v_bias(A_						,A_)

				_SCREAMING_SNAKE_CASE      	=hf_model.load_state_dict(A_						,strict=A_)
				assert len(A_) == 0
				assert unexpected_keys == ["qformer.embeddings.position_ids"]

				_SCREAMING_SNAKE_CASE      	=load_demo_image()
				_SCREAMING_SNAKE_CASE      	=vis_processors['''eval'''](A_).unsqueeze(0).to(A_)
				_SCREAMING_SNAKE_CASE      	=tokenizer(['''\n''']						,return_tensors='''pt''').input_ids.to(A_)

				# create processor
				_SCREAMING_SNAKE_CASE      	=BlipImageProcessor(
				    size={'''height''': image_size, '''width''': image_size}						,image_mean=A_						,image_std=A_)
				_SCREAMING_SNAKE_CASE      	=BlipaProcessor(image_processor=A_						,tokenizer=A_)
				_SCREAMING_SNAKE_CASE      	=processor(images=A_						,return_tensors='''pt''').pixel_values.to(A_)

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

				original_model.to(A_)
				hf_model.to(A_)
				with torch.no_grad():
								if "opt" in model_name:
												_SCREAMING_SNAKE_CASE      	=original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}).logits
												_SCREAMING_SNAKE_CASE      	=hf_model(A_						,A_).logits
								else:
												_SCREAMING_SNAKE_CASE      	=original_model(
												    {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}).logits
												_SCREAMING_SNAKE_CASE      	=input_ids.masked_fill(input_ids == tokenizer.pad_token_id						,-100)
												_SCREAMING_SNAKE_CASE      	=hf_model(A_						,A_						,labels=A_).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":
								_SCREAMING_SNAKE_CASE      	=torch.tensor(
								    [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]						,device=A_)
								assert torch.allclose(logits[0, :3, :3]						,A_						,atol=1e-4)
				elif model_name == "blip2-flan-t5-xl-coco":
								_SCREAMING_SNAKE_CASE      	=torch.tensor(
								    [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]						,device=A_)
				else:
								# cast to same type
								_SCREAMING_SNAKE_CASE      	=logits.dtype
								assert torch.allclose(original_logits.to(A_)						,A_						,atol=1e-2)
				print('''Looks ok!''')

				print('''Generating a caption...''')
				_SCREAMING_SNAKE_CASE      	=''''''
				_SCREAMING_SNAKE_CASE      	=tokenizer(A_						,return_tensors='''pt''').input_ids.to(A_)

				_SCREAMING_SNAKE_CASE      	=original_model.generate({'''image''': original_pixel_values})
				_SCREAMING_SNAKE_CASE      	=hf_model.generate(
				    A_						,A_						,do_sample=A_						,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:'''						,A_)
				_SCREAMING_SNAKE_CASE      	=input_ids.shape[1]
				_SCREAMING_SNAKE_CASE      	=processor.batch_decode(outputs[:, prompt_length:]						,skip_special_tokens=A_)
				_SCREAMING_SNAKE_CASE      	=[text.strip() for text in output_text]
				print('''HF generation:'''						,A_)

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

				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_    :      Optional[Any]				    =				argparse.ArgumentParser()
							snake_case_    :      Tuple				    =				[
							    '''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_    :      Optional[Any]				    =				parser.parse_args()

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


"""simple docstring"""





import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
							import torch

if is_vision_available():
							from PIL import Image

							from transformers import LevitImageProcessor







class 							SCREAMING_SNAKE_CASE   (						unittest.TestCase			):






					"""simple docstring"""



					def __init__(			self     :							Any  ,lowercase_     :							Dict  ,lowercase_     :							List[str]=7  ,lowercase_     :							Tuple=3  ,lowercase_     :							List[str]=1_8  ,lowercase_     :							Optional[Any]=3_0  ,lowercase_     :							List[Any]=4_0_0  ,lowercase_     :							List[Any]=True  ,lowercase_     :							Any=None  ,lowercase_     :							Optional[Any]=True  ,lowercase_     :							str=None  ,lowercase_     :							List[Any]=True  ,lowercase_     :							Dict=[0.5, 0.5, 0.5]  ,lowercase_     :							Dict=[0.5, 0.5, 0.5]  ,):
						lowerCAmelCase__      :		Any    							=      size if size is not None else {'''shortest_edge''': 1_8}
						lowerCAmelCase__      :		Union[str, Any]    							=      crop_size if crop_size is not None else {'''height''': 1_8, '''width''': 1_8}
						lowerCAmelCase__      :		Dict    							=      parent
						lowerCAmelCase__      :		Dict    							=      batch_size
						lowerCAmelCase__      :		List[str]    							=      num_channels
						lowerCAmelCase__      :		Any    							=      image_size
						lowerCAmelCase__      :		Union[str, Any]    							=      min_resolution
						lowerCAmelCase__      :		Dict    							=      max_resolution
						lowerCAmelCase__      :		List[str]    							=      do_resize
						lowerCAmelCase__      :		Optional[Any]    							=      size
						lowerCAmelCase__      :		Tuple    							=      do_center_crop
						lowerCAmelCase__      :		Optional[int]    							=      crop_size
						lowerCAmelCase__      :		List[str]    							=      do_normalize
						lowerCAmelCase__      :		Tuple    							=      image_mean
						lowerCAmelCase__      :		int    							=      image_std







					def 			__lowerCAmelCase					(			self     :							List[str]						):
						return {
						    "image_mean": self.image_mean,
						    "image_std": self.image_std,
						    "do_normalize": self.do_normalize,
						    "do_resize": self.do_resize,
						    "do_center_crop": self.do_center_crop,
						    "size": self.size,
						    "crop_size": self.crop_size,
						}







@require_torch
@require_vision
class 							SCREAMING_SNAKE_CASE   (						a_       ,				unittest.TestCase			):






					"""simple docstring"""



					lowercase__												=			LevitImageProcessor if is_vision_available() else None







					def 			__lowerCAmelCase					(			self     :							List[str]						):
						lowerCAmelCase__      :		Optional[Any]    							=      LevitImageProcessingTester(self						)







					@property
					def 			__lowerCAmelCase					(			self     :							Any						):
						return self.image_processor_tester.prepare_image_processor_dict()







					def 			__lowerCAmelCase					(			self     :							str						):
						lowerCAmelCase__      :		Optional[int]    							=      self.image_processing_class(**self.image_processor_dict						)
						self.assertTrue(hasattr(lowercase_  ,'''image_mean'''						)						)
						self.assertTrue(hasattr(lowercase_  ,'''image_std'''						)						)
						self.assertTrue(hasattr(lowercase_  ,'''do_normalize'''						)						)
						self.assertTrue(hasattr(lowercase_  ,'''do_resize'''						)						)
						self.assertTrue(hasattr(lowercase_  ,'''do_center_crop'''						)						)
						self.assertTrue(hasattr(lowercase_  ,'''size'''						)						)







					def 			__lowerCAmelCase					(			self     :							Any						):
						lowerCAmelCase__      :		int    							=      self.image_processing_class.from_dict(self.image_processor_dict						)
						self.assertEqual(image_processor.size  ,{'''shortest_edge''': 1_8}						)
						self.assertEqual(image_processor.crop_size  ,{'''height''': 1_8, '''width''': 1_8}						)

						lowerCAmelCase__      :		List[str]    							=      self.image_processing_class.from_dict(self.image_processor_dict  ,size=4_2  ,crop_size=8_4						)
						self.assertEqual(image_processor.size  ,{'''shortest_edge''': 4_2}						)
						self.assertEqual(image_processor.crop_size  ,{'''height''': 8_4, '''width''': 8_4}						)







					def 			__lowerCAmelCase					(			self     :							Union[str, Any]						):
						pass







					def 			__lowerCAmelCase					(			self     :							Optional[int]						):
						# Initialize image_processing
						lowerCAmelCase__      :		List[str]    							=      self.image_processing_class(**self.image_processor_dict						)
						# create random PIL images
						lowerCAmelCase__      :		Union[str, Any]    							=      prepare_image_inputs(self.image_processor_tester  ,equal_resolution=lowercase_						)
						for image in image_inputs:
							self.assertIsInstance(lowercase_  ,Image.Image						)

						# Test not batched input
						lowerCAmelCase__      :		Any    							=      image_processing(image_inputs[0]  ,return_tensors='''pt'''						).pixel_values
						self.assertEqual(
						    encoded_images.shape  ,(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['''height'''],
						        self.image_processor_tester.crop_size['''width'''],
						    )  ,)

						# Test batched
						lowerCAmelCase__      :		Optional[int]    							=      image_processing(lowercase_  ,return_tensors='''pt'''						).pixel_values
						self.assertEqual(
						    encoded_images.shape  ,(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['''height'''],
						        self.image_processor_tester.crop_size['''width'''],
						    )  ,)







					def 			__lowerCAmelCase					(			self     :							Union[str, Any]						):
						# Initialize image_processing
						lowerCAmelCase__      :		str    							=      self.image_processing_class(**self.image_processor_dict						)
						# create random numpy tensors
						lowerCAmelCase__      :		Optional[Any]    							=      prepare_image_inputs(self.image_processor_tester  ,equal_resolution=lowercase_  ,numpify=lowercase_						)
						for image in image_inputs:
							self.assertIsInstance(lowercase_  ,np.ndarray						)

						# Test not batched input
						lowerCAmelCase__      :		Tuple    							=      image_processing(image_inputs[0]  ,return_tensors='''pt'''						).pixel_values
						self.assertEqual(
						    encoded_images.shape  ,(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['''height'''],
						        self.image_processor_tester.crop_size['''width'''],
						    )  ,)

						# Test batched
						lowerCAmelCase__      :		Union[str, Any]    							=      image_processing(lowercase_  ,return_tensors='''pt'''						).pixel_values
						self.assertEqual(
						    encoded_images.shape  ,(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['''height'''],
						        self.image_processor_tester.crop_size['''width'''],
						    )  ,)







					def 			__lowerCAmelCase					(			self     :							int						):
						# Initialize image_processing
						lowerCAmelCase__      :		Tuple    							=      self.image_processing_class(**self.image_processor_dict						)
						# create random PyTorch tensors
						lowerCAmelCase__      :		Dict    							=      prepare_image_inputs(self.image_processor_tester  ,equal_resolution=lowercase_  ,torchify=lowercase_						)
						for image in image_inputs:
							self.assertIsInstance(lowercase_  ,torch.Tensor						)

						# Test not batched input
						lowerCAmelCase__      :		Dict    							=      image_processing(image_inputs[0]  ,return_tensors='''pt'''						).pixel_values
						self.assertEqual(
						    encoded_images.shape  ,(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['''height'''],
						        self.image_processor_tester.crop_size['''width'''],
						    )  ,)

						# Test batched
						lowerCAmelCase__      :		Optional[int]    							=      image_processing(lowercase_  ,return_tensors='''pt'''						).pixel_values
						self.assertEqual(
						    encoded_images.shape  ,(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['''height'''],
						        self.image_processor_tester.crop_size['''width'''],
						    )  ,)




