


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       a_ (       ):
       '''simple docstring'''





       _lowerCamelCase			:						List[Any]            ='https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
       _lowerCamelCase			:						List[Any]            =Image.open(requests.get(SCREAMING_SNAKE_CASE__						,     stream=SCREAMING_SNAKE_CASE__	).raw	).convert('RGB'	)

       return image




def       a_ (       SCREAMING_SNAKE_CASE__   : List[Any]	):
       '''simple docstring'''





       _lowerCamelCase			:						Union[str, 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       a_ (       SCREAMING_SNAKE_CASE__   : Union[str, Any]						,     SCREAMING_SNAKE_CASE__   : List[str]						,     SCREAMING_SNAKE_CASE__   : Tuple	):
       '''simple docstring'''





       _lowerCamelCase			:						Optional[int]            =dct.pop(SCREAMING_SNAKE_CASE__	)
       _lowerCamelCase			:						List[Any]            =val




def       a_ (       SCREAMING_SNAKE_CASE__   : List[str]						,     SCREAMING_SNAKE_CASE__   : Tuple	):
       '''simple docstring'''





       for i in range(config.vision_config.num_hidden_layers	):
              # read in original q and v biases
              _lowerCamelCase			:						List[Any]            =state_dict.pop(F'''visual_encoder.blocks.{i}.attn.q_bias'''	)
              _lowerCamelCase			:						Optional[Any]            =state_dict.pop(F'''visual_encoder.blocks.{i}.attn.v_bias'''	)

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




def       a_ (       SCREAMING_SNAKE_CASE__   : Any						,     SCREAMING_SNAKE_CASE__   : int	):
       '''simple docstring'''





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

       _lowerCamelCase			:						Dict            =BlipaConfig(vision_config=SCREAMING_SNAKE_CASE__						,     text_config=SCREAMING_SNAKE_CASE__	)

       return config, image_size




@torch.no_grad()
def       a_ (       SCREAMING_SNAKE_CASE__   : Union[str, Any]						,     SCREAMING_SNAKE_CASE__   : Optional[Any]=None						,     SCREAMING_SNAKE_CASE__   : Tuple=False	):
       '''simple docstring'''





       _lowerCamelCase			:						str            =(
           AutoTokenizer.from_pretrained('facebook/opt-2.7b'	)
           if 'opt' in model_name
           else AutoTokenizer.from_pretrained('google/flan-t5-xl'	)
       )
       _lowerCamelCase			:						Any            =tokenizer('\n'						,     add_special_tokens=SCREAMING_SNAKE_CASE__	).input_ids[0]
       _lowerCamelCase      ,	_lowerCamelCase			:						Tuple            =get_blipa_config(SCREAMING_SNAKE_CASE__						,     eos_token_id=SCREAMING_SNAKE_CASE__	)

       _lowerCamelCase			:						Any            =BlipaForConditionalGeneration(SCREAMING_SNAKE_CASE__	).eval()

       _lowerCamelCase			:						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'),
       }

       _lowerCamelCase      ,	_lowerCamelCase			:						Any            =model_name_to_original[model_name]

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

       # update state dict keys
       _lowerCamelCase			:						Tuple            =original_model.state_dict()
       _lowerCamelCase			:						Any            =create_rename_keys(SCREAMING_SNAKE_CASE__	)
       for src, dest in rename_keys:
              rename_key(SCREAMING_SNAKE_CASE__						,     SCREAMING_SNAKE_CASE__						,     SCREAMING_SNAKE_CASE__	)

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

       # read in qv biases
       read_in_q_v_bias(SCREAMING_SNAKE_CASE__						,     SCREAMING_SNAKE_CASE__	)

       _lowerCamelCase      ,	_lowerCamelCase			:						Tuple            =hf_model.load_state_dict(SCREAMING_SNAKE_CASE__						,     strict=SCREAMING_SNAKE_CASE__	)
       assert len(SCREAMING_SNAKE_CASE__	) == 0
       assert unexpected_keys == ["qformer.embeddings.position_ids"]

       _lowerCamelCase			:						Any            =load_demo_image()
       _lowerCamelCase			:						Any            =vis_processors['eval'](SCREAMING_SNAKE_CASE__	).unsqueeze(0	).to(SCREAMING_SNAKE_CASE__	)
       _lowerCamelCase			:						str            =tokenizer(['\n']						,     return_tensors='pt'	).input_ids.to(SCREAMING_SNAKE_CASE__	)

       # create processor
       _lowerCamelCase			:						str            =BlipImageProcessor(
           size={'height': image_size, 'width': image_size}						,     image_mean=SCREAMING_SNAKE_CASE__						,     image_std=SCREAMING_SNAKE_CASE__	)
       _lowerCamelCase			:						Tuple            =BlipaProcessor(image_processor=SCREAMING_SNAKE_CASE__						,     tokenizer=SCREAMING_SNAKE_CASE__	)
       _lowerCamelCase			:						Optional[int]            =processor(images=SCREAMING_SNAKE_CASE__						,     return_tensors='pt'	).pixel_values.to(SCREAMING_SNAKE_CASE__	)

       # make sure processor creates exact same pixel values
       assert torch.allclose(SCREAMING_SNAKE_CASE__						,     SCREAMING_SNAKE_CASE__	)

       original_model.to(SCREAMING_SNAKE_CASE__	)
       hf_model.to(SCREAMING_SNAKE_CASE__	)
       with torch.no_grad():
              if "opt" in model_name:
                     _lowerCamelCase			:						int            =original_model({'image': original_pixel_values, 'text_input': ['']}	).logits
                     _lowerCamelCase			:						List[Any]            =hf_model(SCREAMING_SNAKE_CASE__						,     SCREAMING_SNAKE_CASE__	).logits
              else:
                     _lowerCamelCase			:						Tuple            =original_model(
                         {'image': original_pixel_values, 'text_input': ['\n'], 'text_output': ['\n']}	).logits
                     _lowerCamelCase			:						Dict            =input_ids.masked_fill(input_ids == tokenizer.pad_token_id						,     -100	)
                     _lowerCamelCase			:						int            =hf_model(SCREAMING_SNAKE_CASE__						,     SCREAMING_SNAKE_CASE__						,     labels=SCREAMING_SNAKE_CASE__	).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			:						Any            =torch.tensor(
                  [[-41.58_50, -4.44_40, -8.99_22], [-47.43_22, -5.91_43, -1.73_40]]						,     device=SCREAMING_SNAKE_CASE__	)
              assert torch.allclose(logits[0, :3, :3]						,     SCREAMING_SNAKE_CASE__						,     atol=1e-4	)
       elif model_name == "blip2-flan-t5-xl-coco":
              _lowerCamelCase			:						int            =torch.tensor(
                  [[-57.01_09, -9.89_67, -12.62_80], [-68.65_78, -12.71_91, -10.50_65]]						,     device=SCREAMING_SNAKE_CASE__	)
       else:
              # cast to same type
              _lowerCamelCase			:						Dict            =logits.dtype
              assert torch.allclose(original_logits.to(SCREAMING_SNAKE_CASE__	)						,     SCREAMING_SNAKE_CASE__						,     atol=1e-2	)
       print('Looks ok!'	)

       print('Generating a caption...'	)
       _lowerCamelCase			:						Union[str, Any]            =''
       _lowerCamelCase			:						Optional[Any]            =tokenizer(SCREAMING_SNAKE_CASE__						,     return_tensors='pt'	).input_ids.to(SCREAMING_SNAKE_CASE__	)

       _lowerCamelCase			:						Optional[Any]            =original_model.generate({'image': original_pixel_values}	)
       _lowerCamelCase			:						List[str]            =hf_model.generate(
           SCREAMING_SNAKE_CASE__						,     SCREAMING_SNAKE_CASE__						,     do_sample=SCREAMING_SNAKE_CASE__						,     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:'						,     SCREAMING_SNAKE_CASE__	)
       _lowerCamelCase			:						Dict            =input_ids.shape[1]
       _lowerCamelCase			:						List[str]            =processor.batch_decode(outputs[:, prompt_length:]						,     skip_special_tokens=SCREAMING_SNAKE_CASE__	)
       _lowerCamelCase			:						List[Any]            =[text.strip() for text in output_text]
       print('HF generation:'						,     SCREAMING_SNAKE_CASE__	)

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

       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		   =			argparse.ArgumentParser()
 lowerCamelCase		   =			[
     '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		   =			parser.parse_args()

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






def       a_ (       SCREAMING_SNAKE_CASE__   : bytes	):
       '''simple docstring'''





       return "".join([hex(SCREAMING_SNAKE_CASE__	)[2:].zfill(2	).upper() for byte in list(SCREAMING_SNAKE_CASE__	)]	)




def       a_ (       SCREAMING_SNAKE_CASE__   : str	):
       '''simple docstring'''





       if (len(SCREAMING_SNAKE_CASE__	) % 2) != 0:
              raise ValueError(
                  'Base16 encoded data is invalid:\nData does not have an even number of hex digits.'	)
       # Check the character set - the standard base16 alphabet
       # is uppercase according to RFC3548 section 6
       if not set(SCREAMING_SNAKE_CASE__	) <= set('0123456789ABCDEF'	):
              raise ValueError(
                  'Base16 encoded data is invalid:\nData is not uppercase hex or it contains invalid characters.'	)
       # For every two hexadecimal digits (= a byte), turn it into an integer.
       # Then, string the result together into bytes, and return it.
       return bytes(int(data[i] + data[i + 1]						,     16	) for i in range(0						,     len(SCREAMING_SNAKE_CASE__	)						,     2	)	)


if __name__ == "__main__":
 import doctest

 doctest.testmod()


