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 		_A		(		_lowercase   )				->   int:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = 'https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'
  
       with open(_lowercase				,     'wb'   ) as f:
            with requests.get(_lowercase				,     stream=_lowercase   ) as req:
                 for chunk in req.iter_content(chunk_size=20_48   ):
                      f.write(_lowercase   )
  
  
  
  
  
  
  def 		_A		(		)				->   VivitConfig:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = VivitConfig()
  
       __UpperCamelCase   = 4_00
       __UpperCamelCase   = 'huggingface/label-files'
       __UpperCamelCase   = 'kinetics400-id2label.json'
  
       __UpperCamelCase   = json.load(open(hf_hub_download(_lowercase				,     _lowercase				,     repo_type='dataset'   )				,     'r'   )   )
       __UpperCamelCase   = {int(_lowercase   ): v for k, v in idalabel.items()}
       __UpperCamelCase   = idalabel
       __UpperCamelCase   = {v: k for k, v in idalabel.items()}
       return config
  
  
  
  
  
  
  def 		_A		(		)				->   Optional[Any]:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = hf_hub_download(
           repo_id='hf-internal-testing/spaghetti-video'				,     filename='eating_spaghetti_32_frames.npy'				,     repo_type='dataset'   )
       __UpperCamelCase   = np.load(_lowercase   )
       return list(_lowercase   )
  
  
  
  
  
  
  def 		_A		(		_lowercase   )				->   List[str]:
  
  
       """simple docstring"""
  
  
       if np.ndim(_lowercase   ) == 2:
            return transform_attention_bias(_lowercase   )
  
       elif np.ndim(_lowercase   ) == 3:
            return transform_attention_kernel(_lowercase   )
  
       else:
            raise Exception(f'''Invalid number of dimesions: {np.ndim(_lowercase   )}'''   )
  
  
  
  
  
  
  def 		_A		(		_lowercase   )				->   str:
  
  
       """simple docstring"""
  
  
       return current.flatten()
  
  
  
  
  
  
  def 		_A		(		_lowercase   )				->   Dict:
  
  
       """simple docstring"""
  
  
       return np.reshape(_lowercase				,     (current.shape[0], current.shape[1] * current.shape[2])   ).T
  
  
  
  
  
  
  def 		_A		(		_lowercase   )				->   Optional[Any]:
  
  
       """simple docstring"""
  
  
       return np.reshape(_lowercase				,     (current.shape[0] * current.shape[1], current.shape[2])   ).T
  
  
  
  
  
  
  def 		_A		(		_lowercase				,     _lowercase   )				->   Any:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = state_dict['optimizer']['target']['Transformer'][f'''encoderblock_{i}''']
  
       __UpperCamelCase   = f'''encoder.layer.{i}.'''
       __UpperCamelCase   = {
           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 		_A		(		_lowercase   )				->   Optional[int]:
  
  
       """simple docstring"""
  
  
       return sum([1 if 'encoderblock_' in k else 0 for k in state_dict['optimizer']['target']['Transformer'].keys()]   )
  
  
  
  
  
  
  def 		_A		(		_lowercase				,     _lowercase=False   )				->   List[str]:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = get_n_layers(_lowercase   )
  
       __UpperCamelCase   = OrderedDict()
  
       __UpperCamelCase   = state_dict['optimizer']['target']['Transformer']['encoder_norm']['bias']
       __UpperCamelCase   = state_dict['optimizer']['target']['Transformer']['encoder_norm']['scale']
  
       __UpperCamelCase   = np.transpose(
           state_dict['optimizer']['target']['embedding']['kernel']				,     (4, 3, 0, 1, 2)   )
       __UpperCamelCase   = state_dict['optimizer']['target']['embedding']['bias']
  
       __UpperCamelCase   = state_dict['optimizer']['target']['cls']
       __UpperCamelCase   = state_dict['optimizer']['target']['Transformer']['posembed_input'][
           'pos_embedding'
       ]
  
       for i in range(_lowercase   ):
            new_state.update(transform_state_encoder_block(_lowercase				,     _lowercase   )   )
  
       if classification_head:
            __UpperCamelCase   = {'vivit.' + k: v for k, v in new_state.items()}
            __UpperCamelCase   = np.transpose(state_dict['optimizer']['target']['output_projection']['kernel']   )
            __UpperCamelCase   = np.transpose(state_dict['optimizer']['target']['output_projection']['bias']   )
  
       return {k: torch.tensor(_lowercase   ) for k, v in new_state.items()}
  
  
  
  
  
  
  def 		_A		(		)				->   VivitImageProcessor:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = 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 		_A		(		_lowercase   )				->   Union[str, Any]:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = 'checkpoint'
  
       if not os.path.exists(_lowercase   ):
            download_checkpoint(_lowercase   )
  
       __UpperCamelCase   = restore_checkpoint(_lowercase				,     _lowercase   )
       __UpperCamelCase   = transform_state(_lowercase				,     classification_head=_lowercase   )
  
       __UpperCamelCase   = get_vivit_config()
  
       assert config.image_size == 2_24
       assert config.num_frames == 32
  
       __UpperCamelCase   = VivitForVideoClassification(_lowercase   )
       model.load_state_dict(_lowercase   )
       model.eval()
  
       __UpperCamelCase   = get_processor()
  
       __UpperCamelCase   = prepare_video()
       __UpperCamelCase   = extractor(_lowercase				,     return_tensors='pt'   )
  
       __UpperCamelCase   = model(**_lowercase   )
  
       __UpperCamelCase   = torch.Size([1, 4_00]   )
       __UpperCamelCase   = torch.tensor([-1.05_43, 2.07_64, -0.21_04, 0.44_39, -0.96_58]   )
  
       assert outputs.logits.shape == expected_shape
       assert torch.allclose(outputs.logits[0, :5]				,     _lowercase				,     atol=1e-4   ), outputs.logits[0, :5]
  
       model.save_pretrained(_lowercase   )
       extractor.save_pretrained(_lowercase   )
  
  
  if __name__ == "__main__":
      __snake_case             =	argparse.ArgumentParser()
  
      parser.add_argument('''--output_model_name''', '''-o''', type=str, help='''Output path for the converted HuggingFace model''')
  
      __snake_case             =	parser.parse_args()
      convert(args.output_model_name)
  
  
  
style_context: 
  
  
  from __future__ import annotations
  
  
  
  
  
  
  def 		_A		(		_lowercase   )				->   int | float:
  
  
       """simple docstring"""
  
  
       if len(_lowercase   ) == 0:
            raise ValueError('find_max() arg is an empty sequence'   )
       __UpperCamelCase   = nums[0]
       for x in nums:
            if x > max_num:
                 __UpperCamelCase   = x
       return max_num
  
  
  if __name__ == "__main__":
      import doctest
  
      doctest.testmod(verbose=True)
  
  
  