



'''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 					UpperCAmelCase   (						):
							SCREAMING_SNAKE_CASE					:      List[Any]     					=			'''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
							SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			Image.open(requests.get(A					,		stream=A				).raw				).convert('''RGB'''				)

							return image
def 					UpperCAmelCase   (						A  :				Union[str, Any]				):
							SCREAMING_SNAKE_CASE					:      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 					UpperCAmelCase   (						A  :				Tuple					,		A  :				Any					,		A  :				Optional[int]				):
							SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			dct.pop(A				)
							SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			val
def 					UpperCAmelCase   (						A  :				str					,		A  :				str				):
							for i in range(config.vision_config.num_hidden_layers				):
														# read in original q and v biases
														SCREAMING_SNAKE_CASE					:      Dict     					=			state_dict.pop(F"""visual_encoder.blocks.{i}.attn.q_bias"""				)
														SCREAMING_SNAKE_CASE					:      int     					=			state_dict.pop(F"""visual_encoder.blocks.{i}.attn.v_bias"""				)

														# next, set bias in the state dict
														SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			torch.cat((q_bias, torch.zeros_like(A					,		requires_grad=A				), v_bias)				)
														SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			qkv_bias
def 					UpperCAmelCase   (						A  :				Any					,		A  :				Union[str, Any]				):
							SCREAMING_SNAKE_CASE					:      Tuple     					=			364 if '''coco''' in model_name else 224
							SCREAMING_SNAKE_CASE					:      Dict     					=			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					:      str     					=			OPTConfig.from_pretrained('''facebook/opt-2.7b'''					,		eos_token_id=A				).to_dict()
							elif "opt-6.7b" in model_name:
														SCREAMING_SNAKE_CASE					:      str     					=			OPTConfig.from_pretrained('''facebook/opt-6.7b'''					,		eos_token_id=A				).to_dict()
							elif "t5-xl" in model_name:
														SCREAMING_SNAKE_CASE					:      Dict     					=			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					:      Optional[Any]     					=			TaConfig.from_pretrained('''google/flan-t5-xxl'''					,		dense_act_fn='''gelu'''					,		bos_token_id=1				).to_dict()

							SCREAMING_SNAKE_CASE					:      Tuple     					=			BlipaConfig(vision_config=A					,		text_config=A				)

							return config, image_size
@torch.no_grad()
def 					UpperCAmelCase   (						A  :				Tuple					,		A  :				List[str]=None					,		A  :				List[str]=False				):
							SCREAMING_SNAKE_CASE					:      str     					=			(
							    AutoTokenizer.from_pretrained('''facebook/opt-2.7b'''				)
							    if '''opt''' in model_name
							    else AutoTokenizer.from_pretrained('''google/flan-t5-xl'''				)
							)
							SCREAMING_SNAKE_CASE					:      Tuple     					=			tokenizer('''\n'''					,		add_special_tokens=A				).input_ids[0]
							SCREAMING_SNAKE_CASE					,SCREAMING_SNAKE_CASE					:      Tuple     					=			get_blipa_config(A					,		eos_token_id=A				)

							SCREAMING_SNAKE_CASE					:      Dict     					=			BlipaForConditionalGeneration(A				).eval()

							SCREAMING_SNAKE_CASE					:      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'''),
							}

							SCREAMING_SNAKE_CASE					,SCREAMING_SNAKE_CASE					:      Optional[int]     					=			model_name_to_original[model_name]

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

							# read in qv biases
							read_in_q_v_bias(A					,		A				)

							SCREAMING_SNAKE_CASE					,SCREAMING_SNAKE_CASE					:      Optional[int]     					=			hf_model.load_state_dict(A					,		strict=A				)
							assert len(A				) == 0
							assert unexpected_keys == ["qformer.embeddings.position_ids"]

							SCREAMING_SNAKE_CASE					:      str     					=			load_demo_image()
							SCREAMING_SNAKE_CASE					:      Optional[int]     					=			vis_processors['''eval'''](A				).unsqueeze(0				).to(A				)
							SCREAMING_SNAKE_CASE					:      int     					=			tokenizer(['''\n''']					,		return_tensors='''pt'''				).input_ids.to(A				)

							# create processor
							SCREAMING_SNAKE_CASE					:      Dict     					=			BlipImageProcessor(
							    size={'''height''': image_size, '''width''': image_size}					,		image_mean=A					,		image_std=A				)
							SCREAMING_SNAKE_CASE					:      int     					=			BlipaProcessor(image_processor=A					,		tokenizer=A				)
							SCREAMING_SNAKE_CASE					:      str     					=			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					:      Any     					=			original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}				).logits
																					SCREAMING_SNAKE_CASE					:      List[Any]     					=			hf_model(A					,		A				).logits
														else:
																					SCREAMING_SNAKE_CASE					:      Any     					=			original_model(
																					    {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}				).logits
																					SCREAMING_SNAKE_CASE					:      str     					=			input_ids.masked_fill(input_ids == tokenizer.pad_token_id					,		-100				)
																					SCREAMING_SNAKE_CASE					:      Tuple     					=			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					:      int     					=			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					:      Any     					=			torch.tensor(
														    [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]					,		device=A				)
							else:
														# cast to same type
														SCREAMING_SNAKE_CASE					:      Tuple     					=			logits.dtype
														assert torch.allclose(original_logits.to(A				)					,		A					,		atol=1e-2				)
							print('''Looks ok!'''				)

							print('''Generating a caption...'''				)
							SCREAMING_SNAKE_CASE					:      Any     					=			''''''
							SCREAMING_SNAKE_CASE					:      Any     					=			tokenizer(A					,		return_tensors='''pt'''				).input_ids.to(A				)

							SCREAMING_SNAKE_CASE					:      Dict     					=			original_model.generate({'''image''': original_pixel_values}				)
							SCREAMING_SNAKE_CASE					:      Optional[int]     					=			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					:      Optional[int]     					=			input_ids.shape[1]
							SCREAMING_SNAKE_CASE					:      Optional[int]     					=			processor.batch_decode(outputs[:, prompt_length:]					,		skip_special_tokens=A				)
							SCREAMING_SNAKE_CASE					:      Tuple     					=			[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__":
							lowerCAmelCase_     :       List[Any]								=						argparse.ArgumentParser()
							lowerCAmelCase_     :       Union[str, 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_     :       Optional[Any]								=						parser.parse_args()

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









'''simple docstring'''



from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase_     :       Optional[int]								=						logging.get_logger(__name__)

lowerCAmelCase_     :       Optional[Any]								=						{
    'microsoft/swinv2-tiny-patch4-window8-256': (
        'https://huggingface.co/microsoft/swinv2-tiny-patch4-window8-256/resolve/main/config.json'
    ),
}
class 	lowerCamelCase_							(					snake_case_			):
						_lowerCAmelCase     :   str	     =				'swinv2'

						_lowerCAmelCase     :   Optional[int]	     =				{
						    'num_attention_heads': 'num_heads',
						    'num_hidden_layers': 'num_layers',
						}







						def __init__(  self   :      Optional[int]				,    lowerCAmelCase__   :      Any=2_24				,    lowerCAmelCase__   :      Optional[Any]=4				,    lowerCAmelCase__   :      Tuple=3				,    lowerCAmelCase__   :      Tuple=96				,    lowerCAmelCase__   :      Optional[int]=[2, 2, 6, 2]				,    lowerCAmelCase__   :      Tuple=[3, 6, 12, 24]				,    lowerCAmelCase__   :      List[str]=7				,    lowerCAmelCase__   :      Tuple=4.0				,    lowerCAmelCase__   :      Dict=True				,    lowerCAmelCase__   :      Any=0.0				,    lowerCAmelCase__   :      Optional[int]=0.0				,    lowerCAmelCase__   :      Optional[Any]=0.1				,    lowerCAmelCase__   :      int="gelu"				,    lowerCAmelCase__   :      Tuple=False				,    lowerCAmelCase__   :      Tuple=0.02				,    lowerCAmelCase__   :      Dict=1e-5				,    lowerCAmelCase__   :      List[str]=32				,    **lowerCAmelCase__   :      Union[str, Any]				,    ):

													"""simple docstring"""
													super().__init__(**lowerCAmelCase__				)

													SCREAMING_SNAKE_CASE					:      int     					=			image_size
													SCREAMING_SNAKE_CASE					:      Optional[int]     					=			patch_size
													SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			num_channels
													SCREAMING_SNAKE_CASE					:      int     					=			embed_dim
													SCREAMING_SNAKE_CASE					:      List[Any]     					=			depths
													SCREAMING_SNAKE_CASE					:      str     					=			len(lowerCAmelCase__				)
													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			num_heads
													SCREAMING_SNAKE_CASE					:      Tuple     					=			window_size
													SCREAMING_SNAKE_CASE					:      Any     					=			mlp_ratio
													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			qkv_bias
													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			hidden_dropout_prob
													SCREAMING_SNAKE_CASE					:      int     					=			attention_probs_dropout_prob
													SCREAMING_SNAKE_CASE					:      List[str]     					=			drop_path_rate
													SCREAMING_SNAKE_CASE					:      Any     					=			hidden_act
													SCREAMING_SNAKE_CASE					:      List[Any]     					=			use_absolute_embeddings
													SCREAMING_SNAKE_CASE					:      List[str]     					=			layer_norm_eps
													SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			initializer_range
													SCREAMING_SNAKE_CASE					:      Dict     					=			encoder_stride
													# we set the hidden_size attribute in order to make Swinv2 work with VisionEncoderDecoderModel
													# this indicates the channel dimension after the last stage of the model
													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			int(embed_dim * 2 ** (len(lowerCAmelCase__				) - 1)				)
													SCREAMING_SNAKE_CASE					:      int     					=			(0, 0, 0, 0)




