



"""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   snake_case__     (	__lowerCamelCase       :					Optional[Any]			):
			"""simple docstring"""
			lowerCamelCase__	:	str				=VideoMAEConfig()

			set_architecture_configs(__lowerCamelCase							,	__lowerCamelCase			)

			if "finetuned" not in model_name:
						lowerCamelCase__	:	int				=False

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

			return config






def   snake_case__     (	__lowerCamelCase       :					List[str]							,	__lowerCamelCase       :					Optional[Any]			):
			"""simple docstring"""
			if "small" in model_name:
						lowerCamelCase__	:	Optional[Any]				=384
						lowerCamelCase__	:	List[Any]				=1536
						lowerCamelCase__	:	int				=12
						lowerCamelCase__	:	Dict				=16
						lowerCamelCase__	:	List[Any]				=12
						lowerCamelCase__	:	Optional[Any]				=3
						lowerCamelCase__	:	Union[str, Any]				=192
						lowerCamelCase__	:	str				=768
			elif "large" in model_name:
						lowerCamelCase__	:	Union[str, Any]				=1024
						lowerCamelCase__	:	str				=4096
						lowerCamelCase__	:	int				=24
						lowerCamelCase__	:	Dict				=16
						lowerCamelCase__	:	Union[str, Any]				=12
						lowerCamelCase__	:	List[Any]				=8
						lowerCamelCase__	:	int				=512
						lowerCamelCase__	:	Optional[Any]				=2048
			elif "huge" in model_name:
						lowerCamelCase__	:	Optional[int]				=1280
						lowerCamelCase__	:	Optional[int]				=5120
						lowerCamelCase__	:	List[Any]				=32
						lowerCamelCase__	:	List[Any]				=16
						lowerCamelCase__	:	Optional[Any]				=12
						lowerCamelCase__	:	Dict				=8
						lowerCamelCase__	:	List[Any]				=640
						lowerCamelCase__	:	Any				=2560
			elif "base" not in model_name:
						raise ValueError('''Model name should include either "small", "base", "large", or "huge"'''			)






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

			return name






def   snake_case__     (	__lowerCamelCase       :					int							,	__lowerCamelCase       :					int			):
			"""simple docstring"""
			for key in orig_state_dict.copy().keys():
						lowerCamelCase__	:	Dict				=orig_state_dict.pop(__lowerCamelCase			)

						if key.startswith('''encoder.'''			):
									lowerCamelCase__	:	Optional[int]				=key.replace('''encoder.'''							,	''''''			)

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

			return orig_state_dict






def   snake_case__     (	):
			"""simple docstring"""
			lowerCamelCase__	:	List[Any]				=hf_hub_download(
			    repo_id='''hf-internal-testing/spaghetti-video'''							,	filename='''eating_spaghetti.npy'''							,	repo_type='''dataset'''			)
			lowerCamelCase__	:	Optional[Any]				=np.load(__lowerCamelCase			)
			return list(__lowerCamelCase			)






def   snake_case__     (	__lowerCamelCase       :					Optional[Any]							,	__lowerCamelCase       :					Union[str, Any]							,	__lowerCamelCase       :					List[Any]							,	__lowerCamelCase       :					Optional[int]			):
			"""simple docstring"""
			lowerCamelCase__	:	str				=get_videomae_config(__lowerCamelCase			)

			if "finetuned" in model_name:
						lowerCamelCase__	:	Tuple				=VideoMAEForVideoClassification(__lowerCamelCase			)
			else:
						lowerCamelCase__	:	int				=VideoMAEForPreTraining(__lowerCamelCase			)

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

			model.load_state_dict(__lowerCamelCase			)
			model.eval()

			# verify model on basic input
			lowerCamelCase__	:	Dict				=VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]							,	image_std=[0.5, 0.5, 0.5]			)
			lowerCamelCase__	:	int				=prepare_video()
			lowerCamelCase__	:	Tuple				=image_processor(__lowerCamelCase							,	return_tensors='''pt'''			)

			if "finetuned" not in model_name:
						lowerCamelCase__	:	Tuple				=hf_hub_download(repo_id='''hf-internal-testing/bool-masked-pos'''							,	filename='''bool_masked_pos.pt'''			)
						lowerCamelCase__	:	Union[str, Any]				=torch.load(__lowerCamelCase			)

			lowerCamelCase__	:	int				=model(**__lowerCamelCase			)
			lowerCamelCase__	:	Dict				=outputs.logits

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

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









"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


_lowercase				:				Optional[Any]   		=		{
    "configuration_clip": [
        "CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "CLIPConfig",
        "CLIPOnnxConfig",
        "CLIPTextConfig",
        "CLIPVisionConfig",
    ],
    "processing_clip": ["CLIPProcessor"],
    "tokenization_clip": ["CLIPTokenizer"],
}

try:
		if not is_tokenizers_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase				:				str   		=		["CLIPTokenizerFast"]

try:
		if not is_vision_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase				:				Any   		=		["CLIPFeatureExtractor"]
		_lowercase				:				int   		=		["CLIPImageProcessor"]

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase				:				Optional[Any]   		=		[
		    "CLIP_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "CLIPModel",
		    "CLIPPreTrainedModel",
		    "CLIPTextModel",
		    "CLIPTextModelWithProjection",
		    "CLIPVisionModel",
		    "CLIPVisionModelWithProjection",
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase				:				Dict   		=		[
		    "TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "TFCLIPModel",
		    "TFCLIPPreTrainedModel",
		    "TFCLIPTextModel",
		    "TFCLIPVisionModel",
		]

try:
		if not is_flax_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase				:				Union[str, Any]   		=		[
		    "FlaxCLIPModel",
		    "FlaxCLIPPreTrainedModel",
		    "FlaxCLIPTextModel",
		    "FlaxCLIPTextPreTrainedModel",
		    "FlaxCLIPVisionModel",
		    "FlaxCLIPVisionPreTrainedModel",
		]


if TYPE_CHECKING:
		from .configuration_clip import (
		    CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    CLIPConfig,
		    CLIPOnnxConfig,
		    CLIPTextConfig,
		    CLIPVisionConfig,
		)
		from .processing_clip import CLIPProcessor
		from .tokenization_clip import CLIPTokenizer

		try:
				if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tokenization_clip_fast import CLIPTokenizerFast

		try:
				if not is_vision_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .feature_extraction_clip import CLIPFeatureExtractor
				from .image_processing_clip import CLIPImageProcessor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_clip import (
				    CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
				    CLIPModel,
				    CLIPPreTrainedModel,
				    CLIPTextModel,
				    CLIPTextModelWithProjection,
				    CLIPVisionModel,
				    CLIPVisionModelWithProjection,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_clip import (
				    TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFCLIPModel,
				    TFCLIPPreTrainedModel,
				    TFCLIPTextModel,
				    TFCLIPVisionModel,
				)

		try:
				if not is_flax_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_flax_clip import (
				    FlaxCLIPModel,
				    FlaxCLIPPreTrainedModel,
				    FlaxCLIPTextModel,
				    FlaxCLIPTextPreTrainedModel,
				    FlaxCLIPVisionModel,
				    FlaxCLIPVisionPreTrainedModel,
				)


else:
		import sys

		_lowercase				:				Union[str, Any]   		=		_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




