



'''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 lowerCAmelCase  (__A):



     """simple docstring"""



     _a      =  VideoMAEConfig()

     set_architecture_configs(_lowerCAmelCase  ,	_lowerCAmelCase)

     if "finetuned" not in model_name:
          _a      =  False

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

     return config




def lowerCAmelCase  (__A  ,	__A):



     """simple docstring"""



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




def lowerCAmelCase  (__A):



     """simple docstring"""



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

     return name




def lowerCAmelCase  (__A  ,	__A):



     """simple docstring"""



     for key in orig_state_dict.copy().keys():
          _a      =  orig_state_dict.pop(_lowerCAmelCase)

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

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

     return orig_state_dict




def lowerCAmelCase  ():



     """simple docstring"""



     _a      =  hf_hub_download(
         repo_id='''hf-internal-testing/spaghetti-video'''  ,	filename='''eating_spaghetti.npy'''  ,	repo_type='''dataset''')
     _a      =  np.load(_lowerCAmelCase)
     return list(_lowerCAmelCase)




def lowerCAmelCase  (__A  ,	__A  ,	__A  ,	__A):



     """simple docstring"""



     _a      =  get_videomae_config(_lowerCAmelCase)

     if "finetuned" in model_name:
          _a      =  VideoMAEForVideoClassification(_lowerCAmelCase)
     else:
          _a      =  VideoMAEForPreTraining(_lowerCAmelCase)

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

     model.load_state_dict(_lowerCAmelCase)
     model.eval()

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

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

     _a      =  model(**_lowerCAmelCase)
     _a      =  outputs.logits

     _a      =  [
         """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":
          _a      =  torch.Size([1, 400])
          _a      =  torch.tensor([-0.92_91, -0.40_61, -0.93_07])
     elif model_name == "videomae-small-finetuned-ssv2":
          _a      =  torch.Size([1, 174])
          _a      =  torch.tensor([0.26_71, -0.46_89, -0.82_35])
     elif model_name == "videomae-base":
          _a      =  torch.Size([1, 1_408, 1_536])
          _a      =  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":
          _a      =  torch.Size([1, 1_408, 1_536])
          _a      =  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
          _a      =  torch.tensor([0.51_42]) if config.norm_pix_loss else torch.tensor([0.64_69])
     elif model_name == "videomae-large":
          _a      =  torch.Size([1, 1_408, 1_536])
          _a      =  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":
          _a      =  torch.Size([1, 400])
          _a      =  torch.tensor([0.07_71, 0.00_11, -0.36_25])
     elif model_name == "videomae-huge-finetuned-kinetics":
          _a      =  torch.Size([1, 400])
          _a      =  torch.tensor([0.24_33, 0.16_32, -0.48_94])
     elif model_name == "videomae-base-short-finetuned-kinetics":
          _a      =  torch.Size([1, 400])
          _a      =  torch.tensor([0.65_88, 0.09_90, -0.24_93])
     elif model_name == "videomae-base-finetuned-kinetics":
          _a      =  torch.Size([1, 400])
          _a      =  torch.tensor([0.36_69, -0.06_88, -0.24_21])
     elif model_name == "videomae-base-short-ssv2":
          _a      =  torch.Size([1, 1_408, 1_536])
          _a      =  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":
          _a      =  torch.Size([1, 174])
          _a      =  torch.tensor([-0.05_37, -0.15_39, -0.32_66])
     elif model_name == "videomae-base-ssv2":
          _a      =  torch.Size([1, 1_408, 1_536])
          _a      =  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":
          _a      =  torch.Size([1, 174])
          _a      =  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":
          _a      =  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__":
  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)








'''simple docstring'''





import importlib
import os
import sys


# This is required to make the module import works (when the python process is running from the root of the repo)
sys.path.append(".")





def 						__snake_case(				_lowerCAmelCase		)		-> int:
      snake_case__  :    Optional[int]    						=      test_file.split(os.path.sep		)
      if components[0:2] != ["tests", "models"]:
            raise ValueError(
                """`test_file` should start with `tests/models/` (with `/` being the OS specific path separator). Got """
                f"{test_file} instead."		)
      snake_case__  :    Dict    						=      components[-1]
      if not test_fn.endswith("""py"""		):
            raise ValueError(f"`test_file` should be a python file. Got {test_fn} instead."		)
      if not test_fn.startswith("""test_modeling_"""		):
            raise ValueError(
                f"`test_file` should point to a file name of the form `test_modeling_*.py`. Got {test_fn} instead."		)

      snake_case__  :    int    						=      components[:-1] + [test_fn.replace(""".py"""     ,   """"""		)]
      snake_case__  :    int    						=      """.""".join(_lowerCAmelCase		)

      return test_module_path





def 						__snake_case(				_lowerCAmelCase		)		-> List[str]:
      snake_case__  :    str    						=      get_module_path(_lowerCAmelCase		)
      snake_case__  :    Union[str, Any]    						=      importlib.import_module(_lowerCAmelCase		)

      return test_module





def 						__snake_case(				_lowerCAmelCase		)		-> int:
      snake_case__  :    List[Any]    						=      []
      snake_case__  :    Optional[int]    						=      get_test_module(_lowerCAmelCase		)
      for attr in dir(_lowerCAmelCase		):
            if attr.endswith("""ModelTester"""		):
                  tester_classes.append(getattr(_lowerCAmelCase     ,   _lowerCAmelCase		)		)

    # sort with class names
      return sorted(_lowerCAmelCase     ,   key=lambda _lowerCAmelCase		: x.__name__		)





def 						__snake_case(				_lowerCAmelCase		)		-> Dict:
      snake_case__  :    List[str]    						=      []
      snake_case__  :    Any    						=      get_test_module(_lowerCAmelCase		)
      for attr in dir(_lowerCAmelCase		):
            snake_case__  :    Dict    						=      getattr(_lowerCAmelCase     ,   _lowerCAmelCase		)
            # (TF/Flax)ModelTesterMixin is also an attribute in specific model test module. Let's exclude them by checking
            # `all_model_classes` is not empty (which also excludes other special classes).
            snake_case__  :    List[str]    						=      getattr(_lowerCAmelCase     ,   """all_model_classes"""     ,   []		)
            if len(_lowerCAmelCase		) > 0:
                  test_classes.append(_lowerCAmelCase		)

    # sort with class names
      return sorted(_lowerCAmelCase     ,   key=lambda _lowerCAmelCase		: x.__name__		)





def 						__snake_case(				_lowerCAmelCase		)		-> Dict:
      snake_case__  :    Any    						=      get_test_classes(_lowerCAmelCase		)
      snake_case__  :    Optional[Any]    						=      set()
      for test_class in test_classes:
            model_classes.update(test_class.all_model_classes		)

      # sort with class names
      return sorted(_lowerCAmelCase     ,   key=lambda _lowerCAmelCase		: x.__name__		)





def 						__snake_case(				_lowerCAmelCase		)		-> Optional[Any]:
      snake_case__  :    Optional[int]    						=      test_class()
      if hasattr(_lowerCAmelCase     ,   """setUp"""		):
            test.setUp()

      snake_case__  :    Any    						=      None
      if hasattr(_lowerCAmelCase     ,   """model_tester"""		):
            # `(TF/Flax)ModelTesterMixin` has this attribute default to `None`. Let's skip this case.
            if test.model_tester is not None:
                  snake_case__  :    Tuple    						=      test.model_tester.__class__

      return model_tester





def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase		)		-> Dict:
      snake_case__  :    Union[str, Any]    						=      get_test_classes(_lowerCAmelCase		)

      snake_case__  :    str    						=      []
      for test_class in test_classes:
            if model_class in test_class.all_model_classes:
                  target_test_classes.append(_lowerCAmelCase		)

    # sort with class names
      return sorted(_lowerCAmelCase     ,   key=lambda _lowerCAmelCase		: x.__name__		)





def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase		)		-> Tuple:
      snake_case__  :    Optional[Any]    						=      get_test_classes_for_model(_lowerCAmelCase     ,   _lowerCAmelCase		)

      snake_case__  :    Union[str, Any]    						=      []
      for test_class in test_classes:
            snake_case__  :    Tuple    						=      get_model_tester_from_test_class(_lowerCAmelCase		)
            if tester_class is not None:
                  tester_classes.append(_lowerCAmelCase		)

    # sort with class names
      return sorted(_lowerCAmelCase     ,   key=lambda _lowerCAmelCase		: x.__name__		)





def 						__snake_case(				_lowerCAmelCase		)		-> Union[str, Any]:
      snake_case__  :    Optional[Any]    						=      get_test_classes(_lowerCAmelCase		)
      snake_case__  :    Union[str, Any]    						=      {test_class: get_model_tester_from_test_class(_lowerCAmelCase		) for test_class in test_classes}
      return test_tester_mapping





def 						__snake_case(				_lowerCAmelCase		)		-> int:
      snake_case__  :    Any    						=      get_model_classes(_lowerCAmelCase		)
      snake_case__  :    Any    						=      {
          model_class: get_test_classes_for_model(_lowerCAmelCase     ,   _lowerCAmelCase		) for model_class in model_classes
      }
      return model_test_mapping





def 						__snake_case(				_lowerCAmelCase		)		-> Optional[int]:
      snake_case__  :    Union[str, Any]    						=      get_model_classes(_lowerCAmelCase		)
      snake_case__  :    str    						=      {
          model_class: get_tester_classes_for_model(_lowerCAmelCase     ,   _lowerCAmelCase		) for model_class in model_classes
      }
      return model_to_tester_mapping





def 						__snake_case(				_lowerCAmelCase		)		-> int:
      if isinstance(_lowerCAmelCase     ,   _lowerCAmelCase		):
            return o
      elif isinstance(_lowerCAmelCase     ,   _lowerCAmelCase		):
            return o.__name__
      elif isinstance(_lowerCAmelCase     ,   (list, tuple)		):
            return [to_json(_lowerCAmelCase		) for x in o]
      elif isinstance(_lowerCAmelCase     ,   _lowerCAmelCase		):
            return {to_json(_lowerCAmelCase		): to_json(_lowerCAmelCase		) for k, v in o.items()}
      else:
            return o
