

import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)
def   a   (    A__			:     Tuple  )					->   Union[str, Any]:





      """simple docstring"""



      _lowercase			 =VideoMAEConfig()

      set_architecture_configs(A__   ,     A__  )

      if "finetuned" not in model_name:
            _lowercase			 =False

      if "finetuned" in model_name:
            _lowercase			 ='huggingface/label-files'
            if "kinetics" in model_name:
                  _lowercase			 =400
                  _lowercase			 ='kinetics400-id2label.json'
            elif "ssv2" in model_name:
                  _lowercase			 =174
                  _lowercase			 ='something-something-v2-id2label.json'
            else:
                  raise ValueError('Model name should either contain \'kinetics\' or \'ssv2\' in case it\'s fine-tuned.'  )
            _lowercase			 =json.load(open(hf_hub_download(A__   ,     A__   ,     repo_type='dataset'  )   ,     'r'  )  )
            _lowercase			 ={int(A__  ): v for k, v in idalabel.items()}
            _lowercase			 =idalabel
            _lowercase			 ={v: k for k, v in idalabel.items()}

      return config
def   a   (    A__			:     str   ,     A__			:     Dict  )					->   Union[str, Any]:





      """simple docstring"""



      if "small" in model_name:
            _lowercase			 =384
            _lowercase			 =1536
            _lowercase			 =12
            _lowercase			 =16
            _lowercase			 =12
            _lowercase			 =3
            _lowercase			 =192
            _lowercase			 =768
      elif "large" in model_name:
            _lowercase			 =1024
            _lowercase			 =4096
            _lowercase			 =24
            _lowercase			 =16
            _lowercase			 =12
            _lowercase			 =8
            _lowercase			 =512
            _lowercase			 =2048
      elif "huge" in model_name:
            _lowercase			 =1280
            _lowercase			 =5120
            _lowercase			 =32
            _lowercase			 =16
            _lowercase			 =12
            _lowercase			 =8
            _lowercase			 =640
            _lowercase			 =2560
      elif "base" not in model_name:
            raise ValueError('Model name should include either "small", "base", "large", or "huge"'  )
def   a   (    A__			:     Optional[int]  )					->   Optional[Any]:





      """simple docstring"""



      if "encoder." in name:
            _lowercase			 =name.replace('encoder.'   ,     ''  )
      if "cls_token" in name:
            _lowercase			 =name.replace('cls_token'   ,     'videomae.embeddings.cls_token'  )
      if "decoder_pos_embed" in name:
            _lowercase			 =name.replace('decoder_pos_embed'   ,     'decoder.decoder_pos_embed'  )
      if "pos_embed" in name and "decoder" not in name:
            _lowercase			 =name.replace('pos_embed'   ,     'videomae.embeddings.position_embeddings'  )
      if "patch_embed.proj" in name:
            _lowercase			 =name.replace('patch_embed.proj'   ,     'videomae.embeddings.patch_embeddings.projection'  )
      if "patch_embed.norm" in name:
            _lowercase			 =name.replace('patch_embed.norm'   ,     'videomae.embeddings.norm'  )
      if "decoder.blocks" in name:
            _lowercase			 =name.replace('decoder.blocks'   ,     'decoder.decoder_layers'  )
      if "blocks" in name:
            _lowercase			 =name.replace('blocks'   ,     'videomae.encoder.layer'  )
      if "attn.proj" in name:
            _lowercase			 =name.replace('attn.proj'   ,     'attention.output.dense'  )
      if "attn" in name and "bias" not in name:
            _lowercase			 =name.replace('attn'   ,     'attention.self'  )
      if "attn" in name:
            _lowercase			 =name.replace('attn'   ,     'attention.attention'  )
      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 "decoder_embed" in name:
            _lowercase			 =name.replace('decoder_embed'   ,     'decoder.decoder_embed'  )
      if "decoder_norm" in name:
            _lowercase			 =name.replace('decoder_norm'   ,     'decoder.decoder_norm'  )
      if "decoder_pred" in name:
            _lowercase			 =name.replace('decoder_pred'   ,     'decoder.decoder_pred'  )
      if "norm.weight" in name and "decoder" not in name and "fc" not in name:
            _lowercase			 =name.replace('norm.weight'   ,     'videomae.layernorm.weight'  )
      if "norm.bias" in name and "decoder" not in name and "fc" not in name:
            _lowercase			 =name.replace('norm.bias'   ,     'videomae.layernorm.bias'  )
      if "head" in name and "decoder" not in name:
            _lowercase			 =name.replace('head'   ,     'classifier'  )

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





      """simple docstring"""



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

            if key.startswith('encoder.'  ):
                  _lowercase			 =key.replace('encoder.'   ,     ''  )

            if "qkv" in key:
                  _lowercase			 =key.split('.'  )
                  if key.startswith('decoder.blocks'  ):
                        _lowercase			 =config.decoder_hidden_size
                        _lowercase			 =int(key_split[2]  )
                        _lowercase			 ='decoder.decoder_layers.'
                        if "weight" in key:
                              _lowercase			 =val[:dim, :]
                              _lowercase			 =val[dim : dim * 2, :]
                              _lowercase			 =val[-dim:, :]
                  else:
                        _lowercase			 =config.hidden_size
                        _lowercase			 =int(key_split[1]  )
                        _lowercase			 ='videomae.encoder.layer.'
                        if "weight" in key:
                              _lowercase			 =val[:dim, :]
                              _lowercase			 =val[dim : dim * 2, :]
                              _lowercase			 =val[-dim:, :]
            else:
                  _lowercase			 =val

      return orig_state_dict
def   a   (    )					->   Optional[Any]:





      """simple docstring"""



      _lowercase			 =hf_hub_download(
          repo_id='hf-internal-testing/spaghetti-video'   ,     filename='eating_spaghetti.npy'   ,     repo_type='dataset'  )
      _lowercase			 =np.load(A__  )
      return list(A__  )
def   a   (    A__			:     List[Any]   ,     A__			:     Optional[Any]   ,     A__			:     int   ,     A__			:     Optional[int]  )					->   Tuple:





      """simple docstring"""



      _lowercase			 =get_videomae_config(A__  )

      if "finetuned" in model_name:
            _lowercase			 =VideoMAEForVideoClassification(A__  )
      else:
            _lowercase			 =VideoMAEForPreTraining(A__  )

      # download original checkpoint, hosted on Google Drive
      _lowercase			 ='pytorch_model.bin'
      gdown.cached_download(A__   ,     A__   ,     quiet=A__  )
      _lowercase			 =torch.load(A__   ,     map_location='cpu'  )
      if "model" in files:
            _lowercase			 =files['model']
      else:
            _lowercase			 =files['module']
      _lowercase			 =convert_state_dict(A__   ,     A__  )

      model.load_state_dict(A__  )
      model.eval()

      # verify model on basic input
      _lowercase			 =VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]   ,     image_std=[0.5, 0.5, 0.5]  )
      _lowercase			 =prepare_video()
      _lowercase			 =image_processor(A__   ,     return_tensors='pt'  )

      if "finetuned" not in model_name:
            _lowercase			 =hf_hub_download(repo_id='hf-internal-testing/bool-masked-pos'   ,     filename='bool_masked_pos.pt'  )
            _lowercase			 =torch.load(A__  )

      _lowercase			 =model(**A__  )
      _lowercase			 =outputs.logits

      _lowercase			 =[
          'videomae-small-finetuned-kinetics',
          'videomae-small-finetuned-ssv2',
          # Kinetics-400 checkpoints (short = pretrained only for 800 epochs instead of 1600)
          'videomae-base-short',
          'videomae-base-short-finetuned-kinetics',
          'videomae-base',
          'videomae-base-finetuned-kinetics',
          'videomae-large',
          'videomae-large-finetuned-kinetics',
          'videomae-huge-finetuned-kinetics',
          # Something-Something-v2 checkpoints (short = pretrained only for 800 epochs instead of 2400)
          'videomae-base-short-ssv2',
          'videomae-base-short-finetuned-ssv2',
          'videomae-base-ssv2',
          'videomae-base-finetuned-ssv2',
      ]

      # NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
      if model_name == "videomae-small-finetuned-kinetics":
            _lowercase			 =torch.Size([1, 400]  )
            _lowercase			 =torch.tensor([-0.9291, -0.4061, -0.9307]  )
      elif model_name == "videomae-small-finetuned-ssv2":
            _lowercase			 =torch.Size([1, 174]  )
            _lowercase			 =torch.tensor([0.2671, -0.4689, -0.8235]  )
      elif model_name == "videomae-base":
            _lowercase			 =torch.Size([1, 1408, 1536]  )
            _lowercase			 =torch.tensor([[0.7739, 0.7968, 0.7089], [0.6701, 0.7487, 0.6209], [0.4287, 0.5158, 0.4773]]  )
      elif model_name == "videomae-base-short":
            _lowercase			 =torch.Size([1, 1408, 1536]  )
            _lowercase			 =torch.tensor([[0.7994, 0.9612, 0.8508], [0.7401, 0.8958, 0.8302], [0.5862, 0.7468, 0.7325]]  )
            # we verified the loss both for normalized and unnormalized targets for this one
            _lowercase			 =torch.tensor([0.5142]  ) if config.norm_pix_loss else torch.tensor([0.6469]  )
      elif model_name == "videomae-large":
            _lowercase			 =torch.Size([1, 1408, 1536]  )
            _lowercase			 =torch.tensor([[0.7149, 0.7997, 0.6966], [0.6768, 0.7869, 0.6948], [0.5139, 0.6221, 0.5605]]  )
      elif model_name == "videomae-large-finetuned-kinetics":
            _lowercase			 =torch.Size([1, 400]  )
            _lowercase			 =torch.tensor([0.0771, 0.0011, -0.3625]  )
      elif model_name == "videomae-huge-finetuned-kinetics":
            _lowercase			 =torch.Size([1, 400]  )
            _lowercase			 =torch.tensor([0.2433, 0.1632, -0.4894]  )
      elif model_name == "videomae-base-short-finetuned-kinetics":
            _lowercase			 =torch.Size([1, 400]  )
            _lowercase			 =torch.tensor([0.6588, 0.0990, -0.2493]  )
      elif model_name == "videomae-base-finetuned-kinetics":
            _lowercase			 =torch.Size([1, 400]  )
            _lowercase			 =torch.tensor([0.3669, -0.0688, -0.2421]  )
      elif model_name == "videomae-base-short-ssv2":
            _lowercase			 =torch.Size([1, 1408, 1536]  )
            _lowercase			 =torch.tensor([[0.4712, 0.5296, 0.5786], [0.2278, 0.2729, 0.4026], [0.0352, 0.0730, 0.2506]]  )
      elif model_name == "videomae-base-short-finetuned-ssv2":
            _lowercase			 =torch.Size([1, 174]  )
            _lowercase			 =torch.tensor([-0.0537, -0.1539, -0.3266]  )
      elif model_name == "videomae-base-ssv2":
            _lowercase			 =torch.Size([1, 1408, 1536]  )
            _lowercase			 =torch.tensor([[0.8131, 0.8727, 0.8546], [0.7366, 0.9377, 0.8870], [0.5935, 0.8874, 0.8564]]  )
      elif model_name == "videomae-base-finetuned-ssv2":
            _lowercase			 =torch.Size([1, 174]  )
            _lowercase			 =torch.tensor([0.1961, -0.8337, -0.6389]  )
      else:
            raise ValueError(F'''Model name not supported. Should be one of {model_names}'''  )

      # verify logits
      assert logits.shape == expected_shape
      if "finetuned" in model_name:
            assert torch.allclose(logits[0, :3]   ,     A__   ,     atol=1e-4  )
      else:
            print('Logits:'   ,     logits[0, :3, :3]  )
            assert torch.allclose(logits[0, :3, :3]   ,     A__   ,     atol=1e-4  )
      print('Logits ok!'  )

      # verify loss, if applicable
      if model_name == "videomae-base-short":
            _lowercase			 =outputs.loss
            assert torch.allclose(A__   ,     A__   ,     atol=1e-4  )
            print('Loss ok!'  )

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

      if push_to_hub:
            print('Pushing to the hub...'  )
            model.push_to_hub(A__   ,     organization='nielsr'  )


if __name__ == "__main__":
  lowercase_             =     argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--checkpoint_url',
      default='https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4',
      type=str,
      help=(
          'URL of the original PyTorch checkpoint (on Google Drive) you\'d like to convert. Should be a direct'
          ' download link.'
      ),
  )
  parser.add_argument(
      '--pytorch_dump_folder_path',
      default='/Users/nielsrogge/Documents/VideoMAE/Test',
      type=str,
      help='Path to the output PyTorch model directory.',
  )
  parser.add_argument('--model_name', default='videomae-base', type=str, help='Name of the model.')
  parser.add_argument(
      '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
  )

  lowercase_             =     parser.parse_args()
  convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, 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)


