




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 		snake_case_   ():





							'''simple docstring'''


							_a     					=    '''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
							_a     					=    Image.open(requests.get(UpperCamelCase			,						stream=UpperCamelCase				).raw				).convert('''RGB'''				)

							return image
def 		snake_case_   (UpperCamelCase	:  Tuple				):





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

							# fmt: on
							return rename_keys
def 		snake_case_   (UpperCamelCase	:  Any			,						UpperCamelCase	:  Tuple			,						UpperCamelCase	:  Optional[int]				):





							'''simple docstring'''


							_a     					=    dct.pop(UpperCamelCase				)
							_a     					=    val
def 		snake_case_   (UpperCamelCase	:  Dict			,						UpperCamelCase	:  int				):





							'''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(UpperCamelCase			,						requires_grad=UpperCamelCase				), v_bias)				)
														_a     					=    qkv_bias
def 		snake_case_   (UpperCamelCase	:  int			,						UpperCamelCase	:  List[Any]				):





							'''simple docstring'''


							_a     					=    364 if '''coco''' in model_name else 224
							_a     					=    BlipaVisionConfig(image_size=UpperCamelCase				).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:
														_a     					=    OPTConfig.from_pretrained('''facebook/opt-2.7b'''			,						eos_token_id=UpperCamelCase				).to_dict()
							elif "opt-6.7b" in model_name:
														_a     					=    OPTConfig.from_pretrained('''facebook/opt-6.7b'''			,						eos_token_id=UpperCamelCase				).to_dict()
							elif "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()

							_a     					=    BlipaConfig(vision_config=UpperCamelCase			,						text_config=UpperCamelCase				)

							return config, image_size
@torch.no_grad()
def 		snake_case_   (UpperCamelCase	:  Dict			,						UpperCamelCase	:  Tuple=None			,						UpperCamelCase	:  List[Any]=False				):





							'''simple docstring'''


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

							_a     					=    BlipaForConditionalGeneration(UpperCamelCase				).eval()

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

							_a					,							_a     					=    model_name_to_original[model_name]

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

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

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

							# read in qv biases
							read_in_q_v_bias(UpperCamelCase			,						UpperCamelCase				)

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

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

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

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

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

							print('''Generating a caption...'''				)
							_a     					=    ''''''
							_a     					=    tokenizer(UpperCamelCase			,						return_tensors='''pt'''				).input_ids.to(UpperCamelCase				)

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

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

							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				:				str       							=		[
				    """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				:				int       							=		parser.parse_args()

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







'''simple docstring'''


import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / 'utils'))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


_snake_case				:				int       							=		get_tests_dir('fixtures')
_snake_case				:				Tuple       							=		get_tests_dir('fixtures/dummy_feature_extractor_config.json')
_snake_case				:				Optional[int]       							=		get_tests_dir('fixtures/dummy-config.json')

class       A   (		unittest.TestCase   ):







							def        __lowerCAmelCase      (  self	:  int )      ->							List[Any]:


														"""simple docstring"""


														_a     					=    0







							def        __lowerCAmelCase      (  self	:  List[str] )      ->							int:


														"""simple docstring"""


														_a     					=    AutoFeatureExtractor.from_pretrained('''facebook/wav2vec2-base-960h''' )
														self.assertIsInstance(lowerCAmelCase_							,      lowerCAmelCase_ )







							def        __lowerCAmelCase      (  self	:  str )      ->							Tuple:


														"""simple docstring"""


														_a     					=    AutoFeatureExtractor.from_pretrained(lowerCAmelCase_ )
														self.assertIsInstance(lowerCAmelCase_							,      lowerCAmelCase_ )







							def        __lowerCAmelCase      (  self	:  List[str] )      ->							Any:


														"""simple docstring"""


														with tempfile.TemporaryDirectory() as tmpdirname:
																					_a     					=    WavaVecaConfig()

																					# remove feature_extractor_type to make sure config.json alone is enough to load feature processor locally
																					_a     					=    AutoFeatureExtractor.from_pretrained(lowerCAmelCase_ ).to_dict()

																					config_dict.pop('''feature_extractor_type''' )
																					_a     					=    WavaVecaFeatureExtractor(**lowerCAmelCase_ )

																					# save in new folder
																					model_config.save_pretrained(lowerCAmelCase_ )
																					config.save_pretrained(lowerCAmelCase_ )

																					_a     					=    AutoFeatureExtractor.from_pretrained(lowerCAmelCase_ )

																					# make sure private variable is not incorrectly saved
																					_a     					=    json.loads(config.to_json_string() )
																					self.assertTrue('''_processor_class''' not in dict_as_saved )

														self.assertIsInstance(lowerCAmelCase_							,      lowerCAmelCase_ )







							def        __lowerCAmelCase      (  self	:  Optional[Any] )      ->							Optional[Any]:


														"""simple docstring"""


														_a     					=    AutoFeatureExtractor.from_pretrained(lowerCAmelCase_ )
														self.assertIsInstance(lowerCAmelCase_							,      lowerCAmelCase_ )







							def        __lowerCAmelCase      (  self	:  Optional[int] )      ->							str:


														"""simple docstring"""


														with self.assertRaisesRegex(
														    lowerCAmelCase_							,      '''bert-base is not a local folder and is not a valid model identifier''' ):
																					_a     					=    AutoFeatureExtractor.from_pretrained('''bert-base''' )







							def        __lowerCAmelCase      (  self	:  Tuple )      ->							Union[str, Any]:


														"""simple docstring"""


														with self.assertRaisesRegex(
														    lowerCAmelCase_							,      R'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)''' ):
																					_a     					=    AutoFeatureExtractor.from_pretrained(lowerCAmelCase_							,      revision='''aaaaaa''' )







							def        __lowerCAmelCase      (  self	:  Any )      ->							Dict:


														"""simple docstring"""


														with self.assertRaisesRegex(
														    lowerCAmelCase_							,      '''hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json.'''							,      ):
																					_a     					=    AutoFeatureExtractor.from_pretrained('''hf-internal-testing/config-no-model''' )







							def        __lowerCAmelCase      (  self	:  List[Any] )      ->							Any:


														"""simple docstring"""


														with self.assertRaises(lowerCAmelCase_ ):
																					_a     					=    AutoFeatureExtractor.from_pretrained(
																					    '''hf-internal-testing/test_dynamic_feature_extractor''' )
														# If remote code is disabled, we can't load this config.
														with self.assertRaises(lowerCAmelCase_ ):
																					_a     					=    AutoFeatureExtractor.from_pretrained(
																					    '''hf-internal-testing/test_dynamic_feature_extractor'''							,      trust_remote_code=lowerCAmelCase_ )

														_a     					=    AutoFeatureExtractor.from_pretrained(
														    '''hf-internal-testing/test_dynamic_feature_extractor'''							,      trust_remote_code=lowerCAmelCase_ )
														self.assertEqual(feature_extractor.__class__.__name__							,      '''NewFeatureExtractor''' )

														# Test feature extractor can be reloaded.
														with tempfile.TemporaryDirectory() as tmp_dir:
																					feature_extractor.save_pretrained(lowerCAmelCase_ )
																					_a     					=    AutoFeatureExtractor.from_pretrained(lowerCAmelCase_							,      trust_remote_code=lowerCAmelCase_ )
														self.assertEqual(reloaded_feature_extractor.__class__.__name__							,      '''NewFeatureExtractor''' )







							def        __lowerCAmelCase      (  self	:  int )      ->							Optional[Any]:


														"""simple docstring"""


														try:
																					AutoConfig.register('''custom'''							,      lowerCAmelCase_ )
																					AutoFeatureExtractor.register(lowerCAmelCase_							,      lowerCAmelCase_ )
																					# Trying to register something existing in the Transformers library will raise an error
																					with self.assertRaises(lowerCAmelCase_ ):
																												AutoFeatureExtractor.register(lowerCAmelCase_							,      lowerCAmelCase_ )

																					# Now that the config is registered, it can be used as any other config with the auto-API
																					_a     					=    CustomFeatureExtractor.from_pretrained(lowerCAmelCase_ )
																					with tempfile.TemporaryDirectory() as tmp_dir:
																												feature_extractor.save_pretrained(lowerCAmelCase_ )
																												_a     					=    AutoFeatureExtractor.from_pretrained(lowerCAmelCase_ )
																												self.assertIsInstance(lowerCAmelCase_							,      lowerCAmelCase_ )

														finally:
																					if "custom" in CONFIG_MAPPING._extra_content:
																												del CONFIG_MAPPING._extra_content["custom"]
																					if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
																												del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]







							def        __lowerCAmelCase      (  self	:  Optional[int] )      ->							Any:


														"""simple docstring"""


														class       A   (		_a   ):
																					lowercase_          =    True

														try:
																					AutoConfig.register('''custom'''							,      lowerCAmelCase_ )
																					AutoFeatureExtractor.register(lowerCAmelCase_							,      lowerCAmelCase_ )
																					# If remote code is not set, the default is to use local
																					_a     					=    AutoFeatureExtractor.from_pretrained(
																					    '''hf-internal-testing/test_dynamic_feature_extractor''' )
																					self.assertEqual(feature_extractor.__class__.__name__							,      '''NewFeatureExtractor''' )
																					self.assertTrue(feature_extractor.is_local )

																					# If remote code is disabled, we load the local one.
																					_a     					=    AutoFeatureExtractor.from_pretrained(
																					    '''hf-internal-testing/test_dynamic_feature_extractor'''							,      trust_remote_code=lowerCAmelCase_ )
																					self.assertEqual(feature_extractor.__class__.__name__							,      '''NewFeatureExtractor''' )
																					self.assertTrue(feature_extractor.is_local )

																					# If remote is enabled, we load from the Hub
																					_a     					=    AutoFeatureExtractor.from_pretrained(
																					    '''hf-internal-testing/test_dynamic_feature_extractor'''							,      trust_remote_code=lowerCAmelCase_ )
																					self.assertEqual(feature_extractor.__class__.__name__							,      '''NewFeatureExtractor''' )
																					self.assertTrue(not hasattr(lowerCAmelCase_							,      '''is_local''' ) )

														finally:
																					if "custom" in CONFIG_MAPPING._extra_content:
																												del CONFIG_MAPPING._extra_content["custom"]
																					if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
																												del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]

