






"""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      SCREAMING_SNAKE_CASE__    (   snake_case    :      Optional[Any]							)->						Optional[Any]:






     '''simple docstring'''



     UpperCAmelCase__		:  List[str]				=				TimesformerConfig()

     if "large" in model_name:
          UpperCAmelCase__		:  Optional[Any]				=				96

     if "hr" in model_name:
          UpperCAmelCase__		:  str				=				16
          UpperCAmelCase__		:  List[Any]				=				448

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

     return config


def      SCREAMING_SNAKE_CASE__    (   snake_case    :      Union[str, Any]							)->						int:






     '''simple docstring'''



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

     return name


def      SCREAMING_SNAKE_CASE__    (   snake_case    :      Any       , snake_case    :      int							)->						int:






     '''simple docstring'''



     for key in orig_state_dict.copy().keys():
          UpperCAmelCase__		:  Optional[Any]				=				orig_state_dict.pop(snake_case							)

          if key.startswith("model."							):
               UpperCAmelCase__		:  List[str]				=				key.replace("model."       , ""							)

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

     return orig_state_dict


def      SCREAMING_SNAKE_CASE__    (   )->						List[str]:






     '''simple docstring'''



     UpperCAmelCase__		:  Union[str, Any]				=				hf_hub_download(
         repo_id="hf-internal-testing/spaghetti-video"       , filename="eating_spaghetti.npy"       , repo_type="dataset"							)
     UpperCAmelCase__		:  Optional[int]				=				np.load(snake_case							)
     return list(snake_case							)


def      SCREAMING_SNAKE_CASE__    (   snake_case    :      Optional[int]       , snake_case    :      Union[str, Any]       , snake_case    :      Tuple       , snake_case    :      Dict							)->						int:






     '''simple docstring'''



     UpperCAmelCase__		:  Dict				=				get_timesformer_config(snake_case							)

     UpperCAmelCase__		:  Dict				=				TimesformerForVideoClassification(snake_case							)

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

     model.load_state_dict(snake_case							)
     model.eval()

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

     UpperCAmelCase__		:  Any				=				model(**snake_case							)
     UpperCAmelCase__		:  Tuple				=				outputs.logits

     UpperCAmelCase__		:  Tuple				=				[
         # 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__		:  Optional[Any]				=				torch.Size([1, 400]							)
          UpperCAmelCase__		:  int				=				torch.tensor([-0.3016, -0.7713, -0.4205]							)
     elif model_name == "timesformer-base-finetuned-k600":
          UpperCAmelCase__		:  Optional[Any]				=				torch.Size([1, 600]							)
          UpperCAmelCase__		:  Tuple				=				torch.tensor([-0.7267, -0.7466, 3.2404]							)
     elif model_name == "timesformer-base-finetuned-ssv2":
          UpperCAmelCase__		:  Tuple				=				torch.Size([1, 174]							)
          UpperCAmelCase__		:  Union[str, Any]				=				torch.tensor([-0.9059, 0.6433, -3.1457]							)
     elif model_name == "timesformer-large-finetuned-k400":
          UpperCAmelCase__		:  Any				=				torch.Size([1, 400]							)
          UpperCAmelCase__		:  Any				=				torch.tensor([0, 0, 0]							)
     elif model_name == "timesformer-large-finetuned-k600":
          UpperCAmelCase__		:  Dict				=				torch.Size([1, 600]							)
          UpperCAmelCase__		:  List[Any]				=				torch.tensor([0, 0, 0]							)
     elif model_name == "timesformer-large-finetuned-ssv2":
          UpperCAmelCase__		:  Dict				=				torch.Size([1, 174]							)
          UpperCAmelCase__		:  Dict				=				torch.tensor([0, 0, 0]							)
     elif model_name == "timesformer-hr-finetuned-k400":
          UpperCAmelCase__		:  Optional[Any]				=				torch.Size([1, 400]							)
          UpperCAmelCase__		:  Dict				=				torch.tensor([-0.9617, -3.7311, -3.7708]							)
     elif model_name == "timesformer-hr-finetuned-k600":
          UpperCAmelCase__		:  Optional[int]				=				torch.Size([1, 600]							)
          UpperCAmelCase__		:  Optional[Any]				=				torch.tensor([2.5273, 0.7127, 1.8848]							)
     elif model_name == "timesformer-hr-finetuned-ssv2":
          UpperCAmelCase__		:  Tuple				=				torch.Size([1, 174]							)
          UpperCAmelCase__		:  Tuple				=				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]       , snake_case       , 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(snake_case							)
          model.save_pretrained(snake_case							)

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


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

 _lowerCAmelCase     :     Dict            =	parser.parse_args()
 convert_timesformer_checkpoint(
     args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
 )










"""simple docstring"""



import unittest

import torch

from diffusers import DDIMScheduler, TextToVideoZeroPipeline
from diffusers.utils import load_pt, require_torch_gpu, slow

from ..test_pipelines_common import assert_mean_pixel_difference




@slow
@require_torch_gpu
class        lowerCAmelCase__       (  unittest.TestCase					):
   def 	__a				(						self	:     Optional[Any]				):


        '''simple docstring'''
        UpperCAmelCase__		:  int				=				"runwayml/stable-diffusion-v1-5"
        UpperCAmelCase__		:  Optional[Any]				=				TextToVideoZeroPipeline.from_pretrained(snake_case__				,       torch_dtype=torch.floataa				).to("cuda"				)
        UpperCAmelCase__		:  Any				=				DDIMScheduler.from_config(pipe.scheduler.config				)
        UpperCAmelCase__		:  Optional[Any]				=				torch.Generator(device="cuda"				).manual_seed(0				)

        UpperCAmelCase__		:  Any				=				"A bear is playing a guitar on Times Square"
        UpperCAmelCase__		:  Dict				=				pipe(prompt=snake_case__				,       generator=snake_case__				).images

        UpperCAmelCase__		:  List[Any]				=				load_pt(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text-to-video/A bear is playing a guitar on Times Square.pt"				)

        assert_mean_pixel_difference(snake_case__				,       snake_case__				)


