




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


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

			return image

def 					_UpperCamelCase    (			UpperCamelCase							)			->			Dict:


			"""simple docstring"""
			__UpperCAmelCase	: Union[str, Any]		=					[]
			# 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    (			UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase							)			->			Dict:


			"""simple docstring"""
			__UpperCAmelCase	: Any		=					dct.pop(UpperCamelCase							)
			__UpperCAmelCase	: Tuple		=					val

def 					_UpperCamelCase    (			UpperCamelCase				,       UpperCamelCase							)			->			Tuple:


			"""simple docstring"""
			for i in range(config.vision_config.num_hidden_layers							):
						# read in original q and v biases
						__UpperCAmelCase	: Optional[int]		=					state_dict.pop(f"visual_encoder.blocks.{i}.attn.q_bias"							)
						__UpperCAmelCase	: Union[str, Any]		=					state_dict.pop(f"visual_encoder.blocks.{i}.attn.v_bias"							)

						# next, set bias in the state dict
						__UpperCAmelCase	: Tuple		=					torch.cat((q_bias, torch.zeros_like(UpperCamelCase				,       requires_grad=UpperCamelCase							), v_bias)							)
						__UpperCAmelCase	: List[str]		=					qkv_bias

def 					_UpperCamelCase    (			UpperCamelCase				,       UpperCamelCase							)			->			List[str]:


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

			__UpperCAmelCase	: Dict		=					BlipaConfig(vision_config=UpperCamelCase				,       text_config=UpperCamelCase							)

			return config, image_size

@torch.no_grad()
def 					_UpperCamelCase    (			UpperCamelCase				,       UpperCamelCase=None				,       UpperCamelCase=False							)			->			Optional[int]:


			"""simple docstring"""
			__UpperCAmelCase	: Union[str, Any]		=					(
			    AutoTokenizer.from_pretrained("facebook/opt-2.7b"							)
			    if "opt" in model_name
			    else AutoTokenizer.from_pretrained("google/flan-t5-xl"							)
			)
			__UpperCAmelCase	: Any		=					tokenizer("\n"				,       add_special_tokens=UpperCamelCase							).input_ids[0]
			__UpperCAmelCase	,     __UpperCAmelCase	: List[str]		=					get_blipa_config(UpperCamelCase				,       eos_token_id=UpperCamelCase							)

			__UpperCAmelCase	: Tuple		=					BlipaForConditionalGeneration(UpperCamelCase							).eval()

			__UpperCAmelCase	: Tuple		=					{
			    "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	: Optional[Any]		=					model_name_to_original[model_name]

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

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

			# read in qv biases
			read_in_q_v_bias(UpperCamelCase				,       UpperCamelCase							)

			__UpperCAmelCase	,     __UpperCAmelCase	: Tuple		=					hf_model.load_state_dict(UpperCamelCase				,       strict=UpperCamelCase							)
			assert len(UpperCamelCase							) == 0
			assert unexpected_keys == ["qformer.embeddings.position_ids"]

			__UpperCAmelCase	: Optional[Any]		=					load_demo_image()
			__UpperCAmelCase	: List[str]		=					vis_processors["eval"](UpperCamelCase							).unsqueeze(0							).to(UpperCamelCase							)
			__UpperCAmelCase	: Dict		=					tokenizer(["\n"]				,       return_tensors="pt"							).input_ids.to(UpperCamelCase							)

			# create processor
			__UpperCAmelCase	: List[Any]		=					BlipImageProcessor(
			    size={"height": image_size, "width": image_size}				,       image_mean=UpperCamelCase				,       image_std=UpperCamelCase							)
			__UpperCAmelCase	: Optional[Any]		=					BlipaProcessor(image_processor=UpperCamelCase				,       tokenizer=UpperCamelCase							)
			__UpperCAmelCase	: Optional[Any]		=					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:
									__UpperCAmelCase	: int		=					original_model({"image": original_pixel_values, "text_input": [""]}							).logits
									__UpperCAmelCase	: Optional[int]		=					hf_model(UpperCamelCase				,       UpperCamelCase							).logits
						else:
									__UpperCAmelCase	: int		=					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				,       -100							)
									__UpperCAmelCase	: List[Any]		=					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":
						__UpperCAmelCase	: Union[str, Any]		=					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":
						__UpperCAmelCase	: List[str]		=					torch.tensor(
						    [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]				,       device=UpperCamelCase							)
			else:
						# cast to same type
						__UpperCAmelCase	: List[str]		=					logits.dtype
						assert torch.allclose(original_logits.to(UpperCamelCase							)				,       UpperCamelCase				,       atol=1e-2							)
			print("Looks ok!"							)

			print("Generating a caption..."							)
			__UpperCAmelCase	: Optional[Any]		=					""
			__UpperCAmelCase	: Union[str, Any]		=					tokenizer(UpperCamelCase				,       return_tensors="pt"							).input_ids.to(UpperCamelCase							)

			__UpperCAmelCase	: Optional[Any]		=					original_model.generate({"image": original_pixel_values}							)
			__UpperCAmelCase	: List[str]		=					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							)
			__UpperCAmelCase	: List[Any]		=					input_ids.shape[1]
			__UpperCAmelCase	: Union[str, Any]		=					processor.batch_decode(outputs[:, prompt_length:]				,       skip_special_tokens=UpperCamelCase							)
			__UpperCAmelCase	: List[str]		=					[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   =       argparse.ArgumentParser()
					A   =       [
					    """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   =       parser.parse_args()

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








"""simple docstring"""





import shutil
import tempfile
import unittest

from transformers import ClapFeatureExtractor, ClapProcessor, RobertaTokenizer, RobertaTokenizerFast
from transformers.testing_utils import require_sentencepiece, require_torchaudio

from .test_feature_extraction_clap import floats_list



@require_torchaudio
@require_sentencepiece
class      a__				(	unittest.TestCase   ):


			def       a_     (      self      :    Any):


						"""simple docstring"""




						__UpperCAmelCase	: Tuple		=					"laion/clap-htsat-unfused"
						__UpperCAmelCase	: List[str]		=					tempfile.mkdtemp()


			def       a_     (      self      :    Optional[Any]      ,  **UpperCamelCase_      :    Union[str, Any]):


						"""simple docstring"""




						return RobertaTokenizer.from_pretrained(self.checkpoint      ,  **UpperCamelCase_)


			def       a_     (      self      :    Union[str, Any]      ,  **UpperCamelCase_      :    Optional[int]):


						"""simple docstring"""




						return ClapFeatureExtractor.from_pretrained(self.checkpoint      ,  **UpperCamelCase_)


			def       a_     (      self      :    List[Any]):


						"""simple docstring"""




						shutil.rmtree(self.tmpdirname)


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					self.get_tokenizer()
						__UpperCAmelCase	: str		=					self.get_feature_extractor()

						__UpperCAmelCase	: List[Any]		=					ClapProcessor(tokenizer=UpperCamelCase_      ,  feature_extractor=UpperCamelCase_)

						processor.save_pretrained(self.tmpdirname)
						__UpperCAmelCase	: Union[str, Any]		=					ClapProcessor.from_pretrained(self.tmpdirname)

						self.assertEqual(processor.tokenizer.get_vocab()      ,  tokenizer.get_vocab())
						self.assertIsInstance(processor.tokenizer      ,  UpperCamelCase_)

						self.assertEqual(processor.feature_extractor.to_json_string()      ,  feature_extractor.to_json_string())
						self.assertIsInstance(processor.feature_extractor      ,  UpperCamelCase_)


			def       a_     (      self      :    Any):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					ClapProcessor(tokenizer=self.get_tokenizer()      ,  feature_extractor=self.get_feature_extractor())
						processor.save_pretrained(self.tmpdirname)

						__UpperCAmelCase	: int		=					self.get_tokenizer(bos_token="(BOS)"      ,  eos_token="(EOS)")
						__UpperCAmelCase	: Any		=					self.get_feature_extractor(do_normalize=UpperCamelCase_      ,  padding_value=1.0)

						__UpperCAmelCase	: Tuple		=					ClapProcessor.from_pretrained(
						    self.tmpdirname      ,  bos_token="(BOS)"      ,  eos_token="(EOS)"      ,  do_normalize=UpperCamelCase_      ,  padding_value=1.0)

						self.assertEqual(processor.tokenizer.get_vocab()      ,  tokenizer_add_kwargs.get_vocab())
						self.assertIsInstance(processor.tokenizer      ,  UpperCamelCase_)

						self.assertEqual(processor.feature_extractor.to_json_string()      ,  feature_extractor_add_kwargs.to_json_string())
						self.assertIsInstance(processor.feature_extractor      ,  UpperCamelCase_)


			def       a_     (      self      :    Any):


						"""simple docstring"""




						__UpperCAmelCase	: Tuple		=					self.get_feature_extractor()
						__UpperCAmelCase	: int		=					self.get_tokenizer()

						__UpperCAmelCase	: Any		=					ClapProcessor(tokenizer=UpperCamelCase_      ,  feature_extractor=UpperCamelCase_)

						__UpperCAmelCase	: Optional[int]		=					floats_list((3, 1000))

						__UpperCAmelCase	: Tuple		=					feature_extractor(UpperCamelCase_      ,  return_tensors="np")
						__UpperCAmelCase	: Any		=					processor(audios=UpperCamelCase_      ,  return_tensors="np")

						for key in input_feat_extract.keys():
									self.assertAlmostEqual(input_feat_extract[key].sum()      ,  input_processor[key].sum()      ,  delta=1e-2)


			def       a_     (      self      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					self.get_feature_extractor()
						__UpperCAmelCase	: Optional[int]		=					self.get_tokenizer()

						__UpperCAmelCase	: str		=					ClapProcessor(tokenizer=UpperCamelCase_      ,  feature_extractor=UpperCamelCase_)

						__UpperCAmelCase	: Any		=					"This is a test string"

						__UpperCAmelCase	: List[Any]		=					processor(text=UpperCamelCase_)

						__UpperCAmelCase	: Dict		=					tokenizer(UpperCamelCase_)

						for key in encoded_tok.keys():
									self.assertListEqual(encoded_tok[key]      ,  encoded_processor[key])


			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					self.get_feature_extractor()
						__UpperCAmelCase	: List[str]		=					self.get_tokenizer()

						__UpperCAmelCase	: List[str]		=					ClapProcessor(tokenizer=UpperCamelCase_      ,  feature_extractor=UpperCamelCase_)

						__UpperCAmelCase	: Tuple		=					[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

						__UpperCAmelCase	: int		=					processor.batch_decode(UpperCamelCase_)
						__UpperCAmelCase	: Tuple		=					tokenizer.batch_decode(UpperCamelCase_)

						self.assertListEqual(UpperCamelCase_      ,  UpperCamelCase_)


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: Any		=					self.get_feature_extractor()
						__UpperCAmelCase	: int		=					self.get_tokenizer()

						__UpperCAmelCase	: int		=					ClapProcessor(tokenizer=UpperCamelCase_      ,  feature_extractor=UpperCamelCase_)

						self.assertListEqual(
						    processor.model_input_names[2:]      ,  feature_extractor.model_input_names      ,  msg="`processor` and `feature_extractor` model input names do not match"      ,  )


