
'''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      lowerCamelCase__       (       ):




			"""simple docstring"""


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

			return image






def      lowerCamelCase__       (       __lowerCAmelCase						:						Optional[int]       ):




			"""simple docstring"""


			lowerCAmelCase_      			=			[]
			# 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						:						int					,							__lowerCAmelCase						:						Optional[int]					,							__lowerCAmelCase						:						str       ):




			"""simple docstring"""


			lowerCAmelCase_      			=			dct.pop(__UpperCAmelCase       )
			lowerCAmelCase_      			=			val






def      lowerCamelCase__       (       __lowerCAmelCase						:						Union[str, Any]					,							__lowerCAmelCase						:						Optional[int]       ):




			"""simple docstring"""


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

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






def      lowerCamelCase__       (       __lowerCAmelCase						:						str					,							__lowerCAmelCase						:						List[str]       ):




			"""simple docstring"""


			lowerCAmelCase_      			=			364 if "coco" in model_name else 224
			lowerCAmelCase_      			=			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:
						lowerCAmelCase_      			=			OPTConfig.from_pretrained("facebook/opt-2.7b"					,							eos_token_id=__UpperCAmelCase       ).to_dict()
			elif "opt-6.7b" in model_name:
						lowerCAmelCase_      			=			OPTConfig.from_pretrained("facebook/opt-6.7b"					,							eos_token_id=__UpperCAmelCase       ).to_dict()
			elif "t5-xl" in model_name:
						lowerCAmelCase_      			=			TaConfig.from_pretrained("google/flan-t5-xl"					,							dense_act_fn="gelu"					,							bos_token_id=1       ).to_dict()
			elif "t5-xxl" in model_name:
						lowerCAmelCase_      			=			TaConfig.from_pretrained("google/flan-t5-xxl"					,							dense_act_fn="gelu"					,							bos_token_id=1       ).to_dict()

			lowerCAmelCase_      			=			BlipaConfig(vision_config=__UpperCAmelCase					,							text_config=__UpperCAmelCase       )

			return config, image_size






@torch.no_grad()
def      lowerCamelCase__       (       __lowerCAmelCase						:						Optional[Any]					,							__lowerCAmelCase						:						List[str]=None					,							__lowerCAmelCase						:						Union[str, Any]=False       ):




			"""simple docstring"""


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

			lowerCAmelCase_      			=			BlipaForConditionalGeneration(__UpperCAmelCase       ).eval()

			lowerCAmelCase_      			=			{
			    "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_      			=			model_name_to_original[model_name]

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

			# update state dict keys
			lowerCAmelCase_      			=			original_model.state_dict()
			lowerCAmelCase_      			=			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():
						lowerCAmelCase_      			=			state_dict.pop(__UpperCAmelCase       )
						if key.startswith("Qformer.bert"       ):
									lowerCAmelCase_      			=			key.replace("Qformer.bert"					,							"qformer"       )
						if "attention.self" in key:
									lowerCAmelCase_      			=			key.replace("self"					,							"attention"       )
						if "opt_proj" in key:
									lowerCAmelCase_      			=			key.replace("opt_proj"					,							"language_projection"       )
						if "t5_proj" in key:
									lowerCAmelCase_      			=			key.replace("t5_proj"					,							"language_projection"       )
						if key.startswith("opt"       ):
									lowerCAmelCase_      			=			key.replace("opt"					,							"language"       )
						if key.startswith("t5"       ):
									lowerCAmelCase_      			=			key.replace("t5"					,							"language"       )
						lowerCAmelCase_      			=			val

			# read in qv biases
			read_in_q_v_bias(__UpperCAmelCase					,							__UpperCAmelCase       )

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

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

			# create processor
			lowerCAmelCase_      			=			BlipImageProcessor(
			    size={"height": image_size, "width": image_size}					,							image_mean=__UpperCAmelCase					,							image_std=__UpperCAmelCase       )
			lowerCAmelCase_      			=			BlipaProcessor(image_processor=__UpperCAmelCase					,							tokenizer=__UpperCAmelCase       )
			lowerCAmelCase_      			=			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:
									lowerCAmelCase_      			=			original_model({"image": original_pixel_values, "text_input": [""]}       ).logits
									lowerCAmelCase_      			=			hf_model(__UpperCAmelCase					,							__UpperCAmelCase       ).logits
						else:
									lowerCAmelCase_      			=			original_model(
									    {"image": original_pixel_values, "text_input": ["\n"], "text_output": ["\n"]}       ).logits
									lowerCAmelCase_      			=			input_ids.masked_fill(input_ids == tokenizer.pad_token_id					,							-100       )
									lowerCAmelCase_      			=			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":
						lowerCAmelCase_      			=			torch.tensor(
						    [[-41.5850, -4.4_440, -8.9_922], [-47.4322, -5.9_143, -1.7_340]]					,							device=__UpperCAmelCase       )
						assert torch.allclose(logits[0, :3, :3]					,							__UpperCAmelCase					,							atol=1e-4       )
			elif model_name == "blip2-flan-t5-xl-coco":
						lowerCAmelCase_      			=			torch.tensor(
						    [[-57.0109, -9.8_967, -12.6280], [-68.6578, -12.7191, -10.5065]]					,							device=__UpperCAmelCase       )
			else:
						# cast to same type
						lowerCAmelCase_      			=			logits.dtype
						assert torch.allclose(original_logits.to(__UpperCAmelCase       )					,							__UpperCAmelCase					,							atol=1e-2       )
			print("Looks ok!"       )

			print("Generating a caption..."       )
			lowerCAmelCase_      			=			""
			lowerCAmelCase_      			=			tokenizer(__UpperCAmelCase					,							return_tensors="pt"       ).input_ids.to(__UpperCAmelCase       )

			lowerCAmelCase_      			=			original_model.generate({"image": original_pixel_values}       )
			lowerCAmelCase_      			=			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       )
			lowerCAmelCase_      			=			input_ids.shape[1]
			lowerCAmelCase_      			=			processor.batch_decode(outputs[:, prompt_length:]					,							skip_special_tokens=__UpperCAmelCase       )
			lowerCAmelCase_      			=			[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)





import json
import os
from collections import Counter

import torch
import torchvision
import torchvision.transforms as transforms
from PIL import Image
from torch import nn
from torch.utils.data import Dataset


_A     =  {1: (1, 1), 2: (2, 1), 3: (3, 1), 4: (2, 2), 5: (5, 1), 6: (3, 2), 7: (7, 1), 8: (4, 2), 9: (3, 3)}





class 			_lowerCAmelCase      (			nn.Module ):




				def __init__(      self       ,       _UpperCamelCase  )       ->   Optional[Any]:
							super().__init__()
							lowerCAmelCase_      			=			torchvision.models.resnetaaa(pretrained=_UpperCamelCase  )
							lowerCAmelCase_      			=			list(model.children()  )[:-2]
							lowerCAmelCase_      			=			nn.Sequential(*_UpperCamelCase  )
							lowerCAmelCase_      			=			nn.AdaptiveAvgPoolad(POOLING_BREAKDOWN[args.num_image_embeds]  )




				def   __a		(      self       ,       _UpperCamelCase  )       ->   Dict:
							# Bx3x224x224 -> Bx2048x7x7 -> Bx2048xN -> BxNx2048
							lowerCAmelCase_      			=			self.pool(self.model(_UpperCamelCase  )  )
							lowerCAmelCase_      			=			torch.flatten(_UpperCamelCase       ,       start_dim=2  )
							lowerCAmelCase_      			=			out.transpose(1       ,       2  ).contiguous()
							return out  # BxNx2048





class 			_lowerCAmelCase      (			__a ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   List[str]:
							lowerCAmelCase_      			=			[json.loads(_UpperCamelCase  ) for l in open(_UpperCamelCase  )]
							lowerCAmelCase_      			=			os.path.dirname(_UpperCamelCase  )
							lowerCAmelCase_      			=			tokenizer
							lowerCAmelCase_      			=			labels
							lowerCAmelCase_      			=			len(_UpperCamelCase  )
							lowerCAmelCase_      			=			max_seq_length

							lowerCAmelCase_      			=			transforms




				def __len__(      self  )       ->   Any:
							return len(self.data  )




				def __getitem__(      self       ,       _UpperCamelCase  )       ->   Any:
							lowerCAmelCase_      			=			torch.LongTensor(self.tokenizer.encode(self.data[index]["text"]       ,       add_special_tokens=_UpperCamelCase  )  )
							lowerCAmelCase_			,						lowerCAmelCase_			,						lowerCAmelCase_      			=			sentence[0], sentence[1:-1], sentence[-1]
							lowerCAmelCase_      			=			sentence[: self.max_seq_length]

							lowerCAmelCase_      			=			torch.zeros(self.n_classes  )
							lowerCAmelCase_      			=			1

							lowerCAmelCase_      			=			Image.open(os.path.join(self.data_dir       ,       self.data[index]["img"]  )  ).convert("RGB"  )
							lowerCAmelCase_      			=			self.transforms(_UpperCamelCase  )

							return {
							    "image_start_token": start_token,
							    "image_end_token": end_token,
							    "sentence": sentence,
							    "image": image,
							    "label": label,
							}




				def   __a		(      self  )       ->   str:
							lowerCAmelCase_      			=			Counter()
							for row in self.data:
										label_freqs.update(row["label"]  )
							return label_freqs






def      lowerCamelCase__       (       __lowerCAmelCase						:						List[str]       ):




			"""simple docstring"""


			lowerCAmelCase_      			=			[len(row["sentence"]       ) for row in batch]
			lowerCAmelCase_			,						lowerCAmelCase_      			=			len(__lowerCAmelCase       ), max(__lowerCAmelCase       )

			lowerCAmelCase_      			=			torch.zeros(__lowerCAmelCase					,							__lowerCAmelCase					,							dtype=torch.long       )
			lowerCAmelCase_      			=			torch.zeros(__lowerCAmelCase					,							__lowerCAmelCase					,							dtype=torch.long       )

			for i_batch, (input_row, length) in enumerate(zip(__lowerCAmelCase					,							__lowerCAmelCase       )       ):
						lowerCAmelCase_      			=			input_row["sentence"]
						lowerCAmelCase_      			=			1

			lowerCAmelCase_      			=			torch.stack([row["image"] for row in batch]       )
			lowerCAmelCase_      			=			torch.stack([row["label"] for row in batch]       )
			lowerCAmelCase_      			=			torch.stack([row["image_start_token"] for row in batch]       )
			lowerCAmelCase_      			=			torch.stack([row["image_end_token"] for row in batch]       )

			return text_tensor, mask_tensor, img_tensor, img_start_token, img_end_token, tgt_tensor






def      lowerCamelCase__       (       ):




			"""simple docstring"""


			return [
			    "Crime",
			    "Drama",
			    "Thriller",
			    "Action",
			    "Comedy",
			    "Romance",
			    "Documentary",
			    "Short",
			    "Mystery",
			    "History",
			    "Family",
			    "Adventure",
			    "Fantasy",
			    "Sci-Fi",
			    "Western",
			    "Horror",
			    "Sport",
			    "War",
			    "Music",
			    "Musical",
			    "Animation",
			    "Biography",
			    "Film-Noir",
			]






def      lowerCamelCase__       (       ):




			"""simple docstring"""


			return transforms.Compose(
			    [
			        transforms.Resize(256       ),
			        transforms.CenterCrop(224       ),
			        transforms.ToTensor(),
			        transforms.Normalize(
			            mean=[0.46_777_044, 0.44_531_429, 0.40_661_017]					,							std=[0.12_221_994, 0.12_145_835, 0.14_380_469]					,							),
			    ]       )



