

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					(      __a :     List[Any]      )     ->	Union[str, Any]:

					"""simple docstring"""

					_a      :       Optional[Any]   =						VideoMAEConfig()

					set_architecture_configs(__a						,__a      )

					if "finetuned" not in model_name:
										_a      :       Optional[Any]   =						False

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

					return config



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

					"""simple docstring"""

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



def   __UpperCAmelCase					(      __a :     Union[str, Any]      )     ->	Dict:

					"""simple docstring"""

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

					return name



def   __UpperCAmelCase					(      __a :     Dict						,__a :     Union[str, Any]      )     ->	Optional[Any]:

					"""simple docstring"""

					for key in orig_state_dict.copy().keys():
										_a      :       List[str]   =						orig_state_dict.pop(__a      )

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

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

					return orig_state_dict



def   __UpperCAmelCase					(      )     ->	Optional[int]:

					"""simple docstring"""

					_a      :       List[Any]   =						hf_hub_download(
					    repo_id='''hf-internal-testing/spaghetti-video'''						,filename='''eating_spaghetti.npy'''						,repo_type='''dataset'''      )
					_a      :       Dict   =						np.load(__a      )
					return list(__a      )



def   __UpperCAmelCase					(      __a :     Optional[int]						,__a :     Any						,__a :     Tuple						,__a :     Union[str, Any]      )     ->	List[Any]:

					"""simple docstring"""

					_a      :       Optional[Any]   =						get_videomae_config(__a      )

					if "finetuned" in model_name:
										_a      :       int   =						VideoMAEForVideoClassification(__a      )
					else:
										_a      :       Tuple   =						VideoMAEForPreTraining(__a      )

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

					model.load_state_dict(__a      )
					model.eval()

					# verify model on basic input
					_a      :       Tuple   =						VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]						,image_std=[0.5, 0.5, 0.5]      )
					_a      :       Union[str, Any]   =						prepare_video()
					_a      :       Any   =						image_processor(__a						,return_tensors='''pt'''      )

					if "finetuned" not in model_name:
										_a      :       Optional[Any]   =						hf_hub_download(repo_id='''hf-internal-testing/bool-masked-pos'''						,filename='''bool_masked_pos.pt'''      )
										_a      :       Dict   =						torch.load(__a      )

					_a      :       Optional[int]   =						model(**__a      )
					_a      :       str   =						outputs.logits

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






import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import is_speech_available, is_vision_available
from transformers.testing_utils import require_torch


if is_vision_available():
			from transformers import TvltImageProcessor

if is_speech_available():
			from transformers import TvltFeatureExtractor

from transformers import TvltProcessor

@require_torch
class 							UpperCAmelCase_						(     unittest.TestCase  ):
							"""simple docstring"""







							def 							__lowercase							(     self			)       ->       int:
												_a      :       Dict   =						'''ZinengTang/tvlt-base'''
												_a      :       List[str]   =						tempfile.mkdtemp()


							def 							__lowercase							(     self			,				**_a			)       ->       int:
												return TvltImageProcessor.from_pretrained(self.checkpoint			,				**_a			)


							def 							__lowercase							(     self			,				**_a			)       ->       List[Any]:
												return TvltFeatureExtractor.from_pretrained(self.checkpoint			,				**_a			)


							def 							__lowercase							(     self			)       ->       Optional[int]:
												shutil.rmtree(self.tmpdirname			)


							def 							__lowercase							(     self			)       ->       Dict:
												_a      :       Union[str, Any]   =						self.get_image_processor()
												_a      :       Dict   =						self.get_feature_extractor()

												_a      :       Optional[int]   =						TvltProcessor(image_processor=_a			,				feature_extractor=_a			)
												processor.save_pretrained(self.tmpdirname			)
												_a      :       Any   =						TvltProcessor.from_pretrained(self.tmpdirname			)

												self.assertIsInstance(processor.feature_extractor			,				_a			)
												self.assertIsInstance(processor.image_processor			,				_a			)


							def 							__lowercase							(     self			)       ->       Any:
												_a      :       Optional[Any]   =						self.get_image_processor()
												_a      :       Dict   =						self.get_feature_extractor()

												_a      :       Dict   =						TvltProcessor(image_processor=_a			,				feature_extractor=_a			)

												_a      :       Union[str, Any]   =						np.ones([1_2_0_0_0]			)

												_a      :       Dict   =						feature_extractor(_a			,				return_tensors='''np'''			)
												_a      :       Tuple   =						processor(audio=_a			,				return_tensors='''np'''			)

												for key in audio_dict.keys():
																	self.assertAlmostEqual(audio_dict[key].sum()			,				input_processor[key].sum()			,				delta=1e-2			)


							def 							__lowercase							(     self			)       ->       int:
												_a      :       Optional[Any]   =						self.get_image_processor()
												_a      :       Union[str, Any]   =						self.get_feature_extractor()

												_a      :       Optional[Any]   =						TvltProcessor(image_processor=_a			,				feature_extractor=_a			)

												_a      :       List[Any]   =						np.ones([3, 2_2_4, 2_2_4]			)

												_a      :       int   =						image_processor(_a			,				return_tensors='''np'''			)
												_a      :       Optional[int]   =						processor(images=_a			,				return_tensors='''np'''			)

												for key in image_dict.keys():
																	self.assertAlmostEqual(image_dict[key].sum()			,				input_processor[key].sum()			,				delta=1e-2			)


							def 							__lowercase							(     self			)       ->       Union[str, Any]:
												_a      :       int   =						self.get_image_processor()
												_a      :       Union[str, Any]   =						self.get_feature_extractor()

												_a      :       Any   =						TvltProcessor(image_processor=_a			,				feature_extractor=_a			)

												_a      :       List[str]   =						np.ones([1_2_0_0_0]			)
												_a      :       Optional[int]   =						np.ones([3, 2_2_4, 2_2_4]			)

												_a      :       int   =						processor(audio=_a			,				images=_a			)

												self.assertListEqual(list(inputs.keys()			)			,				['''audio_values''', '''audio_mask''', '''pixel_values''', '''pixel_mask''']			)

												# test if it raises when no input is passed
												with pytest.raises(_a			):
																	processor()




							def 							__lowercase							(     self			)       ->       Union[str, Any]:
												_a      :       str   =						self.get_image_processor()
												_a      :       Union[str, Any]   =						self.get_feature_extractor()

												_a      :       Dict   =						TvltProcessor(image_processor=_a			,				feature_extractor=_a			)

												self.assertListEqual(
												    processor.model_input_names			,				image_processor.model_input_names + feature_extractor.model_input_names			,				msg='''`processor` and `image_processor`+`feature_extractor` model input names do not match'''			,				)



