






"""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 							A(    snake_case_						):
       """simple docstring"""





       lowercase__:							Union[str, Any]         =     "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=2048						):
                            f.write(snake_case_						)






def 							A(    ):
       """simple docstring"""





       lowercase__:							str         =     VivitConfig()

       lowercase__:							List[str]         =     400
       lowercase__:							str         =     "huggingface/label-files"
       lowercase__:							Dict         =     "kinetics400-id2label.json"

       lowercase__:							int         =     json.load(open(hf_hub_download(snake_case_		,  snake_case_		,  repo_type="dataset"						)		,  "r"						)						)
       lowercase__:							Optional[int]         =     {int(snake_case_						): v for k, v in idalabel.items()}
       lowercase__:							Any         =     idalabel
       lowercase__:							List[Any]         =     {v: k for k, v in idalabel.items()}
       return config






def 							A(    ):
       """simple docstring"""





       lowercase__:							Tuple         =     hf_hub_download(
           repo_id="hf-internal-testing/spaghetti-video"		,  filename="eating_spaghetti_32_frames.npy"		,  repo_type="dataset"						)
       lowercase__:							List[Any]         =     np.load(snake_case_						)
       return list(snake_case_						)






def 							A(    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 							A(    snake_case_						):
       """simple docstring"""





       return current.flatten()






def 							A(    snake_case_						):
       """simple docstring"""





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






def 							A(    snake_case_						):
       """simple docstring"""





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






def 							A(    snake_case_		,  snake_case_						):
       """simple docstring"""





       lowercase__:							Union[str, Any]         =     state_dict["optimizer"]["target"]["Transformer"][F"""encoderblock_{i}"""]

       lowercase__:							Tuple         =     F"""encoder.layer.{i}."""
       lowercase__:							Dict         =     {
           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 							A(    snake_case_						):
       """simple docstring"""





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






def 							A(    snake_case_		,  snake_case_=False						):
       """simple docstring"""





       lowercase__:							Optional[int]         =     get_n_layers(snake_case_						)

       lowercase__:							Tuple         =     OrderedDict()

       lowercase__:							Optional[int]         =     state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["bias"]
       lowercase__:							Union[str, Any]         =     state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["scale"]

       lowercase__:							Any         =     np.transpose(
           state_dict["optimizer"]["target"]["embedding"]["kernel"]		,  (4, 3, 0, 1, 2)						)
       lowercase__:							Union[str, Any]         =     state_dict["optimizer"]["target"]["embedding"]["bias"]

       lowercase__:							str         =     state_dict["optimizer"]["target"]["cls"]
       lowercase__:							Optional[int]         =     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:
              lowercase__:							List[Any]         =     {"vivit." + k: v for k, v in new_state.items()}
              lowercase__:							str         =     np.transpose(state_dict["optimizer"]["target"]["output_projection"]["kernel"]						)
              lowercase__:							List[Any]         =     np.transpose(state_dict["optimizer"]["target"]["output_projection"]["bias"]						)

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






def 							A(    ):
       """simple docstring"""





       lowercase__:							Optional[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 							A(    snake_case_						):
       """simple docstring"""





       lowercase__:							Optional[int]         =     "checkpoint"

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

       lowercase__:							int         =     restore_checkpoint(snake_case_		,  snake_case_						)
       lowercase__:							Any         =     transform_state(snake_case_		,  classification_head=snake_case_						)

       lowercase__:							int         =     get_vivit_config()

       assert config.image_size == 224
       assert config.num_frames == 32

       lowercase__:							Any         =     VivitForVideoClassification(snake_case_						)
       model.load_state_dict(snake_case_						)
       model.eval()

       lowercase__:							int         =     get_processor()

       lowercase__:							Optional[Any]         =     prepare_video()
       lowercase__:							List[str]         =     extractor(snake_case_		,  return_tensors="pt"						)

       lowercase__:							Tuple         =     model(**snake_case_						)

       lowercase__:							Any         =     torch.Size([1, 400]						)
       lowercase__:							int         =     torch.tensor([-1.05_43, 2.07_64, -0.21_04, 0.44_39, -0.96_58]						)

       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__":
  UpperCamelCase					=     argparse.ArgumentParser()

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

  UpperCamelCase					=     parser.parse_args()
  convert(args.output_model_name)








import os
import pathlib
import tempfile
import uuid

import numpy as np

from ..utils import is_soundfile_availble, is_torch_available, is_vision_available, logging


lowerCamelCase_				:							Any       =    logging.get_logger(__name__)

if is_vision_available():
 import PIL.Image
 from PIL import Image
 from PIL.Image import Image as ImageType
else:
 lowerCamelCase_				:							Union[str, Any]       =    object

if is_torch_available():
 import torch

if is_soundfile_availble():
 import soundfile as sf







class  _lowerCamelCase			:

     def __init__(  self							,       SCREAMING_SNAKE_CASE_    ):
            __snake_case						       =   value

     def __str__(  self    ):
            return self.to_string()

     def 						__lowerCamelCase				(  self    ):
            logger.error(
                'This is a raw AgentType of unknown type. Display in notebooks and string conversion will be unreliable'    )
            return self._value

     def 						__lowerCamelCase				(  self    ):
            logger.error(
                'This is a raw AgentType of unknown type. Display in notebooks and string conversion will be unreliable'    )
            return str(self._value    )







class  _lowerCamelCase			(lowerCamelCase							,					lowerCamelCase					):

     def 						__lowerCamelCase				(  self    ):
            return self._value

     def 						__lowerCamelCase				(  self    ):
            return self._value







class  _lowerCamelCase			(lowerCamelCase							,					lowerCamelCase					):

     def __init__(  self							,       SCREAMING_SNAKE_CASE_    ):
            super().__init__(SCREAMING_SNAKE_CASE_    )

            if not is_vision_available():
                   raise ImportError('PIL must be installed in order to handle images.'    )

            __snake_case						       =   None
            __snake_case						       =   None
            __snake_case						       =   None

            if isinstance(SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_    ):
                   __snake_case						       =   value
            elif isinstance(SCREAMING_SNAKE_CASE_							,       (str, pathlib.Path)    ):
                   __snake_case						       =   value
            elif isinstance(SCREAMING_SNAKE_CASE_							,       torch.Tensor    ):
                   __snake_case						       =   value
            else:
                   raise ValueError(f'''Unsupported type for {self.__class__.__name__}: {type(SCREAMING_SNAKE_CASE_    )}'''    )

     def 						__lowerCamelCase				(  self							,       SCREAMING_SNAKE_CASE_=None							,       SCREAMING_SNAKE_CASE_=None    ):
            from IPython.display import Image, display

            display(Image(self.to_string()    )    )

     def 						__lowerCamelCase				(  self    ):
            if self._raw is not None:
                   return self._raw

            if self._path is not None:
                   __snake_case						       =   Image.open(self._path    )
                   return self._raw

     def 						__lowerCamelCase				(  self    ):
            if self._path is not None:
                   return self._path

            if self._raw is not None:
                   __snake_case						       =   tempfile.mkdtemp()
                   __snake_case						       =   os.path.join(SCREAMING_SNAKE_CASE_							,       str(uuid.uuida()    ) + '.png'    )
                   self._raw.save(self._path    )

                   return self._path

            if self._tensor is not None:
                   __snake_case						       =   self._tensor.cpu().detach().numpy()

                   # There is likely simpler than load into image into save
                   __snake_case						       =   Image.fromarray((array * 255).astype(np.uinta    )    )

                   __snake_case						       =   tempfile.mkdtemp()
                   __snake_case						       =   os.path.join(SCREAMING_SNAKE_CASE_							,       str(uuid.uuida()    ) + '.png'    )

                   img.save(self._path    )

                   return self._path







class  _lowerCamelCase			(lowerCamelCase					):

     def __init__(  self							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_=16_000    ):
            super().__init__(SCREAMING_SNAKE_CASE_    )

            if not is_soundfile_availble():
                   raise ImportError('soundfile must be installed in order to handle audio.'    )

            __snake_case						       =   None
            __snake_case						       =   None

            __snake_case						       =   samplerate

            if isinstance(SCREAMING_SNAKE_CASE_							,       (str, pathlib.Path)    ):
                   __snake_case						       =   value
            elif isinstance(SCREAMING_SNAKE_CASE_							,       torch.Tensor    ):
                   __snake_case						       =   value
            else:
                   raise ValueError(f'''Unsupported audio type: {type(SCREAMING_SNAKE_CASE_    )}'''    )

     def 						__lowerCamelCase				(  self							,       SCREAMING_SNAKE_CASE_=None							,       SCREAMING_SNAKE_CASE_=None    ):
            from IPython.display import Audio, display

            display(Audio(self.to_string()							,       rate=self.samplerate    )    )

     def 						__lowerCamelCase				(  self    ):
            if self._tensor is not None:
                   return self._tensor

            if self._path is not None:
                   __snake_case       ,       __snake_case						       =   sf.read(self._path    )
                   __snake_case						       =   torch.tensor(SCREAMING_SNAKE_CASE_    )
                   return self._tensor

     def 						__lowerCamelCase				(  self    ):
            if self._path is not None:
                   return self._path

            if self._tensor is not None:
                   __snake_case						       =   tempfile.mkdtemp()
                   __snake_case						       =   os.path.join(SCREAMING_SNAKE_CASE_							,       str(uuid.uuida()    ) + '.wav'    )
                   sf.write(self._path							,       self._tensor							,       samplerate=self.samplerate    )
                   return self._path





lowerCamelCase_				:							Optional[Any]       =    {"text": AgentText, "image": AgentImage, "audio": AgentAudio}
lowerCamelCase_				:							Any       =    {str: AgentText}

if is_vision_available():
 lowerCamelCase_				:							Optional[int]       =    AgentImage
def        __lowercase(							*__snake_case				:		Optional[Any]    ,**__snake_case				:		str					)       ->   Tuple:
       __snake_case						       =   [(arg.to_raw() if isinstance(__snake_case    ,__snake_case					) else arg) for arg in args]
       __snake_case						       =   {k: (v.to_raw() if isinstance(__snake_case    ,__snake_case					) else v) for k, v in kwargs.items()}
       return args, kwargs
def        __lowercase(							__snake_case				:		Optional[int]    ,__snake_case				:		List[str]=None					)       ->   Tuple:
       if isinstance(__snake_case    ,__snake_case					):
              __snake_case						       =   {}
              for i, (k, v) in enumerate(outputs.items()					):
                     if output_types is not None:
                            # If the class has defined outputs, we can map directly according to the class definition
                            if output_types[i] in AGENT_TYPE_MAPPING:
                                   __snake_case						       =   AGENT_TYPE_MAPPING[output_types[i]](__snake_case					)
                            else:
                                   __snake_case						       =   AgentType(__snake_case					)

                     else:
                            # If the class does not have defined output, then we map according to the type
                            for _k, _v in INSTANCE_TYPE_MAPPING.items():
                                   if isinstance(__snake_case    ,_k					):
                                          __snake_case						       =   _v(__snake_case					)
                            if k not in decoded_outputs:
                                   __snake_case						       =   AgentType[v]

       elif isinstance(__snake_case    ,(list, tuple)					):
              __snake_case						       =   type(__snake_case					)()
              for i, v in enumerate(__snake_case					):
                     if output_types is not None:
                            # If the class has defined outputs, we can map directly according to the class definition
                            if output_types[i] in AGENT_TYPE_MAPPING:
                                   decoded_outputs.append(AGENT_TYPE_MAPPING[output_types[i]](__snake_case					)					)
                            else:
                                   decoded_outputs.append(AgentType(__snake_case					)					)
                     else:
                            # If the class does not have defined output, then we map according to the type
                            __snake_case						       =   False
                            for _k, _v in INSTANCE_TYPE_MAPPING.items():
                                   if isinstance(__snake_case    ,_k					):
                                          decoded_outputs.append(_v(__snake_case					)					)
                                          __snake_case						       =   True

                            if not found:
                                   decoded_outputs.append(AgentType(__snake_case					)					)

       else:
              if output_types[0] in AGENT_TYPE_MAPPING:
                     # If the class has defined outputs, we can map directly according to the class definition
                     __snake_case						       =   AGENT_TYPE_MAPPING[output_types[0]](__snake_case					)

              else:
                     # If the class does not have defined output, then we map according to the type
                     for _k, _v in INSTANCE_TYPE_MAPPING.items():
                            if isinstance(__snake_case    ,_k					):
                                   return _v(__snake_case					)
                     return AgentType(__snake_case					)

       return decoded_outputs
