'''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 __lowerCAmelCase				(							UpperCamelCase__       )			->					List[str]:
    __lowerCamelCase						=						TimesformerConfig()

    if "large" in model_name:
        __lowerCamelCase						=						96

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

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

    return config






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

    return name






def __lowerCAmelCase				(							UpperCamelCase__       ,						UpperCamelCase__       )			->					Optional[int]:
    for key in orig_state_dict.copy().keys():
        __lowerCamelCase						=						orig_state_dict.pop(UpperCamelCase__       )

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

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

    return orig_state_dict






def __lowerCAmelCase				(							)			->					Optional[int]:
    __lowerCamelCase						=						hf_hub_download(
        repo_id='''hf-internal-testing/spaghetti-video'''       ,						filename='''eating_spaghetti.npy'''       ,						repo_type='''dataset'''       )
    __lowerCamelCase						=						np.load(UpperCamelCase__       )
    return list(UpperCamelCase__       )






def __lowerCAmelCase				(							UpperCamelCase__       ,						UpperCamelCase__       ,						UpperCamelCase__       ,						UpperCamelCase__       )			->					Optional[Any]:
    __lowerCamelCase						=						get_timesformer_config(UpperCamelCase__       )

    __lowerCamelCase						=						TimesformerForVideoClassification(UpperCamelCase__       )

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

    model.load_state_dict(UpperCamelCase__       )
    model.eval()

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

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

    __lowerCamelCase						=						[
        # 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":
        __lowerCamelCase						=						torch.Size([1, 4_00]       )
        __lowerCamelCase						=						torch.tensor([-0.3_0_1_6, -0.7_7_1_3, -0.4_2_0_5]       )
    elif model_name == "timesformer-base-finetuned-k600":
        __lowerCamelCase						=						torch.Size([1, 6_00]       )
        __lowerCamelCase						=						torch.tensor([-0.7_2_6_7, -0.7_4_6_6, 3.2_4_0_4]       )
    elif model_name == "timesformer-base-finetuned-ssv2":
        __lowerCamelCase						=						torch.Size([1, 1_74]       )
        __lowerCamelCase						=						torch.tensor([-0.9_0_5_9, 0.6_4_3_3, -3.1_4_5_7]       )
    elif model_name == "timesformer-large-finetuned-k400":
        __lowerCamelCase						=						torch.Size([1, 4_00]       )
        __lowerCamelCase						=						torch.tensor([0, 0, 0]       )
    elif model_name == "timesformer-large-finetuned-k600":
        __lowerCamelCase						=						torch.Size([1, 6_00]       )
        __lowerCamelCase						=						torch.tensor([0, 0, 0]       )
    elif model_name == "timesformer-large-finetuned-ssv2":
        __lowerCamelCase						=						torch.Size([1, 1_74]       )
        __lowerCamelCase						=						torch.tensor([0, 0, 0]       )
    elif model_name == "timesformer-hr-finetuned-k400":
        __lowerCamelCase						=						torch.Size([1, 4_00]       )
        __lowerCamelCase						=						torch.tensor([-0.9_6_1_7, -3.7_3_1_1, -3.7_7_0_8]       )
    elif model_name == "timesformer-hr-finetuned-k600":
        __lowerCamelCase						=						torch.Size([1, 6_00]       )
        __lowerCamelCase						=						torch.tensor([2.5_2_7_3, 0.7_1_2_7, 1.8_8_4_8]       )
    elif model_name == "timesformer-hr-finetuned-ssv2":
        __lowerCamelCase						=						torch.Size([1, 1_74]       )
        __lowerCamelCase						=						torch.tensor([-3.6_7_5_6, -0.7_5_1_3, 0.7_1_8_0]       )
    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__":
  __UpperCAmelCase  							=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."
  )

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

'''simple docstring'''




from __future__ import annotations






def __lowerCAmelCase				(							UpperCamelCase__       ,						UpperCamelCase__       ,						UpperCamelCase__       )			->					tuple[float, float]:
    assert (
        c % greatest_common_divisor(UpperCamelCase__       ,						UpperCamelCase__       ) == 0
    )  # greatest_common_divisor(a,b) function implemented below
    ((__lowerCamelCase)			,						(__lowerCamelCase)			,						(__lowerCamelCase))						=						extended_gcd(UpperCamelCase__       ,						UpperCamelCase__       )  # extended_gcd(a,b) function implemented below
    __lowerCamelCase						=						c / d
    return (r * x, r * y)






def __lowerCAmelCase				(							UpperCamelCase__       ,						UpperCamelCase__       ,						UpperCamelCase__       ,						UpperCamelCase__ = 2       )			->					None:
    ((__lowerCamelCase)			,						(__lowerCamelCase))						=						diophantine(UpperCamelCase__       ,						UpperCamelCase__       ,						UpperCamelCase__       )  # Initial value
    __lowerCamelCase						=						greatest_common_divisor(UpperCamelCase__       ,						UpperCamelCase__       )
    __lowerCamelCase						=						a // d
    __lowerCamelCase						=						b // d

    for i in range(UpperCamelCase__       ):
        __lowerCamelCase						=						xa + i * q
        __lowerCamelCase						=						ya - i * p
        print(UpperCamelCase__       ,						UpperCamelCase__       )






def __lowerCAmelCase				(							UpperCamelCase__       ,						UpperCamelCase__       )			->					int:
    if a < b:
        __lowerCamelCase			,						__lowerCamelCase						=						b, a

    while a % b != 0:
        __lowerCamelCase			,						__lowerCamelCase						=						b, a % b

    return b






def __lowerCAmelCase				(							UpperCamelCase__       ,						UpperCamelCase__       )			->					tuple[int, int, int]:
    assert a >= 0 and b >= 0

    if b == 0:
        __lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase						=						a, 1, 0
    else:
        ((__lowerCamelCase)			,						(__lowerCamelCase)			,						(__lowerCamelCase))						=						extended_gcd(UpperCamelCase__       ,						a % b       )
        __lowerCamelCase						=						q
        __lowerCamelCase						=						p - q * (a // b)

    assert a % d == 0 and b % d == 0
    assert d == a * x + b * y

    return (d, x, y)


if __name__ == "__main__":
  from doctest import testmod

  testmod(name="diophantine", verbose=True)
  testmod(name="diophantine_all_soln", verbose=True)
  testmod(name="extended_gcd", verbose=True)
  testmod(name="greatest_common_divisor", verbose=True)
