


"""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		(		)      ->		Optional[int]:


    """simple docstring"""





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

    return image







def     __lowerCAmelCase		(		lowercase    :    List[str]					)      ->		Optional[Any]:


    """simple docstring"""





    snake_case       : List[str]    =       []
    # 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		(		lowercase    :    Optional[int]		,				lowercase    :    List[str]		,				lowercase    :    List[Any]					)      ->		List[str]:


    """simple docstring"""





    snake_case       : Optional[int]    =       dct.pop(lowercase					)
    snake_case       : Dict    =       val







def     __lowerCAmelCase		(		lowercase    :    Any		,				lowercase    :    List[Any]					)      ->		Optional[Any]:


    """simple docstring"""





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

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







def     __lowerCAmelCase		(		lowercase    :    Any					)      ->		List[str]:


    """simple docstring"""





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

    return config, image_size







@torch.no_grad()
def     __lowerCAmelCase		(		lowercase    :    Union[str, Any]		,				lowercase    :    Any=None		,				lowercase    :    Any=False					)      ->		str:


    """simple docstring"""





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

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

    snake_case   ,snake_case       : Optional[Any]    =       get_blipa_config(lowercase					)
    snake_case       : int    =       InstructBlipForConditionalGeneration(lowercase					).eval()

    snake_case       : Any    =       {
        "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"),
    }

    snake_case   ,snake_case       : List[Any]    =       model_name_to_original[model_name]

    # load original model
    print("Loading original model..."					)
    snake_case       : Optional[int]    =       "cuda:1" if torch.cuda.is_available() else "cpu"
    snake_case       : Union[str, Any]    =       "cuda:2" if torch.cuda.is_available() else "cpu"
    snake_case   ,snake_case   ,snake_case       : Optional[int]    =       load_model_and_preprocess(
        name=lowercase		,				model_type=lowercase		,				is_eval=lowercase		,				device=lowercase					)
    original_model.eval()
    print("Done!"					)

    # update state dict keys
    snake_case       : Dict    =       original_model.state_dict()
    snake_case       : List[Any]    =       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():
        snake_case       : Tuple    =       state_dict.pop(lowercase					)
        if key.startswith("Qformer.bert"					):
            snake_case       : Tuple    =       key.replace("Qformer.bert"		,				"qformer"					)
        if "attention.self" in key:
            snake_case       : List[Any]    =       key.replace("self"		,				"attention"					)
        if "llm_proj" in key:
            snake_case       : Tuple    =       key.replace("llm_proj"		,				"language_projection"					)
        if "t5_proj" in key:
            snake_case       : Dict    =       key.replace("t5_proj"		,				"language_projection"					)
        if key.startswith("llm_model"					):
            snake_case       : Any    =       key.replace("llm_model"		,				"language_model"					)
        if key.startswith("t5"					):
            snake_case       : List[str]    =       key.replace("t5"		,				"language"					)
        snake_case       : List[str]    =       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					)

    snake_case       : int    =       load_demo_image()
    snake_case       : Dict    =       "What is unusual about this image?"

    # create processor
    snake_case       : List[str]    =       BlipImageProcessor(
        size={"height": image_size, "width": image_size}		,				image_mean=lowercase		,				image_std=lowercase					)
    snake_case       : Tuple    =       InstructBlipProcessor(
        image_processor=lowercase		,				tokenizer=lowercase		,				qformer_tokenizer=lowercase		,				)
    snake_case       : Optional[int]    =       processor(images=lowercase		,				text=lowercase		,				return_tensors="pt"					).to(lowercase					)

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

  __snake_case       =       parser.parse_args()

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









"""simple docstring"""



from .imports import is_tqdm_available


if is_tqdm_available():
  from tqdm.auto import tqdm as _tqdm

from ..state import PartialState







def     __lowerCAmelCase		(		lowercase    :    bool = True		,				*lowercase    :    str		,				**lowercase    :    List[Any]					)      ->		List[Any]:


    """simple docstring"""





    if not is_tqdm_available():
        raise ImportError("Accelerate's `tqdm` module requires `tqdm` to be installed. Please run `pip install tqdm`."					)
    snake_case       : Dict    =       False
    if main_process_only:
        snake_case       : int    =       PartialState().local_process_index == 0
    return _tqdm(*lowercase		,				**lowercase		,				disable=lowercase					)





