

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   (    )					->   List[str]:





      """simple docstring"""



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

      return image
def   a   (    A__			:     List[Any]  )					->   Tuple:





      """simple docstring"""



      _lowercase			 =[]
      # 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   (    A__			:     Union[str, Any]   ,     A__			:     Tuple   ,     A__			:     Tuple  )					->   Union[str, Any]:





      """simple docstring"""



      _lowercase			 =dct.pop(A__  )
      _lowercase			 =val
def   a   (    A__			:     List[Any]   ,     A__			:     Union[str, Any]  )					->   List[Any]:





      """simple docstring"""



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

            # next, set bias in the state dict
            _lowercase			 =torch.cat((q_bias, torch.zeros_like(A__   ,     requires_grad=A__  ), v_bias)  )
            _lowercase			 =qkv_bias
def   a   (    A__			:     List[Any]  )					->   Union[str, Any]:





      """simple docstring"""



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

      return config, image_size
@torch.no_grad()
def   a   (    A__			:     int   ,     A__			:     str=None   ,     A__			:     List[str]=False  )					->   Tuple:





      """simple docstring"""



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

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

      _lowercase , _lowercase			 =get_blipa_config(A__  )
      _lowercase			 =InstructBlipForConditionalGeneration(A__  ).eval()

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

      _lowercase , _lowercase			 =model_name_to_original[model_name]

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

      # update state dict keys
      _lowercase			 =original_model.state_dict()
      _lowercase			 =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():
            _lowercase			 =state_dict.pop(A__  )
            if key.startswith('Qformer.bert'  ):
                  _lowercase			 =key.replace('Qformer.bert'   ,     'qformer'  )
            if "attention.self" in key:
                  _lowercase			 =key.replace('self'   ,     'attention'  )
            if "llm_proj" in key:
                  _lowercase			 =key.replace('llm_proj'   ,     'language_projection'  )
            if "t5_proj" in key:
                  _lowercase			 =key.replace('t5_proj'   ,     'language_projection'  )
            if key.startswith('llm_model'  ):
                  _lowercase			 =key.replace('llm_model'   ,     'language_model'  )
            if key.startswith('t5'  ):
                  _lowercase			 =key.replace('t5'   ,     'language'  )
            _lowercase			 =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__  )

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

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

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





import argparse

import torch
from datasets import load_dataset
from donut import DonutModel

from transformers import (
    DonutImageProcessor,
    DonutProcessor,
    DonutSwinConfig,
    DonutSwinModel,
    MBartConfig,
    MBartForCausalLM,
    VisionEncoderDecoderModel,
    XLMRobertaTokenizerFast,
)
def   a   (    A__			:     Tuple  )					->   Any:





      """simple docstring"""



      _lowercase			 =model.config

      _lowercase			 =DonutSwinConfig(
          image_size=original_config.input_size   ,     patch_size=4   ,     depths=original_config.encoder_layer   ,     num_heads=[4, 8, 16, 32]   ,     window_size=original_config.window_size   ,     embed_dim=128   ,     )
      _lowercase			 =MBartConfig(
          is_decoder=A__   ,     is_encoder_decoder=A__   ,     add_cross_attention=A__   ,     decoder_layers=original_config.decoder_layer   ,     max_position_embeddings=original_config.max_position_embeddings   ,     vocab_size=len(
              model.decoder.tokenizer  )   ,     scale_embedding=A__   ,     add_final_layer_norm=A__   ,     )

      return encoder_config, decoder_config
def   a   (    A__			:     Union[str, Any]  )					->   str:





      """simple docstring"""



      if "encoder.model" in name:
            _lowercase			 =name.replace('encoder.model'   ,     'encoder'  )
      if "decoder.model" in name:
            _lowercase			 =name.replace('decoder.model'   ,     'decoder'  )
      if "patch_embed.proj" in name:
            _lowercase			 =name.replace('patch_embed.proj'   ,     'embeddings.patch_embeddings.projection'  )
      if "patch_embed.norm" in name:
            _lowercase			 =name.replace('patch_embed.norm'   ,     'embeddings.norm'  )
      if name.startswith('encoder'  ):
            if "layers" in name:
                  _lowercase			 ='encoder.' + name
            if "attn.proj" in name:
                  _lowercase			 =name.replace('attn.proj'   ,     'attention.output.dense'  )
            if "attn" in name and "mask" not in name:
                  _lowercase			 =name.replace('attn'   ,     'attention.self'  )
            if "norm1" in name:
                  _lowercase			 =name.replace('norm1'   ,     'layernorm_before'  )
            if "norm2" in name:
                  _lowercase			 =name.replace('norm2'   ,     'layernorm_after'  )
            if "mlp.fc1" in name:
                  _lowercase			 =name.replace('mlp.fc1'   ,     'intermediate.dense'  )
            if "mlp.fc2" in name:
                  _lowercase			 =name.replace('mlp.fc2'   ,     'output.dense'  )

            if name == "encoder.norm.weight":
                  _lowercase			 ='encoder.layernorm.weight'
            if name == "encoder.norm.bias":
                  _lowercase			 ='encoder.layernorm.bias'

      return name
def   a   (    A__			:     Optional[Any]   ,     A__			:     List[Any]  )					->   List[str]:





      """simple docstring"""



      for key in orig_state_dict.copy().keys():
            _lowercase			 =orig_state_dict.pop(A__  )

            if "qkv" in key:
                  _lowercase			 =key.split('.'  )
                  _lowercase			 =int(key_split[3]  )
                  _lowercase			 =int(key_split[5]  )
                  _lowercase			 =model.encoder.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

                  if "weight" in key:
                        _lowercase			 =val[:dim, :]
                        _lowercase			 =val[dim : dim * 2, :]
                        _lowercase			 =val[-dim:, :]
                  else:
                        _lowercase			 =val[:dim]
                        _lowercase			 =val[dim : dim * 2]
                        _lowercase			 =val[-dim:]
            elif "attn_mask" in key or key in ["encoder.model.norm.weight", "encoder.model.norm.bias"]:
                  # HuggingFace implementation doesn't use attn_mask buffer
                  # and model doesn't use final LayerNorms for the encoder
                  pass
            else:
                  _lowercase			 =val

      return orig_state_dict
def   a   (    A__			:     str   ,     A__			:     List[str]=None   ,     A__			:     List[Any]=False  )					->   List[str]:





      """simple docstring"""



      _lowercase			 =DonutModel.from_pretrained(A__  ).eval()

      # load HuggingFace model
      _lowercase , _lowercase			 =get_configs(A__  )
      _lowercase			 =DonutSwinModel(A__  )
      _lowercase			 =MBartForCausalLM(A__  )
      _lowercase			 =VisionEncoderDecoderModel(encoder=A__   ,     decoder=A__  )
      model.eval()

      _lowercase			 =original_model.state_dict()
      _lowercase			 =convert_state_dict(A__   ,     A__  )
      model.load_state_dict(A__  )

      # verify results on scanned document
      _lowercase			 =load_dataset('hf-internal-testing/example-documents'  )
      _lowercase			 =dataset['test'][0]['image'].convert('RGB'  )

      _lowercase			 =XLMRobertaTokenizerFast.from_pretrained(A__   ,     from_slow=A__  )
      _lowercase			 =DonutImageProcessor(
          do_align_long_axis=original_model.config.align_long_axis   ,     size=original_model.config.input_size[::-1]  )
      _lowercase			 =DonutProcessor(A__   ,     A__  )
      _lowercase			 =processor(A__   ,     return_tensors='pt'  ).pixel_values

      if model_name == "naver-clova-ix/donut-base-finetuned-docvqa":
            _lowercase			 ='<s_docvqa><s_question>{user_input}</s_question><s_answer>'
            _lowercase			 ='When is the coffee break?'
            _lowercase			 =task_prompt.replace('{user_input}'   ,     A__  )
      elif model_name == "naver-clova-ix/donut-base-finetuned-rvlcdip":
            _lowercase			 ='<s_rvlcdip>'
      elif model_name in [
          "naver-clova-ix/donut-base-finetuned-cord-v1",
          "naver-clova-ix/donut-base-finetuned-cord-v1-2560",
      ]:
            _lowercase			 ='<s_cord>'
      elif model_name == "naver-clova-ix/donut-base-finetuned-cord-v2":
            _lowercase			 ='s_cord-v2>'
      elif model_name == "naver-clova-ix/donut-base-finetuned-zhtrainticket":
            _lowercase			 ='<s_zhtrainticket>'
      elif model_name in ["naver-clova-ix/donut-proto", "naver-clova-ix/donut-base"]:
            # use a random prompt
            _lowercase			 ='hello world'
      else:
            raise ValueError('Model name not supported'  )
      _lowercase			 =original_model.decoder.tokenizer(A__   ,     add_special_tokens=A__   ,     return_tensors='pt'  )[
          'input_ids'
      ]

      _lowercase			 =original_model.encoder.model.patch_embed(A__  )
      _lowercase , _lowercase			 =model.encoder.embeddings(A__  )
      assert torch.allclose(A__   ,     A__   ,     atol=1e-3  )

      # verify encoder hidden states
      _lowercase			 =original_model.encoder(A__  )
      _lowercase			 =model.encoder(A__  ).last_hidden_state
      assert torch.allclose(A__   ,     A__   ,     atol=1e-2  )

      # verify decoder hidden states
      _lowercase			 =original_model(A__   ,     A__   ,     A__  ).logits
      _lowercase			 =model(A__   ,     decoder_input_ids=A__  ).logits
      assert torch.allclose(A__   ,     A__   ,     atol=1e-3  )
      print('Looks ok!'  )

      if pytorch_dump_folder_path is not None:
            print(F'''Saving model and processor to {pytorch_dump_folder_path}'''  )
            model.save_pretrained(A__  )
            processor.save_pretrained(A__  )

      if push_to_hub:
            model.push_to_hub('nielsr/' + model_name.split('/'  )[-1]   ,     commit_message='Update model'  )
            processor.push_to_hub('nielsr/' + model_name.split('/'  )[-1]   ,     commit_message='Update model'  )


if __name__ == "__main__":
  lowercase_             =     argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--model_name',
      default='naver-clova-ix/donut-base-finetuned-docvqa',
      required=False,
      type=str,
      help='Name of the original model you\'d like to convert.',
  )
  parser.add_argument(
      '--pytorch_dump_folder_path',
      default=None,
      required=False,
      type=str,
      help='Path to the output PyTorch model directory.',
  )
  parser.add_argument(
      '--push_to_hub',
      action='store_true',
      help='Whether or not to push the converted model and processor to the 🤗 hub.',
  )

  lowercase_             =     parser.parse_args()
  convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)


