




'''simple docstring'''
from __future__ import annotations

import inspect
import unittest
from typing import List, Tuple

from transformers import ConvNextConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import TFConvNextForImageClassification, TFConvNextModel


if is_vision_available():
  from PIL import Image

  from transformers import ConvNextImageProcessor






class 			SCREAMING_SNAKE_CASE:

       def __init__(       self   ,						lowerCamelCase__   ,						lowerCamelCase__=13   ,						lowerCamelCase__=32   ,						lowerCamelCase__=3   ,						lowerCamelCase__=4   ,						lowerCamelCase__=[10, 20, 30, 40]   ,						lowerCamelCase__=[2, 2, 3, 2]   ,						lowerCamelCase__=True   ,						lowerCamelCase__=True   ,						lowerCamelCase__=37   ,						lowerCamelCase__="gelu"   ,						lowerCamelCase__=10   ,						lowerCamelCase__=0.02   ,						lowerCamelCase__=3   ,						lowerCamelCase__=None   ,						)       ->   Union[str, Any]:

           """simple docstring"""



           __lowercase     							=		parent
           __lowercase     							=		batch_size
           __lowercase     							=		image_size
           __lowercase     							=		num_channels
           __lowercase     							=		num_stages
           __lowercase     							=		hidden_sizes
           __lowercase     							=		depths
           __lowercase     							=		is_training
           __lowercase     							=		use_labels
           __lowercase     							=		intermediate_size
           __lowercase     							=		hidden_act
           __lowercase     							=		type_sequence_label_size
           __lowercase     							=		initializer_range
           __lowercase     							=		scope

       def        snake_case__     (       self	)       ->   str:

           """simple docstring"""



           __lowercase     							=		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

           __lowercase     							=		None
           if self.use_labels:
               __lowercase     							=		ids_tensor([self.batch_size]   ,						self.type_sequence_label_size	)

           __lowercase     							=		self.get_config()

           return config, pixel_values, labels

       def        snake_case__     (       self	)       ->   Tuple:

           """simple docstring"""



           return ConvNextConfig(
               num_channels=self.num_channels   ,						hidden_sizes=self.hidden_sizes   ,						depths=self.depths   ,						num_stages=self.num_stages   ,						hidden_act=self.hidden_act   ,						is_decoder=lowerCamelCase__   ,						initializer_range=self.initializer_range   ,						)

       def        snake_case__     (       self   ,						lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__	)       ->   Tuple:

           """simple docstring"""



           __lowercase     							=		TFConvNextModel(config=lowerCamelCase__	)
           __lowercase     							=		model(lowerCamelCase__   ,						training=lowerCamelCase__	)
           # expected last hidden states: B, C, H // 32, W // 32
           self.parent.assertEqual(
               result.last_hidden_state.shape   ,						(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)   ,						)

       def        snake_case__     (       self   ,						lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__	)       ->   List[Any]:

           """simple docstring"""



           __lowercase     							=		self.type_sequence_label_size
           __lowercase     							=		TFConvNextForImageClassification(lowerCamelCase__	)
           __lowercase     							=		model(lowerCamelCase__   ,						labels=lowerCamelCase__   ,						training=lowerCamelCase__	)
           self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.type_sequence_label_size)	)




       def        snake_case__     (       self	)       ->   Optional[int]:

           """simple docstring"""



           __lowercase     							=		self.prepare_config_and_inputs()
           __lowercase							,__lowercase							,__lowercase     							=		config_and_inputs
           __lowercase     							=		{"""pixel_values""": pixel_values}
           return config, inputs_dict






@require_tf
class 			SCREAMING_SNAKE_CASE(							__A  , __A  , unittest.TestCase	):
       snake_case_					:	List[Any]       =   (TFConvNextModel, TFConvNextForImageClassification) if is_tf_available() else ()
       snake_case_					:	List[Any]       =   (
           {"""feature-extraction""": TFConvNextModel, """image-classification""": TFConvNextForImageClassification}
           if is_tf_available()
           else {}
       )

       snake_case_					:	List[str]       =   False
       snake_case_					:	List[str]       =   False
       snake_case_					:	List[str]       =   False
       snake_case_					:	Tuple       =   False
       snake_case_					:	Tuple       =   False

       def        snake_case__     (       self	)       ->   Dict:

           """simple docstring"""



           __lowercase     							=		TFConvNextModelTester(self	)
           __lowercase     							=		ConfigTester(
               self   ,						config_class=lowerCamelCase__   ,						has_text_modality=lowerCamelCase__   ,						hidden_size=37   ,						)

       @unittest.skip(reason="""ConvNext does not use inputs_embeds"""	)
       def        snake_case__     (       self	)       ->   str:

           """simple docstring"""



           pass

       @unittest.skipIf(
           not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""	)	) == 0   ,						reason="""TF does not support backprop for grouped convolutions on CPU."""   ,						)
       @slow
       def        snake_case__     (       self	)       ->   List[Any]:

           """simple docstring"""



           super().test_keras_fit()

       @unittest.skip(reason="""ConvNext does not support input and output embeddings"""	)
       def        snake_case__     (       self	)       ->   Optional[int]:

           """simple docstring"""



           pass

       def        snake_case__     (       self	)       ->   int:

           """simple docstring"""



           __lowercase							,__lowercase     							=		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               __lowercase     							=		model_class(lowerCamelCase__	)
               __lowercase     							=		inspect.signature(model.call	)
               # signature.parameters is an OrderedDict => so arg_names order is deterministic
               __lowercase     							=		[*signature.parameters.keys()]

               __lowercase     							=		["""pixel_values"""]
               self.assertListEqual(arg_names[:1]   ,						lowerCamelCase__	)

       def        snake_case__     (       self	)       ->   List[Any]:

           """simple docstring"""



           __lowercase     							=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*lowerCamelCase__	)

       @unittest.skipIf(
           not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""	)	) == 0   ,						reason="""TF does not support backprop for grouped convolutions on CPU."""   ,						)
       def        snake_case__     (       self	)       ->   Dict:

           """simple docstring"""



           super().test_dataset_conversion()

       def        snake_case__     (       self	)       ->   Optional[int]:

           """simple docstring"""



           def check_hidden_states_output(lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__	):
               __lowercase     							=		model_class(lowerCamelCase__	)

               __lowercase     							=		model(**self._prepare_for_class(lowerCamelCase__   ,						lowerCamelCase__	)	)
               __lowercase     							=		outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

               __lowercase     							=		self.model_tester.num_stages
               self.assertEqual(len(lowerCamelCase__	)   ,						expected_num_stages + 1	)

               # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
               self.assertListEqual(
                   list(hidden_states[0].shape[-2:]	)   ,						[self.model_tester.image_size // 4, self.model_tester.image_size // 4]   ,						)

           __lowercase							,__lowercase     							=		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               __lowercase     							=		True
               check_hidden_states_output(lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__	)

               # check that output_hidden_states also work using config
               del inputs_dict["output_hidden_states"]
               __lowercase     							=		True

               check_hidden_states_output(lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__	)

       def        snake_case__     (       self	)       ->   List[str]:

           """simple docstring"""



           __lowercase							,__lowercase     							=		self.model_tester.prepare_config_and_inputs_for_common()

           def check_equivalence(lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__={}	):
               __lowercase     							=		model(lowerCamelCase__   ,						return_dict=lowerCamelCase__   ,						**lowerCamelCase__	)
               __lowercase     							=		model(lowerCamelCase__   ,						return_dict=lowerCamelCase__   ,						**lowerCamelCase__	).to_tuple()

               def recursive_check(lowerCamelCase__   ,						lowerCamelCase__	):
                   if isinstance(lowerCamelCase__   ,						(List, Tuple)	):
                       for tuple_iterable_value, dict_iterable_value in zip(lowerCamelCase__   ,						lowerCamelCase__	):
                           recursive_check(lowerCamelCase__   ,						lowerCamelCase__	)
                   elif tuple_object is None:
                       return
                   else:
                       self.assertTrue(
                           all(tf.equal(lowerCamelCase__   ,						lowerCamelCase__	)	)   ,						msg=(
                               """Tuple and dict output are not equal. Difference:"""
                               F' {tf.math.reduce_max(tf.abs(tuple_object - dict_object	)	)}'
                           )   ,						)

                   recursive_check(lowerCamelCase__   ,						lowerCamelCase__	)

           for model_class in self.all_model_classes:
               __lowercase     							=		model_class(lowerCamelCase__	)

               __lowercase     							=		self._prepare_for_class(lowerCamelCase__   ,						lowerCamelCase__	)
               __lowercase     							=		self._prepare_for_class(lowerCamelCase__   ,						lowerCamelCase__	)
               check_equivalence(lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__	)

               __lowercase     							=		self._prepare_for_class(lowerCamelCase__   ,						lowerCamelCase__   ,						return_labels=lowerCamelCase__	)
               __lowercase     							=		self._prepare_for_class(lowerCamelCase__   ,						lowerCamelCase__   ,						return_labels=lowerCamelCase__	)
               check_equivalence(lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__	)

               __lowercase     							=		self._prepare_for_class(lowerCamelCase__   ,						lowerCamelCase__	)
               __lowercase     							=		self._prepare_for_class(lowerCamelCase__   ,						lowerCamelCase__	)
               check_equivalence(lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__   ,						{"""output_hidden_states""": True}	)

               __lowercase     							=		self._prepare_for_class(lowerCamelCase__   ,						lowerCamelCase__   ,						return_labels=lowerCamelCase__	)
               __lowercase     							=		self._prepare_for_class(lowerCamelCase__   ,						lowerCamelCase__   ,						return_labels=lowerCamelCase__	)
               check_equivalence(lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__   ,						{"""output_hidden_states""": True}	)

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

           """simple docstring"""



           __lowercase     							=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*lowerCamelCase__	)




       @slow
       def        snake_case__     (       self	)       ->   Union[str, Any]:

           """simple docstring"""



           __lowercase     							=		TFConvNextModel.from_pretrained("""facebook/convnext-tiny-224"""	)
           self.assertIsNotNone(lowerCamelCase__	)

def 							snake_case_    (  ):



    """simple docstring"""




    __lowercase     							=		Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""     )
    return image






@require_tf
@require_vision
class 			SCREAMING_SNAKE_CASE(							unittest.TestCase	):

       @cached_property
       def        snake_case__     (       self	)       ->   List[Any]:

           """simple docstring"""



           return ConvNextImageProcessor.from_pretrained("""facebook/convnext-tiny-224"""	) if is_vision_available() else None




       @slow
       def        snake_case__     (       self	)       ->   Any:

           """simple docstring"""



           __lowercase     							=		TFConvNextForImageClassification.from_pretrained("""facebook/convnext-tiny-224"""	)

           __lowercase     							=		self.default_image_processor
           __lowercase     							=		prepare_img()
           __lowercase     							=		image_processor(images=lowerCamelCase__   ,						return_tensors="""tf"""	)

           # forward pass
           __lowercase     							=		model(**lowerCamelCase__	)

           # verify the logits
           __lowercase     							=		tf.TensorShape((1, 1000)	)
           self.assertEqual(outputs.logits.shape   ,						lowerCamelCase__	)

           __lowercase     							=		tf.constant([-0.02_60, -0.47_39, 0.19_11]	)

           tf.debugging.assert_near(outputs.logits[0, :3]   ,						lowerCamelCase__   ,						atol=1E-4	)










'''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_    (  a__						:      Any     ):



    """simple docstring"""




    __lowercase     							=		"""https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint"""

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



    """simple docstring"""




    __lowercase     							=		VivitConfig()

    __lowercase     							=		4_00
    __lowercase     							=		"""huggingface/label-files"""
    __lowercase     							=		"""kinetics400-id2label.json"""

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



    """simple docstring"""




    __lowercase     							=		hf_hub_download(
        repo_id="""hf-internal-testing/spaghetti-video"""       ,filename="""eating_spaghetti_32_frames.npy"""       ,repo_type="""dataset"""     )
    __lowercase     							=		np.load(a__     )
    return list(a__     )
def 							snake_case_    (  a__						:      np.ndarray     ):



    """simple docstring"""




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

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

    else:
        raise Exception(f'Invalid number of dimesions: {np.ndim(a__     )}'     )
def 							snake_case_    (  a__						:      np.ndarray     ):



    """simple docstring"""




    return current.flatten()
def 							snake_case_    (  a__						:      np.ndarray     ):



    """simple docstring"""




    return np.reshape(a__       ,(current.shape[0], current.shape[1] * current.shape[2])     ).T
def 							snake_case_    (  a__						:      np.ndarray     ):



    """simple docstring"""




    return np.reshape(a__       ,(current.shape[0] * current.shape[1], current.shape[2])     ).T
def 							snake_case_    (  a__						:      Tuple       ,a__						:      Dict     ):



    """simple docstring"""




    __lowercase     							=		state_dict["""optimizer"""]["""target"""]["""Transformer"""][f'encoderblock_{i}']

    __lowercase     							=		f'encoder.layer.{i}.'
    __lowercase     							=		{
        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_    (  a__						:      str     ):



    """simple docstring"""




    return sum([1 if """encoderblock_""" in k else 0 for k in state_dict["""optimizer"""]["""target"""]["""Transformer"""].keys()]     )
def 							snake_case_    (  a__						:      Any       ,a__						:      int=False     ):



    """simple docstring"""




    __lowercase     							=		get_n_layers(a__     )

    __lowercase     							=		OrderedDict()

    __lowercase     							=		state_dict["""optimizer"""]["""target"""]["""Transformer"""]["""encoder_norm"""]["""bias"""]
    __lowercase     							=		state_dict["""optimizer"""]["""target"""]["""Transformer"""]["""encoder_norm"""]["""scale"""]

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

    __lowercase     							=		state_dict["""optimizer"""]["""target"""]["""cls"""]
    __lowercase     							=		state_dict["""optimizer"""]["""target"""]["""Transformer"""]["""posembed_input"""][
        """pos_embedding"""
    ]

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

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

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



    """simple docstring"""




    __lowercase     							=		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_    (  a__						:      str     ):



    """simple docstring"""




    __lowercase     							=		"""checkpoint"""

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

    __lowercase     							=		restore_checkpoint(a__       ,a__     )
    __lowercase     							=		transform_state(a__       ,classification_head=a__     )

    __lowercase     							=		get_vivit_config()

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

    __lowercase     							=		VivitForVideoClassification(a__     )
    model.load_state_dict(a__     )
    model.eval()

    __lowercase     							=		get_processor()

    __lowercase     							=		prepare_video()
    __lowercase     							=		extractor(a__       ,return_tensors="""pt"""     )

    __lowercase     							=		model(**a__     )

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

    model.save_pretrained(a__     )
    extractor.save_pretrained(a__     )


if __name__ == "__main__":
  A				: Union[str, Any]					      =			argparse.ArgumentParser()

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

  A				: Optional[Any]					      =			parser.parse_args()
  convert(args.output_model_name)




