'''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__							(			_A:					Any )    ->	Optional[Any]:

		'''simple docstring'''


		lowerCAmelCase   					=  VideoMAEConfig()

		set_architecture_configs(_A      ,   _A )

		if "finetuned" not in model_name:
				lowerCAmelCase   					=  False

		if "finetuned" in model_name:
				lowerCAmelCase   					=  """huggingface/label-files"""
				if "kinetics" in model_name:
						lowerCAmelCase   					=  400
						lowerCAmelCase   					=  """kinetics400-id2label.json"""
				elif "ssv2" in model_name:
						lowerCAmelCase   					=  174
						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(_A      ,   _A      ,   repo_type="""dataset""" )      ,   """r""" ) )
				lowerCAmelCase   					=  {int(_A ): v for k, v in idalabel.items()}
				lowerCAmelCase   					=  idalabel
				lowerCAmelCase   					=  {v: k for k, v in idalabel.items()}

		return config






def 				snake_case__							(			_A:					Optional[int]      ,   _A:					Dict )    ->	Optional[int]:

		'''simple docstring'''


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






def 				snake_case__							(			_A:					str )    ->	Any:

		'''simple docstring'''


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

		'''simple docstring'''


		for key in orig_state_dict.copy().keys():
				lowerCAmelCase   					=  orig_state_dict.pop(_A )

				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 				snake_case__							(			)    ->	Dict:

		'''simple docstring'''


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






def 				snake_case__							(			_A:					Tuple      ,   _A:					int      ,   _A:					int      ,   _A:					str )    ->	List[Any]:

		'''simple docstring'''


		lowerCAmelCase   					=  get_videomae_config(_A )

		if "finetuned" in model_name:
				lowerCAmelCase   					=  VideoMAEForVideoClassification(_A )
		else:
				lowerCAmelCase   					=  VideoMAEForPreTraining(_A )

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

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

		lowerCAmelCase   					=  model(**_A )
		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, 400] )
				lowerCAmelCase   					=  torch.tensor([-0.9_291, -0.4_061, -0.9_307] )
		elif model_name == "videomae-small-finetuned-ssv2":
				lowerCAmelCase   					=  torch.Size([1, 174] )
				lowerCAmelCase   					=  torch.tensor([0.2_671, -0.4_689, -0.8_235] )
		elif model_name == "videomae-base":
				lowerCAmelCase   					=  torch.Size([1, 1408, 1536] )
				lowerCAmelCase   					=  torch.tensor([[0.7_739, 0.7_968, 0.7_089], [0.6_701, 0.7_487, 0.6_209], [0.4_287, 0.5_158, 0.4_773]] )
		elif model_name == "videomae-base-short":
				lowerCAmelCase   					=  torch.Size([1, 1408, 1536] )
				lowerCAmelCase   					=  torch.tensor([[0.7_994, 0.9_612, 0.8_508], [0.7_401, 0.8_958, 0.8_302], [0.5_862, 0.7_468, 0.7_325]] )
				# we verified the loss both for normalized and unnormalized targets for this one
				lowerCAmelCase   					=  torch.tensor([0.5_142] ) if config.norm_pix_loss else torch.tensor([0.6_469] )
		elif model_name == "videomae-large":
				lowerCAmelCase   					=  torch.Size([1, 1408, 1536] )
				lowerCAmelCase   					=  torch.tensor([[0.7_149, 0.7_997, 0.6_966], [0.6_768, 0.7_869, 0.6_948], [0.5_139, 0.6_221, 0.5_605]] )
		elif model_name == "videomae-large-finetuned-kinetics":
				lowerCAmelCase   					=  torch.Size([1, 400] )
				lowerCAmelCase   					=  torch.tensor([0.0_771, 0.0_011, -0.3_625] )
		elif model_name == "videomae-huge-finetuned-kinetics":
				lowerCAmelCase   					=  torch.Size([1, 400] )
				lowerCAmelCase   					=  torch.tensor([0.2_433, 0.1_632, -0.4_894] )
		elif model_name == "videomae-base-short-finetuned-kinetics":
				lowerCAmelCase   					=  torch.Size([1, 400] )
				lowerCAmelCase   					=  torch.tensor([0.6_588, 0.0_990, -0.2_493] )
		elif model_name == "videomae-base-finetuned-kinetics":
				lowerCAmelCase   					=  torch.Size([1, 400] )
				lowerCAmelCase   					=  torch.tensor([0.3_669, -0.0_688, -0.2_421] )
		elif model_name == "videomae-base-short-ssv2":
				lowerCAmelCase   					=  torch.Size([1, 1408, 1536] )
				lowerCAmelCase   					=  torch.tensor([[0.4_712, 0.5_296, 0.5_786], [0.2_278, 0.2_729, 0.4_026], [0.0_352, 0.0_730, 0.2_506]] )
		elif model_name == "videomae-base-short-finetuned-ssv2":
				lowerCAmelCase   					=  torch.Size([1, 174] )
				lowerCAmelCase   					=  torch.tensor([-0.0_537, -0.1_539, -0.3_266] )
		elif model_name == "videomae-base-ssv2":
				lowerCAmelCase   					=  torch.Size([1, 1408, 1536] )
				lowerCAmelCase   					=  torch.tensor([[0.8_131, 0.8_727, 0.8_546], [0.7_366, 0.9_377, 0.8_870], [0.5_935, 0.8_874, 0.8_564]] )
		elif model_name == "videomae-base-finetuned-ssv2":
				lowerCAmelCase   					=  torch.Size([1, 174] )
				lowerCAmelCase   					=  torch.tensor([0.1_961, -0.8_337, -0.6_389] )
		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":
				lowerCAmelCase   					=  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__":
						__lowercase          =  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          =  parser.parse_args()
						convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)





'''simple docstring'''




import unittest

import numpy as np

from transformers.file_utils import is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
						import torch

if is_vision_available():
						from PIL import Image

						from transformers import DPTImageProcessor






class 				a__(					unittest.TestCase	):


			'''simple docstring'''






			def __init__(   self       ,    __lowerCAmelCase       ,    __lowerCAmelCase=7       ,    __lowerCAmelCase=3       ,    __lowerCAmelCase=18       ,    __lowerCAmelCase=30       ,    __lowerCAmelCase=400       ,    __lowerCAmelCase=True       ,    __lowerCAmelCase=None       ,    __lowerCAmelCase=True       ,    __lowerCAmelCase=[0.5, 0.5, 0.5]       ,    __lowerCAmelCase=[0.5, 0.5, 0.5]       ,    ):




					"""simple docstring"""


					lowerCAmelCase   					=  size if size is not None else {"""height""": 18, """width""": 18}
					lowerCAmelCase   					=  parent
					lowerCAmelCase   					=  batch_size
					lowerCAmelCase   					=  num_channels
					lowerCAmelCase   					=  image_size
					lowerCAmelCase   					=  min_resolution
					lowerCAmelCase   					=  max_resolution
					lowerCAmelCase   					=  do_resize
					lowerCAmelCase   					=  size
					lowerCAmelCase   					=  do_normalize
					lowerCAmelCase   					=  image_mean
					lowerCAmelCase   					=  image_std







			def 							a_       (   self):




					"""simple docstring"""


					return {
					    "image_mean": self.image_mean,
					    "image_std": self.image_std,
					    "do_normalize": self.do_normalize,
					    "do_resize": self.do_resize,
					    "size": self.size,
					}







@require_torch
@require_vision
class 				a__(					lowerCAmelCase__ ,						unittest.TestCase	):


			'''simple docstring'''






			UpperCAmelCase_			:     Tuple      =       DPTImageProcessor if is_vision_available() else None



			def 							a_       (   self):




					"""simple docstring"""


					lowerCAmelCase   					=  DPTImageProcessingTester(self)



			@property
			def 							a_       (   self):




					"""simple docstring"""


					return self.image_processor_tester.prepare_image_processor_dict()



			def 							a_       (   self):




					"""simple docstring"""


					lowerCAmelCase   					=  self.image_processing_class(**self.image_processor_dict)
					self.assertTrue(hasattr(__lowerCAmelCase       ,    """image_mean"""))
					self.assertTrue(hasattr(__lowerCAmelCase       ,    """image_std"""))
					self.assertTrue(hasattr(__lowerCAmelCase       ,    """do_normalize"""))
					self.assertTrue(hasattr(__lowerCAmelCase       ,    """do_resize"""))
					self.assertTrue(hasattr(__lowerCAmelCase       ,    """size"""))



			def 							a_       (   self):




					"""simple docstring"""


					lowerCAmelCase   					=  self.image_processing_class.from_dict(self.image_processor_dict)
					self.assertEqual(image_processor.size       ,    {"""height""": 18, """width""": 18})

					lowerCAmelCase   					=  self.image_processing_class.from_dict(self.image_processor_dict       ,    size=42)
					self.assertEqual(image_processor.size       ,    {"""height""": 42, """width""": 42})



			def 							a_       (   self):




					"""simple docstring"""


					lowerCAmelCase   					=  self.image_processing_class(**self.image_processor_dict)
					# create random PIL images
					lowerCAmelCase   					=  prepare_image_inputs(self.image_processor_tester       ,    equal_resolution=__lowerCAmelCase)
					for image in image_inputs:
							self.assertIsInstance(__lowerCAmelCase       ,    Image.Image)

					# Test not batched input
					lowerCAmelCase   					=  image_processing(image_inputs[0]       ,    return_tensors="""pt""").pixel_values
					self.assertEqual(
					    encoded_images.shape       ,    (
					        1,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.size["""height"""],
					        self.image_processor_tester.size["""width"""],
					    )       ,    )

					# Test batched
					lowerCAmelCase   					=  image_processing(__lowerCAmelCase       ,    return_tensors="""pt""").pixel_values
					self.assertEqual(
					    encoded_images.shape       ,    (
					        self.image_processor_tester.batch_size,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.size["""height"""],
					        self.image_processor_tester.size["""width"""],
					    )       ,    )



			def 							a_       (   self):




					"""simple docstring"""


					lowerCAmelCase   					=  self.image_processing_class(**self.image_processor_dict)
					# create random numpy tensors
					lowerCAmelCase   					=  prepare_image_inputs(self.image_processor_tester       ,    equal_resolution=__lowerCAmelCase       ,    numpify=__lowerCAmelCase)
					for image in image_inputs:
							self.assertIsInstance(__lowerCAmelCase       ,    np.ndarray)

					# Test not batched input
					lowerCAmelCase   					=  image_processing(image_inputs[0]       ,    return_tensors="""pt""").pixel_values
					self.assertEqual(
					    encoded_images.shape       ,    (
					        1,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.size["""height"""],
					        self.image_processor_tester.size["""width"""],
					    )       ,    )

					# Test batched
					lowerCAmelCase   					=  image_processing(__lowerCAmelCase       ,    return_tensors="""pt""").pixel_values
					self.assertEqual(
					    encoded_images.shape       ,    (
					        self.image_processor_tester.batch_size,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.size["""height"""],
					        self.image_processor_tester.size["""width"""],
					    )       ,    )







			def 							a_       (   self):




					"""simple docstring"""


					lowerCAmelCase   					=  self.image_processing_class(**self.image_processor_dict)
					# create random PyTorch tensors
					lowerCAmelCase   					=  prepare_image_inputs(self.image_processor_tester       ,    equal_resolution=__lowerCAmelCase       ,    torchify=__lowerCAmelCase)
					for image in image_inputs:
							self.assertIsInstance(__lowerCAmelCase       ,    torch.Tensor)

					# Test not batched input
					lowerCAmelCase   					=  image_processing(image_inputs[0]       ,    return_tensors="""pt""").pixel_values
					self.assertEqual(
					    encoded_images.shape       ,    (
					        1,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.size["""height"""],
					        self.image_processor_tester.size["""width"""],
					    )       ,    )

					# Test batched
					lowerCAmelCase   					=  image_processing(__lowerCAmelCase       ,    return_tensors="""pt""").pixel_values
					self.assertEqual(
					    encoded_images.shape       ,    (
					        self.image_processor_tester.batch_size,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.size["""height"""],
					        self.image_processor_tester.size["""width"""],
					    )       ,    )




