


'''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 (      __lowerCamelCase		:			Tuple					)   ->  Union[str, Any]:







     '''simple docstring'''






     _a            =		TimesformerConfig()

     if "large" in model_name:
          _a            =		96

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

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

     return config






def __UpperCamelCase (      __lowerCamelCase		:			int					)   ->  Any:







     '''simple docstring'''






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

     return name






def __UpperCamelCase (      __lowerCamelCase		:			Optional[Any]   ,	__lowerCamelCase		:			Optional[Any]					)   ->  str:







     '''simple docstring'''






     for key in orig_state_dict.copy().keys():
          _a            =		orig_state_dict.pop(__lowerCamelCase					)

          if key.startswith("model."					):
               _a            =		key.replace("model."   ,	""					)

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

     return orig_state_dict






def __UpperCamelCase (      )   ->  Optional[int]:







     '''simple docstring'''






     _a            =		hf_hub_download(
         repo_id="hf-internal-testing/spaghetti-video"   ,	filename="eating_spaghetti.npy"   ,	repo_type="dataset"					)
     _a            =		np.load(__lowerCamelCase					)
     return list(__lowerCamelCase					)






def __UpperCamelCase (      __lowerCamelCase		:			Tuple   ,	__lowerCamelCase		:			Optional[Any]   ,	__lowerCamelCase		:			Optional[Any]   ,	__lowerCamelCase		:			Optional[int]					)   ->  Tuple:







     '''simple docstring'''






     _a            =		get_timesformer_config(__lowerCamelCase					)

     _a            =		TimesformerForVideoClassification(__lowerCamelCase					)

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

     model.load_state_dict(__lowerCamelCase					)
     model.eval()

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

     _a            =		model(**__lowerCamelCase					)
     _a            =		outputs.logits

     _a            =		[
         # 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":
          _a            =		torch.Size([1, 400]					)
          _a            =		torch.tensor([-0.30_16, -0.77_13, -0.42_05]					)
     elif model_name == "timesformer-base-finetuned-k600":
          _a            =		torch.Size([1, 600]					)
          _a            =		torch.tensor([-0.72_67, -0.74_66, 3.24_04]					)
     elif model_name == "timesformer-base-finetuned-ssv2":
          _a            =		torch.Size([1, 174]					)
          _a            =		torch.tensor([-0.90_59, 0.64_33, -3.14_57]					)
     elif model_name == "timesformer-large-finetuned-k400":
          _a            =		torch.Size([1, 400]					)
          _a            =		torch.tensor([0, 0, 0]					)
     elif model_name == "timesformer-large-finetuned-k600":
          _a            =		torch.Size([1, 600]					)
          _a            =		torch.tensor([0, 0, 0]					)
     elif model_name == "timesformer-large-finetuned-ssv2":
          _a            =		torch.Size([1, 174]					)
          _a            =		torch.tensor([0, 0, 0]					)
     elif model_name == "timesformer-hr-finetuned-k400":
          _a            =		torch.Size([1, 400]					)
          _a            =		torch.tensor([-0.96_17, -3.73_11, -3.77_08]					)
     elif model_name == "timesformer-hr-finetuned-k600":
          _a            =		torch.Size([1, 600]					)
          _a            =		torch.tensor([2.52_73, 0.71_27, 1.88_48]					)
     elif model_name == "timesformer-hr-finetuned-ssv2":
          _a            =		torch.Size([1, 174]					)
          _a            =		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]   ,	__lowerCamelCase   ,	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(__lowerCamelCase					)
          model.save_pretrained(__lowerCamelCase					)

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


if __name__ == "__main__":
   lowercase__					  =		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."
   )

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





'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


lowercase__					  =		{
    "configuration_swin2sr": ["SWIN2SR_PRETRAINED_CONFIG_ARCHIVE_MAP", "Swin2SRConfig"],
}


try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowercase__					  =		[
       "SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST",
       "Swin2SRForImageSuperResolution",
       "Swin2SRModel",
       "Swin2SRPreTrainedModel",
   ]


try:
   if not is_vision_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowercase__					  =		["Swin2SRImageProcessor"]


if TYPE_CHECKING:
   from .configuration_swinasr import SWIN2SR_PRETRAINED_CONFIG_ARCHIVE_MAP, SwinaSRConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_swinasr import (
          SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST,
          SwinaSRForImageSuperResolution,
          SwinaSRModel,
          SwinaSRPreTrainedModel,
      )

   try:
      if not is_vision_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .image_processing_swinasr import SwinaSRImageProcessor


else:
   import sys

   lowercase__					  =		_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

