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      )->		Optional[Any]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_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=2048      ):
  																				f.write(UpperCAmelCase      )
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     )->		VivitConfig:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       VivitConfig()
  
  					SCREAMING_SNAKE_CASE_	      =       400
  					SCREAMING_SNAKE_CASE_	      =       '''huggingface/label-files'''
  					SCREAMING_SNAKE_CASE_	      =       '''kinetics400-id2label.json'''
  
  					SCREAMING_SNAKE_CASE_	      =       json.load(open(hf_hub_download(UpperCAmelCase				,UpperCAmelCase				,repo_type='''dataset'''      )				,'''r'''      )      )
  					SCREAMING_SNAKE_CASE_	      =       {int(UpperCAmelCase      ): v for k, v in idalabel.items()}
  					SCREAMING_SNAKE_CASE_	      =       idalabel
  					SCREAMING_SNAKE_CASE_	      =       {v: k for k, v in idalabel.items()}
  					return config
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     )->		Any:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       hf_hub_download(
  					    repo_id='''hf-internal-testing/spaghetti-video'''				,filename='''eating_spaghetti_32_frames.npy'''				,repo_type='''dataset'''      )
  					SCREAMING_SNAKE_CASE_	      =       np.load(UpperCAmelCase      )
  					return list(UpperCAmelCase      )
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase      )->		Any:
  
  
  
  
  
  
  
  					'''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      )->		Any:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					return current.flatten()
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase      )->		Union[str, Any]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					return np.reshape(UpperCAmelCase				,(current.shape[0], current.shape[1] * current.shape[2])      ).T
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase      )->		Optional[int]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					return np.reshape(UpperCAmelCase				,(current.shape[0] * current.shape[1], current.shape[2])      ).T
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		int:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       state_dict['''optimizer''']['''target''']['''Transformer'''][f'''encoderblock_{i}''']
  
  					SCREAMING_SNAKE_CASE_	      =       f'''encoder.layer.{i}.'''
  					SCREAMING_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      )->		List[Any]:
  
  
  
  
  
  
  
  					'''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      )->		Any:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       get_n_layers(UpperCAmelCase      )
  
  					SCREAMING_SNAKE_CASE_	      =       OrderedDict()
  
  					SCREAMING_SNAKE_CASE_	      =       state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''bias''']
  					SCREAMING_SNAKE_CASE_	      =       state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''scale''']
  
  					SCREAMING_SNAKE_CASE_	      =       np.transpose(
  					    state_dict['''optimizer''']['''target''']['''embedding''']['''kernel''']				,(4, 3, 0, 1, 2)      )
  					SCREAMING_SNAKE_CASE_	      =       state_dict['''optimizer''']['''target''']['''embedding''']['''bias''']
  
  					SCREAMING_SNAKE_CASE_	      =       state_dict['''optimizer''']['''target''']['''cls''']
  					SCREAMING_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:
  										SCREAMING_SNAKE_CASE_	      =       {'''vivit.''' + k: v for k, v in new_state.items()}
  										SCREAMING_SNAKE_CASE_	      =       np.transpose(state_dict['''optimizer''']['''target''']['''output_projection''']['''kernel''']      )
  										SCREAMING_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				(     )->		VivitImageProcessor:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       VivitImageProcessor()
  
  					assert extractor.do_resize is True
  					assert extractor.size == {"shortest_edge": 256}
  					assert extractor.do_center_crop is True
  					assert extractor.crop_size == {"width": 224, "height": 224}
  					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 / 255
  
  					# zero-centering = True in original implementation
  					assert extractor.do_zero_centering is True
  
  					return extractor
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase      )->		List[Any]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       '''checkpoint'''
  
  					if not os.path.exists(UpperCAmelCase      ):
  										download_checkpoint(UpperCAmelCase      )
  
  					SCREAMING_SNAKE_CASE_	      =       restore_checkpoint(UpperCAmelCase				,UpperCAmelCase      )
  					SCREAMING_SNAKE_CASE_	      =       transform_state(UpperCAmelCase				,classification_head=UpperCAmelCase      )
  
  					SCREAMING_SNAKE_CASE_	      =       get_vivit_config()
  
  					assert config.image_size == 224
  					assert config.num_frames == 32
  
  					SCREAMING_SNAKE_CASE_	      =       VivitForVideoClassification(UpperCAmelCase      )
  					model.load_state_dict(UpperCAmelCase      )
  					model.eval()
  
  					SCREAMING_SNAKE_CASE_	      =       get_processor()
  
  					SCREAMING_SNAKE_CASE_	      =       prepare_video()
  					SCREAMING_SNAKE_CASE_	      =       extractor(UpperCAmelCase				,return_tensors='''pt'''      )
  
  					SCREAMING_SNAKE_CASE_	      =       model(**UpperCAmelCase      )
  
  					SCREAMING_SNAKE_CASE_	      =       torch.Size([1, 400]      )
  					SCREAMING_SNAKE_CASE_	      =       torch.tensor([-1.0_5_4_3, 2.0_7_6_4, -0.2_1_0_4, 0.4_4_3_9, -0.9_6_5_8]      )
  
  					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__":
  		A_														=      argparse.ArgumentParser()
  
  		parser.add_argument("--output_model_name", "-o", type=str, help="Output path for the converted HuggingFace model")
  
  		A_														=      parser.parse_args()
  		convert(args.output_model_name)
  
  
  
style_context: 
  
  
  import json
  import os
  from collections.abc import Mapping
  from typing import Dict, List, Optional, Union
  
  import numpy as np
  import torch
  
  from .imports import is_safetensors_available
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase				,UpperCAmelCase				,UpperCAmelCase=None      )->		Dict:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       None
  					# Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
  					if str(weight.dtype      ) == "torch.bfloat16":
  										# Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
  										SCREAMING_SNAKE_CASE_	      =       weight.view(torch.intaa      )
  										SCREAMING_SNAKE_CASE_	      =       '''bfloat16'''
  					SCREAMING_SNAKE_CASE_	      =       weight.cpu().numpy()
  					SCREAMING_SNAKE_CASE_	      =       os.path.join(UpperCAmelCase				,f'''{weight_name}.dat'''      )
  					if index is not None:
  										if dtype is None:
  															SCREAMING_SNAKE_CASE_	      =       str(array.dtype      )
  										SCREAMING_SNAKE_CASE_	      =       {'''dtype''': dtype, '''shape''': list(array.shape      )}
  					if array.ndim == 0:
  										SCREAMING_SNAKE_CASE_	      =       array[None]
  					SCREAMING_SNAKE_CASE_	      =       np.memmap(UpperCAmelCase				,dtype=array.dtype				,mode='''w+'''				,shape=array.shape      )
  					SCREAMING_SNAKE_CASE_	      =       array[:]
  					file_array.flush()
  					return index
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		List[str]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       tuple(weight_info['''shape''']      )
  					if shape == ():
  										# NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
  										SCREAMING_SNAKE_CASE_	      =       (1,)
  
  					SCREAMING_SNAKE_CASE_	      =       weight_info['''dtype''']
  					if dtype == "bfloat16":
  										# NumPy does not support bfloat16 so this was saved as a int16
  										SCREAMING_SNAKE_CASE_	      =       '''int16'''
  
  					SCREAMING_SNAKE_CASE_	      =       np.memmap(UpperCAmelCase				,dtype=UpperCAmelCase				,shape=UpperCAmelCase				,mode='''r'''      )
  
  					if len(weight_info['''shape''']      ) == 0:
  										SCREAMING_SNAKE_CASE_	      =       weight[0]
  					SCREAMING_SNAKE_CASE_	      =       torch.tensor(UpperCAmelCase      )
  					if weight_info["dtype"] == "bfloat16":
  										SCREAMING_SNAKE_CASE_	      =       weight.view(torch.bfloataa      )
  
  					return weight
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		str:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					if index is None or len(UpperCAmelCase      ) == 0:
  										# Nothing to save
  										return
  
  					SCREAMING_SNAKE_CASE_	      =       os.path.join(UpperCAmelCase				,'''index.json'''      )
  					if os.path.isfile(UpperCAmelCase      ):
  										with open(UpperCAmelCase				,'''r'''				,encoding='''utf-8'''      ) as f:
  															SCREAMING_SNAKE_CASE_	      =       json.load(UpperCAmelCase      )
  					else:
  										SCREAMING_SNAKE_CASE_	      =       {}
  					current_index.update(UpperCAmelCase      )
  
  					with open(UpperCAmelCase				,'''w'''				,encoding='''utf-8'''      ) as f:
  										json.dump(UpperCAmelCase				,UpperCAmelCase				,indent=2      )
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		Optional[Any]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					os.makedirs(UpperCAmelCase				,exist_ok=UpperCAmelCase      )
  					SCREAMING_SNAKE_CASE_	      =       {}
  					for name, parameter in state_dict.items():
  										SCREAMING_SNAKE_CASE_	      =       offload_weight(UpperCAmelCase				,UpperCAmelCase				,UpperCAmelCase				,index=UpperCAmelCase      )
  
  					# Update index
  					save_offload_index(UpperCAmelCase				,UpperCAmelCase      )
  
  class 						snake_case			(       lowerCAmelCase__	):
  			'''simple docstring'''
  
  			def __init__(	self     :     int							,							lowerCAmelCase_     :     Mapping							,							lowerCAmelCase_     :     str					)					->					Dict:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       dataset
  								SCREAMING_SNAKE_CASE_	      =       prefix
  
  			def __getitem__(	self     :     Optional[int]							,							lowerCAmelCase_     :     Optional[int]					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return self.dataset[F'''{self.prefix}{key}''']
  
  			def __iter__(	self     :     Tuple					)					->					str:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return iter([key for key in self.dataset if key.startswith(self.prefix					)]					)
  
  
  
  
  			def __len__(	self     :     Dict					)					->					Optional[int]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return len(self.dataset					)
  
  class 						snake_case			(       lowerCAmelCase__	):
  			'''simple docstring'''
  
  			def __init__(	self     :     int							,							lowerCAmelCase_     :     Dict[str, torch.Tensor] = None							,							lowerCAmelCase_     :     Optional[Union[str, os.PathLike]] = None							,							lowerCAmelCase_     :     Mapping = None							,							lowerCAmelCase_     :     int=None							,							)					->					Tuple:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								if state_dict is None and save_folder is None:
  													raise ValueError('''Need either a `state_dict` or a `save_folder` containing offloaded weights.'''					)
  
  								SCREAMING_SNAKE_CASE_	      =       {} if state_dict is None else state_dict
  								SCREAMING_SNAKE_CASE_	      =       save_folder
  								if index is None and save_folder is not None:
  													with open(os.path.join(lowerCAmelCase_							,							'''index.json'''					)					) as f:
  																		SCREAMING_SNAKE_CASE_	      =       json.load(lowerCAmelCase_					)
  								SCREAMING_SNAKE_CASE_	      =       {} if index is None else index
  								SCREAMING_SNAKE_CASE_	      =       list(self.state_dict.keys()					)
  								self.all_keys.extend([key for key in self.index if key not in self.all_keys]					)
  								SCREAMING_SNAKE_CASE_	      =       device
  
  			def __getitem__(	self     :     str							,							lowerCAmelCase_     :     str					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								if key in self.state_dict:
  													return self.state_dict[key]
  								SCREAMING_SNAKE_CASE_	      =       self.index[key]
  								if weight_info.get('''safetensors_file'''					) is not None:
  													if not is_safetensors_available():
  																		raise ImportError('''These offloaded weights require the use of safetensors: `pip install safetensors`.'''					)
  
  													from safetensors import safe_open
  
  													SCREAMING_SNAKE_CASE_	      =       '''cpu''' if self.device is None else self.device
  													with safe_open(weight_info['''safetensors_file''']							,							framework='''pt'''							,							device=lowerCAmelCase_					) as f:
  																		SCREAMING_SNAKE_CASE_	      =       f.get_tensor(weight_info.get('''weight_name'''							,							lowerCAmelCase_					)					)
  
  													if "dtype" in weight_info:
  																		return tensor.to(getattr(lowerCAmelCase_							,							weight_info['''dtype''']					)					)
  													else:
  																		return tensor
  
  								SCREAMING_SNAKE_CASE_	      =       os.path.join(self.save_folder							,							F'''{key}.dat'''					)
  								return load_offloaded_weight(lowerCAmelCase_							,							lowerCAmelCase_					)
  
  			def __iter__(	self     :     str					)					->					Optional[int]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return iter(self.all_keys					)
  
  
  
  
  			def __len__(	self     :     Optional[Any]					)					->					Optional[Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return len(self.all_keys					)
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		List[str]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       {}
  					for module_name in submodule_names:
  										# We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
  										# submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
  										result.update(
  										    {
  										        key: param
  										        for key, param in state_dict.items()
  										        if key == module_name or key.startswith(module_name + '''.'''      )
  										    }      )
  					return result
  
  
  