





import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)






def 				_UpperCAmelCase			(    A     ):




   '''simple docstring'''

   UpperCAmelCase__						=VideoMAEConfig()

   set_architecture_configs(A				,							A     )

   if "finetuned" not in model_name:
      UpperCAmelCase__						=False

   if "finetuned" in model_name:
      UpperCAmelCase__						="huggingface/label-files"
      if "kinetics" in model_name:
         UpperCAmelCase__						=400
         UpperCAmelCase__						="kinetics400-id2label.json"
      elif "ssv2" in model_name:
         UpperCAmelCase__						=174
         UpperCAmelCase__						="something-something-v2-id2label.json"
      else:
         raise ValueError("Model name should either contain 'kinetics' or 'ssv2' in case it's fine-tuned."     )
      UpperCAmelCase__						=json.load(open(hf_hub_download(A				,							A				,							repo_type="dataset"     )				,							"r"     )     )
      UpperCAmelCase__						={int(A     ): v for k, v in idalabel.items()}
      UpperCAmelCase__						=idalabel
      UpperCAmelCase__						={v: k for k, v in idalabel.items()}

   return config






def 				_UpperCAmelCase			(    A				,							A     ):




   '''simple docstring'''

   if "small" in model_name:
      UpperCAmelCase__						=384
      UpperCAmelCase__						=1536
      UpperCAmelCase__						=12
      UpperCAmelCase__						=16
      UpperCAmelCase__						=12
      UpperCAmelCase__						=3
      UpperCAmelCase__						=192
      UpperCAmelCase__						=768
   elif "large" in model_name:
      UpperCAmelCase__						=1024
      UpperCAmelCase__						=4096
      UpperCAmelCase__						=24
      UpperCAmelCase__						=16
      UpperCAmelCase__						=12
      UpperCAmelCase__						=8
      UpperCAmelCase__						=512
      UpperCAmelCase__						=2048
   elif "huge" in model_name:
      UpperCAmelCase__						=1280
      UpperCAmelCase__						=5120
      UpperCAmelCase__						=32
      UpperCAmelCase__						=16
      UpperCAmelCase__						=12
      UpperCAmelCase__						=8
      UpperCAmelCase__						=640
      UpperCAmelCase__						=2560
   elif "base" not in model_name:
      raise ValueError("Model name should include either \"small\", \"base\", \"large\", or \"huge\""     )






def 				_UpperCAmelCase			(    A     ):




   '''simple docstring'''

   if "encoder." in name:
      UpperCAmelCase__						=name.replace("encoder."				,							""     )
   if "cls_token" in name:
      UpperCAmelCase__						=name.replace("cls_token"				,							"videomae.embeddings.cls_token"     )
   if "decoder_pos_embed" in name:
      UpperCAmelCase__						=name.replace("decoder_pos_embed"				,							"decoder.decoder_pos_embed"     )
   if "pos_embed" in name and "decoder" not in name:
      UpperCAmelCase__						=name.replace("pos_embed"				,							"videomae.embeddings.position_embeddings"     )
   if "patch_embed.proj" in name:
      UpperCAmelCase__						=name.replace("patch_embed.proj"				,							"videomae.embeddings.patch_embeddings.projection"     )
   if "patch_embed.norm" in name:
      UpperCAmelCase__						=name.replace("patch_embed.norm"				,							"videomae.embeddings.norm"     )
   if "decoder.blocks" in name:
      UpperCAmelCase__						=name.replace("decoder.blocks"				,							"decoder.decoder_layers"     )
   if "blocks" in name:
      UpperCAmelCase__						=name.replace("blocks"				,							"videomae.encoder.layer"     )
   if "attn.proj" in name:
      UpperCAmelCase__						=name.replace("attn.proj"				,							"attention.output.dense"     )
   if "attn" in name and "bias" not in name:
      UpperCAmelCase__						=name.replace("attn"				,							"attention.self"     )
   if "attn" in name:
      UpperCAmelCase__						=name.replace("attn"				,							"attention.attention"     )
   if "norm1" in name:
      UpperCAmelCase__						=name.replace("norm1"				,							"layernorm_before"     )
   if "norm2" in name:
      UpperCAmelCase__						=name.replace("norm2"				,							"layernorm_after"     )
   if "mlp.fc1" in name:
      UpperCAmelCase__						=name.replace("mlp.fc1"				,							"intermediate.dense"     )
   if "mlp.fc2" in name:
      UpperCAmelCase__						=name.replace("mlp.fc2"				,							"output.dense"     )
   if "decoder_embed" in name:
      UpperCAmelCase__						=name.replace("decoder_embed"				,							"decoder.decoder_embed"     )
   if "decoder_norm" in name:
      UpperCAmelCase__						=name.replace("decoder_norm"				,							"decoder.decoder_norm"     )
   if "decoder_pred" in name:
      UpperCAmelCase__						=name.replace("decoder_pred"				,							"decoder.decoder_pred"     )
   if "norm.weight" in name and "decoder" not in name and "fc" not in name:
      UpperCAmelCase__						=name.replace("norm.weight"				,							"videomae.layernorm.weight"     )
   if "norm.bias" in name and "decoder" not in name and "fc" not in name:
      UpperCAmelCase__						=name.replace("norm.bias"				,							"videomae.layernorm.bias"     )
   if "head" in name and "decoder" not in name:
      UpperCAmelCase__						=name.replace("head"				,							"classifier"     )

   return name






def 				_UpperCAmelCase			(    A				,							A     ):




   '''simple docstring'''

   for key in orig_state_dict.copy().keys():
      UpperCAmelCase__						=orig_state_dict.pop(A     )

      if key.startswith("encoder."     ):
         UpperCAmelCase__						=key.replace("encoder."				,							""     )

      if "qkv" in key:
         UpperCAmelCase__						=key.split("."     )
         if key.startswith("decoder.blocks"     ):
            UpperCAmelCase__						=config.decoder_hidden_size
            UpperCAmelCase__						=int(key_split[2]     )
            UpperCAmelCase__						="decoder.decoder_layers."
            if "weight" in key:
               UpperCAmelCase__						=val[:dim, :]
               UpperCAmelCase__						=val[dim : dim * 2, :]
               UpperCAmelCase__						=val[-dim:, :]
         else:
            UpperCAmelCase__						=config.hidden_size
            UpperCAmelCase__						=int(key_split[1]     )
            UpperCAmelCase__						="videomae.encoder.layer."
            if "weight" in key:
               UpperCAmelCase__						=val[:dim, :]
               UpperCAmelCase__						=val[dim : dim * 2, :]
               UpperCAmelCase__						=val[-dim:, :]
      else:
         UpperCAmelCase__						=val

   return orig_state_dict






def 				_UpperCAmelCase			(    ):




   '''simple docstring'''

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






def 				_UpperCAmelCase			(    A				,							A				,							A				,							A     ):




   '''simple docstring'''

   UpperCAmelCase__						=get_videomae_config(A     )

   if "finetuned" in model_name:
      UpperCAmelCase__						=VideoMAEForVideoClassification(A     )
   else:
      UpperCAmelCase__						=VideoMAEForPreTraining(A     )

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

   model.load_state_dict(A     )
   model.eval()

   # verify model on basic input
   UpperCAmelCase__						=VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]				,							image_std=[0.5, 0.5, 0.5]     )
   UpperCAmelCase__						=prepare_video()
   UpperCAmelCase__						=image_processor(A				,							return_tensors="pt"     )

   if "finetuned" not in model_name:
      UpperCAmelCase__						=hf_hub_download(repo_id="hf-internal-testing/bool-masked-pos"				,							filename="bool_masked_pos.pt"     )
      UpperCAmelCase__						=torch.load(A     )

   UpperCAmelCase__						=model(**A     )
   UpperCAmelCase__						=outputs.logits

   UpperCAmelCase__						=[
       "videomae-small-finetuned-kinetics",
       "videomae-small-finetuned-ssv2",
       # Kinetics-400 checkpoints (short = pretrained only for 800 epochs instead of 1600)
       "videomae-base-short",
       "videomae-base-short-finetuned-kinetics",
       "videomae-base",
       "videomae-base-finetuned-kinetics",
       "videomae-large",
       "videomae-large-finetuned-kinetics",
       "videomae-huge-finetuned-kinetics",
       # Something-Something-v2 checkpoints (short = pretrained only for 800 epochs instead of 2400)
       "videomae-base-short-ssv2",
       "videomae-base-short-finetuned-ssv2",
       "videomae-base-ssv2",
       "videomae-base-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 == "videomae-small-finetuned-kinetics":
      UpperCAmelCase__						=torch.Size([1, 400]     )
      UpperCAmelCase__						=torch.tensor([-0.92_91, -0.40_61, -0.93_07]     )
   elif model_name == "videomae-small-finetuned-ssv2":
      UpperCAmelCase__						=torch.Size([1, 174]     )
      UpperCAmelCase__						=torch.tensor([0.26_71, -0.46_89, -0.82_35]     )
   elif model_name == "videomae-base":
      UpperCAmelCase__						=torch.Size([1, 1408, 1536]     )
      UpperCAmelCase__						=torch.tensor([[0.77_39, 0.79_68, 0.70_89], [0.67_01, 0.74_87, 0.62_09], [0.42_87, 0.51_58, 0.47_73]]     )
   elif model_name == "videomae-base-short":
      UpperCAmelCase__						=torch.Size([1, 1408, 1536]     )
      UpperCAmelCase__						=torch.tensor([[0.79_94, 0.96_12, 0.85_08], [0.74_01, 0.89_58, 0.83_02], [0.58_62, 0.74_68, 0.73_25]]     )
      # we verified the loss both for normalized and unnormalized targets for this one
      UpperCAmelCase__						=torch.tensor([0.51_42]     ) if config.norm_pix_loss else torch.tensor([0.64_69]     )
   elif model_name == "videomae-large":
      UpperCAmelCase__						=torch.Size([1, 1408, 1536]     )
      UpperCAmelCase__						=torch.tensor([[0.71_49, 0.79_97, 0.69_66], [0.67_68, 0.78_69, 0.69_48], [0.51_39, 0.62_21, 0.56_05]]     )
   elif model_name == "videomae-large-finetuned-kinetics":
      UpperCAmelCase__						=torch.Size([1, 400]     )
      UpperCAmelCase__						=torch.tensor([0.07_71, 0.00_11, -0.36_25]     )
   elif model_name == "videomae-huge-finetuned-kinetics":
      UpperCAmelCase__						=torch.Size([1, 400]     )
      UpperCAmelCase__						=torch.tensor([0.24_33, 0.16_32, -0.48_94]     )
   elif model_name == "videomae-base-short-finetuned-kinetics":
      UpperCAmelCase__						=torch.Size([1, 400]     )
      UpperCAmelCase__						=torch.tensor([0.65_88, 0.09_90, -0.24_93]     )
   elif model_name == "videomae-base-finetuned-kinetics":
      UpperCAmelCase__						=torch.Size([1, 400]     )
      UpperCAmelCase__						=torch.tensor([0.36_69, -0.06_88, -0.24_21]     )
   elif model_name == "videomae-base-short-ssv2":
      UpperCAmelCase__						=torch.Size([1, 1408, 1536]     )
      UpperCAmelCase__						=torch.tensor([[0.47_12, 0.52_96, 0.57_86], [0.22_78, 0.27_29, 0.40_26], [0.03_52, 0.07_30, 0.25_06]]     )
   elif model_name == "videomae-base-short-finetuned-ssv2":
      UpperCAmelCase__						=torch.Size([1, 174]     )
      UpperCAmelCase__						=torch.tensor([-0.05_37, -0.15_39, -0.32_66]     )
   elif model_name == "videomae-base-ssv2":
      UpperCAmelCase__						=torch.Size([1, 1408, 1536]     )
      UpperCAmelCase__						=torch.tensor([[0.81_31, 0.87_27, 0.85_46], [0.73_66, 0.93_77, 0.88_70], [0.59_35, 0.88_74, 0.85_64]]     )
   elif model_name == "videomae-base-finetuned-ssv2":
      UpperCAmelCase__						=torch.Size([1, 174]     )
      UpperCAmelCase__						=torch.tensor([0.19_61, -0.83_37, -0.63_89]     )
   else:
      raise ValueError(F"""Model name not supported. Should be one of {model_names}"""     )

   # verify logits
   assert logits.shape == expected_shape
   if "finetuned" in model_name:
      assert torch.allclose(logits[0, :3]				,							A				,							atol=1e-4     )
   else:
      print("Logits:"				,							logits[0, :3, :3]     )
      assert torch.allclose(logits[0, :3, :3]				,							A				,							atol=1e-4     )
   print("Logits ok!"     )

   # verify loss, if applicable
   if model_name == "videomae-base-short":
      UpperCAmelCase__						=outputs.loss
      assert torch.allclose(A				,							A				,							atol=1e-4     )
      print("Loss 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(A				,							organization="nielsr"     )


if __name__ == "__main__":
    UpperCamelCase_												=	argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '--checkpoint_url',
        default='https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4',
        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='/Users/nielsrogge/Documents/VideoMAE/Test',
        type=str,
        help='Path to the output PyTorch model directory.',
    )
    parser.add_argument('--model_name', default='videomae-base', 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_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)










import os
from typing import List, Optional, Union

from ...image_processing_utils import BatchFeature
from ...image_utils import ImageInput
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType
from ..auto import AutoTokenizer

class     snake_case_		(   a							):






 '''simple docstring'''





 __UpperCamelCase       =      ['image_processor', 'tokenizer']
 __UpperCamelCase       =      'BlipImageProcessor'
 __UpperCamelCase       =      'AutoTokenizer'

 def __init__(  self,		A_,		A_,		A_		)	->       str:
    super().__init__(A_,		A_		)

    # add QFormer tokenizer
    UpperCAmelCase__						=qformer_tokenizer

 def __call__(  self,		A_ = None,		A_ = None,		A_ = True,		A_ = False,		A_ = None,		A_ = None,		A_ = 0,		A_ = None,		A_ = None,		A_ = False,		A_ = False,		A_ = False,		A_ = False,		A_ = False,		A_ = True,		A_ = None,		**A_,		)	->       BatchFeature:
    if images is None and text is None:
       raise ValueError("You have to specify at least images or text."		)

    UpperCAmelCase__						=BatchFeature()

    if text is not None:
       UpperCAmelCase__						=self.tokenizer(
           text=A_,		add_special_tokens=A_,		padding=A_,		truncation=A_,		max_length=A_,		stride=A_,		pad_to_multiple_of=A_,		return_attention_mask=A_,		return_overflowing_tokens=A_,		return_special_tokens_mask=A_,		return_offsets_mapping=A_,		return_token_type_ids=A_,		return_length=A_,		verbose=A_,		return_tensors=A_,		**A_,		)
       encoding.update(A_		)
       UpperCAmelCase__						=self.qformer_tokenizer(
           text=A_,		add_special_tokens=A_,		padding=A_,		truncation=A_,		max_length=A_,		stride=A_,		pad_to_multiple_of=A_,		return_attention_mask=A_,		return_overflowing_tokens=A_,		return_special_tokens_mask=A_,		return_offsets_mapping=A_,		return_token_type_ids=A_,		return_length=A_,		verbose=A_,		return_tensors=A_,		**A_,		)
       UpperCAmelCase__						=qformer_text_encoding.pop("input_ids"		)
       UpperCAmelCase__						=qformer_text_encoding.pop("attention_mask"		)

    if images is not None:
       UpperCAmelCase__						=self.image_processor(A_,		return_tensors=A_		)
       encoding.update(A_		)

    return encoding

 def  __UpperCAmelCase						(  self,		*A_,		**A_		)	->       List[Any]:
    return self.tokenizer.batch_decode(*A_,		**A_		)

 def  __UpperCAmelCase						(  self,		*A_,		**A_		)	->       int:
    return self.tokenizer.decode(*A_,		**A_		)

 @property
 # Copied from transformers.models.blip.processing_blip.BlipProcessor.model_input_names
 def  __UpperCAmelCase						(  self		)	->       List[str]:
    UpperCAmelCase__						=self.tokenizer.model_input_names
    UpperCAmelCase__						=self.image_processor.model_input_names
    return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names		)		)

 def  __UpperCAmelCase						(  self,		A_,		**A_		)	->       str:
    if os.path.isfile(A_		):
       raise ValueError(f"""Provided path ({save_directory}) should be a directory, not a file"""		)
    os.makedirs(A_,		exist_ok=A_		)
    UpperCAmelCase__						=os.path.join(A_,		"qformer_tokenizer"		)
    self.qformer_tokenizer.save_pretrained(A_		)
    return super().save_pretrained(A_,		**A_		)






 @classmethod
 def  __UpperCAmelCase						(  cls,		A_,		**A_		)	->       int:
    UpperCAmelCase__						=AutoTokenizer.from_pretrained(A_,		subfolder="qformer_tokenizer"		)
    UpperCAmelCase__						=cls._get_arguments_from_pretrained(A_,		**A_		)
    args.append(A_		)
    return cls(*A_		)



