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[int]     )      -> Tuple:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						_lowercase					:  Optional[Any]    			=      "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'''
  
  						_lowercase					:  Tuple    			=      VivitConfig()
  
  						_lowercase					:  Dict    			=      400
  						_lowercase					:  Dict    			=      "huggingface/label-files"
  						_lowercase					:  str    			=      "kinetics400-id2label.json"
  
  						_lowercase					:  Union[str, Any]    			=      json.load(open(hf_hub_download(_UpperCAmelCase   , _UpperCAmelCase   , repo_type="dataset"     )   , "r"     )     )
  						_lowercase					:  str    			=      {int(_UpperCAmelCase     ): v for k, v in idalabel.items()}
  						_lowercase					:  Optional[Any]    			=      idalabel
  						_lowercase					:  int    			=      {v: k for k, v in idalabel.items()}
  						return config
  
  
  
  def 		UpperCamelCase    (					)      -> Any:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						_lowercase					:  Union[str, Any]    			=      hf_hub_download(
  						    repo_id="hf-internal-testing/spaghetti-video"   , filename="eating_spaghetti_32_frames.npy"   , repo_type="dataset"     )
  						_lowercase					:  Any    			=      np.load(_UpperCAmelCase     )
  						return list(_UpperCAmelCase     )
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      np.ndarray     )      -> Optional[int]:
  
  
  
  
  
  
  						'''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     :      np.ndarray     )      -> Tuple:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						return current.flatten()
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      np.ndarray     )      -> Dict:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						return np.reshape(_UpperCAmelCase   , (current.shape[0], current.shape[1] * current.shape[2])     ).T
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      np.ndarray     )      -> Optional[Any]:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						return np.reshape(_UpperCAmelCase   , (current.shape[0] * current.shape[1], current.shape[2])     ).T
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      Union[str, Any]   , _UpperCAmelCase     :      str     )      -> Optional[Any]:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						_lowercase					:  int    			=      state_dict["optimizer"]["target"]["Transformer"][f"""encoderblock_{i}"""]
  
  						_lowercase					:  List[Any]    			=      f"""encoder.layer.{i}."""
  						_lowercase					:  Tuple    			=      {
  						    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     :      Optional[Any]     )      -> List[str]:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						return sum([1 if "encoderblock_" in k else 0 for k in state_dict["optimizer"]["target"]["Transformer"].keys()]     )
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      List[str]   , _UpperCAmelCase     :      Tuple=False     )      -> Tuple:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						_lowercase					:  Optional[Any]    			=      get_n_layers(_UpperCAmelCase     )
  
  						_lowercase					:  Dict    			=      OrderedDict()
  
  						_lowercase					:  Dict    			=      state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["bias"]
  						_lowercase					:  Union[str, Any]    			=      state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["scale"]
  
  						_lowercase					:  Optional[Any]    			=      np.transpose(
  						    state_dict["optimizer"]["target"]["embedding"]["kernel"]   , (4, 3, 0, 1, 2)     )
  						_lowercase					:  Tuple    			=      state_dict["optimizer"]["target"]["embedding"]["bias"]
  
  						_lowercase					:  Tuple    			=      state_dict["optimizer"]["target"]["cls"]
  						_lowercase					:  int    			=      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:
  												_lowercase					:  Optional[Any]    			=      {"vivit." + k: v for k, v in new_state.items()}
  												_lowercase					:  Optional[Any]    			=      np.transpose(state_dict["optimizer"]["target"]["output_projection"]["kernel"]     )
  												_lowercase					:  Union[str, Any]    			=      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'''
  
  						_lowercase					:  Union[str, Any]    			=      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     :      str     )      -> Tuple:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						_lowercase					:  Any    			=      "checkpoint"
  
  						if not os.path.exists(_UpperCAmelCase     ):
  												download_checkpoint(_UpperCAmelCase     )
  
  						_lowercase					:  List[str]    			=      restore_checkpoint(_UpperCAmelCase   , _UpperCAmelCase     )
  						_lowercase					:  Union[str, Any]    			=      transform_state(_UpperCAmelCase   , classification_head=_UpperCAmelCase     )
  
  						_lowercase					:  Dict    			=      get_vivit_config()
  
  						assert config.image_size == 224
  						assert config.num_frames == 32
  
  						_lowercase					:  int    			=      VivitForVideoClassification(_UpperCAmelCase     )
  						model.load_state_dict(_UpperCAmelCase     )
  						model.eval()
  
  						_lowercase					:  List[Any]    			=      get_processor()
  
  						_lowercase					:  Union[str, Any]    			=      prepare_video()
  						_lowercase					:  Optional[int]    			=      extractor(_UpperCAmelCase   , return_tensors="pt"     )
  
  						_lowercase					:  Union[str, Any]    			=      model(**_UpperCAmelCase     )
  
  						_lowercase					:  List[str]    			=      torch.Size([1, 400]     )
  						_lowercase					:  Optional[int]    			=      torch.tensor([-1.0_543, 2.0_764, -0.2_104, 0.4_439, -0.9_658]     )
  
  						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__":
  							UpperCamelCase_	:    Tuple  					=       argparse.ArgumentParser()
  
  							parser.add_argument("""--output_model_name""", """-o""", type=str, help="""Output path for the converted HuggingFace model""")
  
  							UpperCamelCase_	:    int  					=       parser.parse_args()
  							convert(args.output_model_name)
  
  
  
  
  
style_context: 
  import math
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      int     )      -> list:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						if n <= 0:
  												raise ValueError("Only positive integers have prime factors"     )
  						_lowercase					:  Optional[int]    			=      []
  						while n % 2 == 0:
  												pf.append(2     )
  												_lowercase					:  Tuple    			=      int(n / 2     )
  						for i in range(3   , int(math.sqrt(_UpperCAmelCase     )     ) + 1   , 2     ):
  												while n % i == 0:
  																		pf.append(_UpperCAmelCase     )
  																		_lowercase					:  Optional[int]    			=      int(n / i     )
  						if n > 2:
  												pf.append(_UpperCAmelCase     )
  						return pf
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      int     )      -> int:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						if n <= 0:
  												raise ValueError("Only positive numbers are accepted"     )
  						_lowercase					:  str    			=      1
  						_lowercase					:  Tuple    			=      1
  						while n % 2 == 0:
  												temp += 1
  												_lowercase					:  List[Any]    			=      int(n / 2     )
  						div *= temp
  						for i in range(3   , int(math.sqrt(_UpperCAmelCase     )     ) + 1   , 2     ):
  												_lowercase					:  Optional[Any]    			=      1
  												while n % i == 0:
  																		temp += 1
  																		_lowercase					:  Optional[Any]    			=      int(n / i     )
  												div *= temp
  						if n > 1:
  												div *= 2
  						return div
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      int     )      -> int:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						if n <= 0:
  												raise ValueError("Only positive numbers are accepted"     )
  						_lowercase					:  int    			=      1
  						_lowercase					:  List[str]    			=      1
  						while n % 2 == 0:
  												temp += 1
  												_lowercase					:  str    			=      int(n / 2     )
  						if temp > 1:
  												s *= (2**temp - 1) / (2 - 1)
  						for i in range(3   , int(math.sqrt(_UpperCAmelCase     )     ) + 1   , 2     ):
  												_lowercase					:  int    			=      1
  												while n % i == 0:
  																		temp += 1
  																		_lowercase					:  List[Any]    			=      int(n / i     )
  												if temp > 1:
  																		s *= (i**temp - 1) / (i - 1)
  						return int(_UpperCAmelCase     )
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      int     )      -> int:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						_lowercase					:  int    			=      n
  						for x in set(prime_factors(_UpperCAmelCase     )     ):
  												s *= (x - 1) / x
  						return int(_UpperCAmelCase     )
  
  
  if __name__ == "__main__":
  							import doctest
  
  							doctest.testmod()
  
  
  
  
  