
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__    (       A_					):






  '''simple docstring'''





  __magic_name__				     =       VideoMAEConfig()

  set_architecture_configs(A_,     A_					)

  if "finetuned" not in model_name:
    __magic_name__				     =       False

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

  return config

def a__    (       A_,     A_					):






  '''simple docstring'''





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

def a__    (       A_					):






  '''simple docstring'''





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

  return name

def a__    (       A_,     A_					):






  '''simple docstring'''





  for key in orig_state_dict.copy().keys():
    __magic_name__				     =       orig_state_dict.pop(A_					)

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

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

  return orig_state_dict

def a__    (       ):






  '''simple docstring'''





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

def a__    (       A_,     A_,     A_,     A_					):






  '''simple docstring'''





  __magic_name__				     =       get_videomae_config(A_					)

  if "finetuned" in model_name:
    __magic_name__				     =       VideoMAEForVideoClassification(A_					)
  else:
    __magic_name__				     =       VideoMAEForPreTraining(A_					)

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

  model.load_state_dict(A_					)
  model.eval()

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

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

  __magic_name__				     =       model(**A_					)
  __magic_name__				     =       outputs.logits

  __magic_name__				     =       [
      """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":
    __magic_name__				     =       torch.Size([1, 400]					)
    __magic_name__				     =       torch.tensor([-0.9291, -0.4061, -0.9307]					)
  elif model_name == "videomae-small-finetuned-ssv2":
    __magic_name__				     =       torch.Size([1, 174]					)
    __magic_name__				     =       torch.tensor([0.2671, -0.4689, -0.8235]					)
  elif model_name == "videomae-base":
    __magic_name__				     =       torch.Size([1, 1408, 1536]					)
    __magic_name__				     =       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":
    __magic_name__				     =       torch.Size([1, 1408, 1536]					)
    __magic_name__				     =       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
    __magic_name__				     =       torch.tensor([0.5142]					) if config.norm_pix_loss else torch.tensor([0.6469]					)
  elif model_name == "videomae-large":
    __magic_name__				     =       torch.Size([1, 1408, 1536]					)
    __magic_name__				     =       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":
    __magic_name__				     =       torch.Size([1, 400]					)
    __magic_name__				     =       torch.tensor([0.0771, 0.0011, -0.3625]					)
  elif model_name == "videomae-huge-finetuned-kinetics":
    __magic_name__				     =       torch.Size([1, 400]					)
    __magic_name__				     =       torch.tensor([0.2433, 0.1632, -0.4894]					)
  elif model_name == "videomae-base-short-finetuned-kinetics":
    __magic_name__				     =       torch.Size([1, 400]					)
    __magic_name__				     =       torch.tensor([0.6588, 0.0990, -0.2493]					)
  elif model_name == "videomae-base-finetuned-kinetics":
    __magic_name__				     =       torch.Size([1, 400]					)
    __magic_name__				     =       torch.tensor([0.3669, -0.0688, -0.2421]					)
  elif model_name == "videomae-base-short-ssv2":
    __magic_name__				     =       torch.Size([1, 1408, 1536]					)
    __magic_name__				     =       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":
    __magic_name__				     =       torch.Size([1, 174]					)
    __magic_name__				     =       torch.tensor([-0.0537, -0.1539, -0.3266]					)
  elif model_name == "videomae-base-ssv2":
    __magic_name__				     =       torch.Size([1, 1408, 1536]					)
    __magic_name__				     =       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":
    __magic_name__				     =       torch.Size([1, 174]					)
    __magic_name__				     =       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],     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":
    __magic_name__				     =       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__":
    __lowerCAmelCase			:     str           =   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.'
    )

    __lowerCAmelCase			:     Union[str, Any]           =   parser.parse_args()
    convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)



import logging

import numpy as np
import pytest
from scipy.linalg import eigh

logging.basicConfig(level=logging.INFO, format="""%(message)s""")


def 			UpperCamelCase					(    __lowerCamelCase     : np.ndarray    ):
       return input_array.reshape((input_array.size, 1)    )


def 			UpperCamelCase					(    __lowerCamelCase     : np.ndarray				,					__lowerCamelCase     : np.ndarray				,					__lowerCamelCase     : int    ):
       snake_case       :   Any							=      np.nan
       for i in range(__lowerCamelCase    ):
              snake_case       :   List[str]							=      features[:, labels == i]
              snake_case       :   Dict							=      data.mean(1    )
              # Centralize the data of class i
              snake_case       :   Optional[Any]							=      data - column_reshape(__lowerCamelCase    )
              if i > 0:
                     # If covariance_sum is not None
                     covariance_sum += np.dot(__lowerCamelCase				,					centered_data.T    )
              else:
                     # If covariance_sum is np.nan (i.e. first loop)
                     snake_case       :   Optional[Any]							=      np.dot(__lowerCamelCase				,					centered_data.T    )

       return covariance_sum / features.shape[1]


def 			UpperCamelCase					(    __lowerCamelCase     : np.ndarray				,					__lowerCamelCase     : np.ndarray				,					__lowerCamelCase     : int    ):
       snake_case       :   Optional[Any]							=      features.mean(1    )
       snake_case       :   Tuple							=      np.nan
       for i in range(__lowerCamelCase    ):
              snake_case       :   Tuple							=      features[:, labels == i]
              snake_case       :   Tuple							=      data.shape[1]
              snake_case       :   List[str]							=      data.mean(1    )
              if i > 0:
                     # If covariance_sum is not None
                     covariance_sum += device_data * np.dot(
                         column_reshape(__lowerCamelCase    ) - column_reshape(__lowerCamelCase    )				,					(column_reshape(__lowerCamelCase    ) - column_reshape(__lowerCamelCase    )).T				,					)
              else:
                     # If covariance_sum is np.nan (i.e. first loop)
                     snake_case       :   Optional[int]							=      device_data * np.dot(
                         column_reshape(__lowerCamelCase    ) - column_reshape(__lowerCamelCase    )				,					(column_reshape(__lowerCamelCase    ) - column_reshape(__lowerCamelCase    )).T				,					)

       return covariance_sum / features.shape[1]


def 			UpperCamelCase					(    __lowerCamelCase     : np.ndarray				,					__lowerCamelCase     : int    ):

       # Check if the features have been loaded
       if features.any():
              snake_case       :   Tuple							=      features.mean(1    )
              # Center the dataset
              snake_case       :   List[str]							=      features - np.reshape(__lowerCamelCase				,					(data_mean.size, 1)    )
              snake_case       :   Optional[Any]							=      np.dot(__lowerCamelCase				,					centered_data.T    ) / features.shape[1]
              snake_case				,					snake_case       :   Dict							=      np.linalg.eigh(__lowerCamelCase    )
              # Take all the columns in the reverse order (-1), and then takes only the first
              snake_case       :   Optional[Any]							=      eigenvectors[:, ::-1][:, 0:dimensions]
              # Project the database on the new space
              snake_case       :   Union[str, Any]							=      np.dot(filtered_eigenvectors.T				,					__lowerCamelCase    )
              logging.info("Principal Component Analysis computed"    )

              return projected_data
       else:
              logging.basicConfig(level=logging.ERROR				,					format="%(message)s"				,					force=__lowerCamelCase    )
              logging.error("Dataset empty"    )
              raise AssertionError


def 			UpperCamelCase					(    __lowerCamelCase     : np.ndarray				,					__lowerCamelCase     : np.ndarray				,					__lowerCamelCase     : int				,					__lowerCamelCase     : int    ):
       assert classes > dimensions

       # Check if features have been already loaded
       if features.any:
              snake_case				,					snake_case       :   str							=      eigh(
                  covariance_between_classes(__lowerCamelCase				,					__lowerCamelCase				,					__lowerCamelCase    )				,					covariance_within_classes(__lowerCamelCase				,					__lowerCamelCase				,					__lowerCamelCase    )				,					)
              snake_case       :   str							=      eigenvectors[:, ::-1][:, :dimensions]
              snake_case				,					snake_case				,					snake_case       :   int							=      np.linalg.svd(__lowerCamelCase    )
              snake_case       :   List[Any]							=      svd_matrix[:, 0:dimensions]
              snake_case       :   Optional[Any]							=      np.dot(filtered_svd_matrix.T				,					__lowerCamelCase    )
              logging.info("Linear Discriminant Analysis computed"    )

              return projected_data
       else:
              logging.basicConfig(level=logging.ERROR				,					format="%(message)s"				,					force=__lowerCamelCase    )
              logging.error("Dataset empty"    )
              raise AssertionError


def 			UpperCamelCase					(    ):
       # Create dummy dataset with 2 classes and 3 features
       snake_case       :   str							=      np.array([[1, 2, 3, 4, 5], [2, 3, 4, 5, 6], [3, 4, 5, 6, 7]]    )
       snake_case       :   Union[str, Any]							=      np.array([0, 0, 0, 1, 1]    )
       snake_case       :   List[Any]							=      2
       snake_case       :   Any							=      2

       # Assert that the function raises an AssertionError if dimensions > classes
       with pytest.raises(__lowerCamelCase    ) as error_info:
              snake_case       :   str							=      linear_discriminant_analysis(
                  __lowerCamelCase				,					__lowerCamelCase				,					__lowerCamelCase				,					__lowerCamelCase    )
              if isinstance(__lowerCamelCase				,					np.ndarray    ):
                     raise AssertionError(
                         "Did not raise AssertionError for dimensions > classes"    )
              assert error_info.type is AssertionError


def 			UpperCamelCase					(    ):
       snake_case       :   List[str]							=      np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]    )
       snake_case       :   List[str]							=      2
       snake_case       :   int							=      np.array([[6.9282_0323, 8.6602_5404, 10.3923_0485], [3.0, 3.0, 3.0]]    )

       with pytest.raises(__lowerCamelCase    ) as error_info:
              snake_case       :   Union[str, Any]							=      principal_component_analysis(__lowerCamelCase				,					__lowerCamelCase    )
              if not np.allclose(__lowerCamelCase				,					__lowerCamelCase    ):
                     raise AssertionError
              assert error_info.type is AssertionError


if __name__ == "__main__":
   import doctest

   doctest.testmod()





