code: 
  
  
  import argparse
  import json
  import os.path
  from collections import OrderedDict
  
  import numpy as np
  import requests
  import torch
  from flax.training.checkpoints import restore_checkpoint
  from huggingface_hub import hf_hub_download
  
  from transformers import VivitConfig, VivitForVideoClassification, VivitImageProcessor
  from transformers.image_utils import PILImageResampling
  def     UpperCAmelCase__		(UpperCamelCase_       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       snake_case					=       '''https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'''
  
       with open(UpperCamelCase_							,'''wb'''       ) as f:
            with requests.get(UpperCamelCase_							,stream=UpperCamelCase_       ) as req:
                 for chunk in req.iter_content(chunk_size=20_48       ):
                      f.write(UpperCamelCase_       )
  def     UpperCAmelCase__		():
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       snake_case					=       VivitConfig()
  
       snake_case					=       4_00
       snake_case					=       '''huggingface/label-files'''
       snake_case					=       '''kinetics400-id2label.json'''
  
       snake_case					=       json.load(open(hf_hub_download(UpperCamelCase_							,UpperCamelCase_							,repo_type='''dataset'''       )							,'''r'''       )       )
       snake_case					=       {int(UpperCamelCase_       ): v for k, v in idalabel.items()}
       snake_case					=       idalabel
       snake_case					=       {v: k for k, v in idalabel.items()}
       return config
  def     UpperCAmelCase__		():
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       snake_case					=       hf_hub_download(
           repo_id='''hf-internal-testing/spaghetti-video'''							,filename='''eating_spaghetti_32_frames.npy'''							,repo_type='''dataset'''       )
       snake_case					=       np.load(UpperCamelCase_       )
       return list(UpperCamelCase_       )
  def     UpperCAmelCase__		(UpperCamelCase_       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       if np.ndim(UpperCamelCase_       ) == 2:
            return transform_attention_bias(UpperCamelCase_       )
  
       elif np.ndim(UpperCamelCase_       ) == 3:
            return transform_attention_kernel(UpperCamelCase_       )
  
       else:
            raise Exception(F'''Invalid number of dimesions: {np.ndim(UpperCamelCase_       )}'''       )
  def     UpperCAmelCase__		(UpperCamelCase_       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       return current.flatten()
  def     UpperCAmelCase__		(UpperCamelCase_       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       return np.reshape(UpperCamelCase_							,(current.shape[0], current.shape[1] * current.shape[2])       ).T
  def     UpperCAmelCase__		(UpperCamelCase_       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       return np.reshape(UpperCamelCase_							,(current.shape[0] * current.shape[1], current.shape[2])       ).T
  def     UpperCAmelCase__		(UpperCamelCase_							,UpperCamelCase_       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       snake_case					=       state_dict['''optimizer''']['''target''']['''Transformer'''][F'''encoderblock_{i}''']
  
       snake_case					=       F'''encoder.layer.{i}.'''
       snake_case					=       {
           prefix + '''intermediate.dense.bias''': state['''MlpBlock_0''']['''Dense_0''']['''bias'''],
           prefix + '''intermediate.dense.weight''': np.transpose(state['''MlpBlock_0''']['''Dense_0''']['''kernel''']       ),
           prefix + '''output.dense.bias''': state['''MlpBlock_0''']['''Dense_1''']['''bias'''],
           prefix + '''output.dense.weight''': np.transpose(state['''MlpBlock_0''']['''Dense_1''']['''kernel''']       ),
           prefix + '''layernorm_before.bias''': state['''LayerNorm_0''']['''bias'''],
           prefix + '''layernorm_before.weight''': state['''LayerNorm_0''']['''scale'''],
           prefix + '''layernorm_after.bias''': state['''LayerNorm_1''']['''bias'''],
           prefix + '''layernorm_after.weight''': state['''LayerNorm_1''']['''scale'''],
           prefix
           + '''attention.attention.query.bias''': transform_attention(
               state['''MultiHeadDotProductAttention_0''']['''query''']['''bias''']       ),
           prefix
           + '''attention.attention.query.weight''': transform_attention(
               state['''MultiHeadDotProductAttention_0''']['''query''']['''kernel''']       ),
           prefix
           + '''attention.attention.key.bias''': transform_attention(state['''MultiHeadDotProductAttention_0''']['''key''']['''bias''']       ),
           prefix
           + '''attention.attention.key.weight''': transform_attention(
               state['''MultiHeadDotProductAttention_0''']['''key''']['''kernel''']       ),
           prefix
           + '''attention.attention.value.bias''': transform_attention(
               state['''MultiHeadDotProductAttention_0''']['''value''']['''bias''']       ),
           prefix
           + '''attention.attention.value.weight''': transform_attention(
               state['''MultiHeadDotProductAttention_0''']['''value''']['''kernel''']       ),
           prefix + '''attention.output.dense.bias''': state['''MultiHeadDotProductAttention_0''']['''out''']['''bias'''],
           prefix
           + '''attention.output.dense.weight''': transform_attention_output_weight(
               state['''MultiHeadDotProductAttention_0''']['''out''']['''kernel''']       ),
       }
  
       return new_state
  def     UpperCAmelCase__		(UpperCamelCase_       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       return sum([1 if '''encoderblock_''' in k else 0 for k in state_dict['''optimizer''']['''target''']['''Transformer'''].keys()]       )
  def     UpperCAmelCase__		(UpperCamelCase_							,UpperCamelCase_=False       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       snake_case					=       get_n_layers(UpperCamelCase_       )
  
       snake_case					=       OrderedDict()
  
       snake_case					=       state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''bias''']
       snake_case					=       state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''scale''']
  
       snake_case					=       np.transpose(
           state_dict['''optimizer''']['''target''']['''embedding''']['''kernel''']							,(4, 3, 0, 1, 2)       )
       snake_case					=       state_dict['''optimizer''']['''target''']['''embedding''']['''bias''']
  
       snake_case					=       state_dict['''optimizer''']['''target''']['''cls''']
       snake_case					=       state_dict['''optimizer''']['''target''']['''Transformer''']['''posembed_input'''][
           '''pos_embedding'''
       ]
  
       for i in range(UpperCamelCase_       ):
            new_state.update(transform_state_encoder_block(UpperCamelCase_							,UpperCamelCase_       )       )
  
       if classification_head:
            snake_case					=       {'''vivit.''' + k: v for k, v in new_state.items()}
            snake_case					=       np.transpose(state_dict['''optimizer''']['''target''']['''output_projection''']['''kernel''']       )
            snake_case					=       np.transpose(state_dict['''optimizer''']['''target''']['''output_projection''']['''bias''']       )
  
       return {k: torch.tensor(UpperCamelCase_       ) for k, v in new_state.items()}
  def     UpperCAmelCase__		():
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       snake_case					=       VivitImageProcessor()
  
       assert extractor.do_resize is True
       assert extractor.size == {"shortest_edge": 2_56}
       assert extractor.do_center_crop is True
       assert extractor.crop_size == {"width": 2_24, "height": 2_24}
       assert extractor.resample == PILImageResampling.BILINEAR
  
       # here: https://github.com/deepmind/dmvr/blob/master/dmvr/modalities.py
       # one can seen that add_image has default values for normalization_mean and normalization_std set to 0 and 1
       # which effectively means no normalization (and ViViT does not overwrite those when calling this func)
       assert extractor.do_normalize is False
       assert extractor.do_rescale is True
       assert extractor.rescale_factor == 1 / 2_55
  
       # zero-centering = True in original implementation
       assert extractor.do_zero_centering is True
  
       return extractor
  def     UpperCAmelCase__		(UpperCamelCase_       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       snake_case					=       '''checkpoint'''
  
       if not os.path.exists(UpperCamelCase_       ):
            download_checkpoint(UpperCamelCase_       )
  
       snake_case					=       restore_checkpoint(UpperCamelCase_							,UpperCamelCase_       )
       snake_case					=       transform_state(UpperCamelCase_							,classification_head=UpperCamelCase_       )
  
       snake_case					=       get_vivit_config()
  
       assert config.image_size == 2_24
       assert config.num_frames == 32
  
       snake_case					=       VivitForVideoClassification(UpperCamelCase_       )
       model.load_state_dict(UpperCamelCase_       )
       model.eval()
  
       snake_case					=       get_processor()
  
       snake_case					=       prepare_video()
       snake_case					=       extractor(UpperCamelCase_							,return_tensors='''pt'''       )
  
       snake_case					=       model(**UpperCamelCase_       )
  
       snake_case					=       torch.Size([1, 4_00]       )
       snake_case					=       torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658]       )
  
       assert outputs.logits.shape == expected_shape
       assert torch.allclose(outputs.logits[0, :5]							,UpperCamelCase_							,atol=1e-4       ), outputs.logits[0, :5]
  
       model.save_pretrained(UpperCamelCase_       )
       extractor.save_pretrained(UpperCamelCase_       )
  
  
  if __name__ == "__main__":
   _SCREAMING_SNAKE_CASE    :				List[Any]       =					argparse.ArgumentParser()
  
   parser.add_argument("--output_model_name", "-o", type=str, help="Output path for the converted HuggingFace model")
  
   _SCREAMING_SNAKE_CASE    :				Optional[Any]       =					parser.parse_args()
   convert(args.output_model_name)
  
  
  
  
  
style_context: 
  
  
  import torch
  
  from accelerate import Accelerator, DistributedType
  
  class 						A__      :
  
  
  
  
  
  
  
      """simple docstring"""
  
      def __enter__( self					):
           if self.enabled:
                snake_case					=       self.model.no_sync()
                self.model_sync_obj.__enter__()
  
           return self
      def __exit__( self    ,      __snake_case    ,      __snake_case    ,      __snake_case					):
           if self.enabled:
                # Average all models on exit
                self._sync_and_avg_model_params()
                self.model_sync_obj.__exit__(__snake_case    ,      __snake_case    ,      __snake_case					)
      def __init__( self    ,      __snake_case    ,      __snake_case    ,      __snake_case    ,      __snake_case = True					):
           if accelerator.distributed_type not in [
               DistributedType.NO,
               DistributedType.MULTI_CPU,
               DistributedType.MULTI_GPU,
           ]:
                raise NotImplementedError('''LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'''					)
           snake_case					=       enabled and accelerator.distributed_type != DistributedType.NO
           snake_case					=       0
           if self.enabled:
                snake_case					=       accelerator
                snake_case					=       model
                snake_case					=       local_sgd_steps
      def 	a_    ( self					):
           self.num_steps += 1
           if not self.enabled:
                return
  
           if self.num_steps % self.local_sgd_steps == 0:
                self._sync_and_avg_model_params()
  
  
  
  
      def 	a_    ( self					):
           self.accelerator.wait_for_everyone()
           with self.accelerator.autocast():
                for param in self.model.parameters():
                     snake_case					=       self.accelerator.reduce(param.data    ,      reduction='''mean'''					)
  
  
  
  
  