



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       (   snake_case    ):





					"""simple docstring"""

					_lowerCAmelCase      			=					'''https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'''

					with open(snake_case			,		"""wb"""    ) as f:
										with requests.get(snake_case			,		stream=snake_case    ) as req:
															for chunk in req.iter_content(chunk_size=20_48    ):
																				f.write(snake_case    )


def 	_UpperCAmelCase       (   ):





					"""simple docstring"""

					_lowerCAmelCase      			=					VivitConfig()

					_lowerCAmelCase      			=					4_00
					_lowerCAmelCase      			=					'''huggingface/label-files'''
					_lowerCAmelCase      			=					'''kinetics400-id2label.json'''

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


def 	_UpperCAmelCase       (   ):





					"""simple docstring"""

					_lowerCAmelCase      			=					hf_hub_download(
					    repo_id="""hf-internal-testing/spaghetti-video"""			,		filename="""eating_spaghetti_32_frames.npy"""			,		repo_type="""dataset"""    )
					_lowerCAmelCase      			=					np.load(snake_case    )
					return list(snake_case    )


def 	_UpperCAmelCase       (   snake_case    ):





					"""simple docstring"""

					if np.ndim(snake_case    ) == 2:
										return transform_attention_bias(snake_case    )

					elif np.ndim(snake_case    ) == 3:
										return transform_attention_kernel(snake_case    )

					else:
										raise Exception(F'Invalid number of dimesions: {np.ndim(snake_case    )}'    )


def 	_UpperCAmelCase       (   snake_case    ):





					"""simple docstring"""

					return current.flatten()


def 	_UpperCAmelCase       (   snake_case    ):





					"""simple docstring"""

					return np.reshape(snake_case			,		(current.shape[0], current.shape[1] * current.shape[2])    ).T


def 	_UpperCAmelCase       (   snake_case    ):





					"""simple docstring"""

					return np.reshape(snake_case			,		(current.shape[0] * current.shape[1], current.shape[2])    ).T


def 	_UpperCAmelCase       (   snake_case			,		snake_case    ):





					"""simple docstring"""

					_lowerCAmelCase      			=					state_dict['''optimizer''']['''target''']['''Transformer'''][F'encoderblock_{i}']

					_lowerCAmelCase      			=					F'encoder.layer.{i}.'
					_lowerCAmelCase      			=					{
					    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       (   snake_case    ):





					"""simple docstring"""

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


def 	_UpperCAmelCase       (   snake_case			,		snake_case=False    ):





					"""simple docstring"""

					_lowerCAmelCase      			=					get_n_layers(snake_case    )

					_lowerCAmelCase      			=					OrderedDict()

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

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

					_lowerCAmelCase      			=					state_dict['''optimizer''']['''target''']['''cls''']
					_lowerCAmelCase      			=					state_dict['''optimizer''']['''target''']['''Transformer''']['''posembed_input'''][
					    '''pos_embedding'''
					]

					for i in range(snake_case    ):
										new_state.update(transform_state_encoder_block(snake_case			,		snake_case    )    )

					if classification_head:
										_lowerCAmelCase      			=					{'''vivit.''' + k: v for k, v in new_state.items()}
										_lowerCAmelCase      			=					np.transpose(state_dict["""optimizer"""]["""target"""]["""output_projection"""]["""kernel"""]    )
										_lowerCAmelCase      			=					np.transpose(state_dict["""optimizer"""]["""target"""]["""output_projection"""]["""bias"""]    )

					return {k: torch.tensor(snake_case    ) for k, v in new_state.items()}


def 	_UpperCAmelCase       (   ):





					"""simple docstring"""

					_lowerCAmelCase      			=					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 	_UpperCAmelCase       (   snake_case    ):





					"""simple docstring"""

					_lowerCAmelCase      			=					'''checkpoint'''

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

					_lowerCAmelCase      			=					restore_checkpoint(snake_case			,		snake_case    )
					_lowerCAmelCase      			=					transform_state(snake_case			,		classification_head=snake_case    )

					_lowerCAmelCase      			=					get_vivit_config()

					assert config.image_size == 2_24
					assert config.num_frames == 32

					_lowerCAmelCase      			=					VivitForVideoClassification(snake_case    )
					model.load_state_dict(snake_case    )
					model.eval()

					_lowerCAmelCase      			=					get_processor()

					_lowerCAmelCase      			=					prepare_video()
					_lowerCAmelCase      			=					extractor(snake_case			,		return_tensors="""pt"""    )

					_lowerCAmelCase      			=					model(**snake_case    )

					_lowerCAmelCase      			=					torch.Size([1, 4_00]    )
					_lowerCAmelCase      			=					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]			,		snake_case			,		atol=1E-4    ), outputs.logits[0, :5]

					model.save_pretrained(snake_case    )
					extractor.save_pretrained(snake_case    )


if __name__ == "__main__":
	A__     :   Optional[int]			   = argparse.ArgumentParser()

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

	A__     :   str			   = parser.parse_args()
	convert(args.output_model_name)








import itertools
import random
import unittest

import numpy as np

from transformers import EncodecFeatureExtractor
from transformers.testing_utils import require_torch
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_torch_available():
	import torch


A__			   = random.Random()


def 	_UpperCAmelCase       (   snake_case			,		snake_case=1.0			,		snake_case=None			,		snake_case=None    ):





					"""simple docstring"""

					if rng is None:
										_lowerCAmelCase      			=					global_rng

					_lowerCAmelCase      			=					[]
					for batch_idx in range(shape[0]    ):
										values.append([]    )
										for _ in range(shape[1]    ):
															values[-1].append(rng.random() * scale    )

					return values





@require_torch
class       __lowerCAmelCase				(  unittest.TestCase ):




				def __init__(							self		,       _snake_case		,       _snake_case=7		,       _snake_case=400		,       _snake_case=2000		,       _snake_case=1		,       _snake_case=0.0		,       _snake_case=24000		,       _snake_case=True		,       ):



									"""simple docstring"""





									_lowerCAmelCase      			=					parent
									_lowerCAmelCase      			=					batch_size
									_lowerCAmelCase      			=					min_seq_length
									_lowerCAmelCase      			=					max_seq_length
									_lowerCAmelCase      			=					(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
									_lowerCAmelCase      			=					feature_size
									_lowerCAmelCase      			=					padding_value
									_lowerCAmelCase      			=					sampling_rate
									_lowerCAmelCase      			=					return_attention_mask




				def 							snake_case     (							self	):



									"""simple docstring"""





									return {
									    "feature_size": self.feature_size,
									    "padding_value": self.padding_value,
									    "sampling_rate": self.sampling_rate,
									    "return_attention_mask": self.return_attention_mask,
									}




				def 							snake_case     (							self		,       _snake_case=False		,       _snake_case=False	):



									"""simple docstring"""





									def _flatten(_snake_case	):
														return list(itertools.chain(*_snake_case	)	)

									if equal_length:
														_lowerCAmelCase      			=					floats_list((self.batch_size, self.max_seq_length)	)
									else:
														# make sure that inputs increase in size
														_lowerCAmelCase      			=					[
														    _flatten(floats_list((x, self.feature_size)	)	)
														    for x in range(self.min_seq_length		,       self.max_seq_length		,       self.seq_length_diff	)
														]

									if numpify:
														_lowerCAmelCase      			=					[np.asarray(_snake_case	) for x in audio_inputs]

									return audio_inputs





@require_torch
class       __lowerCAmelCase				(  lowerCamelCase__ ,		unittest.TestCase ):
				__lowerCamelCase				=       EncodecFeatureExtractor




				def 							snake_case     (							self	):



									"""simple docstring"""





									_lowerCAmelCase      			=					EnCodecFeatureExtractionTester(self	)




				def 							snake_case     (							self	):



									"""simple docstring"""





									_lowerCAmelCase      			=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()	)
									# create three inputs of length 800, 1000, and 1200
									_lowerCAmelCase      			=					[floats_list((1, x)	)[0] for x in range(800		,       1400		,       200	)]
									_lowerCAmelCase      			=					[np.asarray(_snake_case	) for audio_input in audio_inputs]

									# Test not batched input
									_lowerCAmelCase      			=					feat_extract(audio_inputs[0]		,       return_tensors="""np"""	).input_values
									_lowerCAmelCase      			=					feat_extract(np_audio_inputs[0]		,       return_tensors="""np"""	).input_values
									self.assertTrue(np.allclose(_snake_case		,       _snake_case		,       atol=1e-3	)	)

									# Test batched
									_lowerCAmelCase      			=					feat_extract(_snake_case		,       padding=_snake_case		,       return_tensors="""np"""	).input_values
									_lowerCAmelCase      			=					feat_extract(_snake_case		,       padding=_snake_case		,       return_tensors="""np"""	).input_values
									for enc_seq_a, enc_seq_a in zip(_snake_case		,       _snake_case	):
														self.assertTrue(np.allclose(_snake_case		,       _snake_case		,       atol=1e-3	)	)




				def 							snake_case     (							self	):



									"""simple docstring"""





									_lowerCAmelCase      			=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()	)
									_lowerCAmelCase      			=					np.random.rand(100	).astype(np.floataa	)
									_lowerCAmelCase      			=					np_audio_inputs.tolist()

									for inputs in [py_audio_inputs, np_audio_inputs]:
														_lowerCAmelCase      			=					feature_extractor.pad([{"""input_values""": inputs}]		,       return_tensors="""np"""	)
														self.assertTrue(np_processed.input_values.dtype == np.floataa	)
														_lowerCAmelCase      			=					feature_extractor.pad([{"""input_values""": inputs}]		,       return_tensors="""pt"""	)
														self.assertTrue(pt_processed.input_values.dtype == torch.floataa	)




				def 							snake_case     (							self		,       _snake_case	):



									"""simple docstring"""





									from datasets import load_dataset

									_lowerCAmelCase      			=					load_dataset("""hf-internal-testing/librispeech_asr_dummy"""		,       """clean"""		,       split="""validation"""	)
									# automatic decoding with librispeech
									_lowerCAmelCase      			=					ds.sort("""id"""	).select(range(_snake_case	)	)[:num_samples]["""audio"""]

									return [x["array"] for x in audio_samples]




				def 							snake_case     (							self	):



									"""simple docstring"""





									_lowerCAmelCase      			=					torch.tensor(
									    [2.3804e-03, 2.0752e-03, 1.9836e-03, 2.1057e-03, 1.6174e-03,
									     3.0518e-04, 9.1553e-05, 3.3569e-04, 9.7656e-04, 1.8311e-03,
									     2.0142e-03, 2.1057e-03, 1.7395e-03, 4.5776e-04, -3.9673e-04,
									     4.5776e-04, 1.0071e-03, 9.1553e-05, 4.8828e-04, 1.1597e-03,
									     7.3242e-04, 9.4604e-04, 1.8005e-03, 1.8311e-03, 8.8501e-04,
									     4.2725e-04, 4.8828e-04, 7.3242e-04, 1.0986e-03, 2.1057e-03]	)
									# fmt: on
									_lowerCAmelCase      			=					self._load_datasamples(1	)
									_lowerCAmelCase      			=					EncodecFeatureExtractor()
									_lowerCAmelCase      			=					feature_extractor(_snake_case		,       return_tensors="""pt"""	).input_values
									self.assertEquals(input_values.shape		,       (1, 1, 93680)	)
									self.assertTrue(torch.allclose(input_values[0, 0, :30]		,       _snake_case		,       atol=1e-6	)	)




				def 							snake_case     (							self	):



									"""simple docstring"""





									_lowerCAmelCase      			=					torch.tensor(
									    [2.3804e-03, 2.0752e-03, 1.9836e-03, 2.1057e-03, 1.6174e-03,
									     3.0518e-04, 9.1553e-05, 3.3569e-04, 9.7656e-04, 1.8311e-03,
									     2.0142e-03, 2.1057e-03, 1.7395e-03, 4.5776e-04, -3.9673e-04,
									     4.5776e-04, 1.0071e-03, 9.1553e-05, 4.8828e-04, 1.1597e-03,
									     7.3242e-04, 9.4604e-04, 1.8005e-03, 1.8311e-03, 8.8501e-04,
									     4.2725e-04, 4.8828e-04, 7.3242e-04, 1.0986e-03, 2.1057e-03]	)
									# fmt: on
									_lowerCAmelCase      			=					self._load_datasamples(1	)
									_lowerCAmelCase      			=					[np.tile(input_audio[0][None]		,       reps=(2, 1)	)]
									input_audio[0][1] *= 0.5
									_lowerCAmelCase      			=					EncodecFeatureExtractor(feature_size=2	)
									_lowerCAmelCase      			=					feature_extractor(_snake_case		,       return_tensors="""pt"""	).input_values
									self.assertEquals(input_values.shape		,       (1, 2, 93680)	)
									self.assertTrue(torch.allclose(input_values[0, 0, :30]		,       _snake_case		,       atol=1e-6	)	)
									self.assertTrue(torch.allclose(input_values[0, 1, :30]		,       EXPECTED_INPUT_VALUES * 0.5		,       atol=1e-6	)	)




				def 							snake_case     (							self	):



									"""simple docstring"""





									_lowerCAmelCase      			=					self._load_datasamples(2	)
									# would be easier if the stride was like
									_lowerCAmelCase      			=					EncodecFeatureExtractor(feature_size=1		,       chunk_length_s=1		,       overlap=0.01	)

									# pad and trunc raise an error ?
									with self.assertRaisesRegex(
									    _snake_case		,       """^Both padding and truncation were set. Make sure you only set one.$"""		,       ):
														_lowerCAmelCase      			=					feature_extractor(
														    _snake_case		,       padding="""max_length"""		,       truncation=_snake_case		,       return_tensors="""pt"""	).input_values

									# truncate to chunk
									_lowerCAmelCase      			=					feature_extractor(_snake_case		,       truncation=_snake_case		,       return_tensors="""pt"""	).input_values
									self.assertEquals(truncated_outputs.shape		,       (2, 1, 71520)	)  # 2 chunks

									# force truncate to max_length
									_lowerCAmelCase      			=					feature_extractor(
									    _snake_case		,       truncation=_snake_case		,       max_length=48000		,       return_tensors="""pt"""	).input_values
									self.assertEquals(truncated_outputs.shape		,       (2, 1, 48000)	)

									# pad to chunk
									_lowerCAmelCase      			=					feature_extractor(_snake_case		,       padding=_snake_case		,       return_tensors="""pt"""	).input_values
									self.assertEquals(padded_outputs.shape		,       (2, 1, 95280)	)

									# pad to chunk
									_lowerCAmelCase      			=					feature_extractor(_snake_case		,       return_tensors="""pt"""	).input_values
									self.assertEquals(truncated_outputs.shape		,       (2, 1, 95280)	)

									# force pad to max length
									_lowerCAmelCase      			=					feature_extractor(
									    _snake_case		,       padding="""max_length"""		,       max_length=100000		,       return_tensors="""pt"""	).input_values
									self.assertEquals(truncated_outputs.shape		,       (2, 1, 100000)	)

									# force no pad
									with self.assertRaisesRegex(
									    _snake_case		,       """^Unable to create tensor, you should probably activate padding with 'padding=True' to have batched tensors with the same length.$"""		,       ):
														_lowerCAmelCase      			=					feature_extractor(_snake_case		,       padding=_snake_case		,       return_tensors="""pt"""	).input_values

									_lowerCAmelCase      			=					feature_extractor(input_audio[0]		,       padding=_snake_case		,       return_tensors="""pt"""	).input_values
									self.assertEquals(truncated_outputs.shape		,       (1, 1, 93680)	)

									# no pad if no chunk_length_s
									_lowerCAmelCase      			=					None
									with self.assertRaisesRegex(
									    _snake_case		,       """^Unable to create tensor, you should probably activate padding with 'padding=True' to have batched tensors with the same length.$"""		,       ):
														_lowerCAmelCase      			=					feature_extractor(_snake_case		,       padding=_snake_case		,       return_tensors="""pt"""	).input_values

									_lowerCAmelCase      			=					feature_extractor(input_audio[0]		,       padding=_snake_case		,       return_tensors="""pt"""	).input_values
									self.assertEquals(truncated_outputs.shape		,       (1, 1, 93680)	)

									# no pad if no overlap
									_lowerCAmelCase      			=					2
									_lowerCAmelCase      			=					None
									with self.assertRaisesRegex(
									    _snake_case		,       """^Unable to create tensor, you should probably activate padding with 'padding=True' to have batched tensors with the same length.$"""		,       ):
														_lowerCAmelCase      			=					feature_extractor(_snake_case		,       padding=_snake_case		,       return_tensors="""pt"""	).input_values

									_lowerCAmelCase      			=					feature_extractor(input_audio[0]		,       padding=_snake_case		,       return_tensors="""pt"""	).input_values
									self.assertEquals(truncated_outputs.shape		,       (1, 1, 93680)	)



