def get_latent_frame_prediction_task(config: ml_collections.ConfigDict,
                                     is_train: bool, task_id: str) -> Task:
  """Get latent frame prediction task setup."""
  task_token = SpecialTokens.TASK_FULL_GENERATION
  task_config = ml_collections.FrozenConfigDict(
      config.get(task_id, ml_collections.ConfigDict()))
  weight_mode = task_config.get('weight_mode', config.get('weight_mode'))
  block_shape = task_config.get('mask_block_shape', (1, 1, 1))
  if not is_train:
    block_shape = task_config.get('eval_mask_block_shape', block_shape)
  cond_latent_frames = task_config.get('cond_latent_frames', 1)
  mask_fn = functools.partial(
      mask_utils.random_block_mask,
      condition_mode='cond->cond',
      condition_segment_id=1,
      weight_mode=weight_mode,
      block_shape=block_shape)
  cond_fn = functools.partial(
      cond_utils.latent_frame_prediction_cond,
      cond_latent_frames=cond_latent_frames,
      latent_shape=config.transformer.latent_shape,
  )
  return Task(task_id, task_token, mask_fn, cond_fn, task_config)
Get latent frame prediction task setup.
videogvt/train_lib/task_registry.py