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 						snake_case__	(      UpperCamelCase	) ->      Union[str, Any]:
  		_UpperCamelCase       : Optional[int]    =      '''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=20_48	):
  								f.write(UpperCamelCase	)
  
  
  
  
  
  
  
  def 						snake_case__	(      ) ->      VivitConfig:
  		_UpperCamelCase       : int    =      VivitConfig()
  
  		_UpperCamelCase       : Dict    =      4_00
  		_UpperCamelCase       : Any    =      '''huggingface/label-files'''
  		_UpperCamelCase       : Optional[Any]    =      '''kinetics400-id2label.json'''
  
  		_UpperCamelCase       : Optional[Any]    =      json.load(open(hf_hub_download(UpperCamelCase			,UpperCamelCase			,repo_type='''dataset'''	)			,'''r'''	)	)
  		_UpperCamelCase       : Dict    =      {int(UpperCamelCase	): v for k, v in idalabel.items()}
  		_UpperCamelCase       : Tuple    =      idalabel
  		_UpperCamelCase       : Optional[int]    =      {v: k for k, v in idalabel.items()}
  		return config
  
  
  
  
  
  
  
  def 						snake_case__	(      ) ->      List[Any]:
  		_UpperCamelCase       : Optional[int]    =      hf_hub_download(
  		    repo_id='''hf-internal-testing/spaghetti-video'''			,filename='''eating_spaghetti_32_frames.npy'''			,repo_type='''dataset'''	)
  		_UpperCamelCase       : Any    =      np.load(UpperCamelCase	)
  		return list(UpperCamelCase	)
  
  
  
  
  
  
  
  def 						snake_case__	(      UpperCamelCase	) ->      Optional[Any]:
  		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 						snake_case__	(      UpperCamelCase	) ->      Any:
  		return current.flatten()
  
  
  
  
  
  
  
  def 						snake_case__	(      UpperCamelCase	) ->      str:
  		return np.reshape(UpperCamelCase			,(current.shape[0], current.shape[1] * current.shape[2])	).T
  
  
  
  
  
  
  
  def 						snake_case__	(      UpperCamelCase	) ->      Optional[Any]:
  		return np.reshape(UpperCamelCase			,(current.shape[0] * current.shape[1], current.shape[2])	).T
  
  
  
  
  
  
  
  def 						snake_case__	(      UpperCamelCase			,UpperCamelCase	) ->      Tuple:
  		_UpperCamelCase       : Tuple    =      state_dict['''optimizer''']['''target''']['''Transformer'''][f'''encoderblock_{i}''']
  
  		_UpperCamelCase       : List[str]    =      f'''encoder.layer.{i}.'''
  		_UpperCamelCase       : Union[str, Any]    =      {
  		    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 						snake_case__	(      UpperCamelCase	) ->      List[Any]:
  		return sum([1 if '''encoderblock_''' in k else 0 for k in state_dict['''optimizer''']['''target''']['''Transformer'''].keys()]	)
  
  
  
  
  
  
  
  def 						snake_case__	(      UpperCamelCase			,UpperCamelCase=False	) ->      Optional[int]:
  		_UpperCamelCase       : str    =      get_n_layers(UpperCamelCase	)
  
  		_UpperCamelCase       : List[Any]    =      OrderedDict()
  
  		_UpperCamelCase       : Union[str, Any]    =      state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''bias''']
  		_UpperCamelCase       : List[Any]    =      state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''scale''']
  
  		_UpperCamelCase       : int    =      np.transpose(
  		    state_dict['''optimizer''']['''target''']['''embedding''']['''kernel''']			,(4, 3, 0, 1, 2)	)
  		_UpperCamelCase       : List[str]    =      state_dict['''optimizer''']['''target''']['''embedding''']['''bias''']
  
  		_UpperCamelCase       : int    =      state_dict['''optimizer''']['''target''']['''cls''']
  		_UpperCamelCase       : List[str]    =      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:
  				_UpperCamelCase       : int    =      {'''vivit.''' + k: v for k, v in new_state.items()}
  				_UpperCamelCase       : List[str]    =      np.transpose(state_dict['''optimizer''']['''target''']['''output_projection''']['''kernel''']	)
  				_UpperCamelCase       : str    =      np.transpose(state_dict['''optimizer''']['''target''']['''output_projection''']['''bias''']	)
  
  		return {k: torch.tensor(UpperCamelCase	) for k, v in new_state.items()}
  
  
  
  
  
  
  
  def 						snake_case__	(      ) ->      VivitImageProcessor:
  		_UpperCamelCase       : Tuple    =      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 						snake_case__	(      UpperCamelCase	) ->      Optional[Any]:
  		_UpperCamelCase       : Optional[int]    =      '''checkpoint'''
  
  		if not os.path.exists(UpperCamelCase	):
  				download_checkpoint(UpperCamelCase	)
  
  		_UpperCamelCase       : List[str]    =      restore_checkpoint(UpperCamelCase			,UpperCamelCase	)
  		_UpperCamelCase       : int    =      transform_state(UpperCamelCase			,classification_head=UpperCamelCase	)
  
  		_UpperCamelCase       : Optional[int]    =      get_vivit_config()
  
  		assert config.image_size == 2_24
  		assert config.num_frames == 32
  
  		_UpperCamelCase       : Tuple    =      VivitForVideoClassification(UpperCamelCase	)
  		model.load_state_dict(UpperCamelCase	)
  		model.eval()
  
  		_UpperCamelCase       : Tuple    =      get_processor()
  
  		_UpperCamelCase       : Tuple    =      prepare_video()
  		_UpperCamelCase       : Optional[int]    =      extractor(UpperCamelCase			,return_tensors='''pt'''	)
  
  		_UpperCamelCase       : Dict    =      model(**UpperCamelCase	)
  
  		_UpperCamelCase       : Optional[Any]    =      torch.Size([1, 4_00]	)
  		_UpperCamelCase       : Tuple    =      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]			,UpperCamelCase			,atol=1e-4	), outputs.logits[0, :5]
  
  		model.save_pretrained(UpperCamelCase	)
  		extractor.save_pretrained(UpperCamelCase	)
  
  
  if __name__ == "__main__":
  	_UpperCAmelCase   :      Optional[int] 						=  argparse.ArgumentParser()
  
  	parser.add_argument("""--output_model_name""", """-o""", type=str, help="""Output path for the converted HuggingFace model""")
  
  	_UpperCAmelCase   :      str 						=  parser.parse_args()
  	convert(args.output_model_name)
  
  
  
  
  
  
style_context: 
  '''simple docstring'''
  
  import unittest
  
  from transformers.image_processing_utils import get_size_dict
  
  class 			UpperCAmelCase				(					unittest.TestCase ):
  
  
  
  
  
  		"""simple docstring"""
  
  
  
  
  
  		def _lowercase  (    self     )				->    Dict:
  				# Test a dict with the wrong keys raises an error
  				_UpperCamelCase       : List[Any]    =      {'''wrong_key''': 224}
  				with self.assertRaises(_snake_case     ):
  						get_size_dict(_snake_case     )
  
  				_UpperCamelCase       : Dict    =      {'''height''': 224}
  				with self.assertRaises(_snake_case     ):
  						get_size_dict(_snake_case     )
  
  				_UpperCamelCase       : str    =      {'''width''': 224, '''shortest_edge''': 224}
  				with self.assertRaises(_snake_case     ):
  						get_size_dict(_snake_case     )
  
  				# Test a dict with the correct keys is returned as is
  				_UpperCamelCase       : Any    =      {'''height''': 224, '''width''': 224}
  				_UpperCamelCase       : Optional[int]    =      get_size_dict(_snake_case     )
  				self.assertEqual(_snake_case  ,					_snake_case     )
  
  				_UpperCamelCase       : Optional[Any]    =      {'''shortest_edge''': 224}
  				_UpperCamelCase       : int    =      get_size_dict(_snake_case     )
  				self.assertEqual(_snake_case  ,					{'''shortest_edge''': 224}     )
  
  				_UpperCamelCase       : int    =      {'''longest_edge''': 224, '''shortest_edge''': 224}
  				_UpperCamelCase       : Tuple    =      get_size_dict(_snake_case     )
  				self.assertEqual(_snake_case  ,					{'''longest_edge''': 224, '''shortest_edge''': 224}     )
  
  				# Test a single int value which  represents (size, size)
  				_UpperCamelCase       : Any    =      get_size_dict(224     )
  				self.assertEqual(_snake_case  ,					{'''height''': 224, '''width''': 224}     )
  
  				# Test a single int value which represents the shortest edge
  				_UpperCamelCase       : List[Any]    =      get_size_dict(224  ,					default_to_square=_snake_case     )
  				self.assertEqual(_snake_case  ,					{'''shortest_edge''': 224}     )
  
  				# Test a tuple of ints which represents (height, width)
  				_UpperCamelCase       : Any    =      get_size_dict((150, 200)     )
  				self.assertEqual(_snake_case  ,					{'''height''': 150, '''width''': 200}     )
  
  				# Test a tuple of ints which represents (width, height)
  				_UpperCamelCase       : List[str]    =      get_size_dict((150, 200)  ,					height_width_order=_snake_case     )
  				self.assertEqual(_snake_case  ,					{'''height''': 200, '''width''': 150}     )
  
  				# Test an int representing the shortest edge and max_size which represents the longest edge
  				_UpperCamelCase       : Any    =      get_size_dict(224  ,					max_size=256  ,					default_to_square=_snake_case     )
  				self.assertEqual(_snake_case  ,					{'''shortest_edge''': 224, '''longest_edge''': 256}     )
  
  				# Test int with default_to_square=True and max_size fails
  				with self.assertRaises(_snake_case     ):
  						get_size_dict(224  ,					max_size=256  ,					default_to_square=_snake_case     )
  
  
  
  
  
  