



import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
	from transformers import WhisperFeatureExtractor

if is_torch_available():
	import torch

a      :  str			     =			random.Random()







def  lowerCAmelCase_	(lowerCAmelCase__:		Union[str, Any]		,						lowerCAmelCase__:		Optional[Any]=1.0		,						lowerCAmelCase__:		int=None		,						lowerCAmelCase__:		Optional[int]=None  ):





			"""simple docstring"""


			if rng is None:
						UpperCAmelCase_:		Optional[Any]							      =     global_rng

			UpperCAmelCase_:		Optional[int]							      =     []
			for batch_idx in range(shape[0]  ):
						values.append([]  )
						for _ in range(shape[1]  ):
									values[-1].append(rng.random() * scale  )

			return values







@require_torch
@require_torchaudio
class 				_a				(   unittest.TestCase					):

			def __init__(self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_=7,   SCREAMING_SNAKE_CASE_=400,   SCREAMING_SNAKE_CASE_=2000,   SCREAMING_SNAKE_CASE_=10,   SCREAMING_SNAKE_CASE_=160,   SCREAMING_SNAKE_CASE_=8,   SCREAMING_SNAKE_CASE_=0.0,   SCREAMING_SNAKE_CASE_=4000,   SCREAMING_SNAKE_CASE_=False,   SCREAMING_SNAKE_CASE_=True,   )       ->   Optional[Any]:
						UpperCAmelCase_:		str							      =     parent
						UpperCAmelCase_:		str							      =     batch_size
						UpperCAmelCase_:		List[str]							      =     min_seq_length
						UpperCAmelCase_:		Union[str, Any]							      =     max_seq_length
						UpperCAmelCase_:		Union[str, Any]							      =     (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
						UpperCAmelCase_:		List[Any]							      =     padding_value
						UpperCAmelCase_:		str							      =     sampling_rate
						UpperCAmelCase_:		Any							      =     return_attention_mask
						UpperCAmelCase_:		List[str]							      =     do_normalize
						UpperCAmelCase_:		Optional[int]							      =     feature_size
						UpperCAmelCase_:		int							      =     chunk_length
						UpperCAmelCase_:		int							      =     hop_length

			def      __snake_case       (self					)       ->   Tuple:
						return {
						    "feature_size": self.feature_size,
						    "hop_length": self.hop_length,
						    "chunk_length": self.chunk_length,
						    "padding_value": self.padding_value,
						    "sampling_rate": self.sampling_rate,
						    "return_attention_mask": self.return_attention_mask,
						    "do_normalize": self.do_normalize,
						}

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_=False,   SCREAMING_SNAKE_CASE_=False					)       ->   Dict:
						def _flatten(SCREAMING_SNAKE_CASE_					):
									return list(itertools.chain(*SCREAMING_SNAKE_CASE_					)					)

						if equal_length:
									UpperCAmelCase_:		Any							      =     [floats_list((self.max_seq_length, self.feature_size)					) for _ in range(self.batch_size					)]
						else:
									# make sure that inputs increase in size
									UpperCAmelCase_:		List[str]							      =     [
									    floats_list((x, self.feature_size)					)
									    for x in range(self.min_seq_length,   self.max_seq_length,   self.seq_length_diff					)
									]
						if numpify:
									UpperCAmelCase_:		Optional[int]							      =     [np.asarray(SCREAMING_SNAKE_CASE_					) for x in speech_inputs]
						return speech_inputs







@require_torch
@require_torchaudio
class 				_a				(   _lowerCAmelCase ,		unittest.TestCase					):
			A      		=	WhisperFeatureExtractor if is_speech_available() else None

			def      __snake_case       (self					)       ->   int:
						UpperCAmelCase_:		Tuple							      =     WhisperFeatureExtractionTester(self					)

			def      __snake_case       (self					)       ->   List[Any]:
						UpperCAmelCase_:		str							      =     self.feature_extraction_class(**self.feat_extract_dict					)

						with tempfile.TemporaryDirectory() as tmpdirname:
									UpperCAmelCase_:		Any							      =     feat_extract_first.save_pretrained(SCREAMING_SNAKE_CASE_					)[0]
									check_json_file_has_correct_format(SCREAMING_SNAKE_CASE_					)
									UpperCAmelCase_:		Optional[Any]							      =     self.feature_extraction_class.from_pretrained(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		List[str]							      =     feat_extract_first.to_dict()
						UpperCAmelCase_:		List[str]							      =     feat_extract_second.to_dict()
						UpperCAmelCase_:		List[Any]							      =     feat_extract_first.mel_filters
						UpperCAmelCase_:		Tuple							      =     feat_extract_second.mel_filters
						self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)					)
						self.assertEqual(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Union[str, Any]:
						UpperCAmelCase_:		List[Any]							      =     self.feature_extraction_class(**self.feat_extract_dict					)

						with tempfile.TemporaryDirectory() as tmpdirname:
									UpperCAmelCase_:		int							      =     os.path.join(SCREAMING_SNAKE_CASE_,   """feat_extract.json"""					)
									feat_extract_first.to_json_file(SCREAMING_SNAKE_CASE_					)
									UpperCAmelCase_:		Dict							      =     self.feature_extraction_class.from_json_file(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Dict							      =     feat_extract_first.to_dict()
						UpperCAmelCase_:		Any							      =     feat_extract_second.to_dict()
						UpperCAmelCase_:		Any							      =     feat_extract_first.mel_filters
						UpperCAmelCase_:		List[str]							      =     feat_extract_second.mel_filters
						self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)					)
						self.assertEqual(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Dict:
						# Tests that all call wrap to encode_plus and batch_encode_plus
						UpperCAmelCase_:		Optional[int]							      =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()					)
						# create three inputs of length 800, 1000, and 1200
						UpperCAmelCase_:		str							      =     [floats_list((1, x)					)[0] for x in range(800,   1400,   200					)]
						UpperCAmelCase_:		Any							      =     [np.asarray(SCREAMING_SNAKE_CASE_					) for speech_input in speech_inputs]

						# Test feature size
						UpperCAmelCase_:		Tuple							      =     feature_extractor(SCREAMING_SNAKE_CASE_,   padding="""max_length""",   return_tensors="""np"""					).input_features
						self.assertTrue(input_features.ndim == 3					)
						self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames					)
						self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size					)

						# Test not batched input
						UpperCAmelCase_:		Tuple							      =     feature_extractor(speech_inputs[0],   return_tensors="""np"""					).input_features
						UpperCAmelCase_:		Tuple							      =     feature_extractor(np_speech_inputs[0],   return_tensors="""np"""					).input_features
						self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   atol=1E-3					)					)

						# Test batched
						UpperCAmelCase_:		Optional[int]							      =     feature_extractor(SCREAMING_SNAKE_CASE_,   return_tensors="""np"""					).input_features
						UpperCAmelCase_:		Any							      =     feature_extractor(SCREAMING_SNAKE_CASE_,   return_tensors="""np"""					).input_features
						for enc_seq_a, enc_seq_a in zip(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					):
									self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   atol=1E-3					)					)

						# Test 2-D numpy arrays are batched.
						UpperCAmelCase_:		int							      =     [floats_list((1, x)					)[0] for x in (800, 800, 800)]
						UpperCAmelCase_:		str							      =     np.asarray(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Tuple							      =     feature_extractor(SCREAMING_SNAKE_CASE_,   return_tensors="""np"""					).input_features
						UpperCAmelCase_:		List[str]							      =     feature_extractor(SCREAMING_SNAKE_CASE_,   return_tensors="""np"""					).input_features
						for enc_seq_a, enc_seq_a in zip(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					):
									self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   atol=1E-3					)					)

						# Test truncation required
						UpperCAmelCase_:		List[str]							      =     [floats_list((1, x)					)[0] for x in range(200,   (feature_extractor.n_samples + 500),   200					)]
						UpperCAmelCase_:		Tuple							      =     [np.asarray(SCREAMING_SNAKE_CASE_					) for speech_input in speech_inputs]

						UpperCAmelCase_:		List[Any]							      =     [x[: feature_extractor.n_samples] for x in speech_inputs]
						UpperCAmelCase_:		Optional[Any]							      =     [np.asarray(SCREAMING_SNAKE_CASE_					) for speech_input in speech_inputs_truncated]

						UpperCAmelCase_:		str							      =     feature_extractor(SCREAMING_SNAKE_CASE_,   return_tensors="""np"""					).input_features
						UpperCAmelCase_:		Optional[int]							      =     feature_extractor(SCREAMING_SNAKE_CASE_,   return_tensors="""np"""					).input_features
						for enc_seq_a, enc_seq_a in zip(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					):
									self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   atol=1E-3					)					)

			def      __snake_case       (self					)       ->   Tuple:
						import torch

						UpperCAmelCase_:		Tuple							      =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()					)
						UpperCAmelCase_:		List[str]							      =     np.random.rand(100,   32					).astype(np.floataa					)
						UpperCAmelCase_:		Tuple							      =     np_speech_inputs.tolist()

						for inputs in [py_speech_inputs, np_speech_inputs]:
									UpperCAmelCase_:		Tuple							      =     feature_extractor.pad([{"""input_features""": inputs}],   return_tensors="""np"""					)
									self.assertTrue(np_processed.input_features.dtype == np.floataa					)
									UpperCAmelCase_:		str							      =     feature_extractor.pad([{"""input_features""": inputs}],   return_tensors="""pt"""					)
									self.assertTrue(pt_processed.input_features.dtype == torch.floataa					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_					)       ->   str:
						UpperCAmelCase_:		str							      =     load_dataset("""hf-internal-testing/librispeech_asr_dummy""",   """clean""",   split="""validation"""					)
						# automatic decoding with librispeech
						UpperCAmelCase_:		Tuple							      =     ds.sort("""id"""					).select(range(SCREAMING_SNAKE_CASE_					)					)[:num_samples]["""audio"""]

						return [x["array"] for x in speech_samples]

			def      __snake_case       (self					)       ->   str:
						# fmt: off
						UpperCAmelCase_:		int							      =     torch.tensor(
						    [
						        0.1_1_9_3, -0.0_9_4_6, -0.1_0_9_8, -0.0_1_9_6, 0.0_2_2_5, -0.0_6_9_0, -0.1_7_3_6, 0.0_9_5_1,
						        0.0_9_7_1, -0.0_8_1_7, -0.0_7_0_2, 0.0_1_6_2, 0.0_2_6_0, 0.0_0_1_7, -0.0_1_9_2, -0.1_6_7_8,
						        0.0_7_0_9, -0.1_8_6_7, -0.0_6_5_5, -0.0_2_7_4, -0.0_2_3_4, -0.1_8_8_4, -0.0_5_1_6, -0.0_5_5_4,
						        -0.0_2_7_4, -0.1_4_2_5, -0.1_4_2_3, 0.0_8_3_7, 0.0_3_7_7, -0.0_8_5_4
						    ]					)
						# fmt: on

						UpperCAmelCase_:		Optional[int]							      =     self._load_datasamples(1					)
						UpperCAmelCase_:		Optional[int]							      =     WhisperFeatureExtractor()
						UpperCAmelCase_:		Any							      =     feature_extractor(SCREAMING_SNAKE_CASE_,   return_tensors="""pt"""					).input_features
						self.assertEqual(input_features.shape,   (1, 80, 3000)					)
						self.assertTrue(torch.allclose(input_features[0, 0, :30],   SCREAMING_SNAKE_CASE_,   atol=1E-4					)					)

			def      __snake_case       (self					)       ->   str:
						UpperCAmelCase_:		Optional[Any]							      =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()					)
						UpperCAmelCase_:		Optional[int]							      =     self._load_datasamples(1					)[0]
						UpperCAmelCase_:		Union[str, Any]							      =     ((audio - audio.min()) / (audio.max() - audio.min())) * 65535  # Rescale to [0, 65535] to show issue
						UpperCAmelCase_:		int							      =     feat_extract.zero_mean_unit_var_norm([audio],   attention_mask=SCREAMING_SNAKE_CASE_					)[0]

						self.assertTrue(np.all(np.mean(SCREAMING_SNAKE_CASE_					) < 1E-3					)					)
						self.assertTrue(np.all(np.abs(np.var(SCREAMING_SNAKE_CASE_					) - 1					) < 1E-3					)					)











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





			"""simple docstring"""


			UpperCAmelCase_:		List[Any]							      =     """https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"""
			UpperCAmelCase_:		List[Any]							      =     Image.open(requests.get(lowerCAmelCase__		,						stream=lowerCAmelCase__  ).raw  ).convert("""RGB"""  )

			return image







def  lowerCAmelCase_	(lowerCAmelCase__:		Any  ):





			"""simple docstring"""


			UpperCAmelCase_:		List[str]							      =     []
			# 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  lowerCAmelCase_	(lowerCAmelCase__:		Optional[int]		,						lowerCAmelCase__:		Dict		,						lowerCAmelCase__:		str  ):





			"""simple docstring"""


			UpperCAmelCase_:		int							      =     dct.pop(lowerCAmelCase__  )
			UpperCAmelCase_:		Optional[int]							      =     val







def  lowerCAmelCase_	(lowerCAmelCase__:		List[str]		,						lowerCAmelCase__:		str  ):





			"""simple docstring"""


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

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







def  lowerCAmelCase_	(lowerCAmelCase__:		Optional[Any]		,						lowerCAmelCase__:		Union[str, Any]  ):





			"""simple docstring"""


			UpperCAmelCase_:		List[str]							      =     3_6_4 if """coco""" in model_name else 2_2_4
			UpperCAmelCase_:		List[str]							      =     BlipaVisionConfig(image_size=lowerCAmelCase__  ).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[Any]							      =     OPTConfig.from_pretrained("""facebook/opt-2.7b"""		,						eos_token_id=lowerCAmelCase__  ).to_dict()
			elif "opt-6.7b" in model_name:
						UpperCAmelCase_:		Union[str, Any]							      =     OPTConfig.from_pretrained("""facebook/opt-6.7b"""		,						eos_token_id=lowerCAmelCase__  ).to_dict()
			elif "t5-xl" in model_name:
						UpperCAmelCase_:		Dict							      =     TaConfig.from_pretrained("""google/flan-t5-xl"""		,						dense_act_fn="""gelu"""		,						bos_token_id=1  ).to_dict()
			elif "t5-xxl" in model_name:
						UpperCAmelCase_:		List[Any]							      =     TaConfig.from_pretrained("""google/flan-t5-xxl"""		,						dense_act_fn="""gelu"""		,						bos_token_id=1  ).to_dict()

			UpperCAmelCase_:		Dict							      =     BlipaConfig(vision_config=lowerCAmelCase__		,						text_config=lowerCAmelCase__  )

			return config, image_size







@torch.no_grad()
def  lowerCAmelCase_	(lowerCAmelCase__:		str		,						lowerCAmelCase__:		Optional[Any]=None		,						lowerCAmelCase__:		Any=False  ):





			"""simple docstring"""


			UpperCAmelCase_:		Any							      =     (
			    AutoTokenizer.from_pretrained("""facebook/opt-2.7b"""  )
			    if """opt""" in model_name
			    else AutoTokenizer.from_pretrained("""google/flan-t5-xl"""  )
			)
			UpperCAmelCase_:		Dict							      =     tokenizer("""\n"""		,						add_special_tokens=lowerCAmelCase__  ).input_ids[0]
			UpperCAmelCase_      ,							UpperCAmelCase_:		Optional[Any]							      =     get_blipa_config(lowerCAmelCase__		,						eos_token_id=lowerCAmelCase__  )

			UpperCAmelCase_:		Tuple							      =     BlipaForConditionalGeneration(lowerCAmelCase__  ).eval()

			UpperCAmelCase_:		Optional[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_:		int							      =     load_model_and_preprocess(
			    name=lowerCAmelCase__		,						model_type=lowerCAmelCase__		,						is_eval=lowerCAmelCase__		,						device=lowerCAmelCase__  )
			original_model.eval()
			print("""Done!"""  )

			# update state dict keys
			UpperCAmelCase_:		List[Any]							      =     original_model.state_dict()
			UpperCAmelCase_:		Union[str, Any]							      =     create_rename_keys(lowerCAmelCase__  )
			for src, dest in rename_keys:
						rename_key(lowerCAmelCase__		,						lowerCAmelCase__		,						lowerCAmelCase__  )

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

			# read in qv biases
			read_in_q_v_bias(lowerCAmelCase__		,						lowerCAmelCase__  )

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

			UpperCAmelCase_:		str							      =     load_demo_image()
			UpperCAmelCase_:		str							      =     vis_processors["""eval"""](lowerCAmelCase__  ).unsqueeze(0  ).to(lowerCAmelCase__  )
			UpperCAmelCase_:		Tuple							      =     tokenizer(["""\n"""]		,						return_tensors="""pt"""  ).input_ids.to(lowerCAmelCase__  )

			# create processor
			UpperCAmelCase_:		List[str]							      =     BlipImageProcessor(
			    size={"""height""": image_size, """width""": image_size}		,						image_mean=lowerCAmelCase__		,						image_std=lowerCAmelCase__  )
			UpperCAmelCase_:		Dict							      =     BlipaProcessor(image_processor=lowerCAmelCase__		,						tokenizer=lowerCAmelCase__  )
			UpperCAmelCase_:		Union[str, Any]							      =     processor(images=lowerCAmelCase__		,						return_tensors="""pt"""  ).pixel_values.to(lowerCAmelCase__  )

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

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

			print("""Generating a caption..."""  )
			UpperCAmelCase_:		Tuple							      =     """"""
			UpperCAmelCase_:		int							      =     tokenizer(lowerCAmelCase__		,						return_tensors="""pt"""  ).input_ids.to(lowerCAmelCase__  )

			UpperCAmelCase_:		Tuple							      =     original_model.generate({"""image""": original_pixel_values}  )
			UpperCAmelCase_:		Tuple							      =     hf_model.generate(
			    lowerCAmelCase__		,						lowerCAmelCase__		,						do_sample=lowerCAmelCase__		,						num_beams=5		,						max_length=3_0		,						min_length=1		,						top_p=0.9		,						repetition_penalty=1.0		,						length_penalty=1.0		,						temperature=1		,						)
			print("""Original generation:"""		,						lowerCAmelCase__  )
			UpperCAmelCase_:		str							      =     input_ids.shape[1]
			UpperCAmelCase_:		int							      =     processor.batch_decode(outputs[:, prompt_length:]		,						skip_special_tokens=lowerCAmelCase__  )
			UpperCAmelCase_:		Dict							      =     [text.strip() for text in output_text]
			print("""HF generation:"""		,						lowerCAmelCase__  )

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

			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      :  Union[str, Any]			     =			argparse.ArgumentParser()
	a      :  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      :  int			     =			parser.parse_args()

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






