def reconstruct_from_tokens(segmented_tokens: np.ndarray,
                            tokenizer_dict: Dict[str, nn.Module],
                            overlapping_frames: int):
  """Reconstructs a long video using the segmented tokens.

  Args:
    segmented_tokens: Array of shape [bs, l_t, l_h, l_w] of overlapping tokens.
    tokenizer_dict: A dictionary of tokenizer and detokenizer
    overlapping_frames: How many frames are overlapping between the windows.

  Returns:
    all_reconstructed: RGB video of shape [bs, t, h, w, 3].
  """
  assert (
      segmented_tokens.ndim == 4 or segmented_tokens.ndim == 5
  ), 'Not in [bs, l_t, l_h, l_w] or [bs, l_t, l_h, l_w, rvq]'
  padded_segmented_tokens = pad(segmented_tokens, jax.device_count())

  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
  all_reconstructed = []
  for i in range(padded_segmented_tokens.shape[0]):
    tokens = padded_segmented_tokens[i]
    tokens = tokens[:, None, ...]  # in [#device, 1, l_t, l_h, l_w]
    reconstructed_videos = tokenizer_dict['detokenizer_pmapped'](tokens)
    reconstructed_videos = jax.device_get(reconstructed_videos)
    reconstructed_videos = np.array(reconstructed_videos, copy=False)  # asarray
    logging.log_first_n(logging.INFO, 'Reconstructed video shape %s.', 1,
                        reconstructed_videos.shape)
    all_reconstructed.append(reconstructed_videos)
  all_reconstructed = np.concatenate(all_reconstructed)
  all_reconstructed = all_reconstructed.reshape(
      [-1, *all_reconstructed.shape[2:]])  # collaps the device dim
  all_reconstructed = all_reconstructed[0:segmented_tokens.shape[0]]

  if overlapping_frames > 0:
    all_reconstructed = train_utils.pool_overlapping_segments(
        all_reconstructed, overlapping_frames)
  return all_reconstructed
Reconstructs a long video using the segmented tokens.

  Args:
    segmented_tokens: Array of shape [bs, l_t, l_h, l_w] of overlapping tokens.
    tokenizer_dict: A dictionary of tokenizer and detokenizer
    overlapping_frames: How many frames are overlapping between the windows.

  Returns:
    all_reconstructed: RGB video of shape [bs, t, h, w, 3].
  
videogvt/interfaces/colab_utils.py