





import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch
@require_torch
class      UpperCamelCase_						:




					@property
					def 		lowerCAmelCase     (						self					)						->							int:
										return self.get_dummy_input()




					@property
					def 		lowerCAmelCase     (						self					)						->							Optional[Any]:
										if self.block_type == "down":
															return (4, 32, 16, 16)
										elif self.block_type == "mid":
															return (4, 32, 32, 32)
										elif self.block_type == "up":
															return (4, 32, 64, 64)

										raise ValueError(F'''\'{self.block_type}\' is not a supported block_type. Set it to \'up\', \'mid\', or \'down\'.'''					)




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_=True			,  lowerCAmelCase_=False			,  lowerCAmelCase_=False			,  lowerCAmelCase_=False			,  )						->							List[str]:
										_snake_case													=			4
										_snake_case													=			32
										_snake_case													=			(32, 32)

										_snake_case													=			torch.manual_seed(0					)
										_snake_case													=			torch.device(lowerCAmelCase_					)
										_snake_case													=			(batch_size, num_channels) + sizes
										_snake_case													=			randn_tensor(lowerCAmelCase_			,  generator=lowerCAmelCase_			,  device=lowerCAmelCase_					)
										_snake_case													=			{'hidden_states': hidden_states}

										if include_temb:
															_snake_case													=			128
															_snake_case													=			randn_tensor((batch_size, temb_channels)			,  generator=lowerCAmelCase_			,  device=lowerCAmelCase_					)

										if include_res_hidden_states_tuple:
															_snake_case													=			torch.manual_seed(1					)
															_snake_case													=			(randn_tensor(lowerCAmelCase_			,  generator=lowerCAmelCase_			,  device=lowerCAmelCase_					),)

										if include_encoder_hidden_states:
															_snake_case													=			floats_tensor((batch_size, 32, 32)					).to(lowerCAmelCase_					)

										if include_skip_sample:
															_snake_case													=			randn_tensor(((batch_size, 3) + sizes)			,  generator=lowerCAmelCase_			,  device=lowerCAmelCase_					)

										return dummy_input




					def 		lowerCAmelCase     (						self					)						->							Tuple:
										_snake_case													=			{
										    'in_channels': 32,
										    'out_channels': 32,
										    'temb_channels': 128,
										}
										if self.block_type == "up":
															_snake_case													=			32

										if self.block_type == "mid":
															init_dict.pop('out_channels'					)

										_snake_case													=			self.dummy_input
										return init_dict, inputs_dict




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_					)						->							Optional[int]:
										_snake_case							,       _snake_case													=			self.prepare_init_args_and_inputs_for_common()
										_snake_case													=			self.block_class(**lowerCAmelCase_					)
										unet_block.to(lowerCAmelCase_					)
										unet_block.eval()

										with torch.no_grad():
															_snake_case													=			unet_block(**lowerCAmelCase_					)

										if isinstance(lowerCAmelCase_			,  lowerCAmelCase_					):
															_snake_case													=			output[0]

										self.assertEqual(output.shape			,  self.output_shape					)

										_snake_case													=			output[0, -1, -3:, -3:]
										_snake_case													=			torch.tensor(lowerCAmelCase_					).to(lowerCAmelCase_					)
										assert torch_all_close(output_slice.flatten()			,  lowerCAmelCase_			,  atol=5E-3					)






					@unittest.skipIf(torch_device == 'mps'			,  'Training is not supported in mps'					)
					def 		lowerCAmelCase     (						self					)						->							Tuple:
										_snake_case							,       _snake_case													=			self.prepare_init_args_and_inputs_for_common()
										_snake_case													=			self.block_class(**lowerCAmelCase_					)
										model.to(lowerCAmelCase_					)
										model.train()
										_snake_case													=			model(**lowerCAmelCase_					)

										if isinstance(lowerCAmelCase_			,  lowerCAmelCase_					):
															_snake_case													=			output[0]

										_snake_case													=			torch.device(lowerCAmelCase_					)
										_snake_case													=			randn_tensor(output.shape			,  device=lowerCAmelCase_					)
										_snake_case													=			torch.nn.functional.mse_loss(lowerCAmelCase_			,  lowerCAmelCase_					)
										loss.backward()










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   lowerCamelCase__       (       UpperCamelCase__     :    int  )     ->	List[str]:



					'''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   lowerCamelCase__       (       UpperCamelCase__     :    str    ,      UpperCamelCase__     :    Dict  )     ->	int:



					'''simple docstring'''




					if "small" in model_name:
										_snake_case													=			384
										_snake_case													=			1_536
										_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													=			1_024
										_snake_case													=			4_096
										_snake_case													=			24
										_snake_case													=			16
										_snake_case													=			12
										_snake_case													=			8
										_snake_case													=			512
										_snake_case													=			2_048
					elif "huge" in model_name:
										_snake_case													=			1_280
										_snake_case													=			5_120
										_snake_case													=			32
										_snake_case													=			16
										_snake_case													=			12
										_snake_case													=			8
										_snake_case													=			640
										_snake_case													=			2_560
					elif "base" not in model_name:
										raise ValueError('Model name should include either "small", "base", "large", or "huge"'  )





def   lowerCamelCase__       (       UpperCamelCase__     :    Any  )     ->	Tuple:



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



					'''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   lowerCamelCase__       (       )     ->	Union[str, 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   lowerCamelCase__       (       UpperCamelCase__     :    List[str]    ,      UpperCamelCase__     :    Union[str, Any]    ,      UpperCamelCase__     :    Optional[int]    ,      UpperCamelCase__     :    Optional[int]  )     ->	List[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.9291, -0.4061, -0.9307]  )
					elif model_name == "videomae-small-finetuned-ssv2":
										_snake_case													=			torch.Size([1, 174]  )
										_snake_case													=			torch.tensor([0.2671, -0.4689, -0.8235]  )
					elif model_name == "videomae-base":
										_snake_case													=			torch.Size([1, 1_408, 1_536]  )
										_snake_case													=			torch.tensor([[0.7739, 0.7968, 0.7089], [0.6701, 0.7487, 0.6209], [0.4287, 0.5158, 0.4773]]  )
					elif model_name == "videomae-base-short":
										_snake_case													=			torch.Size([1, 1_408, 1_536]  )
										_snake_case													=			torch.tensor([[0.7994, 0.9612, 0.8508], [0.7401, 0.8958, 0.8302], [0.5862, 0.7468, 0.7325]]  )
										# we verified the loss both for normalized and unnormalized targets for this one
										_snake_case													=			torch.tensor([0.5142]  ) if config.norm_pix_loss else torch.tensor([0.6469]  )
					elif model_name == "videomae-large":
										_snake_case													=			torch.Size([1, 1_408, 1_536]  )
										_snake_case													=			torch.tensor([[0.7149, 0.7997, 0.6966], [0.6768, 0.7869, 0.6948], [0.5139, 0.6221, 0.5605]]  )
					elif model_name == "videomae-large-finetuned-kinetics":
										_snake_case													=			torch.Size([1, 400]  )
										_snake_case													=			torch.tensor([0.0771, 0.0011, -0.3625]  )
					elif model_name == "videomae-huge-finetuned-kinetics":
										_snake_case													=			torch.Size([1, 400]  )
										_snake_case													=			torch.tensor([0.2433, 0.1632, -0.4894]  )
					elif model_name == "videomae-base-short-finetuned-kinetics":
										_snake_case													=			torch.Size([1, 400]  )
										_snake_case													=			torch.tensor([0.6588, 0.0990, -0.2493]  )
					elif model_name == "videomae-base-finetuned-kinetics":
										_snake_case													=			torch.Size([1, 400]  )
										_snake_case													=			torch.tensor([0.3669, -0.0688, -0.2421]  )
					elif model_name == "videomae-base-short-ssv2":
										_snake_case													=			torch.Size([1, 1_408, 1_536]  )
										_snake_case													=			torch.tensor([[0.4712, 0.5296, 0.5786], [0.2278, 0.2729, 0.4026], [0.0352, 0.0730, 0.2506]]  )
					elif model_name == "videomae-base-short-finetuned-ssv2":
										_snake_case													=			torch.Size([1, 174]  )
										_snake_case													=			torch.tensor([-0.0537, -0.1539, -0.3266]  )
					elif model_name == "videomae-base-ssv2":
										_snake_case													=			torch.Size([1, 1_408, 1_536]  )
										_snake_case													=			torch.tensor([[0.8131, 0.8727, 0.8546], [0.7366, 0.9377, 0.8870], [0.5935, 0.8874, 0.8564]]  )
					elif model_name == "videomae-base-finetuned-ssv2":
										_snake_case													=			torch.Size([1, 174]  )
										_snake_case													=			torch.tensor([0.1961, -0.8337, -0.6389]  )
					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__":
					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)



