




'''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     SCREAMING_SNAKE_CASE							(   ):
  __a								=       'https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'
  __a								=       Image.open(requests.get(a_					,       stream=a_      ).raw      ).convert('RGB'      )

  return image


def     SCREAMING_SNAKE_CASE							(   a_						:				int      ):
  __a								=       []
  # 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     SCREAMING_SNAKE_CASE							(   a_						:				Any					,       a_						:				Dict					,       a_						:				Optional[int]      ):
  __a								=       dct.pop(a_      )
  __a								=       val


def     SCREAMING_SNAKE_CASE							(   a_						:				Optional[Any]					,       a_						:				int      ):
  for i in range(config.vision_config.num_hidden_layers      ):
    # read in original q and v biases
    __a								=       state_dict.pop(f"visual_encoder.blocks.{i}.attn.q_bias"      )
    __a								=       state_dict.pop(f"visual_encoder.blocks.{i}.attn.v_bias"      )

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


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

  return config, image_size


@torch.no_grad()
def     SCREAMING_SNAKE_CASE							(   a_						:				Union[str, Any]					,       a_						:				str=None					,       a_						:				List[str]=False      ):
  __a								=       AutoTokenizer.from_pretrained('bert-base-uncased'					,       truncation_side='left'      )
  qformer_tokenizer.add_special_tokens({'bos_token': '[DEC]'}      )

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

  __a						,							__a								=       get_blipa_config(a_      )
  __a								=       InstructBlipForConditionalGeneration(a_      ).eval()

  __a								=       {
      '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								=       model_name_to_original[model_name]

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

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

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

  # read in qv biases
  read_in_q_v_bias(a_					,       a_      )

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

  __a								=       load_demo_image()
  __a								=       'What is unusual about this image?'

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

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

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

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

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

  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'''


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCAmelCase_					=   {
    "configuration_time_series_transformer": [
        "TIME_SERIES_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "TimeSeriesTransformerConfig",
    ],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_					=   [
      "TIME_SERIES_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
      "TimeSeriesTransformerForPrediction",
      "TimeSeriesTransformerModel",
      "TimeSeriesTransformerPreTrainedModel",
  ]


if TYPE_CHECKING:
  from .configuration_time_series_transformer import (
      TIME_SERIES_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
      TimeSeriesTransformerConfig,
  )

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_time_series_transformer import (
        TIME_SERIES_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TimeSeriesTransformerForPrediction,
        TimeSeriesTransformerModel,
        TimeSeriesTransformerPreTrainedModel,
    )

else:
  import sys

  UpperCAmelCase_					=   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



