



"""simple docstring"""






import argparse

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    CLIPTokenizer,
    CLIPTokenizerFast,
    VideoMAEImageProcessor,
    XCLIPConfig,
    XCLIPModel,
    XCLIPProcessor,
    XCLIPTextConfig,
    XCLIPVisionConfig,
)






def   snake_case__     (	__lowerCamelCase       :					str							,	__lowerCamelCase       :					Optional[int]			):
			"""simple docstring"""
			lowerCamelCase__	:	Union[str, Any]				=XCLIPTextConfig()

			# derive patch size from model name
			lowerCamelCase__	:	Any				=model_name.find('''patch'''			)
			lowerCamelCase__	:	str				=int(model_name[start_idx + len('''patch'''			) : start_idx + len('''patch'''			) + 2]			)
			lowerCamelCase__	:	str				=XCLIPVisionConfig(patch_size=__lowerCamelCase							,	num_frames=__lowerCamelCase			)

			if "large" in model_name:
						lowerCamelCase__	:	Union[str, Any]				=768
						lowerCamelCase__	:	Optional[Any]				=3072
						lowerCamelCase__	:	List[Any]				=12

						lowerCamelCase__	:	List[Any]				=1024
						lowerCamelCase__	:	Optional[int]				=4096
						lowerCamelCase__	:	int				=16
						lowerCamelCase__	:	List[Any]				=24
						lowerCamelCase__	:	Optional[Any]				=768
						lowerCamelCase__	:	Any				=3072

			if model_name == "xclip-large-patch14-16-frames":
						lowerCamelCase__	:	Dict				=336

			lowerCamelCase__	:	Any				=XCLIPConfig.from_text_vision_configs(__lowerCamelCase							,	__lowerCamelCase			)

			if "large" in model_name:
						lowerCamelCase__	:	int				=768

			return config






def   snake_case__     (	__lowerCamelCase       :					Dict			):
			"""simple docstring"""
			# text encoder
			if name == "token_embedding.weight":
						lowerCamelCase__	:	str				=name.replace('''token_embedding.weight'''							,	'''text_model.embeddings.token_embedding.weight'''			)
			if name == "positional_embedding":
						lowerCamelCase__	:	Dict				=name.replace('''positional_embedding'''							,	'''text_model.embeddings.position_embedding.weight'''			)
			if "ln_1" in name:
						lowerCamelCase__	:	List[str]				=name.replace('''ln_1'''							,	'''layer_norm1'''			)
			if "ln_2" in name:
						lowerCamelCase__	:	List[Any]				=name.replace('''ln_2'''							,	'''layer_norm2'''			)
			if "c_fc" in name:
						lowerCamelCase__	:	str				=name.replace('''c_fc'''							,	'''fc1'''			)
			if "c_proj" in name:
						lowerCamelCase__	:	Any				=name.replace('''c_proj'''							,	'''fc2'''			)
			if name.startswith('''transformer.resblocks'''			):
						lowerCamelCase__	:	str				=name.replace('''transformer.resblocks'''							,	'''text_model.encoder.layers'''			)
			if "attn.out_proj" in name and "message" not in name:
						lowerCamelCase__	:	Optional[int]				=name.replace('''attn.out_proj'''							,	'''self_attn.out_proj'''			)
			if "ln_final" in name:
						lowerCamelCase__	:	Union[str, Any]				=name.replace('''ln_final'''							,	'''text_model.final_layer_norm'''			)
			# visual encoder
			if name == "visual.class_embedding":
						lowerCamelCase__	:	Optional[Any]				=name.replace('''visual.class_embedding'''							,	'''vision_model.embeddings.class_embedding'''			)
			if name == "visual.positional_embedding":
						lowerCamelCase__	:	str				=name.replace('''visual.positional_embedding'''							,	'''vision_model.embeddings.position_embedding.weight'''			)
			if name.startswith('''visual.transformer.resblocks'''			):
						lowerCamelCase__	:	List[str]				=name.replace('''visual.transformer.resblocks'''							,	'''vision_model.encoder.layers'''			)
			if "visual.conv1" in name:
						lowerCamelCase__	:	Tuple				=name.replace('''visual.conv1'''							,	'''vision_model.embeddings.patch_embedding'''			)
			if "visual.ln_pre" in name:
						lowerCamelCase__	:	Optional[Any]				=name.replace('''visual.ln_pre'''							,	'''vision_model.pre_layernorm'''			)
			if "visual.ln_post" in name:
						lowerCamelCase__	:	List[str]				=name.replace('''visual.ln_post'''							,	'''vision_model.post_layernorm'''			)
			if "visual.proj" in name:
						lowerCamelCase__	:	Dict				=name.replace('''visual.proj'''							,	'''visual_projection.weight'''			)
			if "text_projection" in name:
						lowerCamelCase__	:	str				=name.replace('''text_projection'''							,	'''text_projection.weight'''			)
			# things on top
			if "prompts_visual_proj" in name:
						lowerCamelCase__	:	str				=name.replace('''prompts_visual_proj'''							,	'''prompts_visual_projection'''			)
			if "prompts_visual_ln" in name:
						lowerCamelCase__	:	Union[str, Any]				=name.replace('''prompts_visual_ln'''							,	'''prompts_visual_layernorm'''			)
			# mit
			if name == "mit.positional_embedding":
						lowerCamelCase__	:	str				=name.replace('''positional'''							,	'''position'''			)
			if name.startswith('''mit.resblocks'''			):
						lowerCamelCase__	:	int				=name.replace('''mit.resblocks'''							,	'''mit.encoder.layers'''			)
			# prompts generator
			if name.startswith('''prompts_generator.norm'''			):
						lowerCamelCase__	:	Union[str, Any]				=name.replace('''prompts_generator.norm'''							,	'''prompts_generator.layernorm'''			)

			return name






def   snake_case__     (	__lowerCamelCase       :					Optional[int]							,	__lowerCamelCase       :					Any			):
			"""simple docstring"""
			for key in orig_state_dict.copy().keys():
						lowerCamelCase__	:	Dict				=orig_state_dict.pop(__lowerCamelCase			)

						if "attn.in_proj" in key:
									lowerCamelCase__	:	Union[str, Any]				=key.split('''.'''			)
									if key.startswith('''visual'''			):
												lowerCamelCase__	:	Optional[int]				=key_split[3]
												lowerCamelCase__	:	Tuple				=config.vision_config.hidden_size
												if "message_attn" in key:
															if "weight" in key:
																		lowerCamelCase__	:	Union[str, Any]				=val[
																		    :dim, :
																		]
																		lowerCamelCase__	:	Any				=val[
																		    dim : dim * 2, :
																		]
																		lowerCamelCase__	:	List[str]				=val[
																		    -dim:, :
																		]
															else:
																		lowerCamelCase__	:	Optional[Any]				=val[
																		    :dim
																		]
																		lowerCamelCase__	:	int				=val[
																		    dim : dim * 2
																		]
																		lowerCamelCase__	:	Optional[int]				=val[
																		    -dim:
																		]
												else:
															if "weight" in key:
																		lowerCamelCase__	:	Tuple				=val[
																		    :dim, :
																		]
																		lowerCamelCase__	:	Optional[int]				=val[
																		    dim : dim * 2, :
																		]
																		lowerCamelCase__	:	str				=val[
																		    -dim:, :
																		]
															else:
																		lowerCamelCase__	:	Union[str, Any]				=val[:dim]
																		lowerCamelCase__	:	int				=val[
																		    dim : dim * 2
																		]
																		lowerCamelCase__	:	str				=val[-dim:]
									elif key.startswith('''mit'''			):
												lowerCamelCase__	:	int				=key_split[2]
												lowerCamelCase__	:	List[str]				=config.vision_config.mit_hidden_size
												if "weight" in key:
															lowerCamelCase__	:	Any				=val[:dim, :]
															lowerCamelCase__	:	List[str]				=val[dim : dim * 2, :]
															lowerCamelCase__	:	Optional[int]				=val[-dim:, :]
												else:
															lowerCamelCase__	:	Optional[Any]				=val[:dim]
															lowerCamelCase__	:	Optional[int]				=val[dim : dim * 2]
															lowerCamelCase__	:	str				=val[-dim:]
									else:
												lowerCamelCase__	:	Any				=key_split[2]
												lowerCamelCase__	:	Optional[int]				=config.text_config.hidden_size
												if "weight" in key:
															lowerCamelCase__	:	Union[str, Any]				=val[:dim, :]
															lowerCamelCase__	:	Tuple				=val[
															    dim : dim * 2, :
															]
															lowerCamelCase__	:	List[str]				=val[-dim:, :]
												else:
															lowerCamelCase__	:	Optional[Any]				=val[:dim]
															lowerCamelCase__	:	Union[str, Any]				=val[
															    dim : dim * 2
															]
															lowerCamelCase__	:	Optional[Any]				=val[-dim:]
						else:
									lowerCamelCase__	:	Optional[Any]				=rename_key(__lowerCamelCase			)
									if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
												lowerCamelCase__	:	str				=val.T
									lowerCamelCase__	:	int				=val

			return orig_state_dict






def   snake_case__     (	__lowerCamelCase       :					List[Any]			):
			"""simple docstring"""
			if num_frames == 8:
						lowerCamelCase__	:	List[Any]				='''eating_spaghetti_8_frames.npy'''
			elif num_frames == 16:
						lowerCamelCase__	:	Optional[int]				='''eating_spaghetti.npy'''
			elif num_frames == 32:
						lowerCamelCase__	:	List[str]				='''eating_spaghetti_32_frames.npy'''
			lowerCamelCase__	:	List[str]				=hf_hub_download(
			    repo_id='''hf-internal-testing/spaghetti-video'''							,	filename=__lowerCamelCase							,	repo_type='''dataset'''							,	)
			lowerCamelCase__	:	Tuple				=np.load(__lowerCamelCase			)
			return list(__lowerCamelCase			)






def   snake_case__     (	__lowerCamelCase       :					str							,	__lowerCamelCase       :					List[Any]=None							,	__lowerCamelCase       :					Tuple=False			):
			"""simple docstring"""
			lowerCamelCase__	:	List[Any]				={
			    # fully supervised kinetics-400 checkpoints
			    '''xclip-base-patch32''': '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_8.pth''',
			    '''xclip-base-patch32-16-frames''': (
			        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_16.pth'''
			    ),
			    '''xclip-base-patch16''': '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_8.pth''',
			    '''xclip-base-patch16-16-frames''': (
			        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_16.pth'''
			    ),
			    '''xclip-large-patch14''': '''https://drive.google.com/u/0/uc?id=1NUOImq0o5DlQTST17iIP3vG7DgmHQuCx&amp;export=download&amp;confirm=t&amp;uuid=b26caedc-88e2-473e-830a-9d158b653cdb''',
			    '''xclip-large-patch14-16-frames''': '''https://drive.google.com/u/0/uc?id=1FOYgnJc097OJ4lGwtRCCydQyVPJEOH7d&amp;export=download&amp;confirm=t&amp;uuid=538fa810-e671-4050-b385-9a623f89804f''',
			    # fully supervised kinetics-600 checkpoints
			    '''xclip-base-patch16-kinetics-600''': (
			        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_8.pth'''
			    ),
			    '''xclip-base-patch16-kinetics-600-16-frames''': (
			        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_16.pth'''
			    ),
			    '''xclip-large-patch14-kinetics-600''': '''https://drive.google.com/u/0/uc?id=1FV8C1INuM91sLAN4ImjzePLIlpMSihwV&amp;export=download&amp;confirm=t&amp;uuid=141d4977-4a65-44ae-864f-4b0c19f838be''',
			    # few shot
			    '''xclip-base-patch16-hmdb-2-shot''': (
			        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_2.pth'''
			    ),
			    '''xclip-base-patch16-hmdb-4-shot''': (
			        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_4.pth'''
			    ),
			    '''xclip-base-patch16-hmdb-8-shot''': (
			        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_8.pth'''
			    ),
			    '''xclip-base-patch16-hmdb-16-shot''': (
			        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_16.pth'''
			    ),
			    '''xclip-base-patch16-ucf-2-shot''': (
			        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_2.pth'''
			    ),
			    '''xclip-base-patch16-ucf-4-shot''': (
			        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_4.pth'''
			    ),
			    '''xclip-base-patch16-ucf-8-shot''': (
			        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_8.pth'''
			    ),
			    '''xclip-base-patch16-ucf-16-shot''': (
			        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_16.pth'''
			    ),
			    # zero shot
			    '''xclip-base-patch16-zero-shot''': '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/zero.pth''',
			}

			lowerCamelCase__	:	Union[str, Any]				=model_to_url[model_name]
			lowerCamelCase__	:	str				=8
			if "16-frames" in model_name:
						lowerCamelCase__	:	Tuple				=16
			elif "shot" in model_name:
						lowerCamelCase__	:	List[Any]				=32

			lowerCamelCase__	:	str				=get_xclip_config(__lowerCamelCase							,	__lowerCamelCase			)
			lowerCamelCase__	:	int				=XCLIPModel(__lowerCamelCase			)
			model.eval()

			if "drive" in checkpoint_url:
						lowerCamelCase__	:	Optional[int]				='''pytorch_model.bin'''
						gdown.cached_download(__lowerCamelCase							,	__lowerCamelCase							,	quiet=__lowerCamelCase			)
						lowerCamelCase__	:	str				=torch.load(__lowerCamelCase							,	map_location='''cpu'''			)['''model''']
			else:
						lowerCamelCase__	:	Optional[Any]				=torch.hub.load_state_dict_from_url(__lowerCamelCase			)['''model''']

			lowerCamelCase__	:	Optional[Any]				=convert_state_dict(__lowerCamelCase							,	__lowerCamelCase			)

			lowerCamelCase__	:	Optional[Any]				=XCLIPModel(__lowerCamelCase			)
			lowerCamelCase__	:	Dict				=model.load_state_dict(__lowerCamelCase							,	strict=__lowerCamelCase			)
			assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
			model.eval()

			lowerCamelCase__	:	Tuple				=336 if model_name == '''xclip-large-patch14-16-frames''' else 224
			lowerCamelCase__	:	Optional[int]				=VideoMAEImageProcessor(size=__lowerCamelCase			)
			lowerCamelCase__	:	List[str]				=CLIPTokenizer.from_pretrained('''openai/clip-vit-base-patch32'''			)
			lowerCamelCase__	:	str				=CLIPTokenizerFast.from_pretrained('''openai/clip-vit-base-patch32'''			)
			lowerCamelCase__	:	Any				=XCLIPProcessor(image_processor=__lowerCamelCase							,	tokenizer=__lowerCamelCase			)

			lowerCamelCase__	:	Tuple				=prepare_video(__lowerCamelCase			)
			lowerCamelCase__	:	Dict				=processor(
			    text=['''playing sports''', '''eating spaghetti''', '''go shopping''']							,	videos=__lowerCamelCase							,	return_tensors='''pt'''							,	padding=__lowerCamelCase			)

			print('''Shape of pixel values:'''							,	inputs.pixel_values.shape			)

			with torch.no_grad():
						lowerCamelCase__	:	int				=model(**__lowerCamelCase			)

			# Verify outputs
			lowerCamelCase__	:	str				=outputs.logits_per_video
			lowerCamelCase__	:	Dict				=logits_per_video.softmax(dim=1			)
			print('''Probs:'''							,	__lowerCamelCase			)
			# kinetics-400
			if model_name == "xclip-base-patch32":
						lowerCamelCase__	:	Optional[int]				=torch.tensor([[0.00_19, 0.99_51, 0.00_30]]			)
			elif model_name == "xclip-base-patch32-16-frames":
						lowerCamelCase__	:	Any				=torch.tensor([[7.0_9_9_9e-0_4, 9.9_8_8_3e-0_1, 4.5_5_8_0e-0_4]]			)
			elif model_name == "xclip-base-patch16":
						lowerCamelCase__	:	Tuple				=torch.tensor([[0.00_83, 0.96_81, 0.02_36]]			)
			elif model_name == "xclip-base-patch16-16-frames":
						lowerCamelCase__	:	List[str]				=torch.tensor([[7.6_9_3_7e-0_4, 9.9_7_2_8e-0_1, 1.9_4_7_3e-0_3]]			)
			elif model_name == "xclip-large-patch14":
						lowerCamelCase__	:	str				=torch.tensor([[0.00_62, 0.98_64, 0.00_75]]			)
			elif model_name == "xclip-large-patch14-16-frames":
						lowerCamelCase__	:	int				=torch.tensor([[3.3_8_7_7e-0_4, 9.9_9_3_7e-0_1, 2.8_8_8_8e-0_4]]			)
			# kinetics-600
			elif model_name == "xclip-base-patch16-kinetics-600":
						lowerCamelCase__	:	List[str]				=torch.tensor([[0.05_55, 0.89_14, 0.05_31]]			)
			elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
						lowerCamelCase__	:	Any				=torch.tensor([[3.8_5_5_4e-0_4, 9.9_9_2_9e-0_1, 3.2_7_5_4e-0_4]]			)
			elif model_name == "xclip-large-patch14-kinetics-600":
						lowerCamelCase__	:	List[Any]				=torch.tensor([[0.00_36, 0.99_20, 0.00_45]]			)
			# few shot
			elif model_name == "xclip-base-patch16-hmdb-2-shot":
						lowerCamelCase__	:	Optional[Any]				=torch.tensor([[7.1_8_9_0e-0_6, 9.9_9_9_4e-0_1, 5.6_5_5_9e-0_5]]			)
			elif model_name == "xclip-base-patch16-hmdb-4-shot":
						lowerCamelCase__	:	Any				=torch.tensor([[1.0_3_2_0e-0_5, 9.9_9_9_3e-0_1, 6.2_4_3_5e-0_5]]			)
			elif model_name == "xclip-base-patch16-hmdb-8-shot":
						lowerCamelCase__	:	Optional[int]				=torch.tensor([[4.1_3_7_7e-0_6, 9.9_9_9_0e-0_1, 9.8_3_8_6e-0_5]]			)
			elif model_name == "xclip-base-patch16-hmdb-16-shot":
						lowerCamelCase__	:	str				=torch.tensor([[4.1_3_4_7e-0_5, 9.9_9_6_2e-0_1, 3.3_4_1_1e-0_4]]			)
			elif model_name == "xclip-base-patch16-ucf-2-shot":
						lowerCamelCase__	:	Tuple				=torch.tensor([[8.5_8_5_7e-0_5, 9.9_9_2_8e-0_1, 6.3_2_9_1e-0_4]]			)
			elif model_name == "xclip-base-patch16-ucf-4-shot":
						lowerCamelCase__	:	Optional[Any]				=torch.tensor([[8.5_8_5_7e-0_5, 9.9_9_2_8e-0_1, 6.3_2_9_1e-0_4]]			)
			elif model_name == "xclip-base-patch16-ucf-8-shot":
						lowerCamelCase__	:	List[Any]				=torch.tensor([[0.00_27, 0.99_04, 0.00_70]]			)
			elif model_name == "xclip-base-patch16-ucf-16-shot":
						lowerCamelCase__	:	Dict				=torch.tensor([[9.8_2_1_9e-0_4, 9.9_5_9_3e-0_1, 3.0_8_6_3e-0_3]]			)
			# zero shot
			elif model_name == "xclip-base-patch16-zero-shot":
						lowerCamelCase__	:	Optional[Any]				=torch.tensor([[3.5_0_8_2e-0_4, 9.9_7_8_5e-0_1, 1.7_9_6_6e-0_3]]			)
			else:
						raise ValueError(f'''Model name {model_name} not supported'''			)
			assert torch.allclose(__lowerCamelCase							,	__lowerCamelCase							,	atol=1e-3			)
			print('''Looks ok!'''			)

			if pytorch_dump_folder_path is not None:
						print(f'''Saving model {model_name} to {pytorch_dump_folder_path}'''			)
						model.save_pretrained(__lowerCamelCase			)

			if push_to_hub:
						print('''Pushing model, processor and slow tokenizer files to the hub...'''			)
						model.push_to_hub(__lowerCamelCase							,	organization='''nielsr'''			)
						processor.push_to_hub(__lowerCamelCase							,	organization='''nielsr'''			)
						slow_tokenizer.push_to_hub(__lowerCamelCase							,	organization='''nielsr'''			)


if __name__ == "__main__":
			_lowercase				:				Dict   		=		argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    "--model_name",
			    default="xclip-base-patch32",
			    type=str,
			    help="Name of the model.",
			)
			parser.add_argument(
			    "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
			)
			parser.add_argument(
			    "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
			)

			_lowercase				:				int   		=		parser.parse_args()
			convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









"""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__     (	__lowerCamelCase       :					Optional[Any]			):
			"""simple docstring"""
			lowerCamelCase__	:	str				=VideoMAEConfig()

			set_architecture_configs(__lowerCamelCase							,	__lowerCamelCase			)

			if "finetuned" not in model_name:
						lowerCamelCase__	:	int				=False

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

			return config






def   snake_case__     (	__lowerCamelCase       :					List[str]							,	__lowerCamelCase       :					Optional[Any]			):
			"""simple docstring"""
			if "small" in model_name:
						lowerCamelCase__	:	Optional[Any]				=384
						lowerCamelCase__	:	List[Any]				=1536
						lowerCamelCase__	:	int				=12
						lowerCamelCase__	:	Dict				=16
						lowerCamelCase__	:	List[Any]				=12
						lowerCamelCase__	:	Optional[Any]				=3
						lowerCamelCase__	:	Union[str, Any]				=192
						lowerCamelCase__	:	str				=768
			elif "large" in model_name:
						lowerCamelCase__	:	Union[str, Any]				=1024
						lowerCamelCase__	:	str				=4096
						lowerCamelCase__	:	int				=24
						lowerCamelCase__	:	Dict				=16
						lowerCamelCase__	:	Union[str, Any]				=12
						lowerCamelCase__	:	List[Any]				=8
						lowerCamelCase__	:	int				=512
						lowerCamelCase__	:	Optional[Any]				=2048
			elif "huge" in model_name:
						lowerCamelCase__	:	Optional[int]				=1280
						lowerCamelCase__	:	Optional[int]				=5120
						lowerCamelCase__	:	List[Any]				=32
						lowerCamelCase__	:	List[Any]				=16
						lowerCamelCase__	:	Optional[Any]				=12
						lowerCamelCase__	:	Dict				=8
						lowerCamelCase__	:	List[Any]				=640
						lowerCamelCase__	:	Any				=2560
			elif "base" not in model_name:
						raise ValueError('''Model name should include either "small", "base", "large", or "huge"'''			)






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

			return name






def   snake_case__     (	__lowerCamelCase       :					int							,	__lowerCamelCase       :					int			):
			"""simple docstring"""
			for key in orig_state_dict.copy().keys():
						lowerCamelCase__	:	Dict				=orig_state_dict.pop(__lowerCamelCase			)

						if key.startswith('''encoder.'''			):
									lowerCamelCase__	:	Optional[int]				=key.replace('''encoder.'''							,	''''''			)

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

			return orig_state_dict






def   snake_case__     (	):
			"""simple docstring"""
			lowerCamelCase__	:	List[Any]				=hf_hub_download(
			    repo_id='''hf-internal-testing/spaghetti-video'''							,	filename='''eating_spaghetti.npy'''							,	repo_type='''dataset'''			)
			lowerCamelCase__	:	Optional[Any]				=np.load(__lowerCamelCase			)
			return list(__lowerCamelCase			)






def   snake_case__     (	__lowerCamelCase       :					Optional[Any]							,	__lowerCamelCase       :					Union[str, Any]							,	__lowerCamelCase       :					List[Any]							,	__lowerCamelCase       :					Optional[int]			):
			"""simple docstring"""
			lowerCamelCase__	:	str				=get_videomae_config(__lowerCamelCase			)

			if "finetuned" in model_name:
						lowerCamelCase__	:	Tuple				=VideoMAEForVideoClassification(__lowerCamelCase			)
			else:
						lowerCamelCase__	:	int				=VideoMAEForPreTraining(__lowerCamelCase			)

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

			model.load_state_dict(__lowerCamelCase			)
			model.eval()

			# verify model on basic input
			lowerCamelCase__	:	Dict				=VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]							,	image_std=[0.5, 0.5, 0.5]			)
			lowerCamelCase__	:	int				=prepare_video()
			lowerCamelCase__	:	Tuple				=image_processor(__lowerCamelCase							,	return_tensors='''pt'''			)

			if "finetuned" not in model_name:
						lowerCamelCase__	:	Tuple				=hf_hub_download(repo_id='''hf-internal-testing/bool-masked-pos'''							,	filename='''bool_masked_pos.pt'''			)
						lowerCamelCase__	:	Union[str, Any]				=torch.load(__lowerCamelCase			)

			lowerCamelCase__	:	int				=model(**__lowerCamelCase			)
			lowerCamelCase__	:	Dict				=outputs.logits

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

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

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


if __name__ == "__main__":
		_lowercase				:				Any   		=		argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    "--checkpoint_url",
		    default="https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4",
		    type=str,
		    help=(
		        "URL of the original PyTorch checkpoint (on Google Drive) you'd like to convert. Should be a direct"
		        " download link."
		    ),
		)
		parser.add_argument(
		    "--pytorch_dump_folder_path",
		    default="/Users/nielsrogge/Documents/VideoMAE/Test",
		    type=str,
		    help="Path to the output PyTorch model directory.",
		)
		parser.add_argument("--model_name", default="videomae-base", type=str, help="Name of the model.")
		parser.add_argument(
		    "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
		)

		_lowercase				:				Union[str, Any]   		=		parser.parse_args()
		convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)




