





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       (   __magic_name__	):
   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       (   __magic_name__  , __magic_name__	):
   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       (   __magic_name__	):
   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       (   __magic_name__  , __magic_name__	):
   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       (   ):
   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       (   __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__	):
   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.9_291, -0.4_061, -0.9_307]	)
   elif model_name == "videomae-small-finetuned-ssv2":
      lowercase__			       =  torch.Size([1, 174]	)
      lowercase__			       =  torch.tensor([0.2_671, -0.4_689, -0.8_235]	)
   elif model_name == "videomae-base":
      lowercase__			       =  torch.Size([1, 1408, 1536]	)
      lowercase__			       =  torch.tensor([[0.7_739, 0.7_968, 0.7_089], [0.6_701, 0.7_487, 0.6_209], [0.4_287, 0.5_158, 0.4_773]]	)
   elif model_name == "videomae-base-short":
      lowercase__			       =  torch.Size([1, 1408, 1536]	)
      lowercase__			       =  torch.tensor([[0.7_994, 0.9_612, 0.8_508], [0.7_401, 0.8_958, 0.8_302], [0.5_862, 0.7_468, 0.7_325]]	)
      # we verified the loss both for normalized and unnormalized targets for this one
      lowercase__			       =  torch.tensor([0.5_142]	) if config.norm_pix_loss else torch.tensor([0.6_469]	)
   elif model_name == "videomae-large":
      lowercase__			       =  torch.Size([1, 1408, 1536]	)
      lowercase__			       =  torch.tensor([[0.7_149, 0.7_997, 0.6_966], [0.6_768, 0.7_869, 0.6_948], [0.5_139, 0.6_221, 0.5_605]]	)
   elif model_name == "videomae-large-finetuned-kinetics":
      lowercase__			       =  torch.Size([1, 400]	)
      lowercase__			       =  torch.tensor([0.0_771, 0.0_011, -0.3_625]	)
   elif model_name == "videomae-huge-finetuned-kinetics":
      lowercase__			       =  torch.Size([1, 400]	)
      lowercase__			       =  torch.tensor([0.2_433, 0.1_632, -0.4_894]	)
   elif model_name == "videomae-base-short-finetuned-kinetics":
      lowercase__			       =  torch.Size([1, 400]	)
      lowercase__			       =  torch.tensor([0.6_588, 0.0_990, -0.2_493]	)
   elif model_name == "videomae-base-finetuned-kinetics":
      lowercase__			       =  torch.Size([1, 400]	)
      lowercase__			       =  torch.tensor([0.3_669, -0.0_688, -0.2_421]	)
   elif model_name == "videomae-base-short-ssv2":
      lowercase__			       =  torch.Size([1, 1408, 1536]	)
      lowercase__			       =  torch.tensor([[0.4_712, 0.5_296, 0.5_786], [0.2_278, 0.2_729, 0.4_026], [0.0_352, 0.0_730, 0.2_506]]	)
   elif model_name == "videomae-base-short-finetuned-ssv2":
      lowercase__			       =  torch.Size([1, 174]	)
      lowercase__			       =  torch.tensor([-0.0_537, -0.1_539, -0.3_266]	)
   elif model_name == "videomae-base-ssv2":
      lowercase__			       =  torch.Size([1, 1408, 1536]	)
      lowercase__			       =  torch.tensor([[0.8_131, 0.8_727, 0.8_546], [0.7_366, 0.9_377, 0.8_870], [0.5_935, 0.8_874, 0.8_564]]	)
   elif model_name == "videomae-base-finetuned-ssv2":
      lowercase__			       =  torch.Size([1, 174]	)
      lowercase__			       =  torch.tensor([0.1_961, -0.8_337, -0.6_389]	)
   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__":
  _snake_case           =			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."""
  )

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











import math


def _A       (   __magic_name__	):
   lowercase__			       =  math.loga(math.sqrt(4 * positive_integer + 1	) / 2 + 1 / 2	)

   return exponent == int(__magic_name__	)


def _A       (   __magic_name__ = 1 / 1_2345	):
   lowercase__			       =  0
   lowercase__			       =  0

   lowercase__			       =  3
   while True:
      lowercase__			       =  (integer**2 - 1) / 4
      # if candidate is an integer, then there is a partition for k
      if partition_candidate == int(__magic_name__	):
         lowercase__			       =  int(__magic_name__	)
         total_partitions += 1
         if check_partition_perfect(__magic_name__	):
            perfect_partitions += 1
      if perfect_partitions > 0:
         if perfect_partitions / total_partitions < max_proportion:
            return int(__magic_name__	)
      integer += 1


if __name__ == "__main__":
  print(F"""{solution() = }""")






