fast_spatial ¶

def apply_fast_spatial_upsample(
    frames: Iterable[np.ndarray],
    spatial: FastSpatialPlan,
) -> list[np.ndarray]:
    """Resample decoded stage-1 frames up to the target resolution.

    This runs on decoded RGB frames, *not* on latents: see the module
    docstring for why the latent-space version produced a blurred veil.

    Parameters:
        frames (Iterable[np.ndarray]): Decoded HxWx3 uint8 RGB frames, produced
            by decoding at the stage-one resolution.
        spatial (FastSpatialPlan): Plan defining the target size, interpolation
            kernel, and unsharp strength.

    Returns:
        list[np.ndarray]: Frames at the target resolution. When spatial scaling
            is disabled the frames are returned unchanged, as a list.
    """
    if not spatial.enabled:
        return list(frames)
    return upsample_frames(
        frames,
        width=spatial.target_width,
        height=spatial.target_height,
        mode=spatial.upsample_mode,
        sharpen=spatial.sharpen,
    )

def plan_fast_spatial(
    *,
    height: int,
    width: int,
    num_frames: int,
    spatial_scale: int = 2,
    vae_spatial_compression: int = 8,
    vae_temporal_compression: int = 4,
    patch_size: tuple[int, int, int] = (1, 2, 2),
    upsample_mode: str = DEFAULT_PIXEL_UPSAMPLE_MODE,
    sharpen: float = DEFAULT_FAST_SPATIAL_SHARPEN,
    enabled: bool = True,
) -> FastSpatialPlan:
    """
    Build a plan for reduced-resolution denoising followed by pixel-space upsampling.

    Parameters:
        upsample_mode (str): Pixel interpolation kernel, one of
            :data:`~fastvideo.mlx_runtime.frame_upsample.PIXEL_UPSAMPLE_MODES`.
        sharpen (float): Unsharp strength applied after the resize.

    Returns:
        FastSpatialPlan: The validated spatial-fast processing plan.

    Raises:
        ValueError: If the upsample mode is unsupported or ``sharpen`` is negative.
    """
    if upsample_mode not in PIXEL_UPSAMPLE_MODES:
        raise ValueError(f"Unsupported upsample mode: {upsample_mode!r} "
                         f"(expected one of {', '.join(PIXEL_UPSAMPLE_MODES)})")
    if sharpen < 0.0:
        raise ValueError(f"sharpen must be >= 0, got {sharpen}")
    plan = plan_refine_resolutions(
        height=height,
        width=width,
        num_frames=num_frames,
        spatial_scale=spatial_scale,
        vae_spatial_compression=vae_spatial_compression,
        vae_temporal_compression=vae_temporal_compression,
        patch_size=patch_size,
        enabled=enabled,
        mode_label="fast-spatial",
    )
    if plan.spatial_scale > 1:
        logger.info(
            "[MLX fast-spatial] denoise+decode %dx%d → upsample %dx to %dx%d (%s, sharpen=%.2f)",
            plan.stage1_width,
            plan.stage1_height,
            plan.spatial_scale,
            plan.target_width,
            plan.target_height,
            upsample_mode,
            sharpen,
        )
    return FastSpatialPlan(plan=plan, upsample_mode=upsample_mode, sharpen=sharpen)

def resolve_spatial_mode(
    *,
    refine: bool,
    fast_spatial: bool,
) -> str:
    """Select the active spatial processing mode, with refinement taking precedence.

    Returns:
        str: ``"refine"`` when refinement is enabled, ``"fast_spatial"`` when
            spatial-fast processing is enabled, or ``"off"`` otherwise.
    """
    if refine:
        return "refine"
    if fast_spatial:
        return "fast_spatial"
    return "off"