hao-ai-lab-FastVideo/fastvideo/mlx_runtime/fast_spatial.py

# SPDX-License-Identifier: Apache-2.0
"""Spatial fast mode for the MLX Wan runtime (RIFE's spatial twin).

RIFE ``--fast`` cuts *frames* (temporal). This module cuts *pixels*
(spatial): denoise at ``target // scale``, decode at that size, then
resample the decoded frames up to the target. No second denoise pass —
that is ``--refine`` (quality). The two compose:

* ``--fast-spatial`` alone → speed (≈ scale² fewer tokens)
* ``--refine`` alone → quality two-pass (H3 / LTX-2)
* ``--fast`` + ``--refine`` → fewer frames at base res, full-res refine
* ``--fast`` + ``--fast-spatial`` → fewer frames *and* fewer pixels

The upsample runs in **pixel** space, after the VAE decode. It used to run
in latent space (bilinear over the latent H/W plane, sharing the refine
hand-off primitive) and that is what made spatial fast mode incoherent: an
interpolated Wan latent is off the decoder's manifold, so decode returned
the right silhouette under a smeared veil. ``--refine`` can get away with
the latent-space upsample because a second DMD pass re-denoises the result;
spatial fast mode hands the latent straight to the decoder, so it cannot.
See :mod:`fastvideo.mlx_runtime.frame_upsample` for the full rationale.

MetalFX is intentionally not used: it needs game-engine motion vectors
and depth that diffusion output lacks.
"""

from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass

import numpy as np

from fastvideo.logger import init_logger
from fastvideo.mlx_runtime.frame_upsample import (
    DEFAULT_PIXEL_UPSAMPLE_MODE,
    PIXEL_UPSAMPLE_MODES,
    upsample_frames,
)
from fastvideo.mlx_runtime.refine import RefinePlan, plan_refine_resolutions

logger = init_logger(__name__)

# Resampling from a smaller decode loses high-frequency detail the same way
# RIFE's flow warp does, so spatial fast mode borrows ``--fast``'s remedy: a
# light unsharp pass. 0.4 recovers perceived crispness on Wan2.1 output at 2x
# without the halos that show up by ~0.8.
DEFAULT_FAST_SPATIAL_SHARPEN = 0.4


@dataclass(frozen=True)
class FastSpatialPlan:
    """Resolved geometry for a spatial-fast (upsample-only) run."""

    plan: RefinePlan
    upsample_mode: str
    sharpen: float = DEFAULT_FAST_SPATIAL_SHARPEN

    @property
    def enabled(self) -> bool:
        """
        Determine whether spatial scaling is enabled.

        Returns:
            `true` if the spatial scale is greater than one, `false` otherwise.
        """
        return self.plan.spatial_scale > 1

    @property
    def scale(self) -> int:
        """Provides the configured spatial scaling factor.

        Returns:
            int: The spatial scaling factor.
        """
        return self.plan.spatial_scale

    @property
    def target_height(self) -> int:
        """
        Return the target output height for the spatial plan.

        Returns:
            int: Target output height in pixels.
        """
        return self.plan.target_height

    @property
    def target_width(self) -> int:
        """Return the target image width in pixels.

        Returns:
            int: The target image width.
        """
        return self.plan.target_width

    @property
    def stage1_height(self) -> int:
        """
        Provide the stage-one latent height used for reduced-resolution processing.

        Returns:
            int: The stage-one latent height.
        """
        return self.plan.stage1_height

    @property
    def stage1_width(self) -> int:
        """Get the stage-one latent width.

        Returns:
            int: The stage-one latent width.
        """
        return self.plan.stage1_width


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 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 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"


__all__ = [
    "DEFAULT_FAST_SPATIAL_SHARPEN",
    "FastSpatialPlan",
    "apply_fast_spatial_upsample",
    "plan_fast_spatial",
    "resolve_spatial_mode",
]