





'''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      UpperCamelCase__   (	lowerCAmelCase						):
      _lowerCAmelCase							  =							"""https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint"""

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

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

      _lowerCAmelCase							  =							json.load(open(hf_hub_download(lowerCAmelCase							,			lowerCAmelCase							,			repo_type="""dataset"""						)							,			"""r"""						)						)
      _lowerCAmelCase							  =							{int(lowerCAmelCase						): v for k, v in idalabel.items()}
      _lowerCAmelCase							  =							idalabel
      _lowerCAmelCase							  =							{v: k for k, v in idalabel.items()}
      return config
def      UpperCamelCase__   (	):
      _lowerCAmelCase							  =							hf_hub_download(
          repo_id="""hf-internal-testing/spaghetti-video"""							,			filename="""eating_spaghetti_32_frames.npy"""							,			repo_type="""dataset"""						)
      _lowerCAmelCase							  =							np.load(lowerCAmelCase						)
      return list(lowerCAmelCase						)
def      UpperCamelCase__   (	lowerCAmelCase						):


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

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

      else:
            raise Exception(f"Invalid number of dimesions: {np.ndim(lowerCAmelCase						)}"						)
def      UpperCamelCase__   (	lowerCAmelCase						):
      return current.flatten()
def      UpperCamelCase__   (	lowerCAmelCase						):
      return np.reshape(lowerCAmelCase							,			(current.shape[0], current.shape[1] * current.shape[2])						).T
def      UpperCamelCase__   (	lowerCAmelCase						):
      return np.reshape(lowerCAmelCase							,			(current.shape[0] * current.shape[1], current.shape[2])						).T
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase						):
      _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__   (	lowerCAmelCase						):
      return sum([1 if """encoderblock_""" in k else 0 for k in state_dict["""optimizer"""]["""target"""]["""Transformer"""].keys()]						)
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase=False						):
      _lowerCAmelCase							  =							get_n_layers(lowerCAmelCase						)

      _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(lowerCAmelCase						):
            new_state.update(transform_state_encoder_block(lowerCAmelCase							,			lowerCAmelCase						)						)

      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(lowerCAmelCase						) for k, v in new_state.items()}
def      UpperCamelCase__   (	):
      _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__   (	lowerCAmelCase						):
      _lowerCAmelCase							  =							"""checkpoint"""

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

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

      _lowerCAmelCase							  =							get_vivit_config()

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

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

      _lowerCAmelCase							  =							get_processor()

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

      _lowerCAmelCase							  =							model(**lowerCAmelCase						)

      _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]							,			lowerCAmelCase							,			atol=1e-4						), outputs.logits[0, :5]

      model.save_pretrained(lowerCAmelCase						)
      extractor.save_pretrained(lowerCAmelCase						)


if __name__ == "__main__":
      A__       :     Tuple          =argparse.ArgumentParser()

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

      A__       :     Union[str, Any]          =parser.parse_args()
      convert(args.output_model_name)













'''simple docstring'''






import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
      import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


A__       :     Optional[Any]          =get_logger(__name__)
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase = None							,			):





      """simple docstring"""


      for i, state in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							f"{MODEL_NAME}.bin" if i == 0 else f"{MODEL_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            save(lowerCAmelCase							,			lowerCAmelCase						)
            logger.info(f"Model weights saved in {output_model_file}"						)
      # Optimizer states
      for i, opt in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							opt.state_dict()
            _lowerCAmelCase							  =							f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            save(lowerCAmelCase							,			lowerCAmelCase						)
            logger.info(f"Optimizer state saved in {output_optimizer_file}"						)
      # Scheduler states
      for i, scheduler in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							scheduler.state_dict()
            _lowerCAmelCase							  =							f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            save(lowerCAmelCase							,			lowerCAmelCase						)
            logger.info(f"Scheduler state saved in {output_scheduler_file}"						)
      # GradScaler state
      if scaler is not None:
            _lowerCAmelCase							  =							scaler.state_dict()
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            torch.save(lowerCAmelCase							,			lowerCAmelCase						)
            logger.info(f"Gradient scaler state saved in {output_scaler_file}"						)
      # Random number generator states
      _lowerCAmelCase							  =							{}
      _lowerCAmelCase							  =							f"{RNG_STATE_NAME}_{process_index}.pkl"
      _lowerCAmelCase							  =							random.getstate()
      _lowerCAmelCase							  =							np.random.get_state()
      _lowerCAmelCase							  =							torch.get_rng_state()
      if is_xpu_available():
            _lowerCAmelCase							  =							torch.xpu.get_rng_state_all()
      else:
            _lowerCAmelCase							  =							torch.cuda.get_rng_state_all()
      if is_tpu_available():
            _lowerCAmelCase							  =							xm.get_rng_state()
      _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
      torch.save(lowerCAmelCase							,			lowerCAmelCase						)
      logger.info(f"Random states saved in {output_states_file}"						)
      return output_dir
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase=None							,			lowerCAmelCase=None							,			**lowerCAmelCase							,			):





      """simple docstring"""


      if map_location not in [None, "cpu", "on_device"]:
            raise TypeError(
                """Unsupported optimizer map location passed, please choose one of `None`, `'cpu'`, or `'on_device'`"""						)
      if map_location is None:
            _lowerCAmelCase							  =							"""cpu"""
      elif map_location == "on_device":
            _lowerCAmelCase							  =							PartialState().device
      # Model states
      for i, model in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							f"{MODEL_NAME}.bin" if i == 0 else f"{MODEL_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            models[i].load_state_dict(torch.load(lowerCAmelCase							,			map_location=lowerCAmelCase						)							,			**lowerCAmelCase						)
      logger.info("""All model weights loaded successfully"""						)

      # Optimizer states
      for i, opt in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            _lowerCAmelCase							  =							torch.load(lowerCAmelCase							,			map_location=lowerCAmelCase						)
            optimizers[i].load_state_dict(lowerCAmelCase						)
      logger.info("""All optimizer states loaded successfully"""						)

      # Scheduler states
      for i, scheduler in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            scheduler.load_state_dict(torch.load(lowerCAmelCase						)						)
      logger.info("""All scheduler states loaded successfully"""						)

      # GradScaler state
      if scaler is not None:
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            scaler.load_state_dict(torch.load(lowerCAmelCase						)						)
            logger.info("""GradScaler state loaded successfully"""						)

      # Random states
      try:
            _lowerCAmelCase							  =							torch.load(os.path.join(lowerCAmelCase							,			f"{RNG_STATE_NAME}_{process_index}.pkl"						)						)
            random.setstate(states["""random_state"""]						)
            np.random.set_state(states["""numpy_random_seed"""]						)
            torch.set_rng_state(states["""torch_manual_seed"""]						)
            if is_xpu_available():
                  torch.xpu.set_rng_state_all(states["""torch_xpu_manual_seed"""]						)
            else:
                  torch.cuda.set_rng_state_all(states["""torch_cuda_manual_seed"""]						)
            if is_tpu_available():
                  xm.set_rng_state(states["""xm_seed"""]						)
            logger.info("""All random states loaded successfully"""						)
      except Exception:
            logger.info("""Could not load random states"""						)
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase = 0						):





      """simple docstring"""


      _lowerCAmelCase							  =							Path(lowerCAmelCase						) / f"custom_checkpoint_{index}.pkl"
      logger.info(f"Saving the state of {get_pretty_name(lowerCAmelCase						)} to {save_location}"						)
      torch.save(obj.state_dict()							,			lowerCAmelCase						)





def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase = 0						):





      """simple docstring"""


      _lowerCAmelCase							  =							f"{path}/custom_checkpoint_{index}.pkl"
      logger.info(f"Loading the state of {get_pretty_name(lowerCAmelCase						)} from {load_location}"						)
      obj.load_state_dict(torch.load(lowerCAmelCase							,			map_location="""cpu"""						)						)






