




'''simple docstring'''





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     (       _UpperCamelCase:   List[Any]   )						->			int:

     """simple docstring"""




     _snake_case							=				VideoMAEConfig()

     set_architecture_configs(_UpperCamelCase ,						_UpperCamelCase   )

     if "finetuned" not in model_name:
          _snake_case							=				False

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

     return config





def      __a     (       _UpperCamelCase:   Optional[int] ,						_UpperCamelCase:   Any   )						->			List[Any]:

     """simple docstring"""




     if "small" in model_name:
          _snake_case							=				384
          _snake_case							=				1_536
          _snake_case							=				12
          _snake_case							=				16
          _snake_case							=				12
          _snake_case							=				3
          _snake_case							=				192
          _snake_case							=				768
     elif "large" in model_name:
          _snake_case							=				1_024
          _snake_case							=				4_096
          _snake_case							=				24
          _snake_case							=				16
          _snake_case							=				12
          _snake_case							=				8
          _snake_case							=				512
          _snake_case							=				2_048
     elif "huge" in model_name:
          _snake_case							=				1_280
          _snake_case							=				5_120
          _snake_case							=				32
          _snake_case							=				16
          _snake_case							=				12
          _snake_case							=				8
          _snake_case							=				640
          _snake_case							=				2_560
     elif "base" not in model_name:
          raise ValueError("Model name should include either \"small\", \"base\", \"large\", or \"huge\""   )





def      __a     (       _UpperCamelCase:   Any   )						->			Union[str, Any]:

     """simple docstring"""




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

     return name





def      __a     (       _UpperCamelCase:   Optional[Any] ,						_UpperCamelCase:   int   )						->			int:

     """simple docstring"""




     for key in orig_state_dict.copy().keys():
          _snake_case							=				orig_state_dict.pop(_UpperCamelCase   )

          if key.startswith("encoder."   ):
               _snake_case							=				key.replace("encoder." ,						""   )

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

     return orig_state_dict





def      __a     (       )						->			Union[str, Any]:

     """simple docstring"""




     _snake_case							=				hf_hub_download(
         repo_id="hf-internal-testing/spaghetti-video" ,						filename="eating_spaghetti.npy" ,						repo_type="dataset"   )
     _snake_case							=				np.load(_UpperCamelCase   )
     return list(_UpperCamelCase   )





def      __a     (       _UpperCamelCase:   List[Any] ,						_UpperCamelCase:   Any ,						_UpperCamelCase:   Tuple ,						_UpperCamelCase:   int   )						->			int:

     """simple docstring"""




     _snake_case							=				get_videomae_config(_UpperCamelCase   )

     if "finetuned" in model_name:
          _snake_case							=				VideoMAEForVideoClassification(_UpperCamelCase   )
     else:
          _snake_case							=				VideoMAEForPreTraining(_UpperCamelCase   )

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

     model.load_state_dict(_UpperCamelCase   )
     model.eval()

     # verify model on basic input
     _snake_case							=				VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5] ,						image_std=[0.5, 0.5, 0.5]   )
     _snake_case							=				prepare_video()
     _snake_case							=				image_processor(_UpperCamelCase ,						return_tensors="pt"   )

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

     _snake_case							=				model(**_UpperCamelCase   )
     _snake_case							=				outputs.logits

     _snake_case							=				[
         "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":
          _snake_case							=				torch.Size([1, 400]   )
          _snake_case							=				torch.tensor([-0.9291, -0.4061, -0.9307]   )
     elif model_name == "videomae-small-finetuned-ssv2":
          _snake_case							=				torch.Size([1, 174]   )
          _snake_case							=				torch.tensor([0.2671, -0.4689, -0.8235]   )
     elif model_name == "videomae-base":
          _snake_case							=				torch.Size([1, 1_408, 1_536]   )
          _snake_case							=				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":
          _snake_case							=				torch.Size([1, 1_408, 1_536]   )
          _snake_case							=				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
          _snake_case							=				torch.tensor([0.5142]   ) if config.norm_pix_loss else torch.tensor([0.6469]   )
     elif model_name == "videomae-large":
          _snake_case							=				torch.Size([1, 1_408, 1_536]   )
          _snake_case							=				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":
          _snake_case							=				torch.Size([1, 400]   )
          _snake_case							=				torch.tensor([0.0771, 0.0011, -0.3625]   )
     elif model_name == "videomae-huge-finetuned-kinetics":
          _snake_case							=				torch.Size([1, 400]   )
          _snake_case							=				torch.tensor([0.2433, 0.1632, -0.4894]   )
     elif model_name == "videomae-base-short-finetuned-kinetics":
          _snake_case							=				torch.Size([1, 400]   )
          _snake_case							=				torch.tensor([0.6588, 0.0990, -0.2493]   )
     elif model_name == "videomae-base-finetuned-kinetics":
          _snake_case							=				torch.Size([1, 400]   )
          _snake_case							=				torch.tensor([0.3669, -0.0688, -0.2421]   )
     elif model_name == "videomae-base-short-ssv2":
          _snake_case							=				torch.Size([1, 1_408, 1_536]   )
          _snake_case							=				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":
          _snake_case							=				torch.Size([1, 174]   )
          _snake_case							=				torch.tensor([-0.0537, -0.1539, -0.3266]   )
     elif model_name == "videomae-base-ssv2":
          _snake_case							=				torch.Size([1, 1_408, 1_536]   )
          _snake_case							=				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":
          _snake_case							=				torch.Size([1, 174]   )
          _snake_case							=				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] ,						_UpperCamelCase ,						atol=1E-4   )
     else:
          print("Logits:" ,						logits[0, :3, :3]   )
          assert torch.allclose(logits[0, :3, :3] ,						_UpperCamelCase ,						atol=1E-4   )
     print("Logits ok!"   )

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

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


if __name__ == "__main__":
   UpperCamelCase_						:      List[Any]       =						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.'''
   )

   UpperCamelCase_						:      int       =						parser.parse_args()
   convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)








'''simple docstring'''





from __future__ import annotations

from collections.abc import Callable
from typing import Generic, TypeVar

UpperCamelCase_						:      int       =						TypeVar('''T''')
UpperCamelCase_						:      Dict       =						TypeVar('''U''')

class        _a    (						Generic[T, U]							):






     def __init__(			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )							->      Dict:
          _snake_case							=				key
          _snake_case							=				val
          _snake_case							=				None
          _snake_case							=				None






     def __repr__(			self   )							->      str:
          return (
              f"""Node: key: {self.key}, val: {self.val}, """
              f"""has next: {bool(self.next   )}, has prev: {bool(self.prev   )}"""
          )







class        _a    (						Generic[T, U]							):






     def __init__(			self   )							->      None:
          _snake_case							=				DoubleLinkedListNode(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )
          _snake_case							=				DoubleLinkedListNode(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )
          _snake_case							, _snake_case							=				self.rear, self.head






     def __repr__(			self   )							->      str:
          _snake_case							=				["DoubleLinkedList"]
          _snake_case							=				self.head
          while node.next is not None:
               rep.append(str(_SCREAMING_SNAKE_CASE   )   )
               _snake_case							=				node.next
          rep.append(str(self.rear   )   )
          return ",\n    ".join(_SCREAMING_SNAKE_CASE   )






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE   )							->      None:
          _snake_case							=				self.rear.prev

          # All nodes other than self.head are guaranteed to have non-None previous
          assert previous is not None

          _snake_case							=				node
          _snake_case							=				previous
          _snake_case							=				node
          _snake_case							=				self.rear






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE   )							->      DoubleLinkedListNode[T, U] | None:

          if node.prev is None or node.next is None:
               return None

          _snake_case							=				node.next
          _snake_case							=				node.prev
          _snake_case							=				None
          _snake_case							=				None
          return node







class        _a    (						Generic[T, U]							):
     SCREAMING_SNAKE_CASE_							:   dict[Callable[[T], U], LRUCache[T, U]]       =						{}






     def __init__(			self ,_SCREAMING_SNAKE_CASE   )							->      str:
          _snake_case							=				DoubleLinkedList()
          _snake_case							=				capacity
          _snake_case							=				0
          _snake_case							=				0
          _snake_case							=				0
          _snake_case							=				{}






     def __repr__(			self   )							->      str:
          return (
              f"""CacheInfo(hits={self.hits}, misses={self.miss}, """
              f"""capacity={self.capacity}, current size={self.num_keys})"""
          )






     def __contains__(			self ,_SCREAMING_SNAKE_CASE   )							->      bool:
          return key in self.cache






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE   )							->      U | None:
          # Note: pythonic interface would throw KeyError rather than return None

          if key in self.cache:
               self.hits += 1
               _snake_case							=				self.cache[key]
               _snake_case							=				self.list.remove(self.cache[key]   )
               assert node == value_node

               # node is guaranteed not None because it is in self.cache
               assert node is not None
               self.list.add(_SCREAMING_SNAKE_CASE   )
               return node.val
          self.miss += 1
          return None






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )							->      None:

          if key not in self.cache:
               if self.num_keys >= self.capacity:
                    # delete first node (oldest) when over capacity
                    _snake_case							=				self.list.head.next

                    # guaranteed to have a non-None first node when num_keys > 0
                    # explain to type checker via assertions
                    assert first_node is not None
                    assert first_node.key is not None
                    assert (
                        self.list.remove(_SCREAMING_SNAKE_CASE   ) is not None
                    )  # node guaranteed to be in list assert node.key is not None

                    del self.cache[first_node.key]
                    self.num_keys -= 1
               _snake_case							=				DoubleLinkedListNode(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )
               self.list.add(self.cache[key]   )
               self.num_keys += 1

          else:
               # bump node to the end of the list, update value
               _snake_case							=				self.list.remove(self.cache[key]   )
               assert node is not None  # node guaranteed to be in list
               _snake_case							=				value
               self.list.add(_SCREAMING_SNAKE_CASE   )






     @classmethod
     def     _lowercase     (			cls ,_SCREAMING_SNAKE_CASE = 128   )							->      Callable[[Callable[[T], U]], Callable[..., U]]:

          def cache_decorator_inner(_SCREAMING_SNAKE_CASE   ) -> Callable[..., U]:
               def cache_decorator_wrapper(*_SCREAMING_SNAKE_CASE   ) -> U:
                    if func not in cls.decorator_function_to_instance_map:
                         _snake_case							=				LRUCache(_SCREAMING_SNAKE_CASE   )

                    _snake_case							=				cls.decorator_function_to_instance_map[func].get(args[0]   )
                    if result is None:
                         _snake_case							=				func(*_SCREAMING_SNAKE_CASE   )
                         cls.decorator_function_to_instance_map[func].put(args[0] ,_SCREAMING_SNAKE_CASE   )
                    return result

               def cache_info() -> LRUCache[T, U]:
                    return cls.decorator_function_to_instance_map[func]

               setattr(_SCREAMING_SNAKE_CASE ,"cache_info" ,_SCREAMING_SNAKE_CASE   )  # noqa: B010

               return cache_decorator_wrapper

          return cache_decorator_inner


if __name__ == "__main__":
   import doctest

   doctest.testmod()


