


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 	_a     ( ):







 """simple docstring"""



 UpperCamelCase__    :			Dict      				=      '''https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'''
 UpperCamelCase__    :			Union[str, Any]      				=      Image.open(requests.get(snake_case_			,       stream=snake_case_   ).raw   ).convert('''RGB'''   )

 return image

def 	_a     ( SCREAMING_SNAKE_CASE		: List[str]   ):







 """simple docstring"""



 UpperCamelCase__    :			Tuple      				=      []
 # 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 	_a     ( SCREAMING_SNAKE_CASE		: Dict			,       SCREAMING_SNAKE_CASE		: List[Any]			,       SCREAMING_SNAKE_CASE		: Tuple   ):







 """simple docstring"""



 UpperCamelCase__    :			Union[str, Any]      				=      dct.pop(snake_case_   )
 UpperCamelCase__    :			str      				=      val

def 	_a     ( SCREAMING_SNAKE_CASE		: int			,       SCREAMING_SNAKE_CASE		: Tuple   ):







 """simple docstring"""



 for i in range(config.vision_config.num_hidden_layers   ):
  # read in original q and v biases
  UpperCamelCase__    :			int      				=      state_dict.pop(F"visual_encoder.blocks.{i}.attn.q_bias"   )
  UpperCamelCase__    :			str      				=      state_dict.pop(F"visual_encoder.blocks.{i}.attn.v_bias"   )

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

def 	_a     ( SCREAMING_SNAKE_CASE		: str   ):







 """simple docstring"""



 UpperCamelCase__    :			str      				=      364 if '''coco''' in model_name else 224
 UpperCamelCase__    :			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:
  UpperCamelCase__    :			Optional[Any]      				=      TaConfig.from_pretrained('''google/flan-t5-xl'''			,       dense_act_fn='''gelu'''			,       bos_token_id=1   ).to_dict()
 elif "t5-xxl" in model_name:
  UpperCamelCase__    :			Any      				=      TaConfig.from_pretrained('''google/flan-t5-xxl'''			,       dense_act_fn='''gelu'''			,       bos_token_id=1   ).to_dict()
 elif "vicuna-7b" in model_name:
  UpperCamelCase__    :			Union[str, Any]      				=      LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''			,       vocab_size=32001   ).to_dict()
 elif "vicuna-13b" in model_name:
  UpperCamelCase__    :			Union[str, Any]      				=      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
 UpperCamelCase__    :			Tuple      				=      InstructBlipQFormerConfig(vocab_size=30523   ).to_dict()
 UpperCamelCase__    :			Optional[Any]      				=      InstructBlipConfig(vision_config=snake_case_			,       text_config=snake_case_			,       qformer_config=snake_case_   )

 return config, image_size

@torch.no_grad()
def 	_a     ( SCREAMING_SNAKE_CASE		: Optional[int]			,       SCREAMING_SNAKE_CASE		: Dict=None			,       SCREAMING_SNAKE_CASE		: str=False   ):







 """simple docstring"""



 UpperCamelCase__    :			List[str]      				=      AutoTokenizer.from_pretrained('''bert-base-uncased'''			,       truncation_side='''left'''   )
 qformer_tokenizer.add_special_tokens({'''bos_token''': '''[DEC]'''}   )

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

 UpperCamelCase__				, UpperCamelCase__    :			Dict      				=      get_blipa_config(snake_case_   )
 UpperCamelCase__    :			List[Any]      				=      InstructBlipForConditionalGeneration(snake_case_   ).eval()

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

 UpperCamelCase__				, UpperCamelCase__    :			Optional[int]      				=      model_name_to_original[model_name]

 # load original model
 print('''Loading original model...'''   )
 UpperCamelCase__    :			Any      				=      '''cuda:1''' if torch.cuda.is_available() else '''cpu'''
 UpperCamelCase__    :			Optional[int]      				=      '''cuda:2''' if torch.cuda.is_available() else '''cpu'''
 UpperCamelCase__				, UpperCamelCase__				, UpperCamelCase__    :			List[str]      				=      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
 UpperCamelCase__    :			Tuple      				=      original_model.state_dict()
 UpperCamelCase__    :			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():
  UpperCamelCase__    :			Optional[int]      				=      state_dict.pop(snake_case_   )
  if key.startswith('''Qformer.bert'''   ):
   UpperCamelCase__    :			Tuple      				=      key.replace('''Qformer.bert'''			,       '''qformer'''   )
  if "attention.self" in key:
   UpperCamelCase__    :			int      				=      key.replace('''self'''			,       '''attention'''   )
  if "llm_proj" in key:
   UpperCamelCase__    :			Union[str, Any]      				=      key.replace('''llm_proj'''			,       '''language_projection'''   )
  if "t5_proj" in key:
   UpperCamelCase__    :			Optional[int]      				=      key.replace('''t5_proj'''			,       '''language_projection'''   )
  if key.startswith('''llm_model'''   ):
   UpperCamelCase__    :			List[Any]      				=      key.replace('''llm_model'''			,       '''language_model'''   )
  if key.startswith('''t5'''   ):
   UpperCamelCase__    :			List[Any]      				=      key.replace('''t5'''			,       '''language'''   )
  UpperCamelCase__    :			Any      				=      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_   )

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

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

 # make sure processor creates exact same pixel values
 UpperCamelCase__    :			Optional[Any]      				=      vis_processors['''eval'''](snake_case_   ).unsqueeze(0   ).to(snake_case_   )
 UpperCamelCase__    :			Tuple      				=      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:
   UpperCamelCase__    :			List[str]      				=      original_model({'''image''': original_pixel_values, '''text_input''': [prompt]}   ).logits
   UpperCamelCase__    :			str      				=      hf_model(**snake_case_   ).logits
  else:
   UpperCamelCase__    :			Optional[int]      				=      original_model(
       {'''image''': original_pixel_values, '''text_input''': [prompt], '''text_output''': ['''\n''']}   ).logits
   UpperCamelCase__    :			Union[str, Any]      				=      tokenizer('''\n'''			,       return_tensors='''pt'''   ).input_ids.to(snake_case_   )
   UpperCamelCase__    :			Optional[Any]      				=      label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id			,       -100   )
   UpperCamelCase__    :			Optional[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
 UpperCamelCase__    :			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...'''   )
 UpperCamelCase__    :			Optional[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...'''   )
 UpperCamelCase__    :			Optional[Any]      				=      hf_model.generate(
     **snake_case_			,       do_sample=snake_case_			,       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?
  UpperCamelCase__    :			Dict      				=      2
 print('''Original generation:'''			,       snake_case_   )
 UpperCamelCase__    :			List[Any]      				=      processor.batch_decode(snake_case_			,       skip_special_tokens=snake_case_   )
 UpperCamelCase__    :			Optional[int]      				=      [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__":
       __UpperCamelCase				:				Any  				= argparse.ArgumentParser()
       __UpperCamelCase				:				str  				= [
           '''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				:				Dict  				= parser.parse_args()

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



'''simple docstring'''

import numpy as np
import torch
from torch.utils.data import DataLoader

from accelerate.utils.dataclasses import DistributedType







class 				_lowercase     :


 '''simple docstring'''






 def __init__(      self   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	List[str]=2      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]=3      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]=64      ,       SCREAMING_SNAKE_CASE__   :	int=None							)  ->    Optional[int]:
  __lowerCAmelCase											=						np.random.default_rng(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						length
  __lowerCAmelCase											=						rng.normal(size=(length,)							).astype(np.floataa							)
  __lowerCAmelCase											=						a * self.x + b + rng.normal(scale=0.1      ,       size=(length,)							).astype(np.floataa							)
 def __len__(      self   :	Union[str, Any]							)  ->    Optional[Any]:
  return self.length





 def __getitem__(      self   :	str      ,       SCREAMING_SNAKE_CASE__   :	Dict							)  ->    Optional[int]:
  return {"x": self.x[i], "y": self.y[i]}







class 				_lowercase     (				torch.nn.Module			):


 '''simple docstring'''






 def __init__(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]=0      ,       SCREAMING_SNAKE_CASE__   :	Tuple=0      ,       SCREAMING_SNAKE_CASE__   :	Tuple=False							)  ->    Any:
  super().__init__()
  __lowerCAmelCase											=						torch.nn.Parameter(torch.tensor([2, 3]							).float()							)
  __lowerCAmelCase											=						torch.nn.Parameter(torch.tensor([2, 3]							).float()							)
  __lowerCAmelCase											=						True





 def a					(      self   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	List[Any]=None							)  ->    str:
  if self.first_batch:
   print(f"""Model dtype: {self.a.dtype}, {self.b.dtype}. Input dtype: {x.dtype}"""							)
   __lowerCAmelCase											=						False
  return x * self.a[0] + self.b[0]







class 				_lowercase     (				torch.nn.Module			):


 '''simple docstring'''






 def __init__(      self   :	Any      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]=0      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]=0      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]=False							)  ->    Optional[Any]:
  super().__init__()
  __lowerCAmelCase											=						torch.nn.Parameter(torch.tensor(SCREAMING_SNAKE_CASE__							).float()							)
  __lowerCAmelCase											=						torch.nn.Parameter(torch.tensor(SCREAMING_SNAKE_CASE__							).float()							)
  __lowerCAmelCase											=						True





 def a					(      self   :	List[str]      ,       SCREAMING_SNAKE_CASE__   :	Any=None							)  ->    int:
  if self.first_batch:
   print(f"""Model dtype: {self.a.dtype}, {self.b.dtype}. Input dtype: {x.dtype}"""							)
   __lowerCAmelCase											=						False
  return x * self.a + self.b







def      UpperCamelCase_						(					snake_case_       :       List[str]						,		snake_case_       :       int = 16      )					->     int:






 '''simple docstring'''

 from datasets import load_dataset
 from transformers import AutoTokenizer

 __lowerCAmelCase											=						AutoTokenizer.from_pretrained("""bert-base-cased"""      )
 __lowerCAmelCase											=						{"""train""": """tests/test_samples/MRPC/train.csv""", """validation""": """tests/test_samples/MRPC/dev.csv"""}
 __lowerCAmelCase											=						load_dataset("""csv"""						,		data_files=snake_case_      )
 __lowerCAmelCase											=						datasets["""train"""].unique("""label"""      )

 __lowerCAmelCase											=						{v: i for i, v in enumerate(snake_case_      )}

 def tokenize_function(snake_case_       :       Optional[int]      ):
  # max_length=None => use the model max length (it's actually the default)
  __lowerCAmelCase											=						tokenizer(
      examples["""sentence1"""]						,		examples["""sentence2"""]						,		truncation=snake_case_						,		max_length=snake_case_						,		padding="""max_length"""      )
  if "label" in examples:
   __lowerCAmelCase											=						[label_to_id[l] for l in examples["""label"""]]
  return outputs

 # Apply the method we just defined to all the examples in all the splits of the dataset
 __lowerCAmelCase											=						datasets.map(
     snake_case_						,		batched=snake_case_						,		remove_columns=["""sentence1""", """sentence2""", """label"""]						,		)

 def collate_fn(snake_case_       :       List[Any]      ):
  # On TPU it's best to pad everything to the same length or training will be very slow.
  if accelerator.distributed_type == DistributedType.TPU:
   return tokenizer.pad(snake_case_						,		padding="""max_length"""						,		max_length=1_28						,		return_tensors="""pt"""      )
  return tokenizer.pad(snake_case_						,		padding="""longest"""						,		return_tensors="""pt"""      )

 # Instantiate dataloaders.
 __lowerCAmelCase											=						DataLoader(tokenized_datasets["""train"""]						,		shuffle=snake_case_						,		collate_fn=snake_case_						,		batch_size=2      )
 __lowerCAmelCase											=						DataLoader(tokenized_datasets["""validation"""]						,		shuffle=snake_case_						,		collate_fn=snake_case_						,		batch_size=1      )

 return train_dataloader, eval_dataloader






