
'''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	)->			List[str]:







      """simple docstring"""
      __A       	=			'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=2_0_4_8	):
                        f.write(UpperCAmelCase	)





def        snake_case							(		)->			VivitConfig:







      """simple docstring"""
      __A       	=			VivitConfig()

      __A       	=			4_0_0
      __A       	=			'huggingface/label-files'
      __A       	=			'kinetics400-id2label.json'

      __A       	=			json.load(open(hf_hub_download(UpperCAmelCase      ,  UpperCAmelCase      ,  repo_type='dataset'	)      ,  'r'	)	)
      __A       	=			{int(UpperCAmelCase	): v for k, v in idalabel.items()}
      __A       	=			idalabel
      __A       	=			{v: k for k, v in idalabel.items()}
      return config





def        snake_case							(		)->			List[str]:







      """simple docstring"""
      __A       	=			hf_hub_download(
          repo_id='hf-internal-testing/spaghetti-video'      ,  filename='eating_spaghetti_32_frames.npy'      ,  repo_type='dataset'	)
      __A       	=			np.load(UpperCAmelCase	)
      return list(UpperCAmelCase	)





def        snake_case							(		UpperCAmelCase	)->			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        snake_case							(		UpperCAmelCase	)->			Optional[Any]:







      """simple docstring"""
      return current.flatten()





def        snake_case							(		UpperCAmelCase	)->			Tuple:







      """simple docstring"""
      return np.reshape(UpperCAmelCase      ,  (current.shape[0], current.shape[1] * current.shape[2])	).T





def        snake_case							(		UpperCAmelCase	)->			Dict:







      """simple docstring"""
      return np.reshape(UpperCAmelCase      ,  (current.shape[0] * current.shape[1], current.shape[2])	).T





def        snake_case							(		UpperCAmelCase      ,  UpperCAmelCase	)->			List[str]:







      """simple docstring"""
      __A       	=			state_dict['optimizer']['target']['Transformer'][f'encoderblock_{i}']

      __A       	=			f'encoder.layer.{i}.'
      __A       	=			{
          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]:







      """simple docstring"""
      return sum([1 if 'encoderblock_' in k else 0 for k in state_dict['optimizer']['target']['Transformer'].keys()]	)





def        snake_case							(		UpperCAmelCase      ,  UpperCAmelCase=False	)->			Tuple:







      """simple docstring"""
      __A       	=			get_n_layers(UpperCAmelCase	)

      __A       	=			OrderedDict()

      __A       	=			state_dict['optimizer']['target']['Transformer']['encoder_norm']['bias']
      __A       	=			state_dict['optimizer']['target']['Transformer']['encoder_norm']['scale']

      __A       	=			np.transpose(
          state_dict['optimizer']['target']['embedding']['kernel']      ,  (4, 3, 0, 1, 2)	)
      __A       	=			state_dict['optimizer']['target']['embedding']['bias']

      __A       	=			state_dict['optimizer']['target']['cls']
      __A       	=			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:
            __A       	=			{'vivit.' + k: v for k, v in new_state.items()}
            __A       	=			np.transpose(state_dict['optimizer']['target']['output_projection']['kernel']	)
            __A       	=			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:







      """simple docstring"""
      __A       	=			VivitImageProcessor()

      assert extractor.do_resize is True
      assert extractor.size == {"shortest_edge": 2_5_6}
      assert extractor.do_center_crop is True
      assert extractor.crop_size == {"width": 2_2_4, "height": 2_2_4}
      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_5_5

      # zero-centering = True in original implementation
      assert extractor.do_zero_centering is True

      return extractor





def        snake_case							(		UpperCAmelCase	)->			Union[str, Any]:







      """simple docstring"""
      __A       	=			'checkpoint'

      if not os.path.exists(UpperCAmelCase	):
            download_checkpoint(UpperCAmelCase	)

      __A       	=			restore_checkpoint(UpperCAmelCase      ,  UpperCAmelCase	)
      __A       	=			transform_state(UpperCAmelCase      ,  classification_head=UpperCAmelCase	)

      __A       	=			get_vivit_config()

      assert config.image_size == 2_2_4
      assert config.num_frames == 3_2

      __A       	=			VivitForVideoClassification(UpperCAmelCase	)
      model.load_state_dict(UpperCAmelCase	)
      model.eval()

      __A       	=			get_processor()

      __A       	=			prepare_video()
      __A       	=			extractor(UpperCAmelCase      ,  return_tensors='pt'	)

      __A       	=			model(**UpperCAmelCase	)

      __A       	=			torch.Size([1, 4_0_0]	)
      __A       	=			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__":
 a__		:   str            =       argparse.ArgumentParser()

 parser.add_argument("--output_model_name", "-o", type=str, help="Output path for the converted HuggingFace model")

 a__		:   int            =       parser.parse_args()
 convert(args.output_model_name)








'''simple docstring'''





import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_video_inputs


if is_torch_available():
 import torch

if is_vision_available():
 from PIL import Image

 from transformers import VideoMAEImageProcessor


class  UpperCamelCase__  (    unittest.TestCase):






 def __init__(			self  :Dict	,				_A  :Optional[Any]	,				_A  :Any=7	,				_A  :Dict=3	,				_A  :Optional[Any]=10	,				_A  :List[Any]=18	,				_A  :int=30	,				_A  :List[Any]=400	,				_A  :Any=True	,				_A  :Any=None	,				_A  :Optional[int]=True	,				_A  :Tuple=[0.5, 0.5, 0.5]	,				_A  :int=[0.5, 0.5, 0.5]	,				_A  :List[Any]=None	,				)				->			int:







       '''simple docstring'''



       __A       	=			size if size is not None else {'shortest_edge': 18}
       __A       	=			crop_size if crop_size is not None else {'height': 18, 'width': 18}

       __A       	=			parent
       __A       	=			batch_size
       __A       	=			num_channels
       __A       	=			num_frames
       __A       	=			image_size
       __A       	=			min_resolution
       __A       	=			max_resolution
       __A       	=			do_resize
       __A       	=			size
       __A       	=			do_normalize
       __A       	=			image_mean
       __A       	=			image_std
       __A       	=			crop_size






 def 	lowercase_						(			self  :Any			)				->			Union[str, Any]:







       '''simple docstring'''



       return {
           "image_mean": self.image_mean,
           "image_std": self.image_std,
           "do_normalize": self.do_normalize,
           "do_resize": self.do_resize,
           "size": self.size,
           "crop_size": self.crop_size,
       }


@require_torch
@require_vision
class  UpperCamelCase__  (    SCREAMING_SNAKE_CASE      ,     unittest.TestCase):
 UpperCAmelCase__    :    Tuple										=					VideoMAEImageProcessor if is_vision_available() else None






 def 	lowercase_						(			self  :Dict			)				->			Dict:







       '''simple docstring'''



       __A       	=			VideoMAEImageProcessingTester(self			)






 @property
 def 	lowercase_						(			self  :Dict			)				->			List[str]:







       '''simple docstring'''



       return self.image_processor_tester.prepare_image_processor_dict()






 def 	lowercase_						(			self  :Dict			)				->			Optional[int]:







       '''simple docstring'''



       __A       	=			self.image_processing_class(**self.image_processor_dict			)
       self.assertTrue(hasattr(_A	,				'image_mean'			)			)
       self.assertTrue(hasattr(_A	,				'image_std'			)			)
       self.assertTrue(hasattr(_A	,				'do_normalize'			)			)
       self.assertTrue(hasattr(_A	,				'do_resize'			)			)
       self.assertTrue(hasattr(_A	,				'do_center_crop'			)			)
       self.assertTrue(hasattr(_A	,				'size'			)			)






 def 	lowercase_						(			self  :Optional[int]			)				->			List[Any]:







       '''simple docstring'''



       __A       	=			self.image_processing_class.from_dict(self.image_processor_dict			)
       self.assertEqual(image_processor.size	,				{'shortest_edge': 18}			)
       self.assertEqual(image_processor.crop_size	,				{'height': 18, 'width': 18}			)

       __A       	=			self.image_processing_class.from_dict(self.image_processor_dict	,				size=42	,				crop_size=84			)
       self.assertEqual(image_processor.size	,				{'shortest_edge': 42}			)
       self.assertEqual(image_processor.crop_size	,				{'height': 84, 'width': 84}			)






 def 	lowercase_						(			self  :Any			)				->			Dict:







       '''simple docstring'''



       pass






 def 	lowercase_						(			self  :Optional[int]			)				->			Tuple:







       '''simple docstring'''



       __A       	=			self.image_processing_class(**self.image_processor_dict			)
       # create random PIL videos
       __A       	=			prepare_video_inputs(self.image_processor_tester	,				equal_resolution=_A			)
       for video in video_inputs:
             self.assertIsInstance(_A	,				_A			)
             self.assertIsInstance(video[0]	,				Image.Image			)

       # Test not batched input
       __A       	=			image_processing(video_inputs[0]	,				return_tensors='pt'			).pixel_values
       self.assertEqual(
           encoded_videos.shape	,				(
               1,
               self.image_processor_tester.num_frames,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.crop_size['height'],
               self.image_processor_tester.crop_size['width'],
           )	,				)

       # Test batched
       __A       	=			image_processing(_A	,				return_tensors='pt'			).pixel_values
       self.assertEqual(
           encoded_videos.shape	,				(
               self.image_processor_tester.batch_size,
               self.image_processor_tester.num_frames,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.crop_size['height'],
               self.image_processor_tester.crop_size['width'],
           )	,				)






 def 	lowercase_						(			self  :Dict			)				->			List[Any]:







       '''simple docstring'''



       __A       	=			self.image_processing_class(**self.image_processor_dict			)
       # create random numpy tensors
       __A       	=			prepare_video_inputs(self.image_processor_tester	,				equal_resolution=_A	,				numpify=_A			)
       for video in video_inputs:
             self.assertIsInstance(_A	,				_A			)
             self.assertIsInstance(video[0]	,				np.ndarray			)

       # Test not batched input
       __A       	=			image_processing(video_inputs[0]	,				return_tensors='pt'			).pixel_values
       self.assertEqual(
           encoded_videos.shape	,				(
               1,
               self.image_processor_tester.num_frames,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.crop_size['height'],
               self.image_processor_tester.crop_size['width'],
           )	,				)

       # Test batched
       __A       	=			image_processing(_A	,				return_tensors='pt'			).pixel_values
       self.assertEqual(
           encoded_videos.shape	,				(
               self.image_processor_tester.batch_size,
               self.image_processor_tester.num_frames,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.crop_size['height'],
               self.image_processor_tester.crop_size['width'],
           )	,				)






 def 	lowercase_						(			self  :int			)				->			List[str]:







       '''simple docstring'''



       __A       	=			self.image_processing_class(**self.image_processor_dict			)
       # create random PyTorch tensors
       __A       	=			prepare_video_inputs(self.image_processor_tester	,				equal_resolution=_A	,				torchify=_A			)
       for video in video_inputs:
             self.assertIsInstance(_A	,				_A			)
             self.assertIsInstance(video[0]	,				torch.Tensor			)

       # Test not batched input
       __A       	=			image_processing(video_inputs[0]	,				return_tensors='pt'			).pixel_values
       self.assertEqual(
           encoded_videos.shape	,				(
               1,
               self.image_processor_tester.num_frames,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.crop_size['height'],
               self.image_processor_tester.crop_size['width'],
           )	,				)

       # Test batched
       __A       	=			image_processing(_A	,				return_tensors='pt'			).pixel_values
       self.assertEqual(
           encoded_videos.shape	,				(
               self.image_processor_tester.batch_size,
               self.image_processor_tester.num_frames,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.crop_size['height'],
               self.image_processor_tester.crop_size['width'],
           )	,				)






