code: 
  '''simple docstring'''
  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	(					SCREAMING_SNAKE_CASE__ )     ->       Tuple:
  
  
    '''simple docstring'''
  
  
    snake_case     :					Optional[int]     =	"https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint"
  
    with open(SCREAMING_SNAKE_CASE__   ,					'''wb''' ) as f:
      with requests.get(SCREAMING_SNAKE_CASE__   ,					stream=SCREAMING_SNAKE_CASE__ ) as req:
        for chunk in req.iter_content(chunk_size=2048 ):
          f.write(SCREAMING_SNAKE_CASE__ )
  
  
  
  
  def       _UpperCamelCase	(					)     ->       VivitConfig:
  
  
    '''simple docstring'''
  
  
    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(SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__   ,					repo_type='''dataset''' )   ,					'''r''' ) )
    snake_case     :					Optional[Any]     =	{int(SCREAMING_SNAKE_CASE__ ): 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       _UpperCamelCase	(					)     ->       Union[str, Any]:
  
  
    '''simple docstring'''
  
  
    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(SCREAMING_SNAKE_CASE__ )
    return list(SCREAMING_SNAKE_CASE__ )
  
  
  
  
  def       _UpperCamelCase	(					SCREAMING_SNAKE_CASE__ )     ->       Any:
  
  
    '''simple docstring'''
  
  
    if np.ndim(SCREAMING_SNAKE_CASE__ ) == 2:
      return transform_attention_bias(SCREAMING_SNAKE_CASE__ )
  
    elif np.ndim(SCREAMING_SNAKE_CASE__ ) == 3:
      return transform_attention_kernel(SCREAMING_SNAKE_CASE__ )
  
    else:
      raise Exception(F'Invalid number of dimesions: {np.ndim(SCREAMING_SNAKE_CASE__ )}' )
  
  
  
  
  def       _UpperCamelCase	(					SCREAMING_SNAKE_CASE__ )     ->       Optional[Any]:
  
  
    '''simple docstring'''
  
  
    return current.flatten()
  
  
  
  
  def       _UpperCamelCase	(					SCREAMING_SNAKE_CASE__ )     ->       Any:
  
  
    '''simple docstring'''
  
  
    return np.reshape(SCREAMING_SNAKE_CASE__   ,					(current.shape[0], current.shape[1] * current.shape[2]) ).T
  
  
  
  
  def       _UpperCamelCase	(					SCREAMING_SNAKE_CASE__ )     ->       str:
  
  
    '''simple docstring'''
  
  
    return np.reshape(SCREAMING_SNAKE_CASE__   ,					(current.shape[0] * current.shape[1], current.shape[2]) ).T
  
  
  
  
  def       _UpperCamelCase	(					SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__ )     ->       Optional[Any]:
  
  
    '''simple docstring'''
  
  
    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       _UpperCamelCase	(					SCREAMING_SNAKE_CASE__ )     ->       str:
  
  
    '''simple docstring'''
  
  
    return sum([1 if '''encoderblock_''' in k else 0 for k in state_dict['''optimizer''']['''target''']['''Transformer'''].keys()] )
  
  
  
  
  def       _UpperCamelCase	(					SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__=False )     ->       int:
  
  
    '''simple docstring'''
  
  
    snake_case     :					Union[str, Any]     =	get_n_layers(SCREAMING_SNAKE_CASE__ )
  
    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(SCREAMING_SNAKE_CASE__ ):
      new_state.update(transform_state_encoder_block(SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__ ) )
  
    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(SCREAMING_SNAKE_CASE__ ) for k, v in new_state.items()}
  
  
  
  
  def       _UpperCamelCase	(					)     ->       VivitImageProcessor:
  
  
    '''simple docstring'''
  
  
    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       _UpperCamelCase	(					SCREAMING_SNAKE_CASE__ )     ->       str:
  
  
    '''simple docstring'''
  
  
    snake_case     :					Optional[Any]     =	"checkpoint"
  
    if not os.path.exists(SCREAMING_SNAKE_CASE__ ):
      download_checkpoint(SCREAMING_SNAKE_CASE__ )
  
    snake_case     :					List[str]     =	restore_checkpoint(SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__ )
    snake_case     :					List[Any]     =	transform_state(SCREAMING_SNAKE_CASE__   ,					classification_head=SCREAMING_SNAKE_CASE__ )
  
    snake_case     :					str     =	get_vivit_config()
  
    assert config.image_size == 224
    assert config.num_frames == 32
  
    snake_case     :					str     =	VivitForVideoClassification(SCREAMING_SNAKE_CASE__ )
    model.load_state_dict(SCREAMING_SNAKE_CASE__ )
    model.eval()
  
    snake_case     :					Union[str, Any]     =	get_processor()
  
    snake_case     :					Optional[int]     =	prepare_video()
    snake_case     :					List[Any]     =	extractor(SCREAMING_SNAKE_CASE__   ,					return_tensors='''pt''' )
  
    snake_case     :					Tuple     =	model(**SCREAMING_SNAKE_CASE__ )
  
    snake_case     :					Dict     =	torch.Size([1, 400] )
    snake_case     :					Any     =	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]   ,					SCREAMING_SNAKE_CASE__   ,					atol=1E-4 ), outputs.logits[0, :5]
  
    model.save_pretrained(SCREAMING_SNAKE_CASE__ )
    extractor.save_pretrained(SCREAMING_SNAKE_CASE__ )
  
  
  if __name__ == "__main__":
        lowercase__						=			argparse.ArgumentParser()
  
        parser.add_argument("--output_model_name", "-o", type=str, help="Output path for the converted HuggingFace model")
  
        lowercase__						=			parser.parse_args()
        convert(args.output_model_name)
  
  
  
style_context: 
  '''simple docstring'''
  from typing import TYPE_CHECKING
  
  from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available
  
  
  lowercase__						=			{
      "configuration_flava": [
          "FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP",
          "FlavaConfig",
          "FlavaImageCodebookConfig",
          "FlavaImageConfig",
          "FlavaMultimodalConfig",
          "FlavaTextConfig",
      ],
  }
  
  try:
  						if not is_vision_available():
  												raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
  						pass
  else:
  						lowercase__						=			["FlavaFeatureExtractor"]
  						lowercase__						=			["FlavaImageProcessor"]
  						lowercase__						=			["FlavaProcessor"]
  
  try:
  						if not is_torch_available():
  												raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
  						pass
  else:
  						lowercase__						=			[
  						    "FLAVA_PRETRAINED_MODEL_ARCHIVE_LIST",
  						    "FlavaForPreTraining",
  						    "FlavaImageCodebook",
  						    "FlavaImageModel",
  						    "FlavaModel",
  						    "FlavaMultimodalModel",
  						    "FlavaPreTrainedModel",
  						    "FlavaTextModel",
  						]
  
  if TYPE_CHECKING:
  						from .configuration_flava import (
  						    FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP,
  						    FlavaConfig,
  						    FlavaImageCodebookConfig,
  						    FlavaImageConfig,
  						    FlavaMultimodalConfig,
  						    FlavaTextConfig,
  						)
  
  						try:
  												if not is_vision_available():
  																		raise OptionalDependencyNotAvailable()
  						except OptionalDependencyNotAvailable:
  												pass
  						else:
  												from .feature_extraction_flava import FlavaFeatureExtractor
  												from .image_processing_flava import FlavaImageProcessor
  												from .processing_flava import FlavaProcessor
  
  						try:
  												if not is_torch_available():
  																		raise OptionalDependencyNotAvailable()
  						except OptionalDependencyNotAvailable:
  												pass
  						else:
  												from .modeling_flava import (
  												    FLAVA_PRETRAINED_MODEL_ARCHIVE_LIST,
  												    FlavaForPreTraining,
  												    FlavaImageCodebook,
  												    FlavaImageModel,
  												    FlavaModel,
  												    FlavaMultimodalModel,
  												    FlavaPreTrainedModel,
  												    FlavaTextModel,
  												)
  
  else:
  						import sys
  
  						lowercase__						=			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
  
  
  