




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_ (						):
				__magic_name__						:						Optional[int]										="""https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"""
				__magic_name__						:						Optional[int]										=Image.open(requests.get(lowerCamelCase						,  stream=lowerCamelCase ).raw ).convert("""RGB""" )

				return image




def 		lowerCAmelCase_ (						lowerCamelCase ):
				__magic_name__						:						List[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 		lowerCAmelCase_ (						lowerCamelCase						,  lowerCamelCase						,  lowerCamelCase ):
				__magic_name__						:						List[str]										=dct.pop(lowerCamelCase )
				__magic_name__						:						Optional[int]										=val




def 		lowerCAmelCase_ (						lowerCamelCase						,  lowerCamelCase ):
				for i in range(config.vision_config.num_hidden_layers ):
								# read in original q and v biases
								__magic_name__						:						List[Any]										=state_dict.pop(F"visual_encoder.blocks.{i}.attn.q_bias" )
								__magic_name__						:						List[Any]										=state_dict.pop(F"visual_encoder.blocks.{i}.attn.v_bias" )

								# next, set bias in the state dict
								__magic_name__						:						List[Any]										=torch.cat((q_bias, torch.zeros_like(lowerCamelCase						,  requires_grad=lowerCamelCase ), v_bias) )
								__magic_name__						:						int										=qkv_bias




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

				__magic_name__						:						Dict										=BlipaConfig(vision_config=lowerCamelCase						,  text_config=lowerCamelCase )

				return config, image_size




@torch.no_grad()
def 		lowerCAmelCase_ (						lowerCamelCase						,  lowerCamelCase=None						,  lowerCamelCase=False ):
				__magic_name__						:						List[Any]										=(
				    AutoTokenizer.from_pretrained("""facebook/opt-2.7b""" )
				    if """opt""" in model_name
				    else AutoTokenizer.from_pretrained("""google/flan-t5-xl""" )
				)
				__magic_name__						:						List[str]										=tokenizer("""\n"""						,  add_special_tokens=lowerCamelCase ).input_ids[0]
				__magic_name__					,			__magic_name__						:						Dict										=get_blipa_config(lowerCamelCase						,  eos_token_id=lowerCamelCase )

				__magic_name__						:						Union[str, Any]										=BlipaForConditionalGeneration(lowerCamelCase ).eval()

				__magic_name__						:						List[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"""),
				}

				__magic_name__					,			__magic_name__						:						int										=model_name_to_original[model_name]

				# load original model
				print("""Loading original model...""" )
				__magic_name__						:						List[str]										="""cuda""" if torch.cuda.is_available() else """cpu"""
				__magic_name__					,			__magic_name__					,			__magic_name__						:						str										=load_model_and_preprocess(
				    name=lowerCamelCase						,  model_type=lowerCamelCase						,  is_eval=lowerCamelCase						,  device=lowerCamelCase )
				original_model.eval()
				print("""Done!""" )

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

				# read in qv biases
				read_in_q_v_bias(lowerCamelCase						,  lowerCamelCase )

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

				__magic_name__						:						int										=load_demo_image()
				__magic_name__						:						Tuple										=vis_processors["""eval"""](lowerCamelCase ).unsqueeze(0 ).to(lowerCamelCase )
				__magic_name__						:						List[Any]										=tokenizer(["""\n"""]						,  return_tensors="""pt""" ).input_ids.to(lowerCamelCase )

				# create processor
				__magic_name__						:						Union[str, Any]										=BlipImageProcessor(
				    size={"""height""": image_size, """width""": image_size}						,  image_mean=lowerCamelCase						,  image_std=lowerCamelCase )
				__magic_name__						:						Optional[int]										=BlipaProcessor(image_processor=lowerCamelCase						,  tokenizer=lowerCamelCase )
				__magic_name__						:						Optional[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:
												__magic_name__						:						List[Any]										=original_model({"""image""": original_pixel_values, """text_input""": [""""""]} ).logits
												__magic_name__						:						List[str]										=hf_model(lowerCamelCase						,  lowerCamelCase ).logits
								else:
												__magic_name__						:						Union[str, Any]										=original_model(
												    {"""image""": original_pixel_values, """text_input""": ["""\n"""], """text_output""": ["""\n"""]} ).logits
												__magic_name__						:						Any										=input_ids.masked_fill(input_ids == tokenizer.pad_token_id						,  -100 )
												__magic_name__						:						List[str]										=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":
								__magic_name__						:						Tuple										=torch.tensor(
								    [[-4_1.5_8_5_0, -4.4_4_4_0, -8.9_9_2_2], [-4_7.4_3_2_2, -5.9_1_4_3, -1.7_3_4_0]]						,  device=lowerCamelCase )
								assert torch.allclose(logits[0, :3, :3]						,  lowerCamelCase						,  atol=1E-4 )
				elif model_name == "blip2-flan-t5-xl-coco":
								__magic_name__						:						Union[str, Any]										=torch.tensor(
								    [[-5_7.0_1_0_9, -9.8_9_6_7, -1_2.6_2_8_0], [-6_8.6_5_7_8, -1_2.7_1_9_1, -1_0.5_0_6_5]]						,  device=lowerCamelCase )
				else:
								# cast to same type
								__magic_name__						:						Any										=logits.dtype
								assert torch.allclose(original_logits.to(lowerCamelCase )						,  lowerCamelCase						,  atol=1E-2 )
				print("""Looks ok!""" )

				print("""Generating a caption...""" )
				__magic_name__						:						Dict										=""""""
				__magic_name__						:						List[Any]										=tokenizer(lowerCamelCase						,  return_tensors="""pt""" ).input_ids.to(lowerCamelCase )

				__magic_name__						:						List[str]										=original_model.generate({"""image""": original_pixel_values} )
				__magic_name__						:						int										=hf_model.generate(
				    lowerCamelCase						,  lowerCamelCase						,  do_sample=lowerCamelCase						,  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:"""						,  lowerCamelCase )
				__magic_name__						:						List[Any]										=input_ids.shape[1]
				__magic_name__						:						int										=processor.batch_decode(outputs[:, prompt_length:]						,  skip_special_tokens=lowerCamelCase )
				__magic_name__						:						int										=[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__":
							UpperCAmelCase_		:					Dict          =  argparse.ArgumentParser()
							UpperCAmelCase_		:					Dict          =  [
							    "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",
							)

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

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










import os
from typing import Dict, List, Union

import tensorflow as tf
from keras_nlp.tokenizers import BytePairTokenizer
from tensorflow_text import pad_model_inputs

from .tokenization_gpta import GPTaTokenizer




class 		__A   (      tf.keras.layers.Layer      ):

						def __init__(		self :Optional[int]     ,			__snake_case :Dict[str, int]     ,			__snake_case :List[str]     ,			__snake_case :int = None     ,			__snake_case :int = None			):
										'''simple docstring'''

										super().__init__()
										__magic_name__						:						Optional[int]										=pad_token_id
										__magic_name__						:						List[Any]										=max_length
										__magic_name__						:						Dict										=vocab
										__magic_name__						:						int										=merges
										__magic_name__						:						Optional[int]										=BytePairTokenizer(__snake_case     ,			__snake_case     ,			sequence_length=__snake_case			)

						@classmethod
						def      A__    (		cls :List[Any]     ,			__snake_case :GPTaTokenizer     ,			*__snake_case :int     ,			**__snake_case :Any			):
										'''simple docstring'''

										__magic_name__						:						List[Any]										=[""" """.join(__snake_case			) for m in tokenizer.bpe_ranks.keys()]
										__magic_name__						:						str										=tokenizer.get_vocab()
										return cls(__snake_case     ,			__snake_case     ,			*__snake_case     ,			**__snake_case			)

						@classmethod
						def      A__    (		cls :Dict     ,			__snake_case :Union[str, os.PathLike]     ,			*__snake_case :Union[str, Any]     ,			**__snake_case :int			):
										'''simple docstring'''

										__magic_name__						:						Dict										=GPTaTokenizer.from_pretrained(__snake_case     ,			*__snake_case     ,			**__snake_case			)
										return cls.from_tokenizer(__snake_case     ,			*__snake_case     ,			**__snake_case			)

						@classmethod
						def      A__    (		cls :Optional[Any]     ,			__snake_case :List[Any]			):
										'''simple docstring'''

										return cls(**__snake_case			)

						def      A__    (		self :Union[str, Any]			):
										'''simple docstring'''

										return {
										    "vocab": self.vocab,
										    "merges": self.merges,
										    "max_length": self.max_length,
										    "pad_token_id": self.pad_token_id,
										}


						def      A__    (		self :List[Any]     ,			__snake_case :Dict     ,			__snake_case :int = None			):
										'''simple docstring'''

										__magic_name__						:						Optional[Any]										=self.tf_tokenizer(__snake_case			)
										__magic_name__						:						Tuple										=tf.ones_like(__snake_case			)

										if self.pad_token_id is not None:
														# pad the tokens up to max length
														__magic_name__						:						Tuple										=max_length if max_length is not None else self.max_length

														if max_length is not None:
																		__magic_name__					,			__magic_name__						:						Tuple										=pad_model_inputs(
																		    __snake_case     ,			max_seq_length=__snake_case     ,			pad_value=self.pad_token_id			)

										return {"attention_mask": attention_mask, "input_ids": input_ids}




