def tokenize_long_video_double_fusion(
    long_video: np.ndarray,
    tokenizer_dict: Dict[str, nn.Module],
    temporal_window: int = 16,
    overlapping_frames: int = 8,
    batch_size: int = 8,
    mode: str = 'early') -> Dict[str, Union[np.ndarray, jnp.ndarray]]:
  """Tokenize a long video using fusion of pixel and embeddings."""
  # Experimental. Needs further research to find the pooling layer.
  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 'detokenizer1_pmapped' not in tokenizer_dict:
    tokenizer_dict['detokenizer1_pmapped'] = jax.pmap(
        functools.partial(tokenizer_dict['detokenizer1']),
        axis_name='device')

  if 'detokenizer2_pmapped' not in tokenizer_dict:
    tokenizer_dict['detokenizer2_pmapped'] = jax.pmap(
        functools.partial(tokenizer_dict['detokenizer2']),
        axis_name='device')

  eval_iter = iter(
      train_utils.LongVideoIterable(long_video, batch_size=batch_size,
                                    temporal_window=temporal_window,
                                    overlapping_frames=overlapping_frames))
  all_tokens = []
  all_embeddings = []
  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))
    embeddings = tokenizer_dict['detokenizer1_pmapped'](tokens)
    embeddings = jax.device_get(embeddings)
    embeddings = np.array(embeddings, copy=False)  # asarray
    all_tokens.append(tokens.reshape(-1, *tokens.shape[3:]))
    all_embeddings.append(embeddings)

  all_tokens = np.concatenate(all_tokens)
  all_embeddings = np.concatenate(all_embeddings)
  if overlapping_frames > 0:
    x = all_embeddings
    x = x.reshape([-1, *x.shape[2:]])  # collaps the device dim
    all_embeddings = train_utils.pool_overlapping_segments(
        x, overlapping_frames, alphas=None)

  # decode second stage
  all_embeddings = pad(all_embeddings, jax.device_count())
  all_reconstructed = []
  for i in range(all_embeddings.shape[0]):
    all_reconstructed.append(tokenizer_dict['detokenizer2_pmapped'](
        all_embeddings[i]))

  all_reconstructed = np.array(all_reconstructed)

  if overlapping_frames and mode == 'double':
    x = all_reconstructed
    x = x.reshape([-1, *x.shape[2:]])  # collaps the device dim
    all_reconstructed = train_utils.pool_overlapping_segments(
        x, overlapping_frames)

  all_reconstructed = all_reconstructed.reshape(-1, *long_video.shape[1:])

  result = dict(
      input=long_video[0:all_reconstructed.shape[0]],
      reconstructed=all_reconstructed,
      tokens=all_tokens)

  return result
Tokenize a long video using fusion of pixel and embeddings.
videogvt/interfaces/colab_utils.py