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  __UpperCamelCase    (						A							):
     UpperCamelCase__										=      VideoMAEConfig()

     set_architecture_configs(lowerCAmelCase_ ,							lowerCAmelCase_							)

     if "finetuned" not in model_name:
          UpperCamelCase__										=      False

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

     return config







def  __UpperCamelCase    (						A ,							A							):


     if "small" in model_name:
          UpperCamelCase__										=      384
          UpperCamelCase__										=      1536
          UpperCamelCase__										=      12
          UpperCamelCase__										=      16
          UpperCamelCase__										=      12
          UpperCamelCase__										=      3
          UpperCamelCase__										=      192
          UpperCamelCase__										=      768
     elif "large" in model_name:
          UpperCamelCase__										=      1024
          UpperCamelCase__										=      4096
          UpperCamelCase__										=      24
          UpperCamelCase__										=      16
          UpperCamelCase__										=      12
          UpperCamelCase__										=      8
          UpperCamelCase__										=      512
          UpperCamelCase__										=      2048
     elif "huge" in model_name:
          UpperCamelCase__										=      1280
          UpperCamelCase__										=      5120
          UpperCamelCase__										=      32
          UpperCamelCase__										=      16
          UpperCamelCase__										=      12
          UpperCamelCase__										=      8
          UpperCamelCase__										=      640
          UpperCamelCase__										=      2560
     elif "base" not in model_name:
          raise ValueError('''Model name should include either \"small\", \"base\", \"large\", or \"huge\"'''							)







def  __UpperCamelCase    (						A							):


     if "encoder." in name:
          UpperCamelCase__										=      name.replace('''encoder.''' ,							''''''							)
     if "cls_token" in name:
          UpperCamelCase__										=      name.replace('''cls_token''' ,							'''videomae.embeddings.cls_token'''							)
     if "decoder_pos_embed" in name:
          UpperCamelCase__										=      name.replace('''decoder_pos_embed''' ,							'''decoder.decoder_pos_embed'''							)
     if "pos_embed" in name and "decoder" not in name:
          UpperCamelCase__										=      name.replace('''pos_embed''' ,							'''videomae.embeddings.position_embeddings'''							)
     if "patch_embed.proj" in name:
          UpperCamelCase__										=      name.replace('''patch_embed.proj''' ,							'''videomae.embeddings.patch_embeddings.projection'''							)
     if "patch_embed.norm" in name:
          UpperCamelCase__										=      name.replace('''patch_embed.norm''' ,							'''videomae.embeddings.norm'''							)
     if "decoder.blocks" in name:
          UpperCamelCase__										=      name.replace('''decoder.blocks''' ,							'''decoder.decoder_layers'''							)
     if "blocks" in name:
          UpperCamelCase__										=      name.replace('''blocks''' ,							'''videomae.encoder.layer'''							)
     if "attn.proj" in name:
          UpperCamelCase__										=      name.replace('''attn.proj''' ,							'''attention.output.dense'''							)
     if "attn" in name and "bias" not in name:
          UpperCamelCase__										=      name.replace('''attn''' ,							'''attention.self'''							)
     if "attn" in name:
          UpperCamelCase__										=      name.replace('''attn''' ,							'''attention.attention'''							)
     if "norm1" in name:
          UpperCamelCase__										=      name.replace('''norm1''' ,							'''layernorm_before'''							)
     if "norm2" in name:
          UpperCamelCase__										=      name.replace('''norm2''' ,							'''layernorm_after'''							)
     if "mlp.fc1" in name:
          UpperCamelCase__										=      name.replace('''mlp.fc1''' ,							'''intermediate.dense'''							)
     if "mlp.fc2" in name:
          UpperCamelCase__										=      name.replace('''mlp.fc2''' ,							'''output.dense'''							)
     if "decoder_embed" in name:
          UpperCamelCase__										=      name.replace('''decoder_embed''' ,							'''decoder.decoder_embed'''							)
     if "decoder_norm" in name:
          UpperCamelCase__										=      name.replace('''decoder_norm''' ,							'''decoder.decoder_norm'''							)
     if "decoder_pred" in name:
          UpperCamelCase__										=      name.replace('''decoder_pred''' ,							'''decoder.decoder_pred'''							)
     if "norm.weight" in name and "decoder" not in name and "fc" not in name:
          UpperCamelCase__										=      name.replace('''norm.weight''' ,							'''videomae.layernorm.weight'''							)
     if "norm.bias" in name and "decoder" not in name and "fc" not in name:
          UpperCamelCase__										=      name.replace('''norm.bias''' ,							'''videomae.layernorm.bias'''							)
     if "head" in name and "decoder" not in name:
          UpperCamelCase__										=      name.replace('''head''' ,							'''classifier'''							)

     return name







def  __UpperCamelCase    (						A ,							A							):


     for key in orig_state_dict.copy().keys():
          UpperCamelCase__										=      orig_state_dict.pop(lowerCAmelCase_							)

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

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

     return orig_state_dict







def  __UpperCamelCase    (						):
     UpperCamelCase__										=      hf_hub_download(
         repo_id='''hf-internal-testing/spaghetti-video''' ,							filename='''eating_spaghetti.npy''' ,							repo_type='''dataset'''							)
     UpperCamelCase__										=      np.load(lowerCAmelCase_							)
     return list(lowerCAmelCase_							)







def  __UpperCamelCase    (						A ,							A ,							A ,							A							):
     UpperCamelCase__										=      get_videomae_config(lowerCAmelCase_							)

     if "finetuned" in model_name:
          UpperCamelCase__										=      VideoMAEForVideoClassification(lowerCAmelCase_							)
     else:
          UpperCamelCase__										=      VideoMAEForPreTraining(lowerCAmelCase_							)

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

     model.load_state_dict(lowerCAmelCase_							)
     model.eval()

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

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

     UpperCamelCase__										=      model(**lowerCAmelCase_							)
     UpperCamelCase__										=      outputs.logits

     UpperCamelCase__										=      [
         "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":
          UpperCamelCase__										=      torch.Size([1, 400]							)
          UpperCamelCase__										=      torch.tensor([-0.92_91, -0.40_61, -0.93_07]							)
     elif model_name == "videomae-small-finetuned-ssv2":
          UpperCamelCase__										=      torch.Size([1, 174]							)
          UpperCamelCase__										=      torch.tensor([0.26_71, -0.46_89, -0.82_35]							)
     elif model_name == "videomae-base":
          UpperCamelCase__										=      torch.Size([1, 1408, 1536]							)
          UpperCamelCase__										=      torch.tensor([[0.77_39, 0.79_68, 0.70_89], [0.67_01, 0.74_87, 0.62_09], [0.42_87, 0.51_58, 0.47_73]]							)
     elif model_name == "videomae-base-short":
          UpperCamelCase__										=      torch.Size([1, 1408, 1536]							)
          UpperCamelCase__										=      torch.tensor([[0.79_94, 0.96_12, 0.85_08], [0.74_01, 0.89_58, 0.83_02], [0.58_62, 0.74_68, 0.73_25]]							)
          # we verified the loss both for normalized and unnormalized targets for this one
          UpperCamelCase__										=      torch.tensor([0.51_42]							) if config.norm_pix_loss else torch.tensor([0.64_69]							)
     elif model_name == "videomae-large":
          UpperCamelCase__										=      torch.Size([1, 1408, 1536]							)
          UpperCamelCase__										=      torch.tensor([[0.71_49, 0.79_97, 0.69_66], [0.67_68, 0.78_69, 0.69_48], [0.51_39, 0.62_21, 0.56_05]]							)
     elif model_name == "videomae-large-finetuned-kinetics":
          UpperCamelCase__										=      torch.Size([1, 400]							)
          UpperCamelCase__										=      torch.tensor([0.07_71, 0.00_11, -0.36_25]							)
     elif model_name == "videomae-huge-finetuned-kinetics":
          UpperCamelCase__										=      torch.Size([1, 400]							)
          UpperCamelCase__										=      torch.tensor([0.24_33, 0.16_32, -0.48_94]							)
     elif model_name == "videomae-base-short-finetuned-kinetics":
          UpperCamelCase__										=      torch.Size([1, 400]							)
          UpperCamelCase__										=      torch.tensor([0.65_88, 0.09_90, -0.24_93]							)
     elif model_name == "videomae-base-finetuned-kinetics":
          UpperCamelCase__										=      torch.Size([1, 400]							)
          UpperCamelCase__										=      torch.tensor([0.36_69, -0.06_88, -0.24_21]							)
     elif model_name == "videomae-base-short-ssv2":
          UpperCamelCase__										=      torch.Size([1, 1408, 1536]							)
          UpperCamelCase__										=      torch.tensor([[0.47_12, 0.52_96, 0.57_86], [0.22_78, 0.27_29, 0.40_26], [0.03_52, 0.07_30, 0.25_06]]							)
     elif model_name == "videomae-base-short-finetuned-ssv2":
          UpperCamelCase__										=      torch.Size([1, 174]							)
          UpperCamelCase__										=      torch.tensor([-0.05_37, -0.15_39, -0.32_66]							)
     elif model_name == "videomae-base-ssv2":
          UpperCamelCase__										=      torch.Size([1, 1408, 1536]							)
          UpperCamelCase__										=      torch.tensor([[0.81_31, 0.87_27, 0.85_46], [0.73_66, 0.93_77, 0.88_70], [0.59_35, 0.88_74, 0.85_64]]							)
     elif model_name == "videomae-base-finetuned-ssv2":
          UpperCamelCase__										=      torch.Size([1, 174]							)
          UpperCamelCase__										=      torch.tensor([0.19_61, -0.83_37, -0.63_89]							)
     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] ,							lowerCAmelCase_ ,							atol=1e-4							)
     else:
          print('''Logits:''' ,							logits[0, :3, :3]							)
          assert torch.allclose(logits[0, :3, :3] ,							lowerCAmelCase_ ,							atol=1e-4							)
     print('''Logits ok!'''							)

     # verify loss, if applicable
     if model_name == "videomae-base-short":
          UpperCamelCase__										=      outputs.loss
          assert torch.allclose(lowerCAmelCase_ ,							lowerCAmelCase_ ,							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(lowerCAmelCase_							)
          model.save_pretrained(lowerCAmelCase_							)

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


if __name__ == "__main__":
     __magic_name__								=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.'''
     )

     __magic_name__								=parser.parse_args()
     convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)




from __future__ import annotations







def  __UpperCamelCase    (						A ,							A							):
     UpperCamelCase__										=      get_failure_array(A							)

     # 2) Step through text searching for pattern
     UpperCamelCase__						,       UpperCamelCase__										=      0, 0  # index into text, pattern
     while i < len(A							):
          if pattern[j] == text[i]:
               if j == (len(A							) - 1):
                    return True
               j += 1

          # if this is a prefix in our pattern
          # just go back far enough to continue
          elif j > 0:
               UpperCamelCase__										=      failure[j - 1]
               continue
          i += 1
     return False







def  __UpperCamelCase    (						A							):
     UpperCamelCase__										=      [0]
     UpperCamelCase__										=      0
     UpperCamelCase__										=      1
     while j < len(A							):
          if pattern[i] == pattern[j]:
               i += 1
          elif i > 0:
               UpperCamelCase__										=      failure[i - 1]
               continue
          j += 1
          failure.append(A							)
     return failure


if __name__ == "__main__":
     # Test 1)
     __magic_name__								='''abc1abc12'''
     __magic_name__								='''alskfjaldsabc1abc1abc12k23adsfabcabc'''
     __magic_name__								='''alskfjaldsk23adsfabcabc'''
     assert kmp(pattern, texta) and not kmp(pattern, texta)

     # Test 2)
     __magic_name__								='''ABABX'''
     __magic_name__								='''ABABZABABYABABX'''
     assert kmp(pattern, text)

     # Test 3)
     __magic_name__								='''AAAB'''
     __magic_name__								='''ABAAAAAB'''
     assert kmp(pattern, text)

     # Test 4)
     __magic_name__								='''abcdabcy'''
     __magic_name__								='''abcxabcdabxabcdabcdabcy'''
     assert kmp(pattern, text)

     # Test 5)
     __magic_name__								='''aabaabaaa'''
     assert get_failure_array(pattern) == [0, 1, 0, 1, 2, 3, 4, 5, 2]



