




"""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 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Any:
 A__ =       "https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint"

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







def 						_SCREAMING_SNAKE_CASE     (							)		->	VivitConfig:
 A__ =       VivitConfig()

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

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







def 						_SCREAMING_SNAKE_CASE     (							)		->	Optional[int]:
 A__ =       hf_hub_download(
     repo_id="hf-internal-testing/spaghetti-video"    ,				filename="eating_spaghetti_32_frames.npy"    ,				repo_type="dataset"		)
 A__ =       np.load(lowercase_		)
 return list(lowercase_		)







def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Tuple:
 if np.ndim(lowercase_		) == 2:
  return transform_attention_bias(lowercase_		)

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

 else:
  raise Exception(f"""Invalid number of dimesions: {np.ndim(lowercase_		)}"""		)







def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Any:
 return current.flatten()







def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	str:
 return np.reshape(lowercase_    ,				(current.shape[0], current.shape[1] * current.shape[2])		).T







def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	List[str]:
 return np.reshape(lowercase_    ,				(current.shape[0] * current.shape[1], current.shape[2])		).T







def 						_SCREAMING_SNAKE_CASE     (							lowercase_    ,				lowercase_		)		->	int:
 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 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Any:
 return sum([1 if "encoderblock_" in k else 0 for k in state_dict["optimizer"]["target"]["Transformer"].keys()]		)







def 						_SCREAMING_SNAKE_CASE     (							lowercase_    ,				lowercase_=False		)		->	Union[str, Any]:
 A__ =       get_n_layers(lowercase_		)

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

 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(lowercase_		) for k, v in new_state.items()}







def 						_SCREAMING_SNAKE_CASE     (							)		->	VivitImageProcessor:
 A__ =       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 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Optional[int]:
 A__ =       "checkpoint"

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

 A__ =       restore_checkpoint(lowercase_    ,				lowercase_		)
 A__ =       transform_state(lowercase_    ,				classification_head=lowercase_		)

 A__ =       get_vivit_config()

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

 A__ =       VivitForVideoClassification(lowercase_		)
 model.load_state_dict(lowercase_		)
 model.eval()

 A__ =       get_processor()

 A__ =       prepare_video()
 A__ =       extractor(lowercase_    ,				return_tensors="pt"		)

 A__ =       model(**lowercase_		)

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

 model.save_pretrained(lowercase_		)
 extractor.save_pretrained(lowercase_		)


if __name__ == "__main__":
      SCREAMING_SNAKE_CASE       = argparse.ArgumentParser()

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

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











"""simple docstring"""



from manim import *




class 			UpperCAmelCase_						(		A_							):




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




     '''simple docstring'''

     A__ =       Rectangle(height=0.5			,				width=0.5   )
     A__ =       Rectangle(height=0.46			,				width=0.46   ).set_stroke(width=0   )
     A__ =       Rectangle(height=0.25			,				width=0.25   )

     A__ =       [mem.copy() for i in range(6   )]
     A__ =       [mem.copy() for i in range(6   )]
     A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
     A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
     A__ =       VGroup(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0   )
     A__ =       Text("CPU"			,				font_size=24   )
     A__ =       Group(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0.5			,				aligned_edge=snake_case_   )
     cpu.move_to([-2.5, -0.5, 0]   )
     self.add(snake_case_   )

     A__ =       [mem.copy() for i in range(4   )]
     A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
     A__ =       Text("GPU"			,				font_size=24   )
     A__ =       Group(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0.5			,				aligned_edge=snake_case_   )
     gpu.move_to([-1, -1, 0]   )
     self.add(snake_case_   )

     A__ =       [mem.copy() for i in range(6   )]
     A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )

     A__ =       Text("Model"			,				font_size=24   )
     A__ =       Group(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0.5			,				aligned_edge=snake_case_   )
     model.move_to([3, -1.0, 0]   )
     self.add(snake_case_   )

     A__ =       []
     A__ =       []

     for i, rect in enumerate(snake_case_   ):
      rect.set_stroke(snake_case_   )

      A__ =       Rectangle(height=0.46 / 4			,				width=0.46 / 3   ).set_stroke(width=0.0   ).set_fill(snake_case_			,				opacity=0.7   )

      if i == 0:
       cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT   )			,				buff=0.02			,				direction=snake_case_   )
       cpu_target.set_x(cpu_target.get_x() + 0.1   )
      elif i == 3:
       cpu_target.next_to(model_cpu_arr[0]			,				direction=snake_case_			,				buff=0.0   )
      else:
       cpu_target.next_to(model_cpu_arr[i - 1]			,				direction=snake_case_			,				buff=0.0   )
      self.add(snake_case_   )
      model_cpu_arr.append(snake_case_   )

     self.add(*snake_case_			,				*snake_case_   )

     A__ =       [meta_mem.copy() for i in range(6   )]
     A__ =       [meta_mem.copy() for i in range(6   )]
     A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
     A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
     A__ =       VGroup(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0   )
     A__ =       Text("Disk"			,				font_size=24   )
     A__ =       Group(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0.5			,				aligned_edge=snake_case_   )
     disk.move_to([-4.0, -1.25, 0]   )
     self.add(snake_case_			,				snake_case_   )

     A__ =       []

     for i in range(6   ):
      A__ =       fill.copy().set_fill(snake_case_			,				opacity=0.8   )
      target.move_to(disk_left_col_base[i]   ).scale(0.5   )
      cpu_disk_arr.append(snake_case_   )

     self.add(*snake_case_   )

     A__ =       Square(side_length=2.2   )
     key.move_to([-5, 2, 0]   )

     A__ =       MarkupText(
         F"""<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model"""			,				font_size=18			,				)

     key_text.move_to([-5, 2.4, 0]   )

     self.add(snake_case_			,				snake_case_   )

     A__ =       MarkupText(
         F"""<span fgcolor='{BLUE}'>●</span> Checkpoint"""			,				font_size=18			,				)

     blue_text.next_to(snake_case_			,				DOWN * 2.4			,				aligned_edge=key_text.get_left()   )
     self.add(snake_case_   )

     A__ =       MarkupText(F"""The offloaded weights are all sent to the CPU."""			,				font_size=24   )
     step_a.move_to([2, 2, 0]   )

     self.play(Write(snake_case_			,				run_time=3   )   )

     for i in range(6   ):
      A__ =       cpu_disk_arr[i]
      A__ =       rect.copy().set_fill(snake_case_			,				opacity=0.8   ).scale(2.0   )
      cpa.generate_target()
      cpa.target.move_to(model_base[i]   )

      if i == 0:
       rect.set_fill(snake_case_			,				opacity=0.8   )
       rect.generate_target()
       rect.target.move_to(cpu_left_col_base[0]   ).scale(2.0   )

       self.remove(
           *snake_case_			,				*snake_case_			,				)

      else:
       rect.generate_target()
       rect.target.move_to(cpu_left_col_base[i]   ).scale(2.0   )
      self.play(MoveToTarget(snake_case_   )			,				MoveToTarget(snake_case_   )			,				model_base[i].animate.set_stroke(snake_case_   )   )
     self.play(FadeOut(snake_case_   )   )

     A__ =       MarkupText(
         F"""Finally, hooks are added to each weight in the model\nto transfer the weights from CPU to GPU\n\t\tand back when needed."""			,				font_size=24			,				)
     step_a.move_to([2, 2, 0]   )

     self.play(Write(snake_case_			,				run_time=3   )   )

     A__ =       []
     A__ =       []
     for i in range(6   ):
      A__ =       Arrow(start=snake_case_			,				end=snake_case_			,				color=snake_case_			,				buff=0.5   )
      a.next_to(model_base[i].get_left()			,				snake_case_			,				buff=0.2   )
      arrows.append(snake_case_   )
      animations.append(Write(snake_case_   )   )
     self.play(*snake_case_   )
     self.wait()





