





'''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		(	)		->Union[str, Any]:
    _SCREAMING_SNAKE_CASE   							=     'https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'
    _SCREAMING_SNAKE_CASE   							=     Image.open(requests.get(_snake_case   ,			stream=_snake_case       ).raw       ).convert("""RGB"""       )

    return image





def lowerCamelCase		(	__lowerCamelCase					:       List[str]       )		->Any:
    _SCREAMING_SNAKE_CASE   							=     []
    # 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					:       Optional[Any]   ,			__lowerCamelCase					:       Tuple   ,			__lowerCamelCase					:       str       )		->List[str]:
    _SCREAMING_SNAKE_CASE   							=     dct.pop(_snake_case       )
    _SCREAMING_SNAKE_CASE   							=     val





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

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





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

    return config, image_size





@torch.no_grad()
def lowerCamelCase		(	__lowerCamelCase					:       int   ,			__lowerCamelCase					:       str=None   ,			__lowerCamelCase					:       Union[str, Any]=False       )		->Optional[Any]:
    _SCREAMING_SNAKE_CASE   							=     AutoTokenizer.from_pretrained("""bert-base-uncased"""   ,			truncation_side="""left"""       )
    qformer_tokenizer.add_special_tokens({"""bos_token""": """[DEC]"""}       )

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

    _SCREAMING_SNAKE_CASE   							=     get_blipa_config(_snake_case       )
    _SCREAMING_SNAKE_CASE   							=     InstructBlipForConditionalGeneration(_snake_case       ).eval()

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

    _SCREAMING_SNAKE_CASE   							=     model_name_to_original[model_name]

    # load original model
    print("""Loading original model..."""       )
    _SCREAMING_SNAKE_CASE   							=     'cuda:1' if torch.cuda.is_available() else 'cpu'
    _SCREAMING_SNAKE_CASE   							=     'cuda:2' if torch.cuda.is_available() else 'cpu'
    _SCREAMING_SNAKE_CASE   							=     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
    _SCREAMING_SNAKE_CASE   							=     original_model.state_dict()
    _SCREAMING_SNAKE_CASE   							=     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():
        _SCREAMING_SNAKE_CASE   							=     state_dict.pop(_snake_case       )
        if key.startswith("""Qformer.bert"""       ):
            _SCREAMING_SNAKE_CASE   							=     key.replace("""Qformer.bert"""   ,			"""qformer"""       )
        if "attention.self" in key:
            _SCREAMING_SNAKE_CASE   							=     key.replace("""self"""   ,			"""attention"""       )
        if "llm_proj" in key:
            _SCREAMING_SNAKE_CASE   							=     key.replace("""llm_proj"""   ,			"""language_projection"""       )
        if "t5_proj" in key:
            _SCREAMING_SNAKE_CASE   							=     key.replace("""t5_proj"""   ,			"""language_projection"""       )
        if key.startswith("""llm_model"""       ):
            _SCREAMING_SNAKE_CASE   							=     key.replace("""llm_model"""   ,			"""language_model"""       )
        if key.startswith("""t5"""       ):
            _SCREAMING_SNAKE_CASE   							=     key.replace("""t5"""   ,			"""language"""       )
        _SCREAMING_SNAKE_CASE   							=     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       )

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

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

    # make sure processor creates exact same pixel values
    _SCREAMING_SNAKE_CASE   							=     vis_processors['eval'](_snake_case       ).unsqueeze(0       ).to(_snake_case       )
    _SCREAMING_SNAKE_CASE   							=     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:
            _SCREAMING_SNAKE_CASE   							=     original_model({"""image""": original_pixel_values, """text_input""": [prompt]}       ).logits
            _SCREAMING_SNAKE_CASE   							=     hf_model(**_snake_case       ).logits
        else:
            _SCREAMING_SNAKE_CASE   							=     original_model(
                {"""image""": original_pixel_values, """text_input""": [prompt], """text_output""": ["""\n"""]}       ).logits
            _SCREAMING_SNAKE_CASE   							=     tokenizer("""\n"""   ,			return_tensors="""pt"""       ).input_ids.to(_snake_case       )
            _SCREAMING_SNAKE_CASE   							=     label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id   ,			-100       )
            _SCREAMING_SNAKE_CASE   							=     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
    _SCREAMING_SNAKE_CASE   							=     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..."""       )
    _SCREAMING_SNAKE_CASE   							=     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..."""       )
    _SCREAMING_SNAKE_CASE   							=     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?
        _SCREAMING_SNAKE_CASE   							=     2
    print("""Original generation:"""   ,			_snake_case       )
    _SCREAMING_SNAKE_CASE   							=     processor.batch_decode(_snake_case   ,			skip_special_tokens=_snake_case       )
    _SCREAMING_SNAKE_CASE   							=     [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__":
      lowercase_      							=					argparse.ArgumentParser()
      lowercase_      							=					[
          """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""",
      )

      lowercase_      							=					parser.parse_args()

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





"""simple docstring"""



import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType





class 				a			(					lowercase      ):
      UpperCamelCase							:   Tuple           =     ["""image_processor""", """tokenizer"""]
      UpperCamelCase							:   List[str]           =     """LayoutLMv2ImageProcessor"""
      UpperCamelCase							:   Optional[int]           =     ("""LayoutXLMTokenizer""", """LayoutXLMTokenizerFast""")



      def __init__(   self     ,	UpperCamelCase_=None     ,	UpperCamelCase_=None     ,	**UpperCamelCase_      ):
          if "feature_extractor" in kwargs:
              warnings.warn(
                  'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
                  ' instead.'     ,	UpperCamelCase_     ,	)
              UpperCAmelCase__	:      Tuple     							=					kwargs.pop('feature_extractor'      )

          UpperCAmelCase__	:      Optional[Any]     							=					image_processor if image_processor is not None else feature_extractor
          if image_processor is None:
              raise ValueError('You need to specify an `image_processor`.'      )
          if tokenizer is None:
              raise ValueError('You need to specify a `tokenizer`.'      )

          super().__init__(UpperCamelCase_     ,	UpperCamelCase_      )



      def __call__(   self     ,	UpperCamelCase_     ,	UpperCamelCase_ = None     ,	UpperCamelCase_ = None     ,	UpperCamelCase_ = None     ,	UpperCamelCase_ = None     ,	UpperCamelCase_ = True     ,	UpperCamelCase_ = False     ,	UpperCamelCase_ = None     ,	UpperCamelCase_ = None     ,	UpperCamelCase_ = 0     ,	UpperCamelCase_ = None     ,	UpperCamelCase_ = None     ,	UpperCamelCase_ = None     ,	UpperCamelCase_ = False     ,	UpperCamelCase_ = False     ,	UpperCamelCase_ = False     ,	UpperCamelCase_ = False     ,	UpperCamelCase_ = True     ,	UpperCamelCase_ = None     ,	**UpperCamelCase_     ,	):
          # verify input
          if self.image_processor.apply_ocr and (boxes is not None):
              raise ValueError(
                  'You cannot provide bounding boxes '
                  'if you initialized the image processor with apply_ocr set to True.'      )

          if self.image_processor.apply_ocr and (word_labels is not None):
              raise ValueError(
                  'You cannot provide word labels if you initialized the image processor with apply_ocr set to True.'      )

          if return_overflowing_tokens is True and return_offsets_mapping is False:
              raise ValueError('You cannot return overflowing tokens without returning the offsets mapping.'      )

          # first, apply the image processor
          UpperCAmelCase__	:      List[Any]     							=					self.image_processor(images=UpperCamelCase_     ,	return_tensors=UpperCamelCase_      )

          # second, apply the tokenizer
          if text is not None and self.image_processor.apply_ocr and text_pair is None:
              if isinstance(UpperCamelCase_     ,	UpperCamelCase_      ):
                  UpperCAmelCase__	:      Any     							=					[text]  # add batch dimension (as the image processor always adds a batch dimension)
              UpperCAmelCase__	:      Union[str, Any]     							=					features['words']

          UpperCAmelCase__	:      List[Any]     							=					self.tokenizer(
              text=text if text is not None else features['words']     ,	text_pair=text_pair if text_pair is not None else None     ,	boxes=boxes if boxes is not None else features['boxes']     ,	word_labels=UpperCamelCase_     ,	add_special_tokens=UpperCamelCase_     ,	padding=UpperCamelCase_     ,	truncation=UpperCamelCase_     ,	max_length=UpperCamelCase_     ,	stride=UpperCamelCase_     ,	pad_to_multiple_of=UpperCamelCase_     ,	return_token_type_ids=UpperCamelCase_     ,	return_attention_mask=UpperCamelCase_     ,	return_overflowing_tokens=UpperCamelCase_     ,	return_special_tokens_mask=UpperCamelCase_     ,	return_offsets_mapping=UpperCamelCase_     ,	return_length=UpperCamelCase_     ,	verbose=UpperCamelCase_     ,	return_tensors=UpperCamelCase_     ,	**UpperCamelCase_     ,	)

          # add pixel values
          UpperCAmelCase__	:      List[Any]     							=					features.pop('pixel_values'      )
          if return_overflowing_tokens is True:
              UpperCAmelCase__	:      int     							=					self.get_overflowing_images(UpperCamelCase_     ,	encoded_inputs['overflow_to_sample_mapping']      )
          UpperCAmelCase__	:      Dict     							=					images

          return encoded_inputs



      def  __snake_case							(   self     ,	UpperCamelCase_     ,	UpperCamelCase_      ):
          # in case there's an overflow, ensure each `input_ids` sample is mapped to its corresponding image
          UpperCAmelCase__	:      Dict     							=					[]
          for sample_idx in overflow_to_sample_mapping:
              images_with_overflow.append(images[sample_idx]      )

          if len(UpperCamelCase_      ) != len(UpperCamelCase_      ):
              raise ValueError(
                  'Expected length of images to be the same as the length of `overflow_to_sample_mapping`, but got'
                  F''' {len(UpperCamelCase_      )} and {len(UpperCamelCase_      )}'''      )

          return images_with_overflow



      def  __snake_case							(   self     ,	*UpperCamelCase_     ,	**UpperCamelCase_      ):
          return self.tokenizer.batch_decode(*UpperCamelCase_     ,	**UpperCamelCase_      )



      def  __snake_case							(   self     ,	*UpperCamelCase_     ,	**UpperCamelCase_      ):
          return self.tokenizer.decode(*UpperCamelCase_     ,	**UpperCamelCase_      )



      @property
      def  __snake_case							(   self      ):
          return ["input_ids", "bbox", "attention_mask", "image"]



      @property
      def  __snake_case							(   self      ):
          warnings.warn(
              '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'     ,	UpperCamelCase_     ,	)
          return self.image_processor_class



      @property
      def  __snake_case							(   self      ):
          warnings.warn(
              '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'     ,	UpperCamelCase_     ,	)
          return self.image_processor

