def frame_prediction_cond(input_video: jnp.ndarray,
                          *,
                          cond_frames: int,
                          cond_padding: str,
                          cond_latent_frames: int,
                          latent_shape: tuple[int, int, int],
                          prefix_condition: bool = False) -> CondDict:
  """Frame prediction condition.

  Args:
    input_video: [batch, T, H, W, C] input video.
    cond_frames: number of video frames as condition.
    cond_padding: padding method for masked non-condition video frames.
    cond_latent_frames: number of latent frames as condition.
    latent_shape: 3D shape of latent code.
    prefix_condition: whether condition is prefix or inside latent code. When
      inside, the condition region is reflected in returned masks.

  Returns:
    Dictionary containing the condition video, mask of condition video,
    and mask of condition latent.
  """
  assert input_video.ndim == 5  # bs, t, h, w, c
  bs, t, _, _, _ = input_video.shape
  l_t, _, _ = latent_shape
  assert 0 < cond_frames < t
  cond_video_mask = jnp.zeros((bs, t, 1, 1, 1), dtype=bool)
  cond_latent_mask = jnp.zeros((bs, l_t, 1, 1), dtype=bool)
  if not prefix_condition:
    cond_video_mask = cond_video_mask.at[:, :cond_frames].set(True)
    cond_latent_mask = cond_latent_mask.at[:, :cond_latent_frames].set(True)
  pad = ((0, 0), (0, t - cond_frames), (0, 0), (0, 0), (0, 0))
  cond_video = jnp.pad(input_video[:, :cond_frames], pad, mode=cond_padding)
  assert cond_video.shape == input_video.shape
  return {
      'video': cond_video,
      'video_mask': cond_video_mask,
      'latent_mask': cond_latent_mask
  }
Frame prediction condition.

  Args:
    input_video: [batch, T, H, W, C] input video.
    cond_frames: number of video frames as condition.
    cond_padding: padding method for masked non-condition video frames.
    cond_latent_frames: number of latent frames as condition.
    latent_shape: 3D shape of latent code.
    prefix_condition: whether condition is prefix or inside latent code. When
      inside, the condition region is reflected in returned masks.

  Returns:
    Dictionary containing the condition video, mask of condition video,
    and mask of condition latent.
  
videogvt/train_lib/cond_utils.py