






'''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 						_UpperCAmelCase     (	a		:						Any				)      ->      List[Any]:





						"""simple docstring"""
						lowercase_   : Optional[Any]   	=      TimesformerConfig()

						if "large" in model_name:
												lowercase_   : List[str]   	=      9_6

						if "hr" in model_name:
												lowercase_   : int   	=      1_6
												lowercase_   : Dict   	=      4_4_8

						lowercase_   : Optional[int]   	=      'huggingface/label-files'
						if "k400" in model_name:
												lowercase_   : Optional[Any]   	=      4_0_0
												lowercase_   : List[Any]   	=      'kinetics400-id2label.json'
						elif "k600" in model_name:
												lowercase_   : int   	=      6_0_0
												lowercase_   : Union[str, Any]   	=      'kinetics600-id2label.json'
						elif "ssv2" in model_name:
												lowercase_   : List[str]   	=      1_7_4
												lowercase_   : Optional[int]   	=      'something-something-v2-id2label.json'
						else:
												raise ValueError('Model name should either contain \'k400\', \'k600\' or \'ssv2\'.'				)
						lowercase_   : int   	=      json.load(open(hf_hub_download(a    ,      a    ,      repo_type='dataset'				)    ,      'r'				)				)
						lowercase_   : Any   	=      {int(a				): v for k, v in idalabel.items()}
						lowercase_   : Optional[Any]   	=      idalabel
						lowercase_   : List[Any]   	=      {v: k for k, v in idalabel.items()}

						return config






def 						_UpperCAmelCase     (	a		:						Any				)      ->      List[str]:





						"""simple docstring"""
						if "encoder." in name:
												lowercase_   : Dict   	=      name.replace('encoder.'    ,      ''				)
						if "cls_token" in name:
												lowercase_   : Any   	=      name.replace('cls_token'    ,      'timesformer.embeddings.cls_token'				)
						if "pos_embed" in name:
												lowercase_   : List[str]   	=      name.replace('pos_embed'    ,      'timesformer.embeddings.position_embeddings'				)
						if "time_embed" in name:
												lowercase_   : List[Any]   	=      name.replace('time_embed'    ,      'timesformer.embeddings.time_embeddings'				)
						if "patch_embed.proj" in name:
												lowercase_   : Tuple   	=      name.replace('patch_embed.proj'    ,      'timesformer.embeddings.patch_embeddings.projection'				)
						if "patch_embed.norm" in name:
												lowercase_   : List[Any]   	=      name.replace('patch_embed.norm'    ,      'timesformer.embeddings.norm'				)
						if "blocks" in name:
												lowercase_   : Optional[Any]   	=      name.replace('blocks'    ,      'timesformer.encoder.layer'				)
						if "attn.proj" in name:
												lowercase_   : List[str]   	=      name.replace('attn.proj'    ,      'attention.output.dense'				)
						if "attn" in name and "bias" not in name and "temporal" not in name:
												lowercase_   : int   	=      name.replace('attn'    ,      'attention.self'				)
						if "attn" in name and "temporal" not in name:
												lowercase_   : Optional[int]   	=      name.replace('attn'    ,      'attention.attention'				)
						if "temporal_norm1" in name:
												lowercase_   : Union[str, Any]   	=      name.replace('temporal_norm1'    ,      'temporal_layernorm'				)
						if "temporal_attn.proj" in name:
												lowercase_   : Optional[Any]   	=      name.replace('temporal_attn'    ,      'temporal_attention.output.dense'				)
						if "temporal_fc" in name:
												lowercase_   : Tuple   	=      name.replace('temporal_fc'    ,      'temporal_dense'				)
						if "norm1" in name and "temporal" not in name:
												lowercase_   : Any   	=      name.replace('norm1'    ,      'layernorm_before'				)
						if "norm2" in name:
												lowercase_   : List[str]   	=      name.replace('norm2'    ,      'layernorm_after'				)
						if "mlp.fc1" in name:
												lowercase_   : Any   	=      name.replace('mlp.fc1'    ,      'intermediate.dense'				)
						if "mlp.fc2" in name:
												lowercase_   : Optional[Any]   	=      name.replace('mlp.fc2'    ,      'output.dense'				)
						if "norm.weight" in name and "fc" not in name and "temporal" not in name:
												lowercase_   : List[str]   	=      name.replace('norm.weight'    ,      'timesformer.layernorm.weight'				)
						if "norm.bias" in name and "fc" not in name and "temporal" not in name:
												lowercase_   : Optional[Any]   	=      name.replace('norm.bias'    ,      'timesformer.layernorm.bias'				)
						if "head" in name:
												lowercase_   : Dict   	=      name.replace('head'    ,      'classifier'				)

						return name






def 						_UpperCAmelCase     (	a		:						Optional[int]    ,      a		:						Optional[Any]				)      ->      Any:





						"""simple docstring"""
						for key in orig_state_dict.copy().keys():
												lowercase_   : Dict   	=      orig_state_dict.pop(a				)

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

												if "qkv" in key:
																		lowercase_   : Union[str, Any]   	=      key.split('.'				)
																		lowercase_   : Optional[int]   	=      int(key_split[1]				)
																		lowercase_   : str   	=      'timesformer.encoder.layer.'
																		if "temporal" in key:
																								lowercase_   : Optional[Any]   	=      '.temporal_attention.attention.qkv.'
																		else:
																								lowercase_   : List[str]   	=      '.attention.attention.qkv.'
																		if "weight" in key:
																								lowercase_   : Union[str, Any]   	=      val
																		else:
																								lowercase_   : Union[str, Any]   	=      val
												else:
																		lowercase_   : List[Any]   	=      val

						return orig_state_dict






def 						_UpperCAmelCase     (	)      ->      Any:





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






def 						_UpperCAmelCase     (	a		:						Optional[int]    ,      a		:						List[str]    ,      a		:						List[Any]    ,      a		:						List[Any]				)      ->      List[str]:





						"""simple docstring"""
						lowercase_   : Optional[Any]   	=      get_timesformer_config(a				)

						lowercase_   : Tuple   	=      TimesformerForVideoClassification(a				)

						# download original checkpoint, hosted on Google Drive
						lowercase_   : Tuple   	=      'pytorch_model.bin'
						gdown.cached_download(a    ,      a    ,      quiet=a				)
						lowercase_   : Optional[int]   	=      torch.load(a    ,      map_location='cpu'				)
						if "model" in files:
												lowercase_   : Tuple   	=      files['model']
						elif "module" in files:
												lowercase_   : Any   	=      files['module']
						else:
												lowercase_   : List[str]   	=      files['model_state']
						lowercase_   : str   	=      convert_state_dict(a    ,      a				)

						model.load_state_dict(a				)
						model.eval()

						# verify model on basic input
						lowercase_   : Optional[Any]   	=      VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]    ,      image_std=[0.5, 0.5, 0.5]				)
						lowercase_   : Union[str, Any]   	=      prepare_video()
						lowercase_   : List[str]   	=      image_processor(video[:8]    ,      return_tensors='pt'				)

						lowercase_   : str   	=      model(**a				)
						lowercase_   : Any   	=      outputs.logits

						lowercase_   : Optional[Any]   	=      [
						    # 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_   : List[str]   	=      torch.Size([1, 4_0_0]				)
												lowercase_   : List[Any]   	=      torch.tensor([-0.30_16, -0.77_13, -0.42_05]				)
						elif model_name == "timesformer-base-finetuned-k600":
												lowercase_   : Any   	=      torch.Size([1, 6_0_0]				)
												lowercase_   : Dict   	=      torch.tensor([-0.72_67, -0.74_66, 3.24_04]				)
						elif model_name == "timesformer-base-finetuned-ssv2":
												lowercase_   : str   	=      torch.Size([1, 1_7_4]				)
												lowercase_   : Optional[Any]   	=      torch.tensor([-0.90_59, 0.64_33, -3.14_57]				)
						elif model_name == "timesformer-large-finetuned-k400":
												lowercase_   : int   	=      torch.Size([1, 4_0_0]				)
												lowercase_   : Any   	=      torch.tensor([0, 0, 0]				)
						elif model_name == "timesformer-large-finetuned-k600":
												lowercase_   : List[str]   	=      torch.Size([1, 6_0_0]				)
												lowercase_   : int   	=      torch.tensor([0, 0, 0]				)
						elif model_name == "timesformer-large-finetuned-ssv2":
												lowercase_   : List[str]   	=      torch.Size([1, 1_7_4]				)
												lowercase_   : Any   	=      torch.tensor([0, 0, 0]				)
						elif model_name == "timesformer-hr-finetuned-k400":
												lowercase_   : str   	=      torch.Size([1, 4_0_0]				)
												lowercase_   : int   	=      torch.tensor([-0.96_17, -3.73_11, -3.77_08]				)
						elif model_name == "timesformer-hr-finetuned-k600":
												lowercase_   : Any   	=      torch.Size([1, 6_0_0]				)
												lowercase_   : List[Any]   	=      torch.tensor([2.52_73, 0.71_27, 1.88_48]				)
						elif model_name == "timesformer-hr-finetuned-ssv2":
												lowercase_   : Optional[int]   	=      torch.Size([1, 1_7_4]				)
												lowercase_   : Union[str, Any]   	=      torch.tensor([-3.67_56, -0.75_13, 0.71_80]				)
						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]    ,      a    ,      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(a				)
												model.save_pretrained(a				)

						if push_to_hub:
												print('Pushing to the hub...'				)
												model.push_to_hub(f"fcakyon/{model_name}"				)


if __name__ == "__main__":
						A:   int						=					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:   List[str]						=					parser.parse_args()
						convert_timesformer_checkpoint(
						    args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
						)








'''simple docstring'''

import inspect
import warnings

import torch

from .state import AcceleratorState, GradientState
from .utils import DistributedType, honor_type, is_tpu_available


if is_tpu_available(check_device=False):
						import torch_xla.core.xla_model as xm






def 						_UpperCAmelCase     (	a		:						Tuple    ,      a		:						int				)      ->      List[str]:





						"""simple docstring"""
						if isinstance(a    ,      (list, tuple)				):
												return honor_type(a    ,      (move_to_device(a    ,      a				) for t in state)				)
						elif isinstance(a    ,      a				):
												return type(a				)({k: move_to_device(a    ,      a				) for k, v in state.items()}				)
						elif isinstance(a    ,      torch.Tensor				):
												return state.to(a				)
						return state






class 				__magic_name__  (						torch.optim.Optimizer					):







	"""simple docstring"""



	def __init__(				self	,   _lowercase	,   _lowercase=True	,   _lowercase=None							)    ->	int:
							lowercase_   : List[str]   	=      optimizer
							lowercase_   : Optional[int]   	=      scaler
							lowercase_   : Union[str, Any]   	=      AcceleratorState()
							lowercase_   : Optional[Any]   	=      GradientState()
							lowercase_   : Union[str, Any]   	=      device_placement
							lowercase_   : List[str]   	=      False
							lowercase_   : Optional[Any]   	=      None

							# Handle device placement
							if device_placement:
													lowercase_   : str   	=      self.optimizer.state_dict()
													if self.accelerator_state.distributed_type == DistributedType.TPU:
																			xm.send_cpu_data_to_device(_lowercase	,   self.accelerator_state.device							)
													else:
																			lowercase_   : List[str]   	=      move_to_device(_lowercase	,   self.accelerator_state.device							)
													self.optimizer.load_state_dict(_lowercase							)
	@property
	def 							lowerCamelCase__		(				self							)    ->	Any:
							return self.optimizer.state
	@state.setter
	def 							lowerCamelCase__		(				self	,   _lowercase							)    ->	Union[str, Any]:
							lowercase_   : Optional[Any]   	=      state
	@property
	def 							lowerCamelCase__		(				self							)    ->	Any:
							return self.optimizer.param_groups
	@param_groups.setter
	def 							lowerCamelCase__		(				self	,   _lowercase							)    ->	Optional[Any]:
							lowercase_   : Dict   	=      param_groups
	@property
	def 							lowerCamelCase__		(				self							)    ->	int:
							return self.optimizer.defaults
	@defaults.setter
	def 							lowerCamelCase__		(				self	,   _lowercase							)    ->	Any:
							lowercase_   : List[Any]   	=      defaults
	def 							lowerCamelCase__		(				self	,   _lowercase							)    ->	Tuple:
							self.optimizer.add_param_group(_lowercase							)
	def 							lowerCamelCase__		(				self	,   _lowercase							)    ->	Any:
							if self.accelerator_state.distributed_type == DistributedType.TPU and self.device_placement:
													xm.send_cpu_data_to_device(_lowercase	,   self.accelerator_state.device							)
							self.optimizer.load_state_dict(_lowercase							)
	def 							lowerCamelCase__		(				self							)    ->	Any:
							return self.optimizer.state_dict()
	def 							lowerCamelCase__		(				self	,   _lowercase=None							)    ->	Optional[Any]:
							if self.gradient_state.sync_gradients:
													lowercase_   : Tuple   	=      'set_to_none' in inspect.signature(self.optimizer.zero_grad							).parameters
													if accept_arg:
																			if set_to_none is None:
																									lowercase_   : Optional[int]   	=      False
																			self.optimizer.zero_grad(set_to_none=_lowercase							)
													else:
																			if set_to_none is not None:
																									raise ValueError('`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.'							)
																			self.optimizer.zero_grad()
	def 							lowerCamelCase__		(				self	,   _lowercase=None							)    ->	Optional[Any]:
							if self.gradient_state.sync_gradients:
													if self.accelerator_state.distributed_type == DistributedType.TPU:
																			lowercase_   : Union[str, Any]   	=      {'closure': closure} if closure is not None else {}
																			xm.optimizer_step(self.optimizer	,   optimizer_args=_lowercase							)
													elif self.scaler is not None:
																			lowercase_   : Optional[Any]   	=      False
																			if self._last_scale is None:
																									# `get_scale` is an async operation requiring full synchronization
																									# on CPU and GPUs before finishing. As a result, we store away
																									# the prior one to reduce the call overhead
																									lowercase_   : Tuple   	=      self.scaler.get_scale()
																									lowercase_   : int   	=      True
																			self.scaler.step(self.optimizer	,   _lowercase							)
																			self.scaler.update()
																			lowercase_   : List[str]   	=      self.scaler.get_scale()
																			if not new_scale:
																									# If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
																									lowercase_   : Optional[int]   	=      scale_after < self._last_scale
																			lowercase_   : List[str]   	=      scale_after

													else:
																			self.optimizer.step(_lowercase							)
	def 							lowerCamelCase__		(				self	,   _lowercase							)    ->	Optional[int]:
							for param_group in self.optimizer.param_groups:
													lowercase_   : Dict   	=      [parameters_map.get(_lowercase	,   _lowercase							) for p in param_group['params']]
	@property
	def 							lowerCamelCase__		(				self							)    ->	List[Any]:
							warnings.warn(
							    'The `is_overflow` property is deprecated and will be removed in version 1.0 of Accelerate use '
							    '`optimizer.step_was_skipped` instead.'	,   _lowercase	,   )
							return self._is_overflow
	@property
	def 							lowerCamelCase__		(				self							)    ->	str:
							return self._is_overflow
	def __getstate__(				self							)    ->	str:
							return self.__dict__.copy()







	def __setstate__(				self	,   _lowercase							)    ->	Union[str, Any]:
							self.__dict__.update(_lowercase							)
