



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 		__snake_case						(    _UpperCAmelCase				):
					"""simple docstring"""






					lowercase						   =       TimesformerConfig()

					if "large" in model_name:
										lowercase						   =       96

					if "hr" in model_name:
										lowercase						   =       16
										lowercase						   =       4_48

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

					return config





def 		__snake_case						(    _UpperCAmelCase				):
					"""simple docstring"""






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

					return name





def 		__snake_case						(    _UpperCAmelCase       ,			_UpperCAmelCase				):
					"""simple docstring"""






					for key in orig_state_dict.copy().keys():
										lowercase						   =       orig_state_dict.pop(_UpperCAmelCase				)

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

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

					return orig_state_dict





def 		__snake_case						(    ):
					"""simple docstring"""






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





def 		__snake_case						(    _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase				):
					"""simple docstring"""






					lowercase						   =       get_timesformer_config(_UpperCAmelCase				)

					lowercase						   =       TimesformerForVideoClassification(_UpperCAmelCase				)

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

					model.load_state_dict(_UpperCAmelCase				)
					model.eval()

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

					lowercase						   =       model(**_UpperCAmelCase				)
					lowercase						   =       outputs.logits

					lowercase						   =       [
					    # 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":
										lowercase						   =       torch.Size([1, 4_00]				)
										lowercase						   =       torch.tensor([-0.3_016, -0.7_713, -0.4_205]				)
					elif model_name == "timesformer-base-finetuned-k600":
										lowercase						   =       torch.Size([1, 6_00]				)
										lowercase						   =       torch.tensor([-0.7_267, -0.7_466, 3.2_404]				)
					elif model_name == "timesformer-base-finetuned-ssv2":
										lowercase						   =       torch.Size([1, 1_74]				)
										lowercase						   =       torch.tensor([-0.9_059, 0.6_433, -3.1_457]				)
					elif model_name == "timesformer-large-finetuned-k400":
										lowercase						   =       torch.Size([1, 4_00]				)
										lowercase						   =       torch.tensor([0, 0, 0]				)
					elif model_name == "timesformer-large-finetuned-k600":
										lowercase						   =       torch.Size([1, 6_00]				)
										lowercase						   =       torch.tensor([0, 0, 0]				)
					elif model_name == "timesformer-large-finetuned-ssv2":
										lowercase						   =       torch.Size([1, 1_74]				)
										lowercase						   =       torch.tensor([0, 0, 0]				)
					elif model_name == "timesformer-hr-finetuned-k400":
										lowercase						   =       torch.Size([1, 4_00]				)
										lowercase						   =       torch.tensor([-0.9_617, -3.7_311, -3.7_708]				)
					elif model_name == "timesformer-hr-finetuned-k600":
										lowercase						   =       torch.Size([1, 6_00]				)
										lowercase						   =       torch.tensor([2.5_273, 0.7_127, 1.8_848]				)
					elif model_name == "timesformer-hr-finetuned-ssv2":
										lowercase						   =       torch.Size([1, 1_74]				)
										lowercase						   =       torch.tensor([-3.6_756, -0.7_513, 0.7_180]				)
					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__":
		__magic_name__       =     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.'''
		)

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











from typing import List, Optional, Tuple, Union

from . import logging
from .import_utils import is_torch_available, is_torch_version


if is_torch_available():
		import torch

__magic_name__       =     logging.get_logger(__name__)  # pylint: disable=invalid-name

try:
		from torch._dynamo import allow_in_graph as maybe_allow_in_graph
except (ImportError, ModuleNotFoundError):





		def 		__snake_case						(    cls				):
							"""simple docstring"""






							return cls





def 		__snake_case						(    _UpperCAmelCase       ,			_UpperCAmelCase = None       ,			_UpperCAmelCase = None       ,			_UpperCAmelCase = None       ,			_UpperCAmelCase = None       ,			):
					"""simple docstring"""






					lowercase						   =       device
					lowercase						   =       shape[0]

					lowercase						   =       layout or torch.strided
					lowercase						   =       device or torch.device('cpu'				)

					if generator is not None:
										lowercase						   =       generator.device.type if not isinstance(_UpperCAmelCase       ,			_UpperCAmelCase				) else generator[0].device.type
										if gen_device_type != device.type and gen_device_type == "cpu":
															lowercase						   =       'cpu'
															if device != "mps":
																				logger.info(
																				    f"""The passed generator was created on 'cpu' even though a tensor on {device} was expected."""
																				    f""" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably"""
																				    f""" slighly speed up this function by passing a generator that was created on the {device} device."""				)
										elif gen_device_type != device.type and gen_device_type == "cuda":
															raise ValueError(f"""Cannot generate a {device} tensor from a generator of type {gen_device_type}."""				)

					if isinstance(_UpperCAmelCase       ,			_UpperCAmelCase				):
										lowercase						   =       (1,) + shape[1:]
										lowercase						   =       [
										    torch.randn(_UpperCAmelCase       ,			generator=generator[i]       ,			device=_UpperCAmelCase       ,			dtype=_UpperCAmelCase       ,			layout=_UpperCAmelCase				)
										    for i in range(_UpperCAmelCase				)
										]
										lowercase						   =       torch.cat(_UpperCAmelCase       ,			dim=0				).to(_UpperCAmelCase				)
					else:
										lowercase						   =       torch.randn(_UpperCAmelCase       ,			generator=_UpperCAmelCase       ,			device=_UpperCAmelCase       ,			dtype=_UpperCAmelCase       ,			layout=_UpperCAmelCase				).to(_UpperCAmelCase				)

					return latents





def 		__snake_case						(    _UpperCAmelCase				):
					"""simple docstring"""






					if is_torch_version('<'       ,			'2.0.0'				) or not hasattr(_UpperCAmelCase       ,			'_dynamo'				):
										return False
					return isinstance(_UpperCAmelCase       ,			torch._dynamo.eval_frame.OptimizedModule				)






