




from __future__ import annotations

import inspect
import unittest
from math import floor

import numpy as np

from transformers import CvtConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
					import tensorflow as tf

					from transformers import TFCvtForImageClassification, TFCvtModel
					from transformers.models.cvt.modeling_tf_cvt import TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					from PIL import Image

					from transformers import AutoImageProcessor
class     a__					(      snake_case__   ):






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.config_class(**self.inputs_dict    )
											self.parent.assertTrue(hasattr(_A	,    "embed_dim"    )    )
											self.parent.assertTrue(hasattr(_A	,    "num_heads"    )    )






class     a__					:






						def __init__(     self	,    _A	,    _A=1_3	,    _A=6_4	,    _A=3	,    _A=[1_6, 4_8, 9_6]	,    _A=[1, 3, 6]	,    _A=[1, 2, 1_0]	,    _A=[7, 3, 3]	,    _A=[4, 2, 2]	,    _A=[2, 1, 1]	,    _A=[2, 2, 2]	,    _A=[False, False, True]	,    _A=[0.0, 0.0, 0.0]	,    _A=0.02	,    _A=1E-1_2	,    _A=True	,    _A=True	,    _A=2	,    ):




											"""simple docstring"""



											__lowerCAmelCase            = parent
											__lowerCAmelCase            = batch_size
											__lowerCAmelCase            = image_size
											__lowerCAmelCase            = patch_sizes
											__lowerCAmelCase            = patch_stride
											__lowerCAmelCase            = patch_padding
											__lowerCAmelCase            = is_training
											__lowerCAmelCase            = use_labels
											__lowerCAmelCase            = num_labels
											__lowerCAmelCase            = num_channels
											__lowerCAmelCase            = embed_dim
											__lowerCAmelCase            = num_heads
											__lowerCAmelCase            = stride_kv
											__lowerCAmelCase            = depth
											__lowerCAmelCase            = cls_token
											__lowerCAmelCase            = attention_drop_rate
											__lowerCAmelCase            = initializer_range
											__lowerCAmelCase            = layer_norm_eps






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

											__lowerCAmelCase            = None
											if self.use_labels:
																# create a random int32 tensor of given shape
																__lowerCAmelCase            = ids_tensor([self.batch_size]	,    self.num_labels    )

											__lowerCAmelCase            = self.get_config()
											return config, pixel_values, labels






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											return CvtConfig(
											    image_size=self.image_size	,    num_labels=self.num_labels	,    num_channels=self.num_channels	,    embed_dim=self.embed_dim	,    num_heads=self.num_heads	,    patch_sizes=self.patch_sizes	,    patch_padding=self.patch_padding	,    patch_stride=self.patch_stride	,    stride_kv=self.stride_kv	,    depth=self.depth	,    cls_token=self.cls_token	,    attention_drop_rate=self.attention_drop_rate	,    initializer_range=self.initializer_range	,    )






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFCvtModel(config=_A    )
											__lowerCAmelCase            = model(_A	,    training=_A    )
											__lowerCAmelCase            = (self.image_size, self.image_size)
											__lowerCAmelCase     ,							__lowerCAmelCase            = image_size[0], image_size[1]
											for i in range(len(self.depth    )    ):
																__lowerCAmelCase            = floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1    )
																__lowerCAmelCase            = floor(((width + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1    )
											self.parent.assertEqual(result.last_hidden_state.shape	,    (self.batch_size, self.embed_dim[-1], height, width)    )






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.num_labels
											__lowerCAmelCase            = TFCvtForImageClassification(_A    )
											__lowerCAmelCase            = model(_A	,    labels=_A	,    training=_A    )
											self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.num_labels)    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.prepare_config_and_inputs()
											__lowerCAmelCase     ,							__lowerCAmelCase     ,							__lowerCAmelCase            = config_and_inputs
											__lowerCAmelCase            = {"pixel_values": pixel_values}
											return config, inputs_dict






@require_tf
class     a__					(      snake_case__				,       snake_case__				,       unittest.TestCase   ):
						_a							:					Union[str, Any]      =						(TFCvtModel, TFCvtForImageClassification) if is_tf_available() else ()
						_a							:					Optional[int]      =						(
						    {"""feature-extraction""": TFCvtModel, """image-classification""": TFCvtForImageClassification}
						    if is_tf_available()
						    else {}
						)
						_a							:					Tuple      =						False
						_a							:					Union[str, Any]      =						False
						_a							:					int      =						False
						_a							:					Any      =						False
						_a							:					Tuple      =						False






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFCvtModelTester(self    )
											__lowerCAmelCase            = TFCvtConfigTester(self	,    config_class=_A	,    has_text_modality=_A	,    hidden_size=3_7    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											self.config_tester.create_and_test_config_common_properties()
											self.config_tester.create_and_test_config_to_json_string()
											self.config_tester.create_and_test_config_to_json_file()
											self.config_tester.create_and_test_config_from_and_save_pretrained()
											self.config_tester.create_and_test_config_with_num_labels()
											self.config_tester.check_config_can_be_init_without_params()
											self.config_tester.check_config_arguments_init()






						@unittest.skip(reason="Cvt does not output attentions"    )
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											pass






						@unittest.skip(reason="Cvt does not use inputs_embeds"    )
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											pass






						@unittest.skip(reason="Cvt does not support input and output embeddings"    )
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											pass






						@unittest.skipIf(
						    not is_tf_available() or len(tf.config.list_physical_devices("GPU"    )    ) == 0	,    reason="TF does not support backprop for grouped convolutions on CPU."	,    )
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											super().test_dataset_conversion()






						@unittest.skipIf(
						    not is_tf_available() or len(tf.config.list_physical_devices("GPU"    )    ) == 0	,    reason="TF does not support backprop for grouped convolutions on CPU."	,    )
						@slow
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											super().test_keras_fit()






						@unittest.skip(reason="Get `Failed to determine best cudnn convolution algo.` error after using TF 2.12+cuda 11.8"    )
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = tf.keras.mixed_precision.Policy("mixed_float16"    )
											tf.keras.mixed_precision.set_global_policy(_A    )
											super().test_keras_fit()
											tf.keras.mixed_precision.set_global_policy("float32"    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase     ,							__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																__lowerCAmelCase            = model_class(_A    )
																__lowerCAmelCase            = inspect.signature(model.call    )
																# signature.parameters is an OrderedDict => so arg_names order is deterministic
																__lowerCAmelCase            = [*signature.parameters.keys()]

																__lowerCAmelCase            = ["pixel_values"]
																self.assertListEqual(arg_names[:1]	,    _A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											def check_hidden_states_output(_A	,    _A	,    _A    ):
																__lowerCAmelCase            = model_class(_A    )

																__lowerCAmelCase            = model(**self._prepare_for_class(_A	,    _A    )    )
																__lowerCAmelCase            = outputs.hidden_states

																__lowerCAmelCase            = len(self.model_tester.depth    )
																self.assertEqual(len(_A    )	,    _A    )

																# verify the first hidden states (first block)
																self.assertListEqual(
																    list(hidden_states[0].shape[-3:]    )	,    [
																        self.model_tester.embed_dim[0],
																        self.model_tester.image_size // 4,
																        self.model_tester.image_size // 4,
																    ]	,    )

											__lowerCAmelCase     ,							__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																__lowerCAmelCase            = True
																check_hidden_states_output(_A	,    _A	,    _A    )

																# check that output_hidden_states also work using config
																del inputs_dict["output_hidden_states"]
																__lowerCAmelCase            = True

																check_hidden_states_output(_A	,    _A	,    _A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*_A    )






						@slow
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__lowerCAmelCase            = TFCvtModel.from_pretrained(_A    )
																self.assertIsNotNone(_A    )







def 			_a   (		):
					__lowerCAmelCase            = Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"  )
					return image


@require_tf
@require_vision
class     a__					(      unittest.TestCase   ):






						@cached_property
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]    )






						@slow
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]    )

											__lowerCAmelCase            = self.default_image_processor
											__lowerCAmelCase            = prepare_img()
											__lowerCAmelCase            = image_processor(images=_A	,    return_tensors="tf"    )

											# forward pass
											__lowerCAmelCase            = model(**_A    )

											# verify the logits
											__lowerCAmelCase            = tf.TensorShape((1, 1_0_0_0)    )
											self.assertEqual(outputs.logits.shape	,    _A    )

											__lowerCAmelCase            = tf.constant([0.92_85, 0.90_15, -0.31_50]    )
											self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()	,    _A	,    atol=1E-4    )    )










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   (		SCREAMING_SNAKE_CASE_							:       Tuple  ):
					__lowerCAmelCase            = VideoMAEConfig()

					set_architecture_configs(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  )

					if "finetuned" not in model_name:
										__lowerCAmelCase            = False

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

					return config







def 			_a   (		SCREAMING_SNAKE_CASE_							:       int   ,   SCREAMING_SNAKE_CASE_							:       List[str]  ):
					if "small" in model_name:
										__lowerCAmelCase            = 3_84
										__lowerCAmelCase            = 15_36
										__lowerCAmelCase            = 12
										__lowerCAmelCase            = 16
										__lowerCAmelCase            = 12
										__lowerCAmelCase            = 3
										__lowerCAmelCase            = 1_92
										__lowerCAmelCase            = 7_68
					elif "large" in model_name:
										__lowerCAmelCase            = 10_24
										__lowerCAmelCase            = 40_96
										__lowerCAmelCase            = 24
										__lowerCAmelCase            = 16
										__lowerCAmelCase            = 12
										__lowerCAmelCase            = 8
										__lowerCAmelCase            = 5_12
										__lowerCAmelCase            = 20_48
					elif "huge" in model_name:
										__lowerCAmelCase            = 12_80
										__lowerCAmelCase            = 51_20
										__lowerCAmelCase            = 32
										__lowerCAmelCase            = 16
										__lowerCAmelCase            = 12
										__lowerCAmelCase            = 8
										__lowerCAmelCase            = 6_40
										__lowerCAmelCase            = 25_60
					elif "base" not in model_name:
										raise ValueError("Model name should include either \"small\", \"base\", \"large\", or \"huge\""  )







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

					return name







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Any   ,   SCREAMING_SNAKE_CASE_							:       Tuple  ):
					for key in orig_state_dict.copy().keys():
										__lowerCAmelCase            = orig_state_dict.pop(SCREAMING_SNAKE_CASE_  )

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

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

					return orig_state_dict







def 			_a   (		):
					__lowerCAmelCase            = hf_hub_download(
					    repo_id="hf-internal-testing/spaghetti-video"   ,   filename="eating_spaghetti.npy"   ,   repo_type="dataset"  )
					__lowerCAmelCase            = np.load(SCREAMING_SNAKE_CASE_  )
					return list(SCREAMING_SNAKE_CASE_  )







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Tuple   ,   SCREAMING_SNAKE_CASE_							:       Optional[int]   ,   SCREAMING_SNAKE_CASE_							:       Union[str, Any]   ,   SCREAMING_SNAKE_CASE_							:       List[str]  ):
					__lowerCAmelCase            = get_videomae_config(SCREAMING_SNAKE_CASE_  )

					if "finetuned" in model_name:
										__lowerCAmelCase            = VideoMAEForVideoClassification(SCREAMING_SNAKE_CASE_  )
					else:
										__lowerCAmelCase            = VideoMAEForPreTraining(SCREAMING_SNAKE_CASE_  )

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

					model.load_state_dict(SCREAMING_SNAKE_CASE_  )
					model.eval()

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

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

					__lowerCAmelCase            = model(**SCREAMING_SNAKE_CASE_  )
					__lowerCAmelCase            = outputs.logits

					__lowerCAmelCase            = [
					    "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            = torch.Size([1, 4_00]  )
										__lowerCAmelCase            = torch.tensor([-0.92_91, -0.40_61, -0.93_07]  )
					elif model_name == "videomae-small-finetuned-ssv2":
										__lowerCAmelCase            = torch.Size([1, 1_74]  )
										__lowerCAmelCase            = torch.tensor([0.26_71, -0.46_89, -0.82_35]  )
					elif model_name == "videomae-base":
										__lowerCAmelCase            = torch.Size([1, 14_08, 15_36]  )
										__lowerCAmelCase            = 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            = torch.Size([1, 14_08, 15_36]  )
										__lowerCAmelCase            = 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            = torch.tensor([0.51_42]  ) if config.norm_pix_loss else torch.tensor([0.64_69]  )
					elif model_name == "videomae-large":
										__lowerCAmelCase            = torch.Size([1, 14_08, 15_36]  )
										__lowerCAmelCase            = 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            = torch.Size([1, 4_00]  )
										__lowerCAmelCase            = torch.tensor([0.07_71, 0.00_11, -0.36_25]  )
					elif model_name == "videomae-huge-finetuned-kinetics":
										__lowerCAmelCase            = torch.Size([1, 4_00]  )
										__lowerCAmelCase            = torch.tensor([0.24_33, 0.16_32, -0.48_94]  )
					elif model_name == "videomae-base-short-finetuned-kinetics":
										__lowerCAmelCase            = torch.Size([1, 4_00]  )
										__lowerCAmelCase            = torch.tensor([0.65_88, 0.09_90, -0.24_93]  )
					elif model_name == "videomae-base-finetuned-kinetics":
										__lowerCAmelCase            = torch.Size([1, 4_00]  )
										__lowerCAmelCase            = torch.tensor([0.36_69, -0.06_88, -0.24_21]  )
					elif model_name == "videomae-base-short-ssv2":
										__lowerCAmelCase            = torch.Size([1, 14_08, 15_36]  )
										__lowerCAmelCase            = 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            = torch.Size([1, 1_74]  )
										__lowerCAmelCase            = torch.tensor([-0.05_37, -0.15_39, -0.32_66]  )
					elif model_name == "videomae-base-ssv2":
										__lowerCAmelCase            = torch.Size([1, 14_08, 15_36]  )
										__lowerCAmelCase            = 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            = torch.Size([1, 1_74]  )
										__lowerCAmelCase            = 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]   ,   SCREAMING_SNAKE_CASE_   ,   atol=1E-4  )
					else:
										print("Logits:"   ,   logits[0, :3, :3]  )
										assert torch.allclose(logits[0, :3, :3]   ,   SCREAMING_SNAKE_CASE_   ,   atol=1E-4  )
					print("Logits ok!"  )

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

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


if __name__ == "__main__":
					UpperCamelCase__					    =    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."""
					)

					UpperCamelCase__					    =    parser.parse_args()
					convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)




