


'''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      a (     )    ->      Dict:



					'''simple docstring'''






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

					return image


def      a (     __a  )    ->      str:



					'''simple docstring'''






					UpperCamelCase__  :Tuple				   =					[]
					# 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      a (     __a					,   __a					,   __a  )    ->      int:



					'''simple docstring'''






					UpperCamelCase__  :Tuple				   =					dct.pop(__a  )
					UpperCamelCase__  :str				   =					val


def      a (     __a					,   __a  )    ->      Optional[Any]:



					'''simple docstring'''






					for i in range(config.vision_config.num_hidden_layers  ):
										# read in original q and v biases
										UpperCamelCase__  :int				   =					state_dict.pop(f'''visual_encoder.blocks.{i}.attn.q_bias'''  )
										UpperCamelCase__  :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(__a					,   requires_grad=__a  ), v_bias)  )
										UpperCamelCase__  :Dict				   =					qkv_bias


def      a (     __a					,   __a  )    ->      Union[str, Any]:



					'''simple docstring'''






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

					UpperCamelCase__  :List[str]				   =					BlipaConfig(vision_config=__a					,   text_config=__a  )

					return config, image_size


@torch.no_grad()
def      a (     __a					,   __a=None					,   __a=False  )    ->      Tuple:



					'''simple docstring'''






					UpperCamelCase__  :Optional[Any]				   =					(
					    AutoTokenizer.from_pretrained('''facebook/opt-2.7b'''  )
					    if '''opt''' in model_name
					    else AutoTokenizer.from_pretrained('''google/flan-t5-xl'''  )
					)
					UpperCamelCase__  :int				   =					tokenizer('''\n'''					,   add_special_tokens=__a  ).input_ids[0]
					UpperCamelCase__   ,							UpperCamelCase__  :Tuple				   =					get_blipa_config(__a					,   eos_token_id=__a  )

					UpperCamelCase__  :Dict				   =					BlipaForConditionalGeneration(__a  ).eval()

					UpperCamelCase__  :Union[str, Any]				   =					{
					    '''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__  :List[Any]				   =					model_name_to_original[model_name]

					# load original model
					print('''Loading original model...'''  )
					UpperCamelCase__  :Any				   =					'''cuda''' if torch.cuda.is_available() else '''cpu'''
					UpperCamelCase__   ,							UpperCamelCase__   ,							UpperCamelCase__  :Union[str, Any]				   =					load_model_and_preprocess(
					    name=__a					,   model_type=__a					,   is_eval=__a					,   device=__a  )
					original_model.eval()
					print('''Done!'''  )

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

					# read in qv biases
					read_in_q_v_bias(__a					,   __a  )

					UpperCamelCase__   ,							UpperCamelCase__  :str				   =					hf_model.load_state_dict(__a					,   strict=__a  )
					assert len(__a  ) == 0
					assert unexpected_keys == ["qformer.embeddings.position_ids"]

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

					# create processor
					UpperCamelCase__  :Tuple				   =					BlipImageProcessor(
					    size={'''height''': image_size, '''width''': image_size}					,   image_mean=__a					,   image_std=__a  )
					UpperCamelCase__  :Optional[Any]				   =					BlipaProcessor(image_processor=__a					,   tokenizer=__a  )
					UpperCamelCase__  :int				   =					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:
															UpperCamelCase__  :Dict				   =					original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}  ).logits
															UpperCamelCase__  :Union[str, Any]				   =					hf_model(__a					,   __a  ).logits
										else:
															UpperCamelCase__  :Tuple				   =					original_model(
															    {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}  ).logits
															UpperCamelCase__  :Any				   =					input_ids.masked_fill(input_ids == tokenizer.pad_token_id					,   -100  )
															UpperCamelCase__  :Optional[int]				   =					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":
										UpperCamelCase__  :Dict				   =					torch.tensor(
										    [[-4_1.5_8_5_0, -4.4_4_4_0, -8.9_9_2_2], [-4_7.4_3_2_2, -5.9_1_4_3, -1.7_3_4_0]]					,   device=__a  )
										assert torch.allclose(logits[0, :3, :3]					,   __a					,   atol=1e-4  )
					elif model_name == "blip2-flan-t5-xl-coco":
										UpperCamelCase__  :Union[str, Any]				   =					torch.tensor(
										    [[-5_7.0_1_0_9, -9.8_9_6_7, -1_2.6_2_8_0], [-6_8.6_5_7_8, -1_2.7_1_9_1, -1_0.5_0_6_5]]					,   device=__a  )
					else:
										# cast to same type
										UpperCamelCase__  :int				   =					logits.dtype
										assert torch.allclose(original_logits.to(__a  )					,   __a					,   atol=1e-2  )
					print('''Looks ok!'''  )

					print('''Generating a caption...'''  )
					UpperCamelCase__  :Optional[Any]				   =					''''''
					UpperCamelCase__  :List[Any]				   =					tokenizer(__a					,   return_tensors='''pt'''  ).input_ids.to(__a  )

					UpperCamelCase__  :List[Any]				   =					original_model.generate({'''image''': original_pixel_values}  )
					UpperCamelCase__  :List[Any]				   =					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  )
					UpperCamelCase__  :int				   =					input_ids.shape[1]
					UpperCamelCase__  :List[Any]				   =					processor.batch_decode(outputs[:, prompt_length:]					,   skip_special_tokens=__a  )
					UpperCamelCase__  :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__":
		__snake_case				 =	argparse.ArgumentParser()
		__snake_case				 =	[
		    '''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				 =	parser.parse_args()

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



'''simple docstring'''
import inspect
import os
import unittest

import torch

import accelerate
from accelerate import debug_launcher
from accelerate.test_utils import (
    execute_subprocess_async,
    require_cpu,
    require_huggingface_suite,
    require_multi_gpu,
    require_single_gpu,
)
from accelerate.utils import patch_environment
@require_huggingface_suite
class   lowercase      (				unittest.TestCase   ):

		"""simple docstring"""

		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							UpperCamelCase__  :Any				   =					inspect.getfile(accelerate.test_utils				)
							UpperCamelCase__  :Any				   =					os.path.sep.join(
							    mod_file.split(os.path.sep				)[:-1] + ['''scripts''', '''external_deps''', '''test_metrics.py''']				)

							from accelerate.test_utils.scripts.external_deps import test_metrics  # noqa: F401

							UpperCamelCase__  :List[str]				   =					test_metrics




		@require_cpu
		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							debug_launcher(self.test_metrics.main       ,  num_processes=1				)




		@require_cpu
		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							debug_launcher(self.test_metrics.main				)




		@require_single_gpu
		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							self.test_metrics.main()




		@require_multi_gpu
		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							print(F'''Found {torch.cuda.device_count()} devices.'''				)
							UpperCamelCase__  :Optional[Any]				   =					['''torchrun''', F'''--nproc_per_node={torch.cuda.device_count()}''', self.test_file_path]
							with patch_environment(omp_num_threads=1				):
												execute_subprocess_async(UpperCamelCase_       ,  env=os.environ.copy()				)