






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

  return image







def      lowerCamelCase__       (							_lowerCamelCase       :		Dict )      ->				List[str]:
  lowerCamelCase_         =						[]
  # 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__       (							_lowerCamelCase       :		int		,				_lowerCamelCase       :		Union[str, Any]		,				_lowerCamelCase       :		List[Any] )      ->				Tuple:
  lowerCamelCase_         =						dct.pop(_lowerCamelCase )
  lowerCamelCase_         =						val







def      lowerCamelCase__       (							_lowerCamelCase       :		Optional[int]		,				_lowerCamelCase       :		List[Any] )      ->				str:
  for i in range(config.vision_config.num_hidden_layers ):
    # read in original q and v biases
    lowerCamelCase_         =						state_dict.pop(F'''visual_encoder.blocks.{i}.attn.q_bias''' )
    lowerCamelCase_         =						state_dict.pop(F'''visual_encoder.blocks.{i}.attn.v_bias''' )

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







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

  return config, image_size







@torch.no_grad()
def      lowerCamelCase__       (							_lowerCamelCase       :		int		,				_lowerCamelCase       :		List[str]=None		,				_lowerCamelCase       :		Dict=False )      ->				List[str]:
  lowerCamelCase_         =						AutoTokenizer.from_pretrained('bert-base-uncased'		,				truncation_side='left' )
  qformer_tokenizer.add_special_tokens({'bos_token': '[DEC]'} )

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

  lowerCamelCase_     ,	lowerCamelCase_         =						get_blipa_config(_lowerCamelCase )
  lowerCamelCase_         =						InstructBlipForConditionalGeneration(_lowerCamelCase ).eval()

  lowerCamelCase_         =						{
      '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'),
  }

  lowerCamelCase_     ,	lowerCamelCase_         =						model_name_to_original[model_name]

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

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

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

  # read in qv biases
  read_in_q_v_bias(_lowerCamelCase		,				_lowerCamelCase )

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

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

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

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

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

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

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

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


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

  _SCREAMING_SNAKE_CASE					:				Any					    =						parser.parse_args()

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












"""simple docstring"""





# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

from ..models.auto import AutoModelForVisionaSeq
from ..utils import requires_backends
from .base import PipelineTool


if TYPE_CHECKING:
  from PIL import Image






class      a						(					__snake_case		):
    SCREAMING_SNAKE_CASE						:     Dict						= """Salesforce/blip-image-captioning-base"""
    SCREAMING_SNAKE_CASE						:     int						= (
        """This is a tool that generates a description of an image. It takes an input named `image` which should be the """
        """image to caption, and returns a text that contains the description in English."""
    )
    SCREAMING_SNAKE_CASE						:     Optional[int]						= """image_captioner"""
    SCREAMING_SNAKE_CASE						:     Optional[Any]						= AutoModelForVisionaSeq

    SCREAMING_SNAKE_CASE						:     int						= ["""image"""]
    SCREAMING_SNAKE_CASE						:     Optional[Any]						= ["""text"""]






    def __init__(       self     :							Optional[Any]       ,						*__SCREAMING_SNAKE_CASE     :							Optional[Any]       ,						**__SCREAMING_SNAKE_CASE     :							List[Any]					)     ->		Optional[Any]:
      requires_backends(self       ,						['vision']					)
      super().__init__(*__SCREAMING_SNAKE_CASE       ,						**__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							List[Any]       ,						__SCREAMING_SNAKE_CASE     :							"Image"					)     ->		Dict:
      return self.pre_processor(images=__SCREAMING_SNAKE_CASE       ,						return_tensors='pt'					)






    def    UpperCamelCase       (       self     :							List[str]       ,						__SCREAMING_SNAKE_CASE     :							Tuple					)     ->		str:
      return self.model.generate(**__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							Union[str, Any]       ,						__SCREAMING_SNAKE_CASE     :							List[str]					)     ->		Optional[Any]:
      return self.pre_processor.batch_decode(__SCREAMING_SNAKE_CASE       ,						skip_special_tokens=__SCREAMING_SNAKE_CASE					)[0].strip()




