






'''simple docstring'''
import argparse
import os
from pathlib import Path

import torch
from bark.generation import _load_model as _bark_load_model
from huggingface_hub import hf_hub_download

from transformers import EncodecConfig, EncodecModel, set_seed
from transformers.models.bark.configuration_bark import (
    BarkCoarseConfig,
    BarkConfig,
    BarkFineConfig,
    BarkSemanticConfig,
)
from transformers.models.bark.generation_configuration_bark import (
    BarkCoarseGenerationConfig,
    BarkFineGenerationConfig,
    BarkGenerationConfig,
    BarkSemanticGenerationConfig,
)
from transformers.models.bark.modeling_bark import BarkCoarseModel, BarkFineModel, BarkModel, BarkSemanticModel
from transformers.utils import logging


logging.set_verbosity_info()
a_ 		=       logging.get_logger(__name__)

set_seed(7_7_0)


a_ 		=       {
    '''c_attn''': '''att_proj''',
    '''c_proj''': '''out_proj''',
    '''c_fc''': '''in_proj''',
    '''transformer.''': '''''',
    '''h.''': '''layers.''',
    '''ln_1''': '''layernorm_1''',
    '''ln_2''': '''layernorm_2''',
    '''ln_f''': '''layernorm_final''',
    '''wpe''': '''position_embeds_layer''',
    '''wte''': '''input_embeds_layer''',
}


a_ 		=       {
    '''text_small''': {
        '''repo_id''': '''suno/bark''',
        '''file_name''': '''text.pt''',
    },
    '''coarse_small''': {
        '''repo_id''': '''suno/bark''',
        '''file_name''': '''coarse.pt''',
    },
    '''fine_small''': {
        '''repo_id''': '''suno/bark''',
        '''file_name''': '''fine.pt''',
    },
    '''text''': {
        '''repo_id''': '''suno/bark''',
        '''file_name''': '''text_2.pt''',
    },
    '''coarse''': {
        '''repo_id''': '''suno/bark''',
        '''file_name''': '''coarse_2.pt''',
    },
    '''fine''': {
        '''repo_id''': '''suno/bark''',
        '''file_name''': '''fine_2.pt''',
    },
}

a_ 		=       os.path.dirname(os.path.abspath(__file__))
a_ 		=       os.path.join(os.path.expanduser('~'), '.cache')
a_ 		=       os.path.join(os.getenv('XDG_CACHE_HOME', default_cache_dir), 'suno', 'bark_v0')





def        _a(     UpperCamelCase__ :  Tuple,		UpperCamelCase__ :  Optional[int]=False							):




      '''simple docstring'''






      SCREAMING_SNAKE_CASE__   : Tuple													=model_type
      if use_small:
            key += "_small"
      return os.path.join(__lowerCamelCase,		REMOTE_MODEL_PATHS[key]['''file_name''']							)





def        _a(     UpperCamelCase__ :  Optional[Any],		UpperCamelCase__ :  Dict							):




      '''simple docstring'''






      os.makedirs(__lowerCamelCase,		exist_ok=__lowerCamelCase							)
      hf_hub_download(repo_id=__lowerCamelCase,		filename=__lowerCamelCase,		local_dir=__lowerCamelCase							)





def        _a(     UpperCamelCase__ :  Any,		UpperCamelCase__ :  Optional[Any],		UpperCamelCase__ :  Dict=False,		UpperCamelCase__ :  Optional[int]="text"							):




      '''simple docstring'''






      if model_type == "text":
            SCREAMING_SNAKE_CASE__   : Union[str, Any]													=BarkSemanticModel
            SCREAMING_SNAKE_CASE__   : int													=BarkSemanticConfig
            SCREAMING_SNAKE_CASE__   : List[Any]													=BarkSemanticGenerationConfig
      elif model_type == "coarse":
            SCREAMING_SNAKE_CASE__   : List[Any]													=BarkCoarseModel
            SCREAMING_SNAKE_CASE__   : List[Any]													=BarkCoarseConfig
            SCREAMING_SNAKE_CASE__   : Dict													=BarkCoarseGenerationConfig
      elif model_type == "fine":
            SCREAMING_SNAKE_CASE__   : List[str]													=BarkFineModel
            SCREAMING_SNAKE_CASE__   : Tuple													=BarkFineConfig
            SCREAMING_SNAKE_CASE__   : Optional[Any]													=BarkFineGenerationConfig
      else:
            raise NotImplementedError()
      SCREAMING_SNAKE_CASE__   : Any													=f"{model_type}_small" if use_small else model_type
      SCREAMING_SNAKE_CASE__   : Optional[Any]													=REMOTE_MODEL_PATHS[model_key]
      if not os.path.exists(__lowerCamelCase							):
            logger.info(f"{model_type} model not found, downloading into `{CACHE_DIR}`."							)
            _download(model_info['''repo_id'''],		model_info['''file_name''']							)
      SCREAMING_SNAKE_CASE__   : Dict													=torch.load(__lowerCamelCase,		map_location=__lowerCamelCase							)
      # this is a hack
      SCREAMING_SNAKE_CASE__   : List[Any]													=checkpoint["model_args"]
      if "input_vocab_size" not in model_args:
            SCREAMING_SNAKE_CASE__   : int													=model_args["vocab_size"]
            SCREAMING_SNAKE_CASE__   : List[str]													=model_args["vocab_size"]
            del model_args["vocab_size"]

      # convert Bark model arguments to HF Bark model arguments
      SCREAMING_SNAKE_CASE__   : Any													=model_args.pop('''n_head'''							)
      SCREAMING_SNAKE_CASE__   : Optional[Any]													=model_args.pop('''n_embd'''							)
      SCREAMING_SNAKE_CASE__   : Dict													=model_args.pop('''n_layer'''							)

      SCREAMING_SNAKE_CASE__   : Dict													=ConfigClass(**checkpoint['''model_args''']							)
      SCREAMING_SNAKE_CASE__   : Optional[int]													=ModelClass(config=__lowerCamelCase							)
      SCREAMING_SNAKE_CASE__   : Optional[int]													=GenerationConfigClass()

      SCREAMING_SNAKE_CASE__   : Dict													=model_generation_config
      SCREAMING_SNAKE_CASE__   : Dict													=checkpoint["model"]
      # fixup checkpoint
      SCREAMING_SNAKE_CASE__   : Tuple													="_orig_mod."
      for k, v in list(state_dict.items()							):
            if k.startswith(__lowerCamelCase							):
                  # replace part of the key with corresponding layer name in HF implementation
                  SCREAMING_SNAKE_CASE__   : Tuple													=k[len(__lowerCamelCase							) :]
                  for old_layer_name in new_layer_name_dict:
                        SCREAMING_SNAKE_CASE__   : List[Any]													=new_k.replace(__lowerCamelCase,		new_layer_name_dict[old_layer_name]							)

                  SCREAMING_SNAKE_CASE__   : Dict													=state_dict.pop(__lowerCamelCase							)

      SCREAMING_SNAKE_CASE__   : Union[str, Any]													=set(state_dict.keys()							) - set(model.state_dict().keys()							)
      SCREAMING_SNAKE_CASE__   : Any													={k for k in extra_keys if not k.endswith('''.attn.bias'''							)}
      SCREAMING_SNAKE_CASE__   : Optional[int]													=set(model.state_dict().keys()							) - set(state_dict.keys()							)
      SCREAMING_SNAKE_CASE__   : Any													={k for k in missing_keys if not k.endswith('''.attn.bias'''							)}
      if len(__lowerCamelCase							) != 0:
            raise ValueError(f"extra keys found: {extra_keys}"							)
      if len(__lowerCamelCase							) != 0:
            raise ValueError(f"missing keys: {missing_keys}"							)
      model.load_state_dict(__lowerCamelCase,		strict=__lowerCamelCase							)
      SCREAMING_SNAKE_CASE__   : Tuple													=model.num_parameters(exclude_embeddings=__lowerCamelCase							)
      SCREAMING_SNAKE_CASE__   : Union[str, Any]													=checkpoint["best_val_loss"].item()
      logger.info(f"model loaded: {round(n_params/1e6,		1							)}M params, {round(__lowerCamelCase,		3							)} loss"							)
      model.eval()
      model.to(__lowerCamelCase							)
      del checkpoint, state_dict

      return model





def        _a(     UpperCamelCase__ :  Union[str, Any],		UpperCamelCase__ :  Optional[Any]=False,		UpperCamelCase__ :  Tuple="text"							):




      '''simple docstring'''






      if model_type not in ("text", "coarse", "fine"):
            raise NotImplementedError()

      SCREAMING_SNAKE_CASE__   : Any													="cpu"  # do conversion on cpu

      SCREAMING_SNAKE_CASE__   : Tuple													=_get_ckpt_path(__lowerCamelCase,		use_small=__lowerCamelCase							)
      SCREAMING_SNAKE_CASE__   : Union[str, Any]													=_load_model(__lowerCamelCase,		__lowerCamelCase,		model_type=__lowerCamelCase,		use_small=__lowerCamelCase							)

      # load bark initial model
      SCREAMING_SNAKE_CASE__   : str													=_bark_load_model(__lowerCamelCase,		'''cpu''',		model_type=__lowerCamelCase,		use_small=__lowerCamelCase							)

      if model_type == "text":
            SCREAMING_SNAKE_CASE__   : int													=bark_model["model"]

      if model.num_parameters(exclude_embeddings=__lowerCamelCase							) != bark_model.get_num_params():
            raise ValueError('''initial and new models don\'t have the same number of parameters'''							)

      # check if same output as the bark model
      SCREAMING_SNAKE_CASE__   : Tuple													=5
      SCREAMING_SNAKE_CASE__   : List[Any]													=1_0

      if model_type in ["text", "coarse"]:
            SCREAMING_SNAKE_CASE__   : Union[str, Any]													=torch.randint(2_5_6,		(batch_size, sequence_length),		dtype=torch.int							)
            SCREAMING_SNAKE_CASE__   : Dict													=bark_model(__lowerCamelCase							)[0]

            SCREAMING_SNAKE_CASE__   : List[Any]													=model(__lowerCamelCase							)

            # take last logits
            SCREAMING_SNAKE_CASE__   : List[str]													=output_new_model_total.logits[:, [-1], :]

      else:
            SCREAMING_SNAKE_CASE__   : List[Any]													=3
            SCREAMING_SNAKE_CASE__   : str													=8
            SCREAMING_SNAKE_CASE__   : List[str]													=torch.randint(2_5_6,		(batch_size, sequence_length, n_codes_total),		dtype=torch.int							)

            SCREAMING_SNAKE_CASE__   : Union[str, Any]													=model(__lowerCamelCase,		__lowerCamelCase							)
            SCREAMING_SNAKE_CASE__   : Any													=bark_model(__lowerCamelCase,		__lowerCamelCase							)

            SCREAMING_SNAKE_CASE__   : List[str]													=output_new_model_total.logits

      # output difference should come from the difference of self-attention implementation design
      if output_new_model.shape != output_old_model.shape:
            raise ValueError('''initial and new outputs don\'t have the same shape'''							)
      if (output_new_model - output_old_model).abs().max().item() > 1e-3:
            raise ValueError('''initial and new outputs are not equal'''							)

      Path(__lowerCamelCase							).mkdir(exist_ok=__lowerCamelCase							)
      model.save_pretrained(__lowerCamelCase							)





def        _a(     UpperCamelCase__ :  Dict,		UpperCamelCase__ :  Dict,		UpperCamelCase__ :  Optional[Any],		UpperCamelCase__ :  Any,		UpperCamelCase__ :  Optional[Any],		UpperCamelCase__ :  Dict,		):




      '''simple docstring'''






      SCREAMING_SNAKE_CASE__   : Any													=os.path.join(__lowerCamelCase,		__lowerCamelCase							)

      SCREAMING_SNAKE_CASE__   : Tuple													=BarkSemanticConfig.from_pretrained(os.path.join(__lowerCamelCase,		'''config.json'''							)							)
      SCREAMING_SNAKE_CASE__   : Union[str, Any]													=BarkCoarseConfig.from_pretrained(os.path.join(__lowerCamelCase,		'''config.json'''							)							)
      SCREAMING_SNAKE_CASE__   : Any													=BarkFineConfig.from_pretrained(os.path.join(__lowerCamelCase,		'''config.json'''							)							)
      SCREAMING_SNAKE_CASE__   : Optional[Any]													=EncodecConfig.from_pretrained('''facebook/encodec_24khz'''							)

      SCREAMING_SNAKE_CASE__   : int													=BarkSemanticModel.from_pretrained(__lowerCamelCase							)
      SCREAMING_SNAKE_CASE__   : List[Any]													=BarkCoarseModel.from_pretrained(__lowerCamelCase							)
      SCREAMING_SNAKE_CASE__   : Optional[int]													=BarkFineModel.from_pretrained(__lowerCamelCase							)
      SCREAMING_SNAKE_CASE__   : Tuple													=EncodecModel.from_pretrained('''facebook/encodec_24khz'''							)

      SCREAMING_SNAKE_CASE__   : Optional[int]													=BarkConfig.from_sub_model_configs(
          __lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase							)

      SCREAMING_SNAKE_CASE__   : Optional[Any]													=BarkGenerationConfig.from_sub_model_configs(
          semantic.generation_config,		coarseAcoustic.generation_config,		fineAcoustic.generation_config							)

      SCREAMING_SNAKE_CASE__   : int													=BarkModel(__lowerCamelCase							)

      SCREAMING_SNAKE_CASE__   : Tuple													=semantic
      SCREAMING_SNAKE_CASE__   : int													=coarseAcoustic
      SCREAMING_SNAKE_CASE__   : Any													=fineAcoustic
      SCREAMING_SNAKE_CASE__   : int													=codec

      SCREAMING_SNAKE_CASE__   : Optional[Any]													=bark_generation_config

      Path(__lowerCamelCase							).mkdir(exist_ok=__lowerCamelCase							)
      bark.save_pretrained(__lowerCamelCase,		repo_id=__lowerCamelCase,		push_to_hub=__lowerCamelCase							)


if __name__ == "__main__":
  a_ 		=       argparse.ArgumentParser()
  # Required parameters

  parser.add_argument('model_type', type=str, help='text, coarse or fine.')
  parser.add_argument('pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
  parser.add_argument('--is_small', action='store_true', help='convert the small version instead of the large.')

  a_ 		=       parser.parse_args()

  load_model(args.pytorch_dump_folder_path, model_type=args.model_type, use_small=args.is_small)






'''simple docstring'''


import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VideoMAEConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import (
     MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
     VideoMAEForPreTraining,
     VideoMAEForVideoClassification,
     VideoMAEModel,
 )
 from transformers.models.videomae.modeling_videomae import VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from transformers import VideoMAEImageProcessor
class       a_   :





     def __init__( self		:					Any	,      lowercase		:					Optional[int]	,      lowercase		:					List[Any]=13	,      lowercase		:					int=10	,      lowercase		:					str=3	,      lowercase		:					List[Any]=2	,      lowercase		:					Dict=2	,      lowercase		:					List[str]=2	,      lowercase		:					int=True	,      lowercase		:					List[Any]=True	,      lowercase		:					Union[str, Any]=32	,      lowercase		:					Optional[int]=5	,      lowercase		:					List[Any]=4	,      lowercase		:					List[str]=37	,      lowercase		:					Union[str, Any]="gelu"	,      lowercase		:					List[Any]=0.1	,      lowercase		:					Any=0.1	,      lowercase		:					Optional[Any]=10	,      lowercase		:					Union[str, Any]=0.02	,      lowercase		:					Optional[int]=0.9	,      lowercase		:					List[str]=None	,      ):





        """simple docstring"""




        lowercase_			:Optional[int]    			= parent
        lowercase_			:str    			= batch_size
        lowercase_			:Optional[int]    			= image_size
        lowercase_			:Tuple    			= num_channels
        lowercase_			:Optional[Any]    			= patch_size
        lowercase_			:List[str]    			= tubelet_size
        lowercase_			:List[Any]    			= num_frames
        lowercase_			:Dict    			= is_training
        lowercase_			:Optional[int]    			= use_labels
        lowercase_			:Optional[int]    			= hidden_size
        lowercase_			:List[str]    			= num_hidden_layers
        lowercase_			:List[str]    			= num_attention_heads
        lowercase_			:int    			= intermediate_size
        lowercase_			:Any    			= hidden_act
        lowercase_			:Tuple    			= hidden_dropout_prob
        lowercase_			:str    			= attention_probs_dropout_prob
        lowercase_			:Any    			= type_sequence_label_size
        lowercase_			:int    			= initializer_range
        lowercase_			:Dict    			= mask_ratio
        lowercase_			:Optional[int]    			= scope

        # in VideoMAE, the number of tokens equals num_frames/tubelet_size * num_patches per frame
        lowercase_			:str    			= (image_size // patch_size) ** 2
        lowercase_			:Union[str, Any]    			= (num_frames // tubelet_size) * self.num_patches_per_frame

        # use this variable to define bool_masked_pos
        lowercase_			:Optional[Any]    			= int(mask_ratio * self.seq_length						)





     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= floats_tensor(
            [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]						)

        lowercase_			:Union[str, Any]    			= None
        if self.use_labels:
           lowercase_			:str    			= ids_tensor([self.batch_size]	,      self.type_sequence_label_size						)

        lowercase_			:Optional[int]    			= self.get_config()

        return config, pixel_values, labels





     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        return VideoMAEConfig(
            image_size=self.image_size	,      patch_size=self.patch_size	,      num_channels=self.num_channels	,      num_frames=self.num_frames	,      tubelet_size=self.tubelet_size	,      hidden_size=self.hidden_size	,      num_hidden_layers=self.num_hidden_layers	,      num_attention_heads=self.num_attention_heads	,      intermediate_size=self.intermediate_size	,      hidden_act=self.hidden_act	,      hidden_dropout_prob=self.hidden_dropout_prob	,      attention_probs_dropout_prob=self.attention_probs_dropout_prob	,      is_decoder=lowercase	,      initializer_range=self.initializer_range	,      )





     def 							lowercase__      ( self		:					Optional[int]	,      lowercase		:					Dict	,      lowercase		:					Dict	,      lowercase		:					Optional[int]						):





        """simple docstring"""




        lowercase_			:int    			= VideoMAEModel(config=lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_			:Optional[int]    			= model(lowercase						)
        self.parent.assertEqual(result.last_hidden_state.shape	,      (self.batch_size, self.seq_length, self.hidden_size)						)





     def 							lowercase__      ( self		:					str	,      lowercase		:					str	,      lowercase		:					List[str]	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			:Union[str, Any]    			= VideoMAEForPreTraining(lowercase						)
        model.to(lowercase						)
        model.eval()
        # important: each video needs to have the same number of masked patches
        # hence we define a single mask, which we then repeat for each example in the batch
        lowercase_			:Optional[int]    			= torch.ones((self.num_masks,)						)
        lowercase_			:List[str]    			= torch.cat([mask, torch.zeros(self.seq_length - mask.size(0						)						)]						)
        lowercase_			:Dict    			= mask.expand(self.batch_size	,      -1						).bool()

        lowercase_			:str    			= model(lowercase	,      lowercase						)
        # model only returns predictions for masked patches
        lowercase_			:Any    			= mask.sum().item()
        lowercase_			:Tuple    			= 3 * self.tubelet_size * self.patch_size**2
        self.parent.assertEqual(result.logits.shape	,      (self.batch_size, num_masked_patches, decoder_num_labels)						)







     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        lowercase_			:int    			= self.prepare_config_and_inputs()
        lowercase_			,							lowercase_			,							lowercase_			:Dict    			= config_and_inputs
        lowercase_			:Any    			= {"pixel_values": pixel_values}
        return config, inputs_dict







@require_torch
class       a_   (				_lowerCAmelCase	,						_lowerCAmelCase	,						unittest.TestCase							):
     __A      =			(
         (VideoMAEModel, VideoMAEForPreTraining, VideoMAEForVideoClassification) if is_torch_available() else ()
     )
     __A      =			(
         {"feature-extraction": VideoMAEModel, "video-classification": VideoMAEForVideoClassification}
         if is_torch_available()
         else {}
     )

     __A      =			False
     __A      =			False
     __A      =			False
     __A      =			False





     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        lowercase_			:List[Any]    			= VideoMAEModelTester(self						)
        lowercase_			:Dict    			= ConfigTester(self	,      config_class=lowercase	,      has_text_modality=lowercase	,      hidden_size=37						)





     def 							lowercase__      ( self		:					List[Any]	,      lowercase		:					List[str]	,      lowercase		:					List[str]	,      lowercase		:					List[str]=False						):





        """simple docstring"""




        lowercase_			:Tuple    			= copy.deepcopy(lowercase						)

        if model_class == VideoMAEForPreTraining:
           # important: each video needs to have the same number of masked patches
           # hence we define a single mask, which we then repeat for each example in the batch
           lowercase_			:Tuple    			= torch.ones((self.model_tester.num_masks,)						)
           lowercase_			:Tuple    			= torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0						)						)]						)
           lowercase_			:Optional[int]    			= mask.expand(self.model_tester.batch_size	,      -1						).bool()
           lowercase_			:Dict    			= bool_masked_pos.to(lowercase						)

        if return_labels:
           if model_class in [
               *get_values(lowercase						),
           ]:
              lowercase_			:Union[str, Any]    			= torch.zeros(
                  self.model_tester.batch_size	,      dtype=torch.long	,      device=lowercase						)

        return inputs_dict





     def 							lowercase__      ( self		:					Tuple						):





        """simple docstring"""




        self.config_tester.run_common_tests()





     @unittest.skip(reason="VideoMAE does not use inputs_embeds"						)
     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        pass





     def 							lowercase__      ( self		:					Tuple						):





        """simple docstring"""




        lowercase_			,							lowercase_			:List[Any]    			= self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           lowercase_			:Dict    			= model_class(lowercase						)
           self.assertIsInstance(model.get_input_embeddings()	,      (nn.Module)						)
           lowercase_			:List[Any]    			= model.get_output_embeddings()
           self.assertTrue(x is None or isinstance(lowercase	,      nn.Linear						)						)





     def 							lowercase__      ( self		:					Optional[int]						):





        """simple docstring"""




        lowercase_			,							lowercase_			:Optional[int]    			= self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           lowercase_			:Any    			= model_class(lowercase						)
           lowercase_			:Any    			= inspect.signature(model.forward						)
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           lowercase_			:List[str]    			= [*signature.parameters.keys()]

           lowercase_			:str    			= ["pixel_values"]
           self.assertListEqual(arg_names[:1]	,      lowercase						)





     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        lowercase_			:Tuple    			= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase						)





     def 							lowercase__      ( self		:					Optional[int]						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*lowercase						)





     @slow
     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowercase_			:List[Any]    			= VideoMAEModel.from_pretrained(lowercase						)
           self.assertIsNotNone(lowercase						)





     def 							lowercase__      ( self		:					Union[str, Any]						):





        """simple docstring"""




        if not self.has_attentions:
           pass

        else:
           lowercase_			,							lowercase_			:Any    			= self.model_tester.prepare_config_and_inputs_for_common()
           lowercase_			:Union[str, Any]    			= True

           for model_class in self.all_model_classes:
              lowercase_			:Dict    			= self.model_tester.seq_length - self.model_tester.num_masks
              lowercase_			:Optional[Any]    			= (
                  num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
              )

              lowercase_			:Union[str, Any]    			= True
              lowercase_			:List[Any]    			= False
              lowercase_			:Optional[int]    			= True
              lowercase_			:Union[str, Any]    			= model_class(lowercase						)
              model.to(lowercase						)
              model.eval()
              with torch.no_grad():
                 lowercase_			:List[Any]    			= model(**self._prepare_for_class(lowercase	,      lowercase						)						)
              lowercase_			:str    			= outputs.attentions
              self.assertEqual(len(lowercase						)	,      self.model_tester.num_hidden_layers						)

              # check that output_attentions also work using config
              del inputs_dict["output_attentions"]
              lowercase_			:Union[str, Any]    			= True
              lowercase_			:Optional[Any]    			= model_class(lowercase						)
              model.to(lowercase						)
              model.eval()
              with torch.no_grad():
                 lowercase_			:Optional[Any]    			= model(**self._prepare_for_class(lowercase	,      lowercase						)						)
              lowercase_			:Union[str, Any]    			= outputs.attentions
              self.assertEqual(len(lowercase						)	,      self.model_tester.num_hidden_layers						)

              self.assertListEqual(
                  list(attentions[0].shape[-3:]						)	,      [self.model_tester.num_attention_heads, seq_len, seq_len]	,      )
              lowercase_			:List[str]    			= len(lowercase						)

              # Check attention is always last and order is fine
              lowercase_			:Optional[Any]    			= True
              lowercase_			:Dict    			= True
              lowercase_			:Dict    			= model_class(lowercase						)
              model.to(lowercase						)
              model.eval()
              with torch.no_grad():
                 lowercase_			:List[Any]    			= model(**self._prepare_for_class(lowercase	,      lowercase						)						)

              self.assertEqual(out_len + 1	,      len(lowercase						)						)

              lowercase_			:int    			= outputs.attentions

              self.assertEqual(len(lowercase						)	,      self.model_tester.num_hidden_layers						)
              self.assertListEqual(
                  list(self_attentions[0].shape[-3:]						)	,      [self.model_tester.num_attention_heads, seq_len, seq_len]	,      )





     def 							lowercase__      ( self		:					int						):





        """simple docstring"""




        def check_hidden_states_output(lowercase		:					Union[str, Any]	,      lowercase		:					Dict	,      lowercase		:					Any						):
           lowercase_			:Any    			= model_class(lowercase						)
           model.to(lowercase						)
           model.eval()

           with torch.no_grad():
              lowercase_			:Optional[int]    			= model(**self._prepare_for_class(lowercase	,      lowercase						)						)

           lowercase_			:Optional[int]    			= outputs.hidden_states
           lowercase_			:Any    			= self.model_tester.num_hidden_layers + 1
           self.assertEqual(len(lowercase						)	,      lowercase						)

           lowercase_			:List[str]    			= self.model_tester.seq_length - self.model_tester.num_masks
           lowercase_			:List[Any]    			= num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length

           self.assertListEqual(
               list(hidden_states[0].shape[-2:]						)	,      [seq_length, self.model_tester.hidden_size]	,      )

        lowercase_			,							lowercase_			:List[Any]    			= self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           lowercase_			:Optional[int]    			= True
           check_hidden_states_output(lowercase	,      lowercase	,      lowercase						)

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           lowercase_			:List[Any]    			= True

           check_hidden_states_output(lowercase	,      lowercase	,      lowercase						)







     @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."						)
     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        pass







def 							UpperCAmelCase_  (						):
   lowercase_			:Dict    			= hf_hub_download(
       repo_id="hf-internal-testing/spaghetti-video"	,filename="eating_spaghetti.npy"	,repo_type="dataset" )
   lowercase_			:Optional[Any]    			= np.load(__lowerCamelCase )
   return list(__lowerCamelCase )


@require_torch
@require_vision
class       a_   (				unittest.TestCase							):





     @cached_property
     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        return (
            VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]	,      image_std=[0.5, 0.5, 0.5]						)
            if is_vision_available()
            else None
        )





     @slow
     def 							lowercase__      ( self		:					Optional[int]						):





        """simple docstring"""




        lowercase_			:Union[str, Any]    			= VideoMAEForVideoClassification.from_pretrained("MCG-NJU/videomae-base-finetuned-kinetics"						).to(
            lowercase						)

        lowercase_			:List[str]    			= self.default_image_processor
        lowercase_			:List[str]    			= prepare_video()
        lowercase_			:int    			= image_processor(lowercase	,      return_tensors="pt"						).to(lowercase						)

        # forward pass
        with torch.no_grad():
           lowercase_			:Dict    			= model(**lowercase						)

        # verify the logits
        lowercase_			:Dict    			= torch.Size((1, 400)						)
        self.assertEqual(outputs.logits.shape	,      lowercase						)

        lowercase_			:int    			= torch.tensor([0.36_69, -0.06_88, -0.24_21]						).to(lowercase						)

        self.assertTrue(torch.allclose(outputs.logits[0, :3]	,      lowercase	,      atol=1e-4						)						)







     @slow
     def 							lowercase__      ( self		:					str						):





        """simple docstring"""




        lowercase_			:List[Any]    			= VideoMAEForPreTraining.from_pretrained("MCG-NJU/videomae-base-short"						).to(lowercase						)

        lowercase_			:Dict    			= self.default_image_processor
        lowercase_			:Union[str, Any]    			= prepare_video()
        lowercase_			:List[str]    			= image_processor(lowercase	,      return_tensors="pt"						).to(lowercase						)

        # add boolean mask, indicating which patches to mask
        lowercase_			:int    			= hf_hub_download(repo_id="hf-internal-testing/bool-masked-pos"	,      filename="bool_masked_pos.pt"						)
        lowercase_			:List[str]    			= torch.load(lowercase						)

        # forward pass
        with torch.no_grad():
           lowercase_			:List[Any]    			= model(**lowercase						)

        # verify the logits
        lowercase_			:Union[str, Any]    			= torch.Size([1, 1_408, 1_536]						)
        lowercase_			:List[Any]    			= 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]]	,      device=lowercase						)
        self.assertEqual(outputs.logits.shape	,      lowercase						)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]	,      lowercase	,      atol=1e-4						)						)

        # verify the loss (`config.norm_pix_loss` = `True`)
        lowercase_			:Any    			= torch.tensor([0.51_42]	,      device=lowercase						)
        self.assertTrue(torch.allclose(outputs.loss	,      lowercase	,      atol=1e-4						)						)

        # verify the loss (`config.norm_pix_loss` = `False`)
        lowercase_			:Union[str, Any]    			= VideoMAEForPreTraining.from_pretrained("MCG-NJU/videomae-base-short"	,      norm_pix_loss=lowercase						).to(
            lowercase						)

        with torch.no_grad():
           lowercase_			:Tuple    			= model(**lowercase						)

        lowercase_			:Optional[Any]    			= torch.tensor(torch.tensor([0.64_69]						)	,      device=lowercase						)
        self.assertTrue(torch.allclose(outputs.loss	,      lowercase	,      atol=1e-4						)						)


