

'''simple docstring'''





import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import UperNetForSemanticSegmentation
							from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from PIL import Image

							from transformers import AutoImageProcessor





class 		__SCREAMING_SNAKE_CASE							:







				def __init__(       self     :		Tuple		,					UpperCAmelCase__     :		int		,					UpperCAmelCase__     :		int=13		,					UpperCAmelCase__     :		List[Any]=32		,					UpperCAmelCase__     :		Dict=3		,					UpperCAmelCase__     :		Any=4		,					UpperCAmelCase__     :		str=[10, 20, 30, 40]		,					UpperCAmelCase__     :		Dict=[2, 2, 3, 2]		,					UpperCAmelCase__     :		Union[str, Any]=True		,					UpperCAmelCase__     :		Dict=True		,					UpperCAmelCase__     :		Tuple=37		,					UpperCAmelCase__     :		Tuple="gelu"		,					UpperCAmelCase__     :		str=10		,					UpperCAmelCase__     :		Optional[Any]=0.02		,					UpperCAmelCase__     :		Union[str, Any]=["stage2", "stage3", "stage4"]		,					UpperCAmelCase__     :		str=3		,					UpperCAmelCase__     :		Optional[Any]=None		,					):
										'''simple docstring'''
										lowercase :							Optional[Any]    		=parent
										lowercase :							int    		=batch_size
										lowercase :							Optional[int]    		=image_size
										lowercase :							Any    		=num_channels
										lowercase :							str    		=num_stages
										lowercase :							List[Any]    		=hidden_sizes
										lowercase :							Dict    		=depths
										lowercase :							int    		=is_training
										lowercase :							Any    		=use_labels
										lowercase :							List[str]    		=intermediate_size
										lowercase :							List[str]    		=hidden_act
										lowercase :							Dict    		=type_sequence_label_size
										lowercase :							int    		=initializer_range
										lowercase :							Optional[Any]    		=out_features
										lowercase :							int    		=num_labels
										lowercase :							Dict    		=scope
										lowercase :							Dict    		=num_stages







				def  lowerCamelCase_	(       self     :		Union[str, Any]   ):
										'''simple docstring'''
										lowercase :							List[str]    		=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

										lowercase :							Optional[int]    		=None
										if self.use_labels:
																lowercase :							Tuple    		=ids_tensor([self.batch_size]		,					self.type_sequence_label_size   )

										lowercase :							str    		=self.get_config()

										return config, pixel_values, labels







				def  lowerCamelCase_	(       self     :		Optional[Any]   ):
										'''simple docstring'''
										return ConvNextConfig(
										    num_channels=self.num_channels		,					num_stages=self.num_stages		,					hidden_sizes=self.hidden_sizes		,					depths=self.depths		,					is_training=self.is_training		,					intermediate_size=self.intermediate_size		,					hidden_act=self.hidden_act		,					out_features=self.out_features		,					)







				def  lowerCamelCase_	(       self     :		int   ):
										'''simple docstring'''
										return UperNetConfig(
										    backbone_config=self.get_backbone_config()		,					hidden_size=512		,					pool_scales=[1, 2, 3, 6]		,					use_auxiliary_head=UpperCAmelCase__		,					auxiliary_loss_weight=0.4		,					auxiliary_in_channels=40		,					auxiliary_channels=256		,					auxiliary_num_convs=1		,					auxiliary_concat_input=UpperCAmelCase__		,					loss_ignore_index=255		,					num_labels=self.num_labels		,					)







				def  lowerCamelCase_	(       self     :		Dict		,					UpperCAmelCase__     :		str		,					UpperCAmelCase__     :		Any		,					UpperCAmelCase__     :		Optional[Any]   ):
										'''simple docstring'''
										lowercase :							Tuple    		=UperNetForSemanticSegmentation(config=UpperCAmelCase__   )
										model.to(UpperCAmelCase__   )
										model.eval()
										lowercase :							Any    		=model(UpperCAmelCase__   )
										self.parent.assertEqual(
										    result.logits.shape		,					(self.batch_size, self.num_labels, self.image_size, self.image_size)   )







				def  lowerCamelCase_	(       self     :		Optional[Any]   ):
										'''simple docstring'''
										lowercase :							int    		=self.prepare_config_and_inputs()
										(
										    (
										    lowercase
										)		,	(
										    lowercase
										)		,	(
										    lowercase
										)		,	
										) :							Dict    		=config_and_inputs
										lowercase :							Optional[Any]    		={'''pixel_values''': pixel_values}
										return config, inputs_dict





@require_torch
class 		__SCREAMING_SNAKE_CASE							(			lowercase__						,							lowercase__						,							unittest.TestCase      ):
				lowerCamelCase_		       =    (UperNetForSemanticSegmentation,) if is_torch_available() else ()
				lowerCamelCase_		       =    {'image-segmentation': UperNetForSemanticSegmentation} if is_torch_available() else {}
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    False







				def  lowerCamelCase_	(       self     :		str   ):
										'''simple docstring'''
										lowercase :							int    		=UperNetModelTester(self   )
										lowercase :							Optional[int]    		=ConfigTester(self		,					config_class=UpperCAmelCase__		,					has_text_modality=UpperCAmelCase__		,					hidden_size=37   )







				def  lowerCamelCase_	(       self     :		Union[str, Any]   ):
										'''simple docstring'''
										self.create_and_test_config_common_properties()
										self.config_tester.create_and_test_config_to_json_string()
										self.config_tester.create_and_test_config_to_json_file()
										self.config_tester.create_and_test_config_from_and_save_pretrained()
										self.config_tester.create_and_test_config_with_num_labels()
										self.config_tester.check_config_can_be_init_without_params()
										self.config_tester.check_config_arguments_init()







				def  lowerCamelCase_	(       self     :		Union[str, Any]   ):
										'''simple docstring'''
										return







				def  lowerCamelCase_	(       self     :		Any   ):
										'''simple docstring'''
										lowercase		,	lowercase :							str    		=self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																lowercase :							List[Any]    		=model_class(UpperCAmelCase__   )
																lowercase :							Union[str, Any]    		=inspect.signature(model.forward   )
																# signature.parameters is an OrderedDict => so arg_names order is deterministic
																lowercase :							List[str]    		=[*signature.parameters.keys()]

																lowercase :							str    		=['''pixel_values''']
																self.assertListEqual(arg_names[:1]		,					UpperCAmelCase__   )







				def  lowerCamelCase_	(       self     :		Tuple   ):
										'''simple docstring'''
										lowercase :							Optional[int]    		=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_semantic_segmentation(*UpperCAmelCase__   )







				@unittest.skip(reason='''UperNet does not use inputs_embeds'''   )
				def  lowerCamelCase_	(       self     :		int   ):
										'''simple docstring'''
										pass







				@unittest.skip(reason='''UperNet does not support input and output embeddings'''   )
				def  lowerCamelCase_	(       self     :		Union[str, Any]   ):
										'''simple docstring'''
										pass







				@unittest.skip(reason='''UperNet does not have a base model'''   )
				def  lowerCamelCase_	(       self     :		List[Any]   ):
										'''simple docstring'''
										pass







				@unittest.skip(reason='''UperNet does not have a base model'''   )
				def  lowerCamelCase_	(       self     :		Union[str, Any]   ):
										'''simple docstring'''
										pass







				@require_torch_multi_gpu
				@unittest.skip(reason='''UperNet has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''   )
				def  lowerCamelCase_	(       self     :		List[str]   ):
										'''simple docstring'''
										pass







				@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''   )
				def  lowerCamelCase_	(       self     :		Union[str, Any]   ):
										'''simple docstring'''
										pass







				def  lowerCamelCase_	(       self     :		int   ):
										'''simple docstring'''
										def check_hidden_states_output(UpperCAmelCase__     :		List[Any]		,					UpperCAmelCase__     :		Any		,					UpperCAmelCase__     :		Tuple   ):
																lowercase :							Optional[int]    		=model_class(UpperCAmelCase__   )
																model.to(UpperCAmelCase__   )
																model.eval()

																with torch.no_grad():
																						lowercase :							Union[str, Any]    		=model(**self._prepare_for_class(UpperCAmelCase__		,					UpperCAmelCase__   )   )

																lowercase :							Optional[int]    		=outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																lowercase :							str    		=self.model_tester.num_stages
																self.assertEqual(len(UpperCAmelCase__   )		,					expected_num_stages + 1   )

																# ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
																self.assertListEqual(
																    list(hidden_states[0].shape[-2:]   )		,					[self.model_tester.image_size // 4, self.model_tester.image_size // 4]		,					)

										lowercase		,	lowercase :							Dict    		=self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																lowercase :							Dict    		=True
																check_hidden_states_output(UpperCAmelCase__		,					UpperCAmelCase__		,					UpperCAmelCase__   )

																# check that output_hidden_states also work using config
																del inputs_dict["output_hidden_states"]
																lowercase :							int    		=True

																check_hidden_states_output(UpperCAmelCase__		,					UpperCAmelCase__		,					UpperCAmelCase__   )







				def  lowerCamelCase_	(       self     :		Optional[int]   ):
										'''simple docstring'''
										lowercase		,	lowercase :							Optional[int]    		=self.model_tester.prepare_config_and_inputs_for_common()

										lowercase :							int    		=_config_zero_init(UpperCAmelCase__   )
										lowercase :							List[Any]    		=_config_zero_init(configs_no_init.backbone_config   )
										for model_class in self.all_model_classes:
																lowercase :							List[str]    		=model_class(config=UpperCAmelCase__   )
																for name, param in model.named_parameters():
																						if param.requires_grad:
																												self.assertIn(
																												    ((param.data.mean() * 1E9).round() / 1E9).item()		,					[0.0, 1.0]		,					msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''		,					)







				@unittest.skip(reason='''UperNet does not have tied weights'''   )
				def  lowerCamelCase_	(       self     :		Dict   ):
										'''simple docstring'''
										pass







				@slow
				def  lowerCamelCase_	(       self     :		Any   ):
										'''simple docstring'''
										for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																lowercase :							str    		=UperNetForSemanticSegmentation.from_pretrained(UpperCAmelCase__   )
																self.assertIsNotNone(UpperCAmelCase__   )






def 				_lowerCAmelCase			(   )				->  str:
						lowercase :							str    		=hf_hub_download(
						    repo_id='''hf-internal-testing/fixtures_ade20k'''   , repo_type='''dataset'''   , filename='''ADE_val_00000001.jpg'''						)
						lowercase :							str    		=Image.open(__magic_name__						).convert('''RGB'''						)
						return image







@require_torch
@require_vision
@slow
class 		__SCREAMING_SNAKE_CASE							(			unittest.TestCase      ):







				def  lowerCamelCase_	(       self     :		Any   ):
										'''simple docstring'''
										lowercase :							Tuple    		=AutoImageProcessor.from_pretrained('''openmmlab/upernet-swin-tiny'''   )
										lowercase :							int    		=UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-swin-tiny'''   ).to(UpperCAmelCase__   )

										lowercase :							Dict    		=prepare_img()
										lowercase :							List[str]    		=processor(images=UpperCAmelCase__		,					return_tensors='''pt'''   ).to(UpperCAmelCase__   )

										with torch.no_grad():
																lowercase :							Union[str, Any]    		=model(**UpperCAmelCase__   )

										lowercase :							Tuple    		=torch.Size((1, model.config.num_labels, 512, 512)   )
										self.assertEqual(outputs.logits.shape		,					UpperCAmelCase__   )

										lowercase :							int    		=torch.tensor(
										    [[-7.59_58, -7.59_58, -7.43_02], [-7.59_58, -7.59_58, -7.43_02], [-7.47_97, -7.47_97, -7.30_68]]   ).to(UpperCAmelCase__   )
										self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]		,					UpperCAmelCase__		,					atol=1E-4   )   )







				def  lowerCamelCase_	(       self     :		str   ):
										'''simple docstring'''
										lowercase :							Dict    		=AutoImageProcessor.from_pretrained('''openmmlab/upernet-convnext-tiny'''   )
										lowercase :							Optional[int]    		=UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-convnext-tiny'''   ).to(UpperCAmelCase__   )

										lowercase :							Optional[int]    		=prepare_img()
										lowercase :							Dict    		=processor(images=UpperCAmelCase__		,					return_tensors='''pt'''   ).to(UpperCAmelCase__   )

										with torch.no_grad():
																lowercase :							Any    		=model(**UpperCAmelCase__   )

										lowercase :							int    		=torch.Size((1, model.config.num_labels, 512, 512)   )
										self.assertEqual(outputs.logits.shape		,					UpperCAmelCase__   )

										lowercase :							List[str]    		=torch.tensor(
										    [[-8.81_10, -8.81_10, -8.65_21], [-8.81_10, -8.81_10, -8.65_21], [-8.77_46, -8.77_46, -8.61_30]]   ).to(UpperCAmelCase__   )
										self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]		,					UpperCAmelCase__		,					atol=1E-4   )   )










"""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 		__snake_case		( )    ->  Tuple:





     '''simple docstring'''

     _UpperCAmelCase     :      Optional[int]														=       "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"
     _UpperCAmelCase     :      int														=       Image.open(requests.get(SCREAMING_SNAKE_CASE__		,    stream=SCREAMING_SNAKE_CASE__							).raw							).convert("RGB"							)

     return image

def 		__snake_case		( SCREAMING_SNAKE_CASE__      :		Optional[int]							)    ->  Dict:





     '''simple docstring'''

     _UpperCAmelCase     :      List[Any]														=       []
     # 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		( SCREAMING_SNAKE_CASE__      :		List[str]		,    SCREAMING_SNAKE_CASE__      :		Dict		,    SCREAMING_SNAKE_CASE__      :		Union[str, Any]							)    ->  Dict:





     '''simple docstring'''

     _UpperCAmelCase     :      Optional[int]														=       dct.pop(SCREAMING_SNAKE_CASE__							)
     _UpperCAmelCase     :      Tuple														=       val

def 		__snake_case		( SCREAMING_SNAKE_CASE__      :		Dict		,    SCREAMING_SNAKE_CASE__      :		int							)    ->  Optional[Any]:





     '''simple docstring'''

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

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

def 		__snake_case		( SCREAMING_SNAKE_CASE__      :		str		,    SCREAMING_SNAKE_CASE__      :		Any							)    ->  str:





     '''simple docstring'''

     _UpperCAmelCase     :      int														=       364 if "coco" in model_name else 224
     _UpperCAmelCase     :      Dict														=       BlipaVisionConfig(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 "opt-2.7b" in model_name:
          _UpperCAmelCase     :      Any														=       OPTConfig.from_pretrained("facebook/opt-2.7b"		,    eos_token_id=SCREAMING_SNAKE_CASE__							).to_dict()
     elif "opt-6.7b" in model_name:
          _UpperCAmelCase     :      List[Any]														=       OPTConfig.from_pretrained("facebook/opt-6.7b"		,    eos_token_id=SCREAMING_SNAKE_CASE__							).to_dict()
     elif "t5-xl" in model_name:
          _UpperCAmelCase     :      Tuple														=       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()

     _UpperCAmelCase     :      Optional[Any]														=       BlipaConfig(vision_config=SCREAMING_SNAKE_CASE__		,    text_config=SCREAMING_SNAKE_CASE__							)

     return config, image_size

@torch.no_grad()
def 		__snake_case		( SCREAMING_SNAKE_CASE__      :		Tuple		,    SCREAMING_SNAKE_CASE__      :		Tuple=None		,    SCREAMING_SNAKE_CASE__      :		Union[str, Any]=False							)    ->  Optional[Any]:





     '''simple docstring'''

     _UpperCAmelCase     :      str														=       (
         AutoTokenizer.from_pretrained("facebook/opt-2.7b"							)
         if "opt" in model_name
         else AutoTokenizer.from_pretrained("google/flan-t5-xl"							)
     )
     _UpperCAmelCase     :      Optional[int]														=       tokenizer("\n"		,    add_special_tokens=SCREAMING_SNAKE_CASE__							).input_ids[0]
     _UpperCAmelCase							,			_UpperCAmelCase     :      int														=       get_blipa_config(SCREAMING_SNAKE_CASE__		,    eos_token_id=SCREAMING_SNAKE_CASE__							)

     _UpperCAmelCase     :      int														=       BlipaForConditionalGeneration(SCREAMING_SNAKE_CASE__							).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     :      str														=       model_name_to_original[model_name]

     # load original model
     print("Loading original model..."							)
     _UpperCAmelCase     :      List[Any]														=       "cuda" if torch.cuda.is_available() else "cpu"
     _UpperCAmelCase							,			_UpperCAmelCase							,			_UpperCAmelCase     :      Optional[Any]														=       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
     _UpperCAmelCase     :      Optional[Any]														=       original_model.state_dict()
     _UpperCAmelCase     :      Union[str, Any]														=       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():
          _UpperCAmelCase     :      Tuple														=       state_dict.pop(SCREAMING_SNAKE_CASE__							)
          if key.startswith("Qformer.bert"							):
               _UpperCAmelCase     :      List[Any]														=       key.replace("Qformer.bert"		,    "qformer"							)
          if "attention.self" in key:
               _UpperCAmelCase     :      Dict														=       key.replace("self"		,    "attention"							)
          if "opt_proj" in key:
               _UpperCAmelCase     :      Union[str, Any]														=       key.replace("opt_proj"		,    "language_projection"							)
          if "t5_proj" in key:
               _UpperCAmelCase     :      Optional[Any]														=       key.replace("t5_proj"		,    "language_projection"							)
          if key.startswith("opt"							):
               _UpperCAmelCase     :      Dict														=       key.replace("opt"		,    "language"							)
          if key.startswith("t5"							):
               _UpperCAmelCase     :      Union[str, Any]														=       key.replace("t5"		,    "language"							)
          _UpperCAmelCase     :      List[str]														=       val

     # read in qv biases
     read_in_q_v_bias(SCREAMING_SNAKE_CASE__		,    SCREAMING_SNAKE_CASE__							)

     _UpperCAmelCase							,			_UpperCAmelCase     :      Optional[Any]														=       hf_model.load_state_dict(SCREAMING_SNAKE_CASE__		,    strict=SCREAMING_SNAKE_CASE__							)
     assert len(SCREAMING_SNAKE_CASE__							) == 0
     assert unexpected_keys == ["qformer.embeddings.position_ids"]

     _UpperCAmelCase     :      Dict														=       load_demo_image()
     _UpperCAmelCase     :      List[Any]														=       vis_processors["eval"](SCREAMING_SNAKE_CASE__							).unsqueeze(0							).to(SCREAMING_SNAKE_CASE__							)
     _UpperCAmelCase     :      Optional[Any]														=       tokenizer(["\n"]		,    return_tensors="pt"							).input_ids.to(SCREAMING_SNAKE_CASE__							)

     # create processor
     _UpperCAmelCase     :      Any														=       BlipImageProcessor(
         size={"height": image_size, "width": image_size}		,    image_mean=SCREAMING_SNAKE_CASE__		,    image_std=SCREAMING_SNAKE_CASE__							)
     _UpperCAmelCase     :      Dict														=       BlipaProcessor(image_processor=SCREAMING_SNAKE_CASE__		,    tokenizer=SCREAMING_SNAKE_CASE__							)
     _UpperCAmelCase     :      Dict														=       processor(images=SCREAMING_SNAKE_CASE__		,    return_tensors="pt"							).pixel_values.to(SCREAMING_SNAKE_CASE__							)

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

     original_model.to(SCREAMING_SNAKE_CASE__							)
     hf_model.to(SCREAMING_SNAKE_CASE__							)
     with torch.no_grad():
          if "opt" in model_name:
               _UpperCAmelCase     :      str														=       original_model({"image": original_pixel_values, "text_input": [""]}							).logits
               _UpperCAmelCase     :      Optional[int]														=       hf_model(SCREAMING_SNAKE_CASE__		,    SCREAMING_SNAKE_CASE__							).logits
          else:
               _UpperCAmelCase     :      List[Any]														=       original_model(
                   {"image": original_pixel_values, "text_input": ["\n"], "text_output": ["\n"]}							).logits
               _UpperCAmelCase     :      List[str]														=       input_ids.masked_fill(input_ids == tokenizer.pad_token_id		,    -100							)
               _UpperCAmelCase     :      Union[str, Any]														=       hf_model(SCREAMING_SNAKE_CASE__		,    SCREAMING_SNAKE_CASE__		,    labels=SCREAMING_SNAKE_CASE__							).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     :      List[str]														=       torch.tensor(
              [[-41.5_850, -4.4_440, -8.9_922], [-47.4_322, -5.9_143, -1.7_340]]		,    device=SCREAMING_SNAKE_CASE__							)
          assert torch.allclose(logits[0, :3, :3]		,    SCREAMING_SNAKE_CASE__		,    atol=1E-4							)
     elif model_name == "blip2-flan-t5-xl-coco":
          _UpperCAmelCase     :      Any														=       torch.tensor(
              [[-57.0_109, -9.8_967, -12.6_280], [-68.6_578, -12.7_191, -10.5_065]]		,    device=SCREAMING_SNAKE_CASE__							)
     else:
          # cast to same type
          _UpperCAmelCase     :      List[str]														=       logits.dtype
          assert torch.allclose(original_logits.to(SCREAMING_SNAKE_CASE__							)		,    SCREAMING_SNAKE_CASE__		,    atol=1E-2							)
     print("Looks ok!"							)

     print("Generating a caption..."							)
     _UpperCAmelCase     :      int														=       ""
     _UpperCAmelCase     :      Any														=       tokenizer(SCREAMING_SNAKE_CASE__		,    return_tensors="pt"							).input_ids.to(SCREAMING_SNAKE_CASE__							)

     _UpperCAmelCase     :      List[Any]														=       original_model.generate({"image": original_pixel_values}							)
     _UpperCAmelCase     :      int														=       hf_model.generate(
         SCREAMING_SNAKE_CASE__		,    SCREAMING_SNAKE_CASE__		,    do_sample=SCREAMING_SNAKE_CASE__		,    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:"		,    SCREAMING_SNAKE_CASE__							)
     _UpperCAmelCase     :      str														=       input_ids.shape[1]
     _UpperCAmelCase     :      List[Any]														=       processor.batch_decode(outputs[:, prompt_length:]		,    skip_special_tokens=SCREAMING_SNAKE_CASE__							)
     _UpperCAmelCase     :      int														=       [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'nielsr/{model_name}'							)
          hf_model.push_to_hub(f'nielsr/{model_name}'							)


if __name__ == "__main__":
 _lowerCAmelCase							:   Union[str, Any]				=				argparse.ArgumentParser()
 _lowerCAmelCase							:   List[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",
 )

 _lowerCAmelCase							:   str				=				parser.parse_args()

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