






'''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"""






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

				return image







def 							UpperCAmelCase_    (	lowerCAmelCase_							):

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

				"""simple docstring"""






				lowercase         = dct.pop(__A							)
				lowercase         = val







def 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_							):

				"""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(__A		,							requires_grad=__A							), v_bias)							)
								lowercase         = qkv_bias







def 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_							):

				"""simple docstring"""






				lowercase         = 364 if "coco" in model_name else 224
				lowercase         = 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:
								lowercase         = OPTConfig.from_pretrained("facebook/opt-2.7b"		,							eos_token_id=__A							).to_dict()
				elif "opt-6.7b" in model_name:
								lowercase         = OPTConfig.from_pretrained("facebook/opt-6.7b"		,							eos_token_id=__A							).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=__A		,							text_config=__A							)

				return config, image_size







@torch.no_grad()
def 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_=None		,							lowerCAmelCase_=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=__A							).input_ids[0]
				lowercase   ,							lowercase         = get_blipa_config(__A		,							eos_token_id=__A							)

				lowercase         = BlipaForConditionalGeneration(__A							).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=__A		,							model_type=__A		,							is_eval=__A		,							device=__A							)
				original_model.eval()
				print("Done!"							)

				# update state dict keys
				lowercase         = original_model.state_dict()
				lowercase         = 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():
								lowercase         = state_dict.pop(__A							)
								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(__A		,							__A							)

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

				lowercase         = load_demo_image()
				lowercase         = vis_processors["eval"](__A							).unsqueeze(0							).to(__A							)
				lowercase         = tokenizer(["\n"]		,							return_tensors="pt"							).input_ids.to(__A							)

				# create processor
				lowercase         = BlipImageProcessor(
				    size={"height": image_size, "width": image_size}		,							image_mean=__A		,							image_std=__A							)
				lowercase         = BlipaProcessor(image_processor=__A		,							tokenizer=__A							)
				lowercase         = 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:
												lowercase         = original_model({"image": original_pixel_values, "text_input": [""]}							).logits
												lowercase         = hf_model(__A		,							__A							).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		,							-100							)
												lowercase         = 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":
								lowercase         = torch.tensor(
								    [[-41.58_50, -4.44_40, -8.99_22], [-47.43_22, -5.91_43, -1.73_40]]		,							device=__A							)
								assert torch.allclose(logits[0, :3, :3]		,							__A		,							atol=1E-4							)
				elif model_name == "blip2-flan-t5-xl-coco":
								lowercase         = torch.tensor(
								    [[-57.01_09, -9.89_67, -12.62_80], [-68.65_78, -12.71_91, -10.50_65]]		,							device=__A							)
				else:
								# cast to same type
								lowercase         = logits.dtype
								assert torch.allclose(original_logits.to(__A							)		,							__A		,							atol=1E-2							)
				print("Looks ok!"							)

				print("Generating a caption..."							)
				lowercase         = ""
				lowercase         = tokenizer(__A		,							return_tensors="pt"							).input_ids.to(__A							)

				lowercase         = original_model.generate({"image": original_pixel_values}							)
				lowercase         = 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							)
				lowercase         = input_ids.shape[1]
				lowercase         = processor.batch_decode(outputs[:, prompt_length:]		,							skip_special_tokens=__A							)
				lowercase         = [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__":
			__lowerCamelCase : Optional[Any]											=			argparse.ArgumentParser()
			__lowerCamelCase : List[str]											=			[
			    "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",
			)

			__lowerCamelCase : Optional[int]											=			parser.parse_args()

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











'''simple docstring'''

import warnings
from contextlib import contextmanager

from ....processing_utils import ProcessorMixin

class       UpperCAmelCase (						_lowercase  ):
		UpperCAmelCase					:	Optional[Any] 			=		'''MCTCTFeatureExtractor'''
		UpperCAmelCase					:	Tuple 			=		'''AutoTokenizer'''




		def __init__(self      :	int							,			A__      :	Tuple							,			A__      :	Union[str, Any]				)					-> Dict:
						super().__init__(A__							,			A__				)
						lowercase         = self.feature_extractor
						lowercase         = False




		def __call__(self      :	Tuple							,			*A__      :	str							,			**A__      :	Dict				)					-> int:
						# For backward compatibility
						if self._in_target_context_manager:
										return self.current_processor(*A__							,			**A__				)

						if "raw_speech" in kwargs:
										warnings.warn("Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead."				)
										lowercase         = kwargs.pop("raw_speech"				)
						else:
										lowercase         = kwargs.pop("audio"							,			A__				)
						lowercase         = kwargs.pop("sampling_rate"							,			A__				)
						lowercase         = kwargs.pop("text"							,			A__				)
						if len(A__				) > 0:
										lowercase         = args[0]
										lowercase         = args[1:]

						if audio is None and text is None:
										raise ValueError("You need to specify either an `audio` or `text` input to process."				)

						if audio is not None:
										lowercase         = self.feature_extractor(A__							,			*A__							,			sampling_rate=A__							,			**A__				)
						if text is not None:
										lowercase         = self.tokenizer(A__							,			**A__				)

						if text is None:
										return inputs
						elif audio is None:
										return encodings
						else:
										lowercase         = encodings["input_ids"]
										return inputs




		def 				UpperCAmelCase__   (self      :	Tuple							,			*A__      :	str							,			**A__      :	str				)					-> str:
						return self.tokenizer.batch_decode(*A__							,			**A__				)




		def 				UpperCAmelCase__   (self      :	Any							,			*A__      :	List[Any]							,			**A__      :	List[str]				)					-> Tuple:
						# For backward compatibility
						if self._in_target_context_manager:
										return self.current_processor.pad(*A__							,			**A__				)

						lowercase         = kwargs.pop("input_features"							,			A__				)
						lowercase         = kwargs.pop("labels"							,			A__				)
						if len(A__				) > 0:
										lowercase         = args[0]
										lowercase         = args[1:]

						if input_features is not None:
										lowercase         = self.feature_extractor.pad(A__							,			*A__							,			**A__				)
						if labels is not None:
										lowercase         = self.tokenizer.pad(A__							,			**A__				)

						if labels is None:
										return input_features
						elif input_features is None:
										return labels
						else:
										lowercase         = labels["input_ids"]
										return input_features




		def 				UpperCAmelCase__   (self      :	Tuple							,			*A__      :	Optional[int]							,			**A__      :	Optional[int]				)					-> Tuple:
						return self.tokenizer.decode(*A__							,			**A__				)





		@contextmanager
		def 				UpperCAmelCase__   (self      :	Optional[Any]				)					-> Union[str, Any]:
						warnings.warn(
						    "`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your "
						    "labels by using the argument `text` of the regular `__call__` method (either in the same call as "
						    "your audio inputs, or in a separate call."				)
						lowercase         = True
						lowercase         = self.tokenizer
						yield
						lowercase         = self.feature_extractor
						lowercase         = False




