




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_float32 (there's also the fix_lavis branch)
# also note: to convert Vicuna checkpoints, we had to include /home/niels/python_projects/checkpoints/FastChat/vicuna-7b in lavis/configs/models/blip2/blip2_instruct_vicuna7b.yaml
# same for Vicuna-13b
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipImageProcessor,
    InstructBlipConfig,
    InstructBlipForConditionalGeneration,
    InstructBlipProcessor,
    InstructBlipQFormerConfig,
    InstructBlipVisionConfig,
    LlamaConfig,
    LlamaTokenizerFast,
    TaConfig,
    TaTokenizerFast,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD




def 					lowerCamelCase__       (						):





       '''simple docstring'''



       __lowerCamelCase   =						"""https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"""
       __lowerCamelCase   =						Image.open(requests.get(A__  ,			stream=A__	).raw	).convert("""RGB"""	)

       return image




def 					lowerCamelCase__       (						A__       :					Tuple	):





       '''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.embeddings.layernorm.weight""")	)
       rename_keys.append(("""Qformer.bert.embeddings.LayerNorm.bias""", """qformer.embeddings.layernorm.bias""")	)

       # fmt: on
       return rename_keys




def 					lowerCamelCase__       (						A__       :					Optional[Any]  ,			A__       :					Optional[Any]  ,			A__       :					List[Any]	):





       '''simple docstring'''



       __lowerCamelCase   =						dct.pop(A__	)
       __lowerCamelCase   =						val




def 					lowerCamelCase__       (						A__       :					str  ,			A__       :					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(A__  ,			requires_grad=A__	), v_bias)	)
              __lowerCamelCase   =						qkv_bias




def 					lowerCamelCase__       (						A__       :					str	):





       '''simple docstring'''



       __lowerCamelCase   =						364 if """coco""" in model_name else 224
       __lowerCamelCase   =						InstructBlipVisionConfig(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 "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()
       elif "vicuna-7b" in model_name:
              __lowerCamelCase   =						LlamaConfig.from_pretrained("""decapoda-research/llama-7b-hf"""  ,			vocab_size=32001	).to_dict()
       elif "vicuna-13b" in model_name:
              __lowerCamelCase   =						LlamaConfig.from_pretrained("""decapoda-research/llama-13b-hf"""  ,			vocab_size=32001	).to_dict()
       else:
              raise ValueError("""Model name not supported"""	)

       # the authors add one special "[DEC]" token to the vocab of Q-Former, hence vocab size = 30522 + 1
       __lowerCamelCase   =						InstructBlipQFormerConfig(vocab_size=30523	).to_dict()
       __lowerCamelCase   =						InstructBlipConfig(vision_config=A__  ,			text_config=A__  ,			qformer_config=A__	)

       return config, image_size




@torch.no_grad()
def 					lowerCamelCase__       (						A__       :					Any  ,			A__       :					str=None  ,			A__       :					Any=False	):





       '''simple docstring'''



       __lowerCamelCase   =						AutoTokenizer.from_pretrained("""bert-base-uncased"""  ,			truncation_side="""left"""	)
       qformer_tokenizer.add_special_tokens({"""bos_token""": """[DEC]"""}	)

       if "t5" in model_name:
              __lowerCamelCase   =						TaTokenizerFast.from_pretrained("""google/flan-t5-xl"""  ,			truncation_side="""left"""	)
       elif "vicuna" in model_name:
              # the following was used in the original implementation:
              # tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b", use_fast=False, truncation_side="left")
              # tokenizer.add_special_tokens({"pad_token": "[PAD]"})
              # tokenizer.add_special_tokens({"bos_token": "</s>"})
              # tokenizer.add_special_tokens({"eos_token": "</s>"})
              # tokenizer.add_special_tokens({"unk_token": "</s>"})
              __lowerCamelCase   =						LlamaTokenizerFast.from_pretrained(
                  """huggyllama/llama-7b"""  ,			truncation_side="""left"""  ,			bos_token="""</s>"""  ,			unk_token="""</s>"""	)
              tokenizer.add_special_tokens({"""pad_token""": """[PAD]"""}	)

       __lowerCamelCase,   __lowerCamelCase   =						get_blipa_config(A__	)
       __lowerCamelCase   =						InstructBlipForConditionalGeneration(A__	).eval()

       __lowerCamelCase   =						{
           """instructblip-vicuna-7b""": ("""blip2_vicuna_instruct""", """vicuna7b"""),
           """instructblip-vicuna-13b""": ("""blip2_vicuna_instruct""", """vicuna13b"""),
           """instructblip-flan-t5-xl""": ("""blip2_t5_instruct""", """flant5xl"""),
           """instructblip-flan-t5-xxl""": ("""blip2_t5_instruct""", """flant5xxl"""),
       }

       __lowerCamelCase,   __lowerCamelCase   =						model_name_to_original[model_name]

       # load original model
       print("""Loading original model..."""	)
       __lowerCamelCase   =						"""cuda:1""" if torch.cuda.is_available() else """cpu"""
       __lowerCamelCase   =						"""cuda:2""" if torch.cuda.is_available() else """cpu"""
       __lowerCamelCase,   __lowerCamelCase,   __lowerCamelCase   =						load_model_and_preprocess(
           name=A__  ,			model_type=A__  ,			is_eval=A__  ,			device=A__	)
       original_model.eval()
       print("""Done!"""	)

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

       # read in qv biases
       read_in_q_v_bias(A__  ,			A__	)

       # note: weights get loaded in torch.float32 by default
       hf_model.load_state_dict(A__  ,			strict=A__	)

       __lowerCamelCase   =						load_demo_image()
       __lowerCamelCase   =						"""What is unusual about this image?"""

       # create processor
       __lowerCamelCase   =						BlipImageProcessor(
           size={"""height""": image_size, """width""": image_size}  ,			image_mean=A__  ,			image_std=A__	)
       __lowerCamelCase   =						InstructBlipProcessor(
           image_processor=A__  ,			tokenizer=A__  ,			qformer_tokenizer=A__  ,			)
       __lowerCamelCase   =						processor(images=A__  ,			text=A__  ,			return_tensors="""pt"""	).to(A__	)

       # make sure processor creates exact same pixel values
       __lowerCamelCase   =						vis_processors["""eval"""](A__	).unsqueeze(0	).to(A__	)
       __lowerCamelCase   =						inputs.pixel_values
       assert torch.allclose(original_pixel_values.to(pixel_values.device	)  ,			A__	)

       original_model.to(A__	)
       hf_model.to(A__	)
       with torch.no_grad():
              if "vicuna" in model_name:
                     __lowerCamelCase   =						original_model({"""image""": original_pixel_values, """text_input""": [prompt]}	).logits
                     __lowerCamelCase   =						hf_model(**A__	).logits
              else:
                     __lowerCamelCase   =						original_model(
                         {"""image""": original_pixel_values, """text_input""": [prompt], """text_output""": ["""\n"""]}	).logits
                     __lowerCamelCase   =						tokenizer("""\n"""  ,			return_tensors="""pt"""	).input_ids.to(A__	)
                     __lowerCamelCase   =						label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id  ,			-100	)
                     __lowerCamelCase   =						hf_model(**A__  ,			labels=A__	).logits

       print("""First values of original logits:"""  ,			original_logits[0, :3, :3]	)
       print("""First values of HF logits:"""  ,			logits[0, :3, :3]	)

       # assert values
       assert original_logits.shape == logits.shape
       __lowerCamelCase   =						1E-4 if """vicuna""" in model_name else 1E-5
       assert torch.allclose(original_logits.to(logits.device	)  ,			A__  ,			atol=A__	)
       print("""Looks ok!"""	)

       print("""Generating with original model..."""	)
       __lowerCamelCase   =						original_model.generate({"""image""": original_pixel_values, """prompt""": prompt}  ,			num_beams=5	)

       # important: we need to cast the weights of the HF model to the appropriate type
       print("""Generating with HF model..."""	)
       __lowerCamelCase   =						hf_model.generate(
           **A__  ,			do_sample=A__  ,			num_beams=5  ,			max_length=256  ,			min_length=1  ,			top_p=0.9  ,			repetition_penalty=1.5  ,			length_penalty=1.0  ,			temperature=1  ,			)
       if "vicuna" in model_name:
              # convert output id 0 to 2 (eos_token_id)
              # TODO add this in the generate method?
              __lowerCamelCase   =						2
       print("""Original generation:"""  ,			A__	)
       __lowerCamelCase   =						processor.batch_decode(A__  ,			skip_special_tokens=A__	)
       __lowerCamelCase   =						[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'Salesforce/{model_name}'	)
              hf_model.push_to_hub(f'Salesforce/{model_name}'	)


if __name__ == "__main__":
       UpperCAmelCase_									=	argparse.ArgumentParser()
       UpperCAmelCase_									=	[
           'instructblip-vicuna-7b',
           'instructblip-vicuna-13b',
           'instructblip-flan-t5-xl',
           'instructblip-flan-t5-xxl',
       ]
       parser.add_argument(
           '--model_name',
           default='instructblip-flan-t5-xl',
           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_									=	parser.parse_args()

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







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


UpperCAmelCase_									=	{
    'configuration_vision_encoder_decoder': ['VisionEncoderDecoderConfig', 'VisionEncoderDecoderOnnxConfig']
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCAmelCase_									=	['VisionEncoderDecoderModel']

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCAmelCase_									=	['TFVisionEncoderDecoderModel']

try:
       if not is_flax_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCAmelCase_									=	['FlaxVisionEncoderDecoderModel']

if TYPE_CHECKING:
       from .configuration_vision_encoder_decoder import VisionEncoderDecoderConfig, VisionEncoderDecoderOnnxConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_vision_encoder_decoder import VisionEncoderDecoderModel

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_vision_encoder_decoder import TFVisionEncoderDecoderModel

       try:
              if not is_flax_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_flax_vision_encoder_decoder import FlaxVisionEncoderDecoderModel

else:
       import sys

       UpperCAmelCase_									=	_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)

