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       lowerCAmelCase__(	lowercase       :    Dict				) ->    Tuple:
    __snake_case   :				Optional[int]        =       "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=2048				):
          f.write(lowercase				)
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	) ->    VivitConfig:
    __snake_case   :				int        =       VivitConfig()
  
    __snake_case   :				int        =       400
    __snake_case   :				Tuple        =       "huggingface/label-files"
    __snake_case   :				List[str]        =       "kinetics400-id2label.json"
  
    __snake_case   :				Any        =       json.load(open(hf_hub_download(lowercase   ,     lowercase   ,     repo_type="dataset"				)   ,     "r"				)				)
    __snake_case   :				Optional[Any]        =       {int(lowercase				): v for k, v in idalabel.items()}
    __snake_case   :				List[Any]        =       idalabel
    __snake_case   :				str        =       {v: k for k, v in idalabel.items()}
    return config
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	) ->    Union[str, Any]:
    __snake_case   :				Union[str, Any]        =       hf_hub_download(
        repo_id="hf-internal-testing/spaghetti-video"   ,     filename="eating_spaghetti_32_frames.npy"   ,     repo_type="dataset"				)
    __snake_case   :				str        =       np.load(lowercase				)
    return list(lowercase				)
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    np.ndarray				) ->    Any:
    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       lowerCAmelCase__(	lowercase       :    np.ndarray				) ->    Optional[Any]:
    return current.flatten()
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    np.ndarray				) ->    Any:
    return np.reshape(lowercase   ,     (current.shape[0], current.shape[1] * current.shape[2])				).T
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    np.ndarray				) ->    str:
    return np.reshape(lowercase   ,     (current.shape[0] * current.shape[1], current.shape[2])				).T
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    int   ,     lowercase       :    List[str]				) ->    Optional[Any]:
    __snake_case   :				Dict        =       state_dict["optimizer"]["target"]["Transformer"][f"""encoderblock_{i}"""]
  
    __snake_case   :				Optional[int]        =       f"""encoder.layer.{i}."""
    __snake_case   :				str        =       {
        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       lowerCAmelCase__(	lowercase       :    List[str]				) ->    str:
    return sum([1 if "encoderblock_" in k else 0 for k in state_dict["optimizer"]["target"]["Transformer"].keys()]				)
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    Tuple   ,     lowercase       :    Union[str, Any]=False				) ->    int:
    __snake_case   :				Union[str, Any]        =       get_n_layers(lowercase				)
  
    __snake_case   :				Optional[Any]        =       OrderedDict()
  
    __snake_case   :				Union[str, Any]        =       state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["bias"]
    __snake_case   :				Dict        =       state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["scale"]
  
    __snake_case   :				Optional[int]        =       np.transpose(
        state_dict["optimizer"]["target"]["embedding"]["kernel"]   ,     (4, 3, 0, 1, 2)				)
    __snake_case   :				Dict        =       state_dict["optimizer"]["target"]["embedding"]["bias"]
  
    __snake_case   :				List[Any]        =       state_dict["optimizer"]["target"]["cls"]
    __snake_case   :				Optional[int]        =       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:
      __snake_case   :				Optional[int]        =       {"vivit." + k: v for k, v in new_state.items()}
      __snake_case   :				List[str]        =       np.transpose(state_dict["optimizer"]["target"]["output_projection"]["kernel"]				)
      __snake_case   :				Union[str, Any]        =       np.transpose(state_dict["optimizer"]["target"]["output_projection"]["bias"]				)
  
    return {k: torch.tensor(lowercase				) for k, v in new_state.items()}
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	) ->    VivitImageProcessor:
    __snake_case   :				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       lowerCAmelCase__(	lowercase       :    str				) ->    str:
    __snake_case   :				Optional[Any]        =       "checkpoint"
  
    if not os.path.exists(lowercase				):
      download_checkpoint(lowercase				)
  
    __snake_case   :				List[str]        =       restore_checkpoint(lowercase   ,     lowercase				)
    __snake_case   :				List[Any]        =       transform_state(lowercase   ,     classification_head=lowercase				)
  
    __snake_case   :				str        =       get_vivit_config()
  
    assert config.image_size == 224
    assert config.num_frames == 32
  
    __snake_case   :				str        =       VivitForVideoClassification(lowercase				)
    model.load_state_dict(lowercase				)
    model.eval()
  
    __snake_case   :				Union[str, Any]        =       get_processor()
  
    __snake_case   :				Optional[int]        =       prepare_video()
    __snake_case   :				List[Any]        =       extractor(lowercase   ,     return_tensors="pt"				)
  
    __snake_case   :				Tuple        =       model(**lowercase				)
  
    __snake_case   :				Dict        =       torch.Size([1, 400]				)
    __snake_case   :				Any        =       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]   ,     lowercase   ,     atol=1E-4				), outputs.logits[0, :5]
  
    model.save_pretrained(lowercase				)
    extractor.save_pretrained(lowercase				)
  
  
  if __name__ == "__main__":
       _UpperCamelCase       = argparse.ArgumentParser()
  
       parser.add_argument('''--output_model_name''', '''-o''', type=str, help='''Output path for the converted HuggingFace model''')
  
       _UpperCamelCase       = parser.parse_args()
       convert(args.output_model_name)
  
  
style_context: 
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    str				) ->    int:
    __snake_case   :				List[str]        =       str(lowercase				).strip()
    if not oct_string:
      raise ValueError("Empty string was passed to the function"				)
    __snake_case   :				int        =       oct_string[0] == "-"
    if is_negative:
      __snake_case   :				Union[str, Any]        =       oct_string[1:]
    if not oct_string.isdigit() or not all(0 <= int(lowercase				) <= 7 for char in oct_string				):
      raise ValueError("Non-octal value was passed to the function"				)
    __snake_case   :				List[Any]        =       0
    for char in oct_string:
      __snake_case   :				Optional[Any]        =       8 * decimal_number + int(lowercase				)
    if is_negative:
      __snake_case   :				Any        =       -decimal_number
    return decimal_number
  
  
  if __name__ == "__main__":
       from doctest import testmod
  
       testmod()
  
  