





'''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 							UpperCAmelCase_    (	lowerCAmelCase_							):

				"""simple docstring"""






				lowercase         = VideoMAEConfig()

				set_architecture_configs(lowerCAmelCase_		,							lowerCAmelCase_							)

				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(lowerCAmelCase_		,							lowerCAmelCase_		,							repo_type="dataset"							)		,							"r"							)							)
								lowercase         = {int(lowerCAmelCase_							): v for k, v in idalabel.items()}
								lowercase         = idalabel
								lowercase         = {v: k for k, v in idalabel.items()}

				return config







def 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_							):

				"""simple docstring"""






				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 							UpperCAmelCase_    (	lowerCAmelCase_							):

				"""simple docstring"""






				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 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_							):

				"""simple docstring"""






				for key in orig_state_dict.copy().keys():
								lowercase         = orig_state_dict.pop(lowerCAmelCase_							)

								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 							UpperCAmelCase_    (	):

				"""simple docstring"""






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







def 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_		,							lowerCAmelCase_		,							lowerCAmelCase_							):

				"""simple docstring"""






				lowercase         = get_videomae_config(lowerCAmelCase_							)

				if "finetuned" in model_name:
								lowercase         = VideoMAEForVideoClassification(lowerCAmelCase_							)
				else:
								lowercase         = VideoMAEForPreTraining(lowerCAmelCase_							)

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

				model.load_state_dict(lowerCAmelCase_							)
				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(lowerCAmelCase_		,							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(lowerCAmelCase_							)

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

			__lowerCamelCase : Dict											=			parser.parse_args()
			convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)











'''simple docstring'''

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCamelCase : Optional[Any]											=			logging.get_logger(__name__)

__lowerCamelCase : Optional[Any]											=			{
    "facebook/vit-mae-base": "https://huggingface.co/facebook/vit-mae-base/resolve/main/config.json",
    # See all ViT MAE models at https://huggingface.co/models?filter=vit-mae
}

class       UpperCAmelCase (						_lowercase  ):
		UpperCAmelCase					:	Union[str, Any] 			=		'''vit_mae'''




		def __init__(self      :	Dict							,			A__      :	Tuple=7_6_8							,			A__      :	str=1_2							,			A__      :	Dict=1_2							,			A__      :	Optional[Any]=3_0_7_2							,			A__      :	Union[str, Any]="gelu"							,			A__      :	Optional[int]=0.0							,			A__      :	str=0.0							,			A__      :	Optional[Any]=0.0_2							,			A__      :	Optional[Any]=1e-12							,			A__      :	Tuple=2_2_4							,			A__      :	str=1_6							,			A__      :	Optional[int]=3							,			A__      :	List[Any]=True							,			A__      :	Dict=1_6							,			A__      :	Dict=5_1_2							,			A__      :	Dict=8							,			A__      :	Union[str, Any]=2_0_4_8							,			A__      :	str=0.7_5							,			A__      :	str=False							,			**A__      :	Any							,			)					-> Union[str, Any]:
						super().__init__(**A__				)

						lowercase         = hidden_size
						lowercase         = num_hidden_layers
						lowercase         = num_attention_heads
						lowercase         = intermediate_size
						lowercase         = hidden_act
						lowercase         = hidden_dropout_prob
						lowercase         = attention_probs_dropout_prob
						lowercase         = initializer_range
						lowercase         = layer_norm_eps
						lowercase         = image_size
						lowercase         = patch_size
						lowercase         = num_channels
						lowercase         = qkv_bias
						lowercase         = decoder_num_attention_heads
						lowercase         = decoder_hidden_size
						lowercase         = decoder_num_hidden_layers
						lowercase         = decoder_intermediate_size
						lowercase         = mask_ratio
						lowercase         = norm_pix_loss




