





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  __UpperCAmelCase(   ):
       _lowerCamelCase				:  List[Any]						=		'''https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'''
       _lowerCamelCase				:  Tuple						=		Image.open(requests.get(lowercase_			,     stream=lowercase_   ).raw   ).convert('''RGB'''   )

       return image




def  __UpperCAmelCase(   lowercase_   ):
       _lowerCamelCase				:  Optional[int]						=		[]
       # 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  __UpperCAmelCase(   lowercase_			,     lowercase_			,     lowercase_   ):
       _lowerCamelCase				:  Optional[int]						=		dct.pop(lowercase_   )
       _lowerCamelCase				:  Optional[Any]						=		val




def  __UpperCAmelCase(   lowercase_			,     lowercase_   ):
       for i in range(config.vision_config.num_hidden_layers   ):
              # read in original q and v biases
              _lowerCamelCase				:  Tuple						=		state_dict.pop(F"""visual_encoder.blocks.{i}.attn.q_bias"""   )
              _lowerCamelCase				:  Union[str, Any]						=		state_dict.pop(F"""visual_encoder.blocks.{i}.attn.v_bias"""   )

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




def  __UpperCAmelCase(   lowercase_   ):
       _lowerCamelCase				:  Optional[int]						=		3_64 if '''coco''' in model_name else 2_24
       _lowerCamelCase				:  Tuple						=		InstructBlipVisionConfig(image_size=lowercase_   ).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				:  Tuple						=		TaConfig.from_pretrained('''google/flan-t5-xl'''			,     dense_act_fn='''gelu'''			,     bos_token_id=1   ).to_dict()
       elif "t5-xxl" in model_name:
              _lowerCamelCase				:  str						=		TaConfig.from_pretrained('''google/flan-t5-xxl'''			,     dense_act_fn='''gelu'''			,     bos_token_id=1   ).to_dict()
       elif "vicuna-7b" in model_name:
              _lowerCamelCase				:  List[Any]						=		LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''			,     vocab_size=3_20_01   ).to_dict()
       elif "vicuna-13b" in model_name:
              _lowerCamelCase				:  Optional[Any]						=		LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf'''			,     vocab_size=3_20_01   ).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				:  Tuple						=		InstructBlipQFormerConfig(vocab_size=3_05_23   ).to_dict()
       _lowerCamelCase				:  Dict						=		InstructBlipConfig(vision_config=lowercase_			,     text_config=lowercase_			,     qformer_config=lowercase_   )

       return config, image_size




@torch.no_grad()
def  __UpperCAmelCase(   lowercase_			,     lowercase_=None			,     lowercase_=False   ):
       _lowerCamelCase				:  Tuple						=		AutoTokenizer.from_pretrained('''bert-base-uncased'''			,     truncation_side='''left'''   )
       qformer_tokenizer.add_special_tokens({'''bos_token''': '''[DEC]'''}   )

       if "t5" in model_name:
              _lowerCamelCase				:  Any						=		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				:  List[Any]						=		LlamaTokenizerFast.from_pretrained(
                  '''huggyllama/llama-7b'''			,     truncation_side='''left'''			,     bos_token='''</s>'''			,     unk_token='''</s>'''   )
              tokenizer.add_special_tokens({'''pad_token''': '''[PAD]'''}   )

       _lowerCamelCase,	_lowerCamelCase				:  Optional[int]						=		get_blipa_config(lowercase_   )
       _lowerCamelCase				:  List[Any]						=		InstructBlipForConditionalGeneration(lowercase_   ).eval()

       _lowerCamelCase				:  str						=		{
           '''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				:  Any						=		model_name_to_original[model_name]

       # load original model
       print('''Loading original model...'''   )
       _lowerCamelCase				:  List[Any]						=		'''cuda:1''' if torch.cuda.is_available() else '''cpu'''
       _lowerCamelCase				:  Dict						=		'''cuda:2''' if torch.cuda.is_available() else '''cpu'''
       _lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase				:  int						=		load_model_and_preprocess(
           name=lowercase_			,     model_type=lowercase_			,     is_eval=lowercase_			,     device=lowercase_   )
       original_model.eval()
       print('''Done!'''   )

       # update state dict keys
       _lowerCamelCase				:  Optional[Any]						=		original_model.state_dict()
       _lowerCamelCase				:  Optional[int]						=		create_rename_keys(lowercase_   )
       for src, dest in rename_keys:
              rename_key(lowercase_			,     lowercase_			,     lowercase_   )

       # some keys can be renamed efficiently
       for key, val in state_dict.copy().items():
              _lowerCamelCase				:  Dict						=		state_dict.pop(lowercase_   )
              if key.startswith('''Qformer.bert'''   ):
                     _lowerCamelCase				:  int						=		key.replace('''Qformer.bert'''			,     '''qformer'''   )
              if "attention.self" in key:
                     _lowerCamelCase				:  Tuple						=		key.replace('''self'''			,     '''attention'''   )
              if "llm_proj" in key:
                     _lowerCamelCase				:  List[Any]						=		key.replace('''llm_proj'''			,     '''language_projection'''   )
              if "t5_proj" in key:
                     _lowerCamelCase				:  Optional[int]						=		key.replace('''t5_proj'''			,     '''language_projection'''   )
              if key.startswith('''llm_model'''   ):
                     _lowerCamelCase				:  Tuple						=		key.replace('''llm_model'''			,     '''language_model'''   )
              if key.startswith('''t5'''   ):
                     _lowerCamelCase				:  List[Any]						=		key.replace('''t5'''			,     '''language'''   )
              _lowerCamelCase				:  Union[str, Any]						=		val

       # read in qv biases
       read_in_q_v_bias(lowercase_			,     lowercase_   )

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

       _lowerCamelCase				:  List[Any]						=		load_demo_image()
       _lowerCamelCase				:  Optional[Any]						=		'''What is unusual about this image?'''

       # create processor
       _lowerCamelCase				:  List[str]						=		BlipImageProcessor(
           size={'''height''': image_size, '''width''': image_size}			,     image_mean=lowercase_			,     image_std=lowercase_   )
       _lowerCamelCase				:  Any						=		InstructBlipProcessor(
           image_processor=lowercase_			,     tokenizer=lowercase_			,     qformer_tokenizer=lowercase_			,     )
       _lowerCamelCase				:  Dict						=		processor(images=lowercase_			,     text=lowercase_			,     return_tensors='''pt'''   ).to(lowercase_   )

       # make sure processor creates exact same pixel values
       _lowerCamelCase				:  Any						=		vis_processors['''eval'''](lowercase_   ).unsqueeze(0   ).to(lowercase_   )
       _lowerCamelCase				:  List[str]						=		inputs.pixel_values
       assert torch.allclose(original_pixel_values.to(pixel_values.device   )			,     lowercase_   )

       original_model.to(lowercase_   )
       hf_model.to(lowercase_   )
       with torch.no_grad():
              if "vicuna" in model_name:
                     _lowerCamelCase				:  Optional[Any]						=		original_model({'''image''': original_pixel_values, '''text_input''': [prompt]}   ).logits
                     _lowerCamelCase				:  Union[str, Any]						=		hf_model(**lowercase_   ).logits
              else:
                     _lowerCamelCase				:  Dict						=		original_model(
                         {'''image''': original_pixel_values, '''text_input''': [prompt], '''text_output''': ['''\n''']}   ).logits
                     _lowerCamelCase				:  Union[str, Any]						=		tokenizer('''\n'''			,     return_tensors='''pt'''   ).input_ids.to(lowercase_   )
                     _lowerCamelCase				:  int						=		label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id			,     -1_00   )
                     _lowerCamelCase				:  List[Any]						=		hf_model(**lowercase_			,     labels=lowercase_   ).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				:  Tuple						=		1e-4 if '''vicuna''' in model_name else 1e-5
       assert torch.allclose(original_logits.to(logits.device   )			,     lowercase_			,     atol=lowercase_   )
       print('''Looks ok!'''   )

       print('''Generating with original model...'''   )
       _lowerCamelCase				:  int						=		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				:  Union[str, Any]						=		hf_model.generate(
           **lowercase_			,     do_sample=lowercase_			,     num_beams=5			,     max_length=2_56			,     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				:  Dict						=		2
       print('''Original generation:'''			,     lowercase_   )
       _lowerCamelCase				:  int						=		processor.batch_decode(lowercase_			,     skip_special_tokens=lowercase_   )
       _lowerCamelCase				:  Dict						=		[text.strip() for text in output_text]
       print('''HF generation:'''			,     lowercase_   )

       if pytorch_dump_folder_path is not None:
              processor.save_pretrained(lowercase_   )
              hf_model.save_pretrained(lowercase_   )

       if push_to_hub:
              processor.push_to_hub(F"""Salesforce/{model_name}"""   )
              hf_model.push_to_hub(F"""Salesforce/{model_name}"""   )


if __name__ == "__main__":
       _lowerCamelCase      							=   argparse.ArgumentParser()
       _lowerCamelCase      							=   [
           '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',
       )

       _lowerCamelCase      							=   parser.parse_args()

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








import copy
import random

from transformers import CLIPTokenizer







class       __A							(   lowerCamelCase__	):




      """simple docstring"""







      def __init__(		self   ,				*a__   ,				**a__):



             """simple docstring"""





             super().__init__(*a__   ,				**a__)
             _lowerCamelCase				:  Optional[Any]						=		{}


      def 					__snake_case      (		self   ,				a__   ,				*a__   ,				**a__):



             """simple docstring"""





             _lowerCamelCase				:  Optional[int]						=		super().add_tokens(a__   ,				*a__   ,				**a__)
             if num_added_tokens == 0:
                    raise ValueError(
                        F"""The tokenizer already contains the token {placeholder_token}. Please pass a different"""
                        ''' `placeholder_token` that is not already in the tokenizer.''')


      def 					__snake_case      (		self   ,				a__   ,				*a__   ,				a__=1   ,				**a__):



             """simple docstring"""





             _lowerCamelCase				:  Tuple						=		[]
             if num_vec_per_token == 1:
                    self.try_adding_tokens(a__   ,				*a__   ,				**a__)
                    output.append(a__)
             else:
                    _lowerCamelCase				:  Dict						=		[]
                    for i in range(a__):
                           _lowerCamelCase				:  Optional[Any]						=		placeholder_token + F"""_{i}"""
                           self.try_adding_tokens(a__   ,				*a__   ,				**a__)
                           output.append(a__)
        # handle cases where there is a new placeholder token that contains the current placeholder token but is larger
             for token in self.token_map:
                    if token in placeholder_token:
                           raise ValueError(
                               F"""The tokenizer already has placeholder token {token} that can get confused with"""
                               F""" {placeholder_token}keep placeholder tokens independent""")
             _lowerCamelCase				:  int						=		output


      def 					__snake_case      (		self   ,				a__   ,				a__=False   ,				a__=1.0):



             """simple docstring"""





             if isinstance(a__   ,				a__):
                    _lowerCamelCase				:  Tuple						=		[]
                    for i in range(len(a__)):
                           output.append(self.replace_placeholder_tokens_in_text(text[i]   ,				vector_shuffle=a__))
                    return output
             for placeholder_token in self.token_map:
                    if placeholder_token in text:
                           _lowerCamelCase				:  Union[str, Any]						=		self.token_map[placeholder_token]
                           _lowerCamelCase				:  Optional[Any]						=		tokens[: 1 + int(len(a__) * prop_tokens_to_load)]
                           if vector_shuffle:
                                  _lowerCamelCase				:  int						=		copy.copy(a__)
                                  random.shuffle(a__)
                           _lowerCamelCase				:  List[str]						=		text.replace(a__   ,				''' '''.join(a__))
             return text


      def __call__(		self   ,				a__   ,				*a__   ,				a__=False   ,				a__=1.0   ,				**a__):



             """simple docstring"""





             return super().__call__(
                 self.replace_placeholder_tokens_in_text(
                     a__   ,				vector_shuffle=a__   ,				prop_tokens_to_load=a__)   ,				*a__   ,				**a__   ,				)



      def 					__snake_case      (		self   ,				a__   ,				*a__   ,				a__=False   ,				a__=1.0   ,				**a__):



             """simple docstring"""





             return super().encode(
                 self.replace_placeholder_tokens_in_text(
                     a__   ,				vector_shuffle=a__   ,				prop_tokens_to_load=a__)   ,				*a__   ,				**a__   ,				)

