


"""simple docstring"""







import argparse
import json

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

from transformers import TimesformerConfig, TimesformerForVideoClassification, VideoMAEImageProcessor






def 					a__ (       _UpperCamelCase ):





						"""simple docstring"""

						UpperCAmelCase__		=    TimesformerConfig()

						if "large" in model_name:
												UpperCAmelCase__		=    96

						if "hr" in model_name:
												UpperCAmelCase__		=    16
												UpperCAmelCase__		=    448

						UpperCAmelCase__		=    '''huggingface/label-files'''
						if "k400" in model_name:
												UpperCAmelCase__		=    400
												UpperCAmelCase__		=    '''kinetics400-id2label.json'''
						elif "k600" in model_name:
												UpperCAmelCase__		=    600
												UpperCAmelCase__		=    '''kinetics600-id2label.json'''
						elif "ssv2" in model_name:
												UpperCAmelCase__		=    174
												UpperCAmelCase__		=    '''something-something-v2-id2label.json'''
						else:
												raise ValueError('''Model name should either contain \'k400\', \'k600\' or \'ssv2\'.''' )
						UpperCAmelCase__		=    json.load(open(hf_hub_download(_UpperCamelCase	,			_UpperCamelCase	,			repo_type='''dataset''' )	,			'''r''' ) )
						UpperCAmelCase__		=    {int(_UpperCamelCase ): v for k, v in idalabel.items()}
						UpperCAmelCase__		=    idalabel
						UpperCAmelCase__		=    {v: k for k, v in idalabel.items()}

						return config






def 					a__ (       _UpperCamelCase ):





						"""simple docstring"""

						if "encoder." in name:
												UpperCAmelCase__		=    name.replace('''encoder.'''	,			'''''' )
						if "cls_token" in name:
												UpperCAmelCase__		=    name.replace('''cls_token'''	,			'''timesformer.embeddings.cls_token''' )
						if "pos_embed" in name:
												UpperCAmelCase__		=    name.replace('''pos_embed'''	,			'''timesformer.embeddings.position_embeddings''' )
						if "time_embed" in name:
												UpperCAmelCase__		=    name.replace('''time_embed'''	,			'''timesformer.embeddings.time_embeddings''' )
						if "patch_embed.proj" in name:
												UpperCAmelCase__		=    name.replace('''patch_embed.proj'''	,			'''timesformer.embeddings.patch_embeddings.projection''' )
						if "patch_embed.norm" in name:
												UpperCAmelCase__		=    name.replace('''patch_embed.norm'''	,			'''timesformer.embeddings.norm''' )
						if "blocks" in name:
												UpperCAmelCase__		=    name.replace('''blocks'''	,			'''timesformer.encoder.layer''' )
						if "attn.proj" in name:
												UpperCAmelCase__		=    name.replace('''attn.proj'''	,			'''attention.output.dense''' )
						if "attn" in name and "bias" not in name and "temporal" not in name:
												UpperCAmelCase__		=    name.replace('''attn'''	,			'''attention.self''' )
						if "attn" in name and "temporal" not in name:
												UpperCAmelCase__		=    name.replace('''attn'''	,			'''attention.attention''' )
						if "temporal_norm1" in name:
												UpperCAmelCase__		=    name.replace('''temporal_norm1'''	,			'''temporal_layernorm''' )
						if "temporal_attn.proj" in name:
												UpperCAmelCase__		=    name.replace('''temporal_attn'''	,			'''temporal_attention.output.dense''' )
						if "temporal_fc" in name:
												UpperCAmelCase__		=    name.replace('''temporal_fc'''	,			'''temporal_dense''' )
						if "norm1" in name and "temporal" not in name:
												UpperCAmelCase__		=    name.replace('''norm1'''	,			'''layernorm_before''' )
						if "norm2" in name:
												UpperCAmelCase__		=    name.replace('''norm2'''	,			'''layernorm_after''' )
						if "mlp.fc1" in name:
												UpperCAmelCase__		=    name.replace('''mlp.fc1'''	,			'''intermediate.dense''' )
						if "mlp.fc2" in name:
												UpperCAmelCase__		=    name.replace('''mlp.fc2'''	,			'''output.dense''' )
						if "norm.weight" in name and "fc" not in name and "temporal" not in name:
												UpperCAmelCase__		=    name.replace('''norm.weight'''	,			'''timesformer.layernorm.weight''' )
						if "norm.bias" in name and "fc" not in name and "temporal" not in name:
												UpperCAmelCase__		=    name.replace('''norm.bias'''	,			'''timesformer.layernorm.bias''' )
						if "head" in name:
												UpperCAmelCase__		=    name.replace('''head'''	,			'''classifier''' )

						return name






def 					a__ (       _UpperCamelCase	,			_UpperCamelCase ):





						"""simple docstring"""

						for key in orig_state_dict.copy().keys():
												UpperCAmelCase__		=    orig_state_dict.pop(_UpperCamelCase )

												if key.startswith('''model.''' ):
																		UpperCAmelCase__		=    key.replace('''model.'''	,			'''''' )

												if "qkv" in key:
																		UpperCAmelCase__		=    key.split('''.''' )
																		UpperCAmelCase__		=    int(key_split[1] )
																		UpperCAmelCase__		=    '''timesformer.encoder.layer.'''
																		if "temporal" in key:
																								UpperCAmelCase__		=    '''.temporal_attention.attention.qkv.'''
																		else:
																								UpperCAmelCase__		=    '''.attention.attention.qkv.'''
																		if "weight" in key:
																								UpperCAmelCase__		=    val
																		else:
																								UpperCAmelCase__		=    val
												else:
																		UpperCAmelCase__		=    val

						return orig_state_dict






def 					a__ (       ):





						"""simple docstring"""

						UpperCAmelCase__		=    hf_hub_download(
						    repo_id='''hf-internal-testing/spaghetti-video'''	,			filename='''eating_spaghetti.npy'''	,			repo_type='''dataset''' )
						UpperCAmelCase__		=    np.load(_UpperCamelCase )
						return list(_UpperCamelCase )






def 					a__ (       _UpperCamelCase	,			_UpperCamelCase	,			_UpperCamelCase	,			_UpperCamelCase ):





						"""simple docstring"""

						UpperCAmelCase__		=    get_timesformer_config(_UpperCamelCase )

						UpperCAmelCase__		=    TimesformerForVideoClassification(_UpperCamelCase )

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

						model.load_state_dict(_UpperCamelCase )
						model.eval()

						# verify model on basic input
						UpperCAmelCase__		=    VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]	,			image_std=[0.5, 0.5, 0.5] )
						UpperCAmelCase__		=    prepare_video()
						UpperCAmelCase__		=    image_processor(video[:8]	,			return_tensors='''pt''' )

						UpperCAmelCase__		=    model(**_UpperCamelCase )
						UpperCAmelCase__		=    outputs.logits

						UpperCAmelCase__		=    [
						    # Kinetics-400 checkpoints (hr = high resolution input of 448px instead of 224px)
						    '''timesformer-base-finetuned-k400''',
						    '''timesformer-large-finetuned-k400''',
						    '''timesformer-hr-finetuned-k400''',
						    # Kinetics-600 checkpoints (hr = high resolution input of 448px instead of 224px)
						    '''timesformer-base-finetuned-k600''',
						    '''timesformer-large-finetuned-k600''',
						    '''timesformer-hr-finetuned-k600''',
						    # Something-Something-v2 checkpoints (hr = high resolution input of 448px instead of 224px)
						    '''timesformer-base-finetuned-ssv2''',
						    '''timesformer-large-finetuned-ssv2''',
						    '''timesformer-hr-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 == "timesformer-base-finetuned-k400":
												UpperCAmelCase__		=    torch.Size([1, 400] )
												UpperCAmelCase__		=    torch.tensor([-0.3016, -0.7713, -0.4205] )
						elif model_name == "timesformer-base-finetuned-k600":
												UpperCAmelCase__		=    torch.Size([1, 600] )
												UpperCAmelCase__		=    torch.tensor([-0.7267, -0.7466, 3.2404] )
						elif model_name == "timesformer-base-finetuned-ssv2":
												UpperCAmelCase__		=    torch.Size([1, 174] )
												UpperCAmelCase__		=    torch.tensor([-0.9059, 0.6433, -3.1457] )
						elif model_name == "timesformer-large-finetuned-k400":
												UpperCAmelCase__		=    torch.Size([1, 400] )
												UpperCAmelCase__		=    torch.tensor([0, 0, 0] )
						elif model_name == "timesformer-large-finetuned-k600":
												UpperCAmelCase__		=    torch.Size([1, 600] )
												UpperCAmelCase__		=    torch.tensor([0, 0, 0] )
						elif model_name == "timesformer-large-finetuned-ssv2":
												UpperCAmelCase__		=    torch.Size([1, 174] )
												UpperCAmelCase__		=    torch.tensor([0, 0, 0] )
						elif model_name == "timesformer-hr-finetuned-k400":
												UpperCAmelCase__		=    torch.Size([1, 400] )
												UpperCAmelCase__		=    torch.tensor([-0.9617, -3.7311, -3.7708] )
						elif model_name == "timesformer-hr-finetuned-k600":
												UpperCAmelCase__		=    torch.Size([1, 600] )
												UpperCAmelCase__		=    torch.tensor([2.5273, 0.7127, 1.8848] )
						elif model_name == "timesformer-hr-finetuned-ssv2":
												UpperCAmelCase__		=    torch.Size([1, 174] )
												UpperCAmelCase__		=    torch.tensor([-3.6756, -0.7513, 0.7180] )
						else:
												raise ValueError(F"""Model name not supported. Should be one of {model_names}""" )

						# verify logits
						assert logits.shape == expected_shape
						assert torch.allclose(logits[0, :3]	,			_UpperCamelCase	,			atol=1E-4 )
						print('''Logits 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(F"""fcakyon/{model_name}""" )


if __name__ == "__main__":
						__a			 =   argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    "--checkpoint_url",
						    default="https://drive.google.com/u/1/uc?id=17yvuYp9L4mn-HpIcK5Zo6K3UoOy1kA5l&export=download",
						    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="",
						    type=str,
						    help="Path to the output PyTorch model directory.",
						)
						parser.add_argument("--model_name", default="timesformer-base-finetuned-k400", type=str, help="Name of the model.")
						parser.add_argument(
						    "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
						)

						__a			 =   parser.parse_args()
						convert_timesformer_checkpoint(
						    args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
						)



"""simple docstring"""







import safs
class 							lowerCamelCase       (			safs.SaFileSystem ):


							"""simple docstring"""