





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

     return image






def        _lowerCAmelCase			(  __snake_case  : int			)     ->     List[Any]:
     __A						: 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        _lowerCAmelCase			(  __snake_case  : Optional[int] ,							__snake_case  : Optional[Any] ,							__snake_case  : Optional[int]			)     ->     List[str]:
     __A						: List[Any]      =       dct.pop(__snake_case			)
     __A						: Dict      =       val






def        _lowerCAmelCase			(  __snake_case  : List[Any] ,							__snake_case  : Tuple			)     ->     Any:
     for i in range(config.vision_config.num_hidden_layers			):
          # read in original q and v biases
          __A						: int      =       state_dict.pop(f'visual_encoder.blocks.{i}.attn.q_bias'			)
          __A						: Any      =       state_dict.pop(f'visual_encoder.blocks.{i}.attn.v_bias'			)

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






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

     return config, image_size






@torch.no_grad()
def        _lowerCAmelCase			(  __snake_case  : Dict ,							__snake_case  : Optional[Any]=None ,							__snake_case  : List[Any]=False			)     ->     int:
     __A						: Union[str, Any]      =       AutoTokenizer.from_pretrained('bert-base-uncased' ,							truncation_side='left'			)
     qformer_tokenizer.add_special_tokens({'bos_token': '[DEC]'}			)

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

     __A ,__A						: str      =       get_blipa_config(__snake_case			)
     __A						: Optional[Any]      =       InstructBlipForConditionalGeneration(__snake_case			).eval()

     __A						: int      =       {
         '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'),
     }

     __A ,__A						: int      =       model_name_to_original[model_name]

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

     # update state dict keys
     __A						: Optional[int]      =       original_model.state_dict()
     __A						: List[str]      =       create_rename_keys(__snake_case			)
     for src, dest in rename_keys:
          rename_key(__snake_case ,							__snake_case ,							__snake_case			)

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

     # read in qv biases
     read_in_q_v_bias(__snake_case ,							__snake_case			)

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

     __A						: List[str]      =       load_demo_image()
     __A						: Optional[int]      =       'What is unusual about this image?'

     # create processor
     __A						: Any      =       BlipImageProcessor(
         size={'height': image_size, 'width': image_size} ,							image_mean=__snake_case ,							image_std=__snake_case			)
     __A						: Tuple      =       InstructBlipProcessor(
         image_processor=__snake_case ,							tokenizer=__snake_case ,							qformer_tokenizer=__snake_case ,							)
     __A						: int      =       processor(images=__snake_case ,							text=__snake_case ,							return_tensors='pt'			).to(__snake_case			)

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

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

     print('Generating with original model...'			)
     __A						: Tuple      =       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...'			)
     __A						: Any      =       hf_model.generate(
         **__snake_case ,							do_sample=__snake_case ,							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?
          __A						: str      =       2
     print('Original generation:' ,							__snake_case			)
     __A						: Optional[Any]      =       processor.batch_decode(__snake_case ,							skip_special_tokens=__snake_case			)
     __A						: Tuple      =       [text.strip() for text in output_text]
     print('HF generation:' ,							__snake_case			)

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

     if push_to_hub:
          processor.push_to_hub(f'Salesforce/{model_name}'			)
          hf_model.push_to_hub(f'Salesforce/{model_name}'			)


if __name__ == "__main__":
  lowercase__    :      Optional[int]						=    argparse.ArgumentParser()
  lowercase__    :      Any						=    [
      '''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''',
  )

  lowercase__    :      Tuple						=    parser.parse_args()

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






'''simple docstring'''




lowercase__    :      Dict						=    {
    '''Pillow''': '''Pillow''',
    '''accelerate''': '''accelerate>=0.11.0''',
    '''compel''': '''compel==0.1.8''',
    '''black''': '''black~=23.1''',
    '''datasets''': '''datasets''',
    '''filelock''': '''filelock''',
    '''flax''': '''flax>=0.4.1''',
    '''hf-doc-builder''': '''hf-doc-builder>=0.3.0''',
    '''huggingface-hub''': '''huggingface-hub>=0.13.2''',
    '''requests-mock''': '''requests-mock==1.10.0''',
    '''importlib_metadata''': '''importlib_metadata''',
    '''invisible-watermark''': '''invisible-watermark''',
    '''isort''': '''isort>=5.5.4''',
    '''jax''': '''jax>=0.2.8,!=0.3.2''',
    '''jaxlib''': '''jaxlib>=0.1.65''',
    '''Jinja2''': '''Jinja2''',
    '''k-diffusion''': '''k-diffusion>=0.0.12''',
    '''torchsde''': '''torchsde''',
    '''note_seq''': '''note_seq''',
    '''librosa''': '''librosa''',
    '''numpy''': '''numpy''',
    '''omegaconf''': '''omegaconf''',
    '''parameterized''': '''parameterized''',
    '''protobuf''': '''protobuf>=3.20.3,<4''',
    '''pytest''': '''pytest''',
    '''pytest-timeout''': '''pytest-timeout''',
    '''pytest-xdist''': '''pytest-xdist''',
    '''ruff''': '''ruff>=0.0.241''',
    '''safetensors''': '''safetensors''',
    '''sentencepiece''': '''sentencepiece>=0.1.91,!=0.1.92''',
    '''scipy''': '''scipy''',
    '''onnx''': '''onnx''',
    '''regex''': '''regex!=2019.12.17''',
    '''requests''': '''requests''',
    '''tensorboard''': '''tensorboard''',
    '''torch''': '''torch>=1.4''',
    '''torchvision''': '''torchvision''',
    '''transformers''': '''transformers>=4.25.1''',
    '''urllib3''': '''urllib3<=2.0.0''',
}