



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

			return image






def 					_UpperCAmelCase       (		__lowerCamelCase    :   Any     ):
			"""simple docstring"""
			lowerCamelCase__		:		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 					_UpperCAmelCase       (		__lowerCamelCase    :   Optional[int]							,		__lowerCamelCase    :   int							,		__lowerCamelCase    :   int     ):
			"""simple docstring"""
			lowerCamelCase__		:		str		     =dct.pop(snake_case_     )
			lowerCamelCase__		:		Any		     =val






def 					_UpperCAmelCase       (		__lowerCamelCase    :   Optional[Any]							,		__lowerCamelCase    :   Optional[Any]     ):
			"""simple docstring"""




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

						# next, set bias in the state dict
						lowerCamelCase__		:		Optional[int]		     =torch.cat((q_bias, torch.zeros_like(snake_case_							,		requires_grad=snake_case_     ), v_bias)     )
						lowerCamelCase__		:		Optional[Any]		     =qkv_bias






def 					_UpperCAmelCase       (		__lowerCamelCase    :   List[str]							,		__lowerCamelCase    :   Optional[int]     ):
			"""simple docstring"""
			lowerCamelCase__		:		int		     =364 if '''coco''' in model_name else 224
			lowerCamelCase__		:		List[Any]		     =BlipaVisionConfig(image_size=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:
						lowerCamelCase__		:		List[Any]		     =OPTConfig.from_pretrained('''facebook/opt-2.7b'''							,		eos_token_id=snake_case_     ).to_dict()
			elif "opt-6.7b" in model_name:
						lowerCamelCase__		:		Dict		     =OPTConfig.from_pretrained('''facebook/opt-6.7b'''							,		eos_token_id=snake_case_     ).to_dict()
			elif "t5-xl" in model_name:
						lowerCamelCase__		:		Tuple		     =TaConfig.from_pretrained('''google/flan-t5-xl'''							,		dense_act_fn='''gelu'''							,		bos_token_id=1     ).to_dict()
			elif "t5-xxl" in model_name:
						lowerCamelCase__		:		int		     =TaConfig.from_pretrained('''google/flan-t5-xxl'''							,		dense_act_fn='''gelu'''							,		bos_token_id=1     ).to_dict()

			lowerCamelCase__		:		Optional[Any]		     =BlipaConfig(vision_config=snake_case_							,		text_config=snake_case_     )

			return config, image_size






@torch.no_grad()
def 					_UpperCAmelCase       (		__lowerCamelCase    :   int							,		__lowerCamelCase    :   Optional[int]=None							,		__lowerCamelCase    :   Optional[Any]=False     ):
			"""simple docstring"""
			lowerCamelCase__		:		int		     =(
			    AutoTokenizer.from_pretrained('''facebook/opt-2.7b'''     )
			    if '''opt''' in model_name
			    else AutoTokenizer.from_pretrained('''google/flan-t5-xl'''     )
			)
			lowerCamelCase__		:		str		     =tokenizer('''\n'''							,		add_special_tokens=snake_case_     ).input_ids[0]
			lowerCamelCase__     ,       lowerCamelCase__		:		Dict		     =get_blipa_config(snake_case_							,		eos_token_id=snake_case_     )

			lowerCamelCase__		:		Dict		     =BlipaForConditionalGeneration(snake_case_     ).eval()

			lowerCamelCase__		:		str		     ={
			    '''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'''),
			}

			lowerCamelCase__     ,       lowerCamelCase__		:		Tuple		     =model_name_to_original[model_name]

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

			# update state dict keys
			lowerCamelCase__		:		Optional[Any]		     =original_model.state_dict()
			lowerCamelCase__		:		Tuple		     =create_rename_keys(snake_case_     )
			for src, dest in rename_keys:
						rename_key(snake_case_							,		snake_case_							,		snake_case_     )

			# some keys can be renamed efficiently
			for key, val in state_dict.copy().items():
						lowerCamelCase__		:		Optional[int]		     =state_dict.pop(snake_case_     )
						if key.startswith('''Qformer.bert'''     ):
									lowerCamelCase__		:		Optional[int]		     =key.replace('''Qformer.bert'''							,		'''qformer'''     )
						if "attention.self" in key:
									lowerCamelCase__		:		Any		     =key.replace('''self'''							,		'''attention'''     )
						if "opt_proj" in key:
									lowerCamelCase__		:		int		     =key.replace('''opt_proj'''							,		'''language_projection'''     )
						if "t5_proj" in key:
									lowerCamelCase__		:		List[str]		     =key.replace('''t5_proj'''							,		'''language_projection'''     )
						if key.startswith('''opt'''     ):
									lowerCamelCase__		:		Dict		     =key.replace('''opt'''							,		'''language'''     )
						if key.startswith('''t5'''     ):
									lowerCamelCase__		:		List[str]		     =key.replace('''t5'''							,		'''language'''     )
						lowerCamelCase__		:		int		     =val

			# read in qv biases
			read_in_q_v_bias(snake_case_							,		snake_case_     )

			lowerCamelCase__     ,       lowerCamelCase__		:		Union[str, Any]		     =hf_model.load_state_dict(snake_case_							,		strict=snake_case_     )
			assert len(snake_case_     ) == 0
			assert unexpected_keys == ["qformer.embeddings.position_ids"]

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

			# create processor
			lowerCamelCase__		:		Union[str, Any]		     =BlipImageProcessor(
			    size={'''height''': image_size, '''width''': image_size}							,		image_mean=snake_case_							,		image_std=snake_case_     )
			lowerCamelCase__		:		str		     =BlipaProcessor(image_processor=snake_case_							,		tokenizer=snake_case_     )
			lowerCamelCase__		:		Union[str, Any]		     =processor(images=snake_case_							,		return_tensors='''pt'''     ).pixel_values.to(snake_case_     )

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

			original_model.to(snake_case_     )
			hf_model.to(snake_case_     )
			with torch.no_grad():
						if "opt" in model_name:
									lowerCamelCase__		:		str		     =original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}     ).logits
									lowerCamelCase__		:		Any		     =hf_model(snake_case_							,		snake_case_     ).logits
						else:
									lowerCamelCase__		:		Union[str, Any]		     =original_model(
									    {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}     ).logits
									lowerCamelCase__		:		Union[str, Any]		     =input_ids.masked_fill(input_ids == tokenizer.pad_token_id							,		-100     )
									lowerCamelCase__		:		str		     =hf_model(snake_case_							,		snake_case_							,		labels=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":
						lowerCamelCase__		:		str		     =torch.tensor(
						    [[-41.58_50, -4.44_40, -8.99_22], [-47.43_22, -5.91_43, -1.73_40]]							,		device=snake_case_     )
						assert torch.allclose(logits[0, :3, :3]							,		snake_case_							,		atol=1e-4     )
			elif model_name == "blip2-flan-t5-xl-coco":
						lowerCamelCase__		:		Union[str, Any]		     =torch.tensor(
						    [[-57.01_09, -9.89_67, -12.62_80], [-68.65_78, -12.71_91, -10.50_65]]							,		device=snake_case_     )
			else:
						# cast to same type
						lowerCamelCase__		:		str		     =logits.dtype
						assert torch.allclose(original_logits.to(snake_case_     )							,		snake_case_							,		atol=1e-2     )
			print('''Looks ok!'''     )

			print('''Generating a caption...'''     )
			lowerCamelCase__		:		str		     =''''''
			lowerCamelCase__		:		Tuple		     =tokenizer(snake_case_							,		return_tensors='''pt'''     ).input_ids.to(snake_case_     )

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

			if pytorch_dump_folder_path is not None:
						processor.save_pretrained(snake_case_     )
						hf_model.save_pretrained(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__":
		_lowercase			:			Any										=						argparse.ArgumentParser()
		_lowercase			:			Optional[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",
		)

		_lowercase			:			List[str]										=						parser.parse_args()

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









"""simple docstring"""






import math

import flax.linen as nn
import jax.numpy as jnp






def 					snake_case__       (		__lowerCamelCase    :   jnp.ndarray							,		__lowerCamelCase    :   int							,		__lowerCamelCase    :   float = 1							,		__lowerCamelCase    :   float = 1							,		__lowerCamelCase    :   float = 1.0e4							,		__lowerCamelCase    :   bool = False							,		__lowerCamelCase    :   float = 1.0							,		):
			"""simple docstring"""
			assert timesteps.ndim == 1, "Timesteps should be a 1d-array"
			assert embedding_dim % 2 == 0, f'''Embedding dimension {embedding_dim} should be even'''
			lowerCamelCase__		:		Any		     =float(embedding_dim // 2     )
			lowerCamelCase__		:		List[str]		     =math.log(max_timescale / min_timescale     ) / (num_timescales - freq_shift)
			lowerCamelCase__		:		int		     =min_timescale * jnp.exp(jnp.arange(__lowerCamelCase							,		dtype=jnp.floataa     ) * -log_timescale_increment     )
			lowerCamelCase__		:		Tuple		     =jnp.expand_dims(__lowerCamelCase							,		1     ) * jnp.expand_dims(__lowerCamelCase							,		0     )

			# scale embeddings
			lowerCamelCase__		:		List[str]		     =scale * emb

			if flip_sin_to_cos:
						lowerCamelCase__		:		int		     =jnp.concatenate([jnp.cos(__lowerCamelCase     ), jnp.sin(__lowerCamelCase     )]							,		axis=1     )
			else:
						lowerCamelCase__		:		List[str]		     =jnp.concatenate([jnp.sin(__lowerCamelCase     ), jnp.cos(__lowerCamelCase     )]							,		axis=1     )
			lowerCamelCase__		:		str		     =jnp.reshape(__lowerCamelCase							,		[jnp.shape(__lowerCamelCase     )[0], embedding_dim]     )
			return signal

class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					_a             =							3_2
					_a             =							jnp.floataa



					@nn.compact
					def __call__(	self							:  Optional[Any],  lowerCamelCase							:  int					)-> Any:
								lowerCamelCase__		:		Optional[Any]		     =nn.Dense(self.time_embed_dim,  dtype=self.dtype,  name='''linear_1'''					)(lowerCamelCase					)
								lowerCamelCase__		:		List[str]		     =nn.silu(lowerCamelCase					)
								lowerCamelCase__		:		Any		     =nn.Dense(self.time_embed_dim,  dtype=self.dtype,  name='''linear_2'''					)(lowerCamelCase					)
								return temb





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					_a             =							3_2
					_a             =							False
					_a             =							1



					@nn.compact
					def __call__(	self							:  Any,  lowerCamelCase							:  int					)-> int:
								return get_sinusoidal_embeddings(
								    lowerCamelCase,  embedding_dim=self.dim,  flip_sin_to_cos=self.flip_sin_to_cos,  freq_shift=self.freq_shift					)




