



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

		'''simple docstring'''


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

		return image




def 	UpperCAmelCase     (		A	:						Any       ):

		'''simple docstring'''


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

		'''simple docstring'''


		_UpperCAmelCase							 =		dct.pop(A       )
		_UpperCAmelCase							 =		val




def 	UpperCAmelCase     (		A	:						Optional[Any]     ,				A	:						Optional[Any]       ):

		'''simple docstring'''


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

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




def 	UpperCAmelCase     (		A	:						Optional[Any]     ,				A	:						List[Any]       ):

		'''simple docstring'''


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

		_UpperCAmelCase							 =		BlipaConfig(vision_config=A     ,				text_config=A       )

		return config, image_size




@torch.no_grad()
def 	UpperCAmelCase     (		A	:						str     ,				A	:						List[Any]=None     ,				A	:						Any=False       ):

		'''simple docstring'''


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

		_UpperCAmelCase							 =		BlipaForConditionalGeneration(A       ).eval()

		_UpperCAmelCase							 =		{
		    '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							 =		model_name_to_original[model_name]

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

		# update state dict keys
		_UpperCAmelCase							 =		original_model.state_dict()
		_UpperCAmelCase							 =		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							 =		state_dict.pop(A       )
				if key.startswith('Qformer.bert'       ):
						_UpperCAmelCase							 =		key.replace('Qformer.bert'     ,				'qformer'       )
				if "attention.self" in key:
						_UpperCAmelCase							 =		key.replace('self'     ,				'attention'       )
				if "opt_proj" in key:
						_UpperCAmelCase							 =		key.replace('opt_proj'     ,				'language_projection'       )
				if "t5_proj" in key:
						_UpperCAmelCase							 =		key.replace('t5_proj'     ,				'language_projection'       )
				if key.startswith('opt'       ):
						_UpperCAmelCase							 =		key.replace('opt'     ,				'language'       )
				if key.startswith('t5'       ):
						_UpperCAmelCase							 =		key.replace('t5'     ,				'language'       )
				_UpperCAmelCase							 =		val

		# read in qv biases
		read_in_q_v_bias(A     ,				A       )

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

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

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

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

		_UpperCAmelCase							 =		original_model.generate({'image': original_pixel_values}       )
		_UpperCAmelCase							 =		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							 =		input_ids.shape[1]
		_UpperCAmelCase							 =		processor.batch_decode(outputs[:, prompt_length:]     ,				skip_special_tokens=A       )
		_UpperCAmelCase							 =		[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__":
					lowercase  							=					argparse.ArgumentParser()
					lowercase  							=					[
					    '''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''',
					)

					lowercase  							=					parser.parse_args()

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







"""simple docstring"""





def 	UpperCAmelCase     (		A	:						int       ):

		'''simple docstring'''


		_UpperCAmelCase							 =		abs(A       )
		_UpperCAmelCase							 =		0
		while n > 0:
				res += n % 10
				n //= 10
		return res




def 	UpperCAmelCase     (		A	:						int       ):

		'''simple docstring'''


		_UpperCAmelCase							 =		abs(A       )
		return n if n < 10 else n % 10 + sum_of_digits(n // 10       )




def 	UpperCAmelCase     (		A	:						int       ):

		'''simple docstring'''


		return sum(int(A       ) for c in str(abs(A       )       )       )




def 	UpperCAmelCase     (		):

		'''simple docstring'''


		from collections.abc import Callable
		from timeit import timeit

		def benchmark_a_function(A	:						Callable     ,				A	:						int       ) -> None:
				_UpperCAmelCase							 =		f'{func.__name__}({value})'
				_UpperCAmelCase							 =		timeit(f'__main__.{call}'     ,				setup='import __main__'       )
				print(f'{call:56} = {func(A       )} -- {timing:.4f} seconds'       )

		for value in (26_2144, 1125_8999_0684_2624, 126_7650_6002_2822_9401_4967_0320_5376):
				for func in (sum_of_digits, sum_of_digits_recursion, sum_of_digits_compact):
						benchmark_a_function(A     ,				A       )
				print()


if __name__ == "__main__":
					import doctest

					doctest.testmod()
					benchmark()


