






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







      """simple docstring"""






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

      return image

def _a			(		_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	[]
      # 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			(		_snake_case				,	_snake_case				,	_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	dct.pop(_UpperCAmelCase							)
      UpperCAmelCase          =	val

def _a			(		_snake_case				,	_snake_case							):







      """simple docstring"""






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

            # next, set bias in the state dict
            UpperCAmelCase          =	torch.cat((q_bias, torch.zeros_like(_UpperCAmelCase				,	requires_grad=_UpperCAmelCase							), v_bias)							)
            UpperCAmelCase          =	qkv_bias

def _a			(		_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	364 if "coco" in model_name else 224
      UpperCAmelCase          =	InstructBlipVisionConfig(image_size=_UpperCAmelCase							).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          =	TaConfig.from_pretrained("""google/flan-t5-xl"""				,	dense_act_fn="""gelu"""				,	bos_token_id=1							).to_dict()
      elif "t5-xxl" in model_name:
            UpperCAmelCase          =	TaConfig.from_pretrained("""google/flan-t5-xxl"""				,	dense_act_fn="""gelu"""				,	bos_token_id=1							).to_dict()
      elif "vicuna-7b" in model_name:
            UpperCAmelCase          =	LlamaConfig.from_pretrained("""decapoda-research/llama-7b-hf"""				,	vocab_size=3_2001							).to_dict()
      elif "vicuna-13b" in model_name:
            UpperCAmelCase          =	LlamaConfig.from_pretrained("""decapoda-research/llama-13b-hf"""				,	vocab_size=3_2001							).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          =	InstructBlipQFormerConfig(vocab_size=3_0523							).to_dict()
      UpperCAmelCase          =	InstructBlipConfig(vision_config=_UpperCAmelCase				,	text_config=_UpperCAmelCase				,	qformer_config=_UpperCAmelCase							)

      return config, image_size

@torch.no_grad()
def _a			(		_snake_case				,	_snake_case=None				,	_snake_case=False							):







      """simple docstring"""






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

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

      UpperCAmelCase          =	get_blipa_config(_UpperCAmelCase							)
      UpperCAmelCase          =	InstructBlipForConditionalGeneration(_UpperCAmelCase							).eval()

      UpperCAmelCase          =	{
          "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          =	model_name_to_original[model_name]

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

      # update state dict keys
      UpperCAmelCase          =	original_model.state_dict()
      UpperCAmelCase          =	create_rename_keys(_UpperCAmelCase							)
      for src, dest in rename_keys:
            rename_key(_UpperCAmelCase				,	_UpperCAmelCase				,	_UpperCAmelCase							)

      # some keys can be renamed efficiently
      for key, val in state_dict.copy().items():
            UpperCAmelCase          =	state_dict.pop(_UpperCAmelCase							)
            if key.startswith("""Qformer.bert"""							):
                  UpperCAmelCase          =	key.replace("""Qformer.bert"""				,	"""qformer"""							)
            if "attention.self" in key:
                  UpperCAmelCase          =	key.replace("""self"""				,	"""attention"""							)
            if "llm_proj" in key:
                  UpperCAmelCase          =	key.replace("""llm_proj"""				,	"""language_projection"""							)
            if "t5_proj" in key:
                  UpperCAmelCase          =	key.replace("""t5_proj"""				,	"""language_projection"""							)
            if key.startswith("""llm_model"""							):
                  UpperCAmelCase          =	key.replace("""llm_model"""				,	"""language_model"""							)
            if key.startswith("""t5"""							):
                  UpperCAmelCase          =	key.replace("""t5"""				,	"""language"""							)
            UpperCAmelCase          =	val

      # read in qv biases
      read_in_q_v_bias(_UpperCAmelCase				,	_UpperCAmelCase							)

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

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

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

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

      original_model.to(_UpperCAmelCase							)
      hf_model.to(_UpperCAmelCase							)
      with torch.no_grad():
            if "vicuna" in model_name:
                  UpperCAmelCase          =	original_model({"""image""": original_pixel_values, """text_input""": [prompt]}							).logits
                  UpperCAmelCase          =	hf_model(**_UpperCAmelCase							).logits
            else:
                  UpperCAmelCase          =	original_model(
                      {"""image""": original_pixel_values, """text_input""": [prompt], """text_output""": ["""\n"""]}							).logits
                  UpperCAmelCase          =	tokenizer("""\n"""				,	return_tensors="""pt"""							).input_ids.to(_UpperCAmelCase							)
                  UpperCAmelCase          =	label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id				,	-100							)
                  UpperCAmelCase          =	hf_model(**_UpperCAmelCase				,	labels=_UpperCAmelCase							).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          =	1E-4 if "vicuna" in model_name else 1E-5
      assert torch.allclose(original_logits.to(logits.device							)				,	_UpperCAmelCase				,	atol=_UpperCAmelCase							)
      print("""Looks ok!"""							)

      print("""Generating with original model..."""							)
      UpperCAmelCase          =	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          =	hf_model.generate(
          **_UpperCAmelCase				,	do_sample=_UpperCAmelCase				,	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          =	2
      print("""Original generation:"""				,	_UpperCAmelCase							)
      UpperCAmelCase          =	processor.batch_decode(_UpperCAmelCase				,	skip_special_tokens=_UpperCAmelCase							)
      UpperCAmelCase          =	[text.strip() for text in output_text]
      print("""HF generation:"""				,	_UpperCAmelCase							)

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

      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)











"""simple docstring"""







import argparse

import torch

from transformers import (
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaForAudioFrameClassification,
    WavaVecaForSequenceClassification,
    WavaVecaForXVector,
    logging,
)


logging.set_verbosity_info()
_UpperCamelCase					      =  logging.get_logger(__name__)

def _a			(		_snake_case				,	_snake_case				,	_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	WavaVecaForSequenceClassification.from_pretrained(_snake_case				,	config=_snake_case							)
      UpperCAmelCase          =	downstream_dict["""projector.weight"""]
      UpperCAmelCase          =	downstream_dict["""projector.bias"""]
      UpperCAmelCase          =	downstream_dict["""model.post_net.linear.weight"""]
      UpperCAmelCase          =	downstream_dict["""model.post_net.linear.bias"""]
      return model

def _a			(		_snake_case				,	_snake_case				,	_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	WavaVecaForAudioFrameClassification.from_pretrained(_snake_case				,	config=_snake_case							)
      UpperCAmelCase          =	downstream_dict["""model.linear.weight"""]
      UpperCAmelCase          =	downstream_dict["""model.linear.bias"""]
      return model

def _a			(		_snake_case				,	_snake_case				,	_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	WavaVecaForXVector.from_pretrained(_snake_case				,	config=_snake_case							)
      UpperCAmelCase          =	downstream_dict["""connector.weight"""]
      UpperCAmelCase          =	downstream_dict["""connector.bias"""]
      for i, kernel_size in enumerate(hf_config.tdnn_kernel							):
            UpperCAmelCase          =	downstream_dict[
                F'''model.framelevel_feature_extractor.module.{i}.kernel.weight'''
            ]
            UpperCAmelCase          =	downstream_dict[F'''model.framelevel_feature_extractor.module.{i}.kernel.bias''']

      UpperCAmelCase          =	downstream_dict["""model.utterancelevel_feature_extractor.linear1.weight"""]
      UpperCAmelCase          =	downstream_dict["""model.utterancelevel_feature_extractor.linear1.bias"""]
      UpperCAmelCase          =	downstream_dict["""model.utterancelevel_feature_extractor.linear2.weight"""]
      UpperCAmelCase          =	downstream_dict["""model.utterancelevel_feature_extractor.linear2.bias"""]
      UpperCAmelCase          =	downstream_dict["""objective.W"""]
      return model

@torch.no_grad()
def _a			(		_snake_case				,	_snake_case				,	_snake_case				,	_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	torch.load(_snake_case				,	map_location="""cpu"""							)

      UpperCAmelCase          =	checkpoint["""Downstream"""]

      UpperCAmelCase          =	WavaVecaConfig.from_pretrained(_snake_case							)
      UpperCAmelCase          =	WavaVecaFeatureExtractor.from_pretrained(
          _snake_case				,	return_attention_mask=_snake_case				,	do_normalize=_snake_case							)

      UpperCAmelCase          =	hf_config.architectures[0]
      if arch.endswith("""ForSequenceClassification"""							):
            UpperCAmelCase          =	convert_classification(_snake_case				,	_snake_case				,	_snake_case							)
      elif arch.endswith("""ForAudioFrameClassification"""							):
            UpperCAmelCase          =	convert_diarization(_snake_case				,	_snake_case				,	_snake_case							)
      elif arch.endswith("""ForXVector"""							):
            UpperCAmelCase          =	convert_xvector(_snake_case				,	_snake_case				,	_snake_case							)
      else:
            raise NotImplementedError(F'''S3PRL weights conversion is not supported for {arch}'''							)

      if hf_config.use_weighted_layer_sum:
            UpperCAmelCase          =	checkpoint["""Featurizer"""]["""weights"""]

      hf_feature_extractor.save_pretrained(_snake_case							)
      hf_model.save_pretrained(_snake_case							)


if __name__ == "__main__":
  _UpperCamelCase					      =  argparse.ArgumentParser()
  parser.add_argument(
      """--base_model_name""", default=None, type=str, help="""Name of the huggingface pretrained base model."""
  )
  parser.add_argument("""--config_path""", default=None, type=str, help="""Path to the huggingface classifier config.""")
  parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to the s3prl checkpoint.""")
  parser.add_argument("""--model_dump_path""", default=None, type=str, help="""Path to the final converted model.""")
  _UpperCamelCase					      =  parser.parse_args()
  convert_saprl_checkpoint(args.base_model_name, args.config_path, args.checkpoint_path, args.model_dump_path)



