




'''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 		__magic_name__    (	__UpperCAmelCase	)		->					Union[str, Any]:

					'''simple docstring'''

					snake_case_									=				VideoMAEConfig()

					set_architecture_configs(__UpperCAmelCase,      __UpperCAmelCase	)

					if "finetuned" not in model_name:
										snake_case_									=				False

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

					return config


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase	)		->					List[str]:

					'''simple docstring'''

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


def 		__magic_name__    (	__UpperCAmelCase	)		->					List[Any]:

					'''simple docstring'''

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

					return name


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase	)		->					Optional[int]:

					'''simple docstring'''

					for key in orig_state_dict.copy().keys():
										snake_case_									=				orig_state_dict.pop(__UpperCAmelCase	)

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

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

					return orig_state_dict


def 		__magic_name__    (	)		->					List[Any]:

					'''simple docstring'''

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


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase,      __UpperCAmelCase,      __UpperCAmelCase	)		->					Any:

					'''simple docstring'''

					snake_case_									=				get_videomae_config(__UpperCAmelCase	)

					if "finetuned" in model_name:
										snake_case_									=				VideoMAEForVideoClassification(__UpperCAmelCase	)
					else:
										snake_case_									=				VideoMAEForPreTraining(__UpperCAmelCase	)

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

					model.load_state_dict(__UpperCAmelCase	)
					model.eval()

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

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

					snake_case_									=				model(**__UpperCAmelCase	)
					snake_case_									=				outputs.logits

					snake_case_									=				[
					    '''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_									=				torch.Size([1, 400]	)
										snake_case_									=				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_									=				torch.Size([1, 174]	)
										snake_case_									=				torch.tensor([0.2_6_7_1, -0.4_6_8_9, -0.8_2_3_5]	)
					elif model_name == "videomae-base":
										snake_case_									=				torch.Size([1, 1408, 1536]	)
										snake_case_									=				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_									=				torch.Size([1, 1408, 1536]	)
										snake_case_									=				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_									=				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_									=				torch.Size([1, 1408, 1536]	)
										snake_case_									=				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_									=				torch.Size([1, 400]	)
										snake_case_									=				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_									=				torch.Size([1, 400]	)
										snake_case_									=				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_									=				torch.Size([1, 400]	)
										snake_case_									=				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_									=				torch.Size([1, 400]	)
										snake_case_									=				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_									=				torch.Size([1, 1408, 1536]	)
										snake_case_									=				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_									=				torch.Size([1, 174]	)
										snake_case_									=				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_									=				torch.Size([1, 1408, 1536]	)
										snake_case_									=				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_									=				torch.Size([1, 174]	)
										snake_case_									=				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],      __UpperCAmelCase,      atol=1e-4	)
					else:
										print('''Logits:''',      logits[0, :3, :3]	)
										assert torch.allclose(logits[0, :3, :3],      __UpperCAmelCase,      atol=1e-4	)
					print('''Logits ok!'''	)

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

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


if __name__ == "__main__":
							a  :				Optional[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.'
							)

							a  :				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'''







import warnings

from ...utils import logging
from .image_processing_videomae import VideoMAEImageProcessor


__lowerCAmelCase         =    logging.get_logger(__name__)







class UpperCAmelCase__	( lowercase__	):





  """simple docstring"""

  def __init__(		self      :     Tuple							,*_a      :     List[str]							,**_a      :     Any       ):


     '''simple docstring'''






     warnings.warn(
         'The class VideoMAEFeatureExtractor is deprecated and will be removed in version 5 of Transformers.'
         ' Please use VideoMAEImageProcessor instead.'							,_a							,)
     super().__init__(*_a							,**_a       )



