def tokenize_long_video(
    long_video: np.ndarray,
    tokenizer_dict: Dict[str, nn.Module],
    temporal_window: int = 16,
    overlapping_frames: int = 0,
    batch_size: int = 8) -> Dict[str, Union[np.ndarray, jnp.ndarray]]:
  """Tokenize and detokenize a long video using the tokenizer_dict.

  This is supposed to be called only once and hence not optimized for
  efficiency.

  Args:
    long_video: A single long video of shape [frames, h, w, c]
    tokenizer_dict: A dictionary of tokenizer and detokenizer
    temporal_window: Temporal window to chunk the input video for 3D Conv.
    overlapping_frames: How many frames are overlapping between the windows.
    batch_size: Batch size to run the tokenizer. Must be divisiable by
      jax.device_count().

  Returns:
    A dictionary of original video, detokenized/reconstructed video,
      and corresponding visual codes. It will also return pmaped tokenizer,
      and detokenizer if return_pmapped_func=True.
  """
  assert long_video.ndim == 4, 'Video format should be [#frames, h, w, c].'
  assert np.min(long_video) >= 0 and np.max(
      long_video) < 1 + 1e-1, 'Video value range is wrong.'
  assert batch_size % jax.device_count() == 0, 'Not divisiable by device_count'

  if 'tokenizer_pmapped' not in tokenizer_dict:
    tokenizer_pmapped = jax.pmap(
        functools.partial(mm_vq.tokenize, tokenizer_dict=tokenizer_dict),
        axis_name='device')
    tokenizer_dict['tokenizer_pmapped'] = tokenizer_pmapped

  if 'detokenizer_pmapped' not in tokenizer_dict:
    detokenizer_pmapped = jax.pmap(
        functools.partial(mm_vq.detokenize, tokenizer_dict=tokenizer_dict),
        axis_name='device')
    tokenizer_dict['detokenizer_pmapped'] = detokenizer_pmapped

  eval_iter = iter(
      train_utils.LongVideoIterable(long_video, batch_size=batch_size,
                                    temporal_window=temporal_window,
                                    overlapping_frames=overlapping_frames))
  all_tokens = []
  # Uses np not jnp as much as possible to save TPU memory.
  local_batch_size = batch_size // jax.device_count()
  for train_batch in eval_iter:
    train_batch = jax.tree_util.tree_map(lambda x: np.reshape(  # pylint: disable=g-long-lambda
        x, [-1, local_batch_size, *x.shape[1:]]), train_batch)
    tokens = jax.device_get(tokenizer_dict['tokenizer_pmapped'](train_batch))

    logging.log_first_n(logging.INFO, 'Input video shape %s.', 1,
                        train_batch['inputs'].shape)
    logging.log_first_n(logging.INFO, 'Tokens shape %s.', 1, tokens.shape)
    all_tokens.append(tokens.reshape(-1, *tokens.shape[2:]))
  all_tokens = np.concatenate(all_tokens)

  all_reconstructed = reconstruct_from_tokens(all_tokens, tokenizer_dict,
                                              overlapping_frames)
  all_reconstructed = all_reconstructed.reshape(-1, *long_video.shape[1:])
  video_len = min(all_reconstructed.shape[0], long_video.shape[0])
  result = dict(
      input=long_video[0:video_len],
      reconstructed=all_reconstructed[0:video_len],
      tokens=all_tokens)
  return result
Tokenize and detokenize a long video using the tokenizer_dict.

  This is supposed to be called only once and hence not optimized for
  efficiency.

  Args:
    long_video: A single long video of shape [frames, h, w, c]
    tokenizer_dict: A dictionary of tokenizer and detokenizer
    temporal_window: Temporal window to chunk the input video for 3D Conv.
    overlapping_frames: How many frames are overlapping between the windows.
    batch_size: Batch size to run the tokenizer. Must be divisiable by
      jax.device_count().

  Returns:
    A dictionary of original video, detokenized/reconstructed video,
      and corresponding visual codes. It will also return pmaped tokenizer,
      and detokenizer if return_pmapped_func=True.
  
videogvt/interfaces/colab_utils.py