


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


			'''simple docstring'''







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

			return image


def 		lowercase_	(  _UpperCamelCase	):


			'''simple docstring'''







			__lowercase			       = []
			# 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 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			__lowercase			       = dct.pop(_UpperCamelCase	)
			__lowercase			       = val


def 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







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

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


def 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			__lowercase			       = 3_64 if '''coco''' in model_name else 2_24
			__lowercase			       = BlipaVisionConfig(image_size=_UpperCamelCase	).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:
						__lowercase			       = OPTConfig.from_pretrained('''facebook/opt-2.7b'''							,      eos_token_id=_UpperCamelCase	).to_dict()
			elif "opt-6.7b" in model_name:
						__lowercase			       = OPTConfig.from_pretrained('''facebook/opt-6.7b'''							,      eos_token_id=_UpperCamelCase	).to_dict()
			elif "t5-xl" in model_name:
						__lowercase			       = TaConfig.from_pretrained('''google/flan-t5-xl'''							,      dense_act_fn='''gelu'''							,      bos_token_id=1	).to_dict()
			elif "t5-xxl" in model_name:
						__lowercase			       = TaConfig.from_pretrained('''google/flan-t5-xxl'''							,      dense_act_fn='''gelu'''							,      bos_token_id=1	).to_dict()

			__lowercase			       = BlipaConfig(vision_config=_UpperCamelCase							,      text_config=_UpperCamelCase	)

			return config, image_size


@torch.no_grad()
def 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase=None							,      _UpperCamelCase=False	):


			'''simple docstring'''







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

			__lowercase			       = BlipaForConditionalGeneration(_UpperCamelCase	).eval()

			__lowercase			       = {
			    '''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'''),
			}

			__lowercase       ,			__lowercase			       = model_name_to_original[model_name]

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

			# update state dict keys
			__lowercase			       = original_model.state_dict()
			__lowercase			       = create_rename_keys(_UpperCamelCase	)
			for src, dest in rename_keys:
						rename_key(_UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	)

			# some keys can be renamed efficiently
			for key, val in state_dict.copy().items():
						__lowercase			       = state_dict.pop(_UpperCamelCase	)
						if key.startswith('''Qformer.bert'''	):
									__lowercase			       = key.replace('''Qformer.bert'''							,      '''qformer'''	)
						if "attention.self" in key:
									__lowercase			       = key.replace('''self'''							,      '''attention'''	)
						if "opt_proj" in key:
									__lowercase			       = key.replace('''opt_proj'''							,      '''language_projection'''	)
						if "t5_proj" in key:
									__lowercase			       = key.replace('''t5_proj'''							,      '''language_projection'''	)
						if key.startswith('''opt'''	):
									__lowercase			       = key.replace('''opt'''							,      '''language'''	)
						if key.startswith('''t5'''	):
									__lowercase			       = key.replace('''t5'''							,      '''language'''	)
						__lowercase			       = val

			# read in qv biases
			read_in_q_v_bias(_UpperCamelCase							,      _UpperCamelCase	)

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

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

			# create processor
			__lowercase			       = BlipImageProcessor(
			    size={'''height''': image_size, '''width''': image_size}							,      image_mean=_UpperCamelCase							,      image_std=_UpperCamelCase	)
			__lowercase			       = BlipaProcessor(image_processor=_UpperCamelCase							,      tokenizer=_UpperCamelCase	)
			__lowercase			       = processor(images=_UpperCamelCase							,      return_tensors='''pt'''	).pixel_values.to(_UpperCamelCase	)

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

			original_model.to(_UpperCamelCase	)
			hf_model.to(_UpperCamelCase	)
			with torch.no_grad():
						if "opt" in model_name:
									__lowercase			       = original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}	).logits
									__lowercase			       = hf_model(_UpperCamelCase							,      _UpperCamelCase	).logits
						else:
									__lowercase			       = original_model(
									    {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}	).logits
									__lowercase			       = input_ids.masked_fill(input_ids == tokenizer.pad_token_id							,      -1_00	)
									__lowercase			       = hf_model(_UpperCamelCase							,      _UpperCamelCase							,      labels=_UpperCamelCase	).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":
						__lowercase			       = torch.tensor(
						    [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]							,      device=_UpperCamelCase	)
						assert torch.allclose(logits[0, :3, :3]							,      _UpperCamelCase							,      atol=1E-4	)
			elif model_name == "blip2-flan-t5-xl-coco":
						__lowercase			       = torch.tensor(
						    [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]							,      device=_UpperCamelCase	)
			else:
						# cast to same type
						__lowercase			       = logits.dtype
						assert torch.allclose(original_logits.to(_UpperCamelCase	)							,      _UpperCamelCase							,      atol=1E-2	)
			print('''Looks ok!'''	)

			print('''Generating a caption...'''	)
			__lowercase			       = ''''''
			__lowercase			       = tokenizer(_UpperCamelCase							,      return_tensors='''pt'''	).input_ids.to(_UpperCamelCase	)

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

			if pytorch_dump_folder_path is not None:
						processor.save_pretrained(_UpperCamelCase	)
						hf_model.save_pretrained(_UpperCamelCase	)

			if push_to_hub:
						processor.push_to_hub(F'nielsr/{model_name}'	)
						hf_model.push_to_hub(F'nielsr/{model_name}'	)


if __name__ == "__main__":
	a			:			int												=					argparse.ArgumentParser()
	a			:			Optional[int]												=					[
	    '''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''',
	)

	a			:			str												=					parser.parse_args()

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






import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Audio, Features, Value
from .base import TaskTemplate
@dataclass(frozen=lowerCAmelCase__					)
class 						lowerCamelCase_							(					lowerCAmelCase__					):



					'''simple docstring'''







					__UpperCAmelCase							  =					field(default="automatic-speech-recognition"  ,  metadata={"include_in_asdict_even_if_is_default": True}					)
					__UpperCAmelCase							  =					Features({"audio": Audio()}					)
					__UpperCAmelCase							  =					Features({"transcription": Value("string"					)}					)
					__UpperCAmelCase							  =					"audio"
					__UpperCAmelCase							  =					"transcription"






					def A   (					self			,							snake_case_    ) ->      Any:







								'''simple docstring'''

								if self.audio_column not in features:
											raise ValueError(F'Column {self.audio_column} is not present in features.'    )
								if not isinstance(features[self.audio_column]			,							snake_case_    ):
											raise ValueError(F'Column {self.audio_column} is not an Audio type.'    )
								__lowercase			       = copy.deepcopy(self    )
								__lowercase			       = self.input_schema.copy()
								__lowercase			       = features[self.audio_column]
								__lowercase			       = input_schema
								return task_template






					@property
					def A   (					self    ) ->      Dict[str, str]:







								'''simple docstring'''

								return {self.audio_column: "audio", self.transcription_column: "transcription"}


