def frame_interpolation_cond(input_video: jnp.ndarray, *, cond_frames: int,
                             cond_padding: str, cond_latent_frames: int,
                             latent_shape: tuple[int, int, int]) -> CondDict:
  """Frame interpolation condition.

  Args:
    input_video: [batch, T, H, W, C] input video.
    cond_frames: number of video frames as condition at both ends.
    cond_padding: padding method for masked non-condition video frames.
    cond_latent_frames: number of latent frames as condition at both ends.
    latent_shape: 3D shape of latent code.

  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 < 2 * cond_frames < t
  cond_video_mask = jnp.ones((bs, t, 1, 1, 1), dtype=bool)
  cond_latent_mask = jnp.ones((bs, l_t, 1, 1), dtype=bool)
  cond_video_mask = cond_video_mask.at[:, cond_frames:-cond_frames].set(False)
  cond_latent_mask = cond_latent_mask.at[:, cond_latent_frames:
                                         -cond_latent_frames].set(False)
  if cond_padding == 'interpolate':
    start = input_video[:, cond_frames - 1:cond_frames]
    end = input_video[:, t - cond_frames:t - cond_frames + 1]
    length = t - 2 * cond_frames
    fractions = jnp.arange(1, length + 1) / (length + 1)
    middle = start + (end - start) * fractions[:, None, None, None]
    cond_video = input_video.at[:, cond_frames:-cond_frames].set(middle)
  else:
    raise NotImplementedError(f'Unsupported cond_padding: {cond_padding}')
  assert cond_video.shape == input_video.shape
  return {
      'video': cond_video,
      'video_mask': cond_video_mask,
      'latent_mask': cond_latent_mask
  }
Frame interpolation condition.

  Args:
    input_video: [batch, T, H, W, C] input video.
    cond_frames: number of video frames as condition at both ends.
    cond_padding: padding method for masked non-condition video frames.
    cond_latent_frames: number of latent frames as condition at both ends.
    latent_shape: 3D shape of latent code.

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