



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       lowercase       (   SCREAMING_SNAKE_CASE__       :	int				)  ->     Optional[int]:
 _snake_case    :  List[Any]					=					VideoMAEConfig()

 set_architecture_configs(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				)

 if "finetuned" not in model_name:
  _snake_case    :  str					=					False

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

 return config



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	int				,       SCREAMING_SNAKE_CASE__       :	str				)  ->     List[str]:
 if "small" in model_name:
  _snake_case    :  Union[str, Any]					=					384
  _snake_case    :  List[Any]					=					1_536
  _snake_case    :  Optional[int]					=					12
  _snake_case    :  str					=					16
  _snake_case    :  int					=					12
  _snake_case    :  List[Any]					=					3
  _snake_case    :  Optional[int]					=					192
  _snake_case    :  List[Any]					=					768
 elif "large" in model_name:
  _snake_case    :  Tuple					=					1_024
  _snake_case    :  int					=					4_096
  _snake_case    :  Union[str, Any]					=					24
  _snake_case    :  str					=					16
  _snake_case    :  Dict					=					12
  _snake_case    :  Any					=					8
  _snake_case    :  Union[str, Any]					=					512
  _snake_case    :  Tuple					=					2_048
 elif "huge" in model_name:
  _snake_case    :  Tuple					=					1_280
  _snake_case    :  Dict					=					5_120
  _snake_case    :  List[str]					=					32
  _snake_case    :  Tuple					=					16
  _snake_case    :  int					=					12
  _snake_case    :  int					=					8
  _snake_case    :  str					=					640
  _snake_case    :  List[Any]					=					2_560
 elif "base" not in model_name:
  raise ValueError("""Model name should include either \"small\", \"base\", \"large\", or \"huge\""""				)



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	Optional[int]				)  ->     Optional[int]:
 if "encoder." in name:
  _snake_case    :  Optional[Any]					=					name.replace("""encoder."""				,       """"""				)
 if "cls_token" in name:
  _snake_case    :  Optional[Any]					=					name.replace("""cls_token"""				,       """videomae.embeddings.cls_token"""				)
 if "decoder_pos_embed" in name:
  _snake_case    :  Optional[Any]					=					name.replace("""decoder_pos_embed"""				,       """decoder.decoder_pos_embed"""				)
 if "pos_embed" in name and "decoder" not in name:
  _snake_case    :  Dict					=					name.replace("""pos_embed"""				,       """videomae.embeddings.position_embeddings"""				)
 if "patch_embed.proj" in name:
  _snake_case    :  str					=					name.replace("""patch_embed.proj"""				,       """videomae.embeddings.patch_embeddings.projection"""				)
 if "patch_embed.norm" in name:
  _snake_case    :  List[str]					=					name.replace("""patch_embed.norm"""				,       """videomae.embeddings.norm"""				)
 if "decoder.blocks" in name:
  _snake_case    :  List[str]					=					name.replace("""decoder.blocks"""				,       """decoder.decoder_layers"""				)
 if "blocks" in name:
  _snake_case    :  Union[str, Any]					=					name.replace("""blocks"""				,       """videomae.encoder.layer"""				)
 if "attn.proj" in name:
  _snake_case    :  List[str]					=					name.replace("""attn.proj"""				,       """attention.output.dense"""				)
 if "attn" in name and "bias" not in name:
  _snake_case    :  str					=					name.replace("""attn"""				,       """attention.self"""				)
 if "attn" in name:
  _snake_case    :  Union[str, Any]					=					name.replace("""attn"""				,       """attention.attention"""				)
 if "norm1" in name:
  _snake_case    :  Optional[Any]					=					name.replace("""norm1"""				,       """layernorm_before"""				)
 if "norm2" in name:
  _snake_case    :  Tuple					=					name.replace("""norm2"""				,       """layernorm_after"""				)
 if "mlp.fc1" in name:
  _snake_case    :  Dict					=					name.replace("""mlp.fc1"""				,       """intermediate.dense"""				)
 if "mlp.fc2" in name:
  _snake_case    :  Tuple					=					name.replace("""mlp.fc2"""				,       """output.dense"""				)
 if "decoder_embed" in name:
  _snake_case    :  Optional[int]					=					name.replace("""decoder_embed"""				,       """decoder.decoder_embed"""				)
 if "decoder_norm" in name:
  _snake_case    :  List[Any]					=					name.replace("""decoder_norm"""				,       """decoder.decoder_norm"""				)
 if "decoder_pred" in name:
  _snake_case    :  Optional[int]					=					name.replace("""decoder_pred"""				,       """decoder.decoder_pred"""				)
 if "norm.weight" in name and "decoder" not in name and "fc" not in name:
  _snake_case    :  Tuple					=					name.replace("""norm.weight"""				,       """videomae.layernorm.weight"""				)
 if "norm.bias" in name and "decoder" not in name and "fc" not in name:
  _snake_case    :  Union[str, Any]					=					name.replace("""norm.bias"""				,       """videomae.layernorm.bias"""				)
 if "head" in name and "decoder" not in name:
  _snake_case    :  List[str]					=					name.replace("""head"""				,       """classifier"""				)

 return name



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	Dict				,       SCREAMING_SNAKE_CASE__       :	List[Any]				)  ->     Optional[Any]:
 for key in orig_state_dict.copy().keys():
  _snake_case    :  str					=					orig_state_dict.pop(SCREAMING_SNAKE_CASE__				)

  if key.startswith("""encoder."""				):
   _snake_case    :  Union[str, Any]					=					key.replace("""encoder."""				,       """"""				)

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

 return orig_state_dict



def       lowercase       (   )  ->     Dict:
 _snake_case    :  Any					=					hf_hub_download(
     repo_id="""hf-internal-testing/spaghetti-video"""				,       filename="""eating_spaghetti.npy"""				,       repo_type="""dataset"""				)
 _snake_case    :  str					=					np.load(SCREAMING_SNAKE_CASE__				)
 return list(SCREAMING_SNAKE_CASE__				)



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	str				,       SCREAMING_SNAKE_CASE__       :	int				,       SCREAMING_SNAKE_CASE__       :	int				,       SCREAMING_SNAKE_CASE__       :	Optional[Any]				)  ->     List[str]:
 _snake_case    :  str					=					get_videomae_config(SCREAMING_SNAKE_CASE__				)

 if "finetuned" in model_name:
  _snake_case    :  List[Any]					=					VideoMAEForVideoClassification(SCREAMING_SNAKE_CASE__				)
 else:
  _snake_case    :  int					=					VideoMAEForPreTraining(SCREAMING_SNAKE_CASE__				)

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

 model.load_state_dict(SCREAMING_SNAKE_CASE__				)
 model.eval()

 # verify model on basic input
 _snake_case    :  Any					=					VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]				,       image_std=[0.5, 0.5, 0.5]				)
 _snake_case    :  int					=					prepare_video()
 _snake_case    :  Optional[Any]					=					image_processor(SCREAMING_SNAKE_CASE__				,       return_tensors="""pt"""				)

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

 _snake_case    :  Dict					=					model(**SCREAMING_SNAKE_CASE__				)
 _snake_case    :  Dict					=					outputs.logits

 _snake_case    :  Optional[Any]					=					[
     """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    :  List[Any]					=					torch.Size([1, 400]				)
  _snake_case    :  int					=					torch.tensor([-0.9_2_9_1, -0.4_0_6_1, -0.9_3_0_7]				)
 elif model_name == "videomae-small-finetuned-ssv2":
  _snake_case    :  Tuple					=					torch.Size([1, 174]				)
  _snake_case    :  Tuple					=					torch.tensor([0.2_6_7_1, -0.4_6_8_9, -0.8_2_3_5]				)
 elif model_name == "videomae-base":
  _snake_case    :  Dict					=					torch.Size([1, 1_408, 1_536]				)
  _snake_case    :  str					=					torch.tensor([[0.7_7_3_9, 0.7_9_6_8, 0.7_0_8_9], [0.6_7_0_1, 0.7_4_8_7, 0.6_2_0_9], [0.4_2_8_7, 0.5_1_5_8, 0.4_7_7_3]]				)
 elif model_name == "videomae-base-short":
  _snake_case    :  Tuple					=					torch.Size([1, 1_408, 1_536]				)
  _snake_case    :  str					=					torch.tensor([[0.7_9_9_4, 0.9_6_1_2, 0.8_5_0_8], [0.7_4_0_1, 0.8_9_5_8, 0.8_3_0_2], [0.5_8_6_2, 0.7_4_6_8, 0.7_3_2_5]]				)
  # we verified the loss both for normalized and unnormalized targets for this one
  _snake_case    :  Optional[Any]					=					torch.tensor([0.5_1_4_2]				) if config.norm_pix_loss else torch.tensor([0.6_4_6_9]				)
 elif model_name == "videomae-large":
  _snake_case    :  List[Any]					=					torch.Size([1, 1_408, 1_536]				)
  _snake_case    :  Optional[Any]					=					torch.tensor([[0.7_1_4_9, 0.7_9_9_7, 0.6_9_6_6], [0.6_7_6_8, 0.7_8_6_9, 0.6_9_4_8], [0.5_1_3_9, 0.6_2_2_1, 0.5_6_0_5]]				)
 elif model_name == "videomae-large-finetuned-kinetics":
  _snake_case    :  Optional[Any]					=					torch.Size([1, 400]				)
  _snake_case    :  Any					=					torch.tensor([0.0_7_7_1, 0.0_0_1_1, -0.3_6_2_5]				)
 elif model_name == "videomae-huge-finetuned-kinetics":
  _snake_case    :  Dict					=					torch.Size([1, 400]				)
  _snake_case    :  Tuple					=					torch.tensor([0.2_4_3_3, 0.1_6_3_2, -0.4_8_9_4]				)
 elif model_name == "videomae-base-short-finetuned-kinetics":
  _snake_case    :  str					=					torch.Size([1, 400]				)
  _snake_case    :  Tuple					=					torch.tensor([0.6_5_8_8, 0.0_9_9_0, -0.2_4_9_3]				)
 elif model_name == "videomae-base-finetuned-kinetics":
  _snake_case    :  List[Any]					=					torch.Size([1, 400]				)
  _snake_case    :  Optional[Any]					=					torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1]				)
 elif model_name == "videomae-base-short-ssv2":
  _snake_case    :  Optional[int]					=					torch.Size([1, 1_408, 1_536]				)
  _snake_case    :  str					=					torch.tensor([[0.4_7_1_2, 0.5_2_9_6, 0.5_7_8_6], [0.2_2_7_8, 0.2_7_2_9, 0.4_0_2_6], [0.0_3_5_2, 0.0_7_3_0, 0.2_5_0_6]]				)
 elif model_name == "videomae-base-short-finetuned-ssv2":
  _snake_case    :  str					=					torch.Size([1, 174]				)
  _snake_case    :  List[str]					=					torch.tensor([-0.0_5_3_7, -0.1_5_3_9, -0.3_2_6_6]				)
 elif model_name == "videomae-base-ssv2":
  _snake_case    :  Dict					=					torch.Size([1, 1_408, 1_536]				)
  _snake_case    :  Dict					=					torch.tensor([[0.8_1_3_1, 0.8_7_2_7, 0.8_5_4_6], [0.7_3_6_6, 0.9_3_7_7, 0.8_8_7_0], [0.5_9_3_5, 0.8_8_7_4, 0.8_5_6_4]]				)
 elif model_name == "videomae-base-finetuned-ssv2":
  _snake_case    :  Any					=					torch.Size([1, 174]				)
  _snake_case    :  List[Any]					=					torch.tensor([0.1_9_6_1, -0.8_3_3_7, -0.6_3_8_9]				)
 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]				,       SCREAMING_SNAKE_CASE__				,       atol=1e-4				)
 else:
  print("""Logits:"""				,       logits[0, :3, :3]				)
  assert torch.allclose(logits[0, :3, :3]				,       SCREAMING_SNAKE_CASE__				,       atol=1e-4				)
 print("""Logits ok!"""				)

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

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


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

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





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


a__    =	{
    """configuration_canine""": ["""CANINE_PRETRAINED_CONFIG_ARCHIVE_MAP""", """CanineConfig"""],
    """tokenization_canine""": ["""CanineTokenizer"""],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  a__    =	[
      """CANINE_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """CanineForMultipleChoice""",
      """CanineForQuestionAnswering""",
      """CanineForSequenceClassification""",
      """CanineForTokenClassification""",
      """CanineLayer""",
      """CanineModel""",
      """CaninePreTrainedModel""",
      """load_tf_weights_in_canine""",
  ]


if TYPE_CHECKING:
  from .configuration_canine import CANINE_PRETRAINED_CONFIG_ARCHIVE_MAP, CanineConfig
  from .tokenization_canine import CanineTokenizer

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_canine import (
        CANINE_PRETRAINED_MODEL_ARCHIVE_LIST,
        CanineForMultipleChoice,
        CanineForQuestionAnswering,
        CanineForSequenceClassification,
        CanineForTokenClassification,
        CanineLayer,
        CanineModel,
        CaninePreTrainedModel,
        load_tf_weights_in_canine,
    )


else:
  import sys

  a__    =	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
