src/fpgm/training/dataset.py

"""``torch.utils.data.Dataset`` over ``vace/window_*`` bundles, in the exact
tensor format DiffSynth-Studio's ``WanTrainingModule`` consumes.

**The tensor-format contract, read from the trainer's own code, not assumed:**
``examples/wanvideo/model_training/train.py``'s ``WanTrainingModule.get_pipeline_inputs``
pulls ``data["video"]`` and ``data["prompt"]`` unconditionally, and (via
``--extra_inputs``) ``data["vace_video"]``, optionally ``data["vace_video_mask"]``,
and ``data["vace_reference_image"][0]``. Each of those is a **list of
``PIL.Image.Image`` (RGB)**, one entry per frame (length 1 for the reference
image) -- confirmed by ``diffsynth/core/data/operators.py``'s ``LoadVideo``/
``LoadImage`` (the stock loaders these fields ordinarily come from) and by
``diffsynth/diffusion/base_pipeline.py``'s ``preprocess_video``/``preprocess_image``,
which convert straight from ``np.array(PIL.Image)`` with **no resize** --
every image handed to the pipeline must already be the target (H, W). That
last point is exactly why :func:`fpgm.training.bundle_io.letterbox_to_white`
exists in this package rather than being left to the trainer: DiffSynth does
none of the resize/crop/letterbox work ``third_party/wan2.1``'s own
``VaceVideoProcessor`` does for the native inference CLI.

``torch`` is imported at module scope only for the ``Dataset`` base class
(no GPU touched at import time) -- matching ``fpgm.utils.gpu``'s existing
precedent for CPU-safe torch imports elsewhere in this codebase.
"""

from __future__ import annotations

from collections.abc import Sequence
from pathlib import Path

import numpy as np
import torch
from PIL import Image

from fpgm.training.augment import apply_augmentations
from fpgm.training.bundle_io import letterbox_to_white, load_window_arrays
from fpgm.training.manifest import discover_windows
from fpgm.training.types import (
    AugmentConfig,
    BundleAssemblyConfig,
    GateFilterConfig,
    ManifestEmptyError,
    WindowSample,
)


def _to_pil_list(frames_rgb: np.ndarray) -> list[Image.Image]:
    """``(T, H, W, 3)`` uint8 -> list of RGB ``PIL.Image``, one per frame."""
    return [Image.fromarray(frames_rgb[t]) for t in range(frames_rgb.shape[0])]


def _largest_4k_plus_1(cap: int) -> int:
    """Largest ``n = 4k+1`` (``k >= 0``) with ``n <= cap``, matching the
    Wan2.1-VACE frame-count contract every exported window already satisfies."""
    if cap < 1:
        raise ValueError(f"cap must be >= 1, got {cap}")
    return cap - ((cap - 1) % 4)


class WindowBundleDataset(torch.utils.data.Dataset):
    """One item = one ``vace/window_*`` bundle, decoded + (optionally) augmented.

    Deliberately **not** built on DiffSynth's ``UnifiedDataset``/metadata.csv
    path: that loader's ``LoadVideo`` re-decodes through a generic
    ``imageio``/ffmpeg reader with resize baked in, which is exactly the
    "just re-encode losslessly-decoded ids through a second lossy hop"
    failure mode :mod:`fpgm.training.bundle_io` exists to avoid. This class
    decodes every channel itself, composites, augments, and hands the
    trainer already-correctly-shaped PIL images -- see the module docstring.

    Batch size is always effectively 1: DiffSynth's own training loop
    (``diffsynth/diffusion/runner.py::launch_training_task``) uses
    ``collate_fn=lambda x: x[0]``, i.e. no batching/padding across samples of
    different length -- consistent with windows all being exactly 81 frames
    here, but not assumed by this class either way.
    """

    def __init__(
        self,
        samples: Sequence[WindowSample],
        assembly_cfg: BundleAssemblyConfig | None = None,
        augment_cfg: AugmentConfig | None = None,
        max_num_frames: int | None = None,
    ):
        """
        Args:
            samples: kept windows, e.g. ``discover_windows(...).kept`` --
                already gate-filtered; this class does not filter further.
            assembly_cfg: how to composite the control channels / letterbox
                the reference image. Defaults to
                :class:`~fpgm.training.types.BundleAssemblyConfig`'s defaults.
            augment_cfg: augmentation magnitudes/toggles. ``None`` disables
                every augmentation (e.g. for eval/validation datasets) --
                distinct from an all-``enable_*=False`` config only in that
                it also skips deriving a per-sample RNG.
            max_num_frames: if set, truncate every window to its first
                ``N`` frames, where ``N`` is the largest value of the form
                ``4k+1`` not exceeding this cap (the contract every window
                already satisfies at 81 -- see
                ``fpgm.datagen.export_vace.compute_windows``). Exists for the
                VRAM-vs-frame-count measurement in the training plan's
                feasibility job, not for normal training runs (``None``
                keeps the full 81-frame window).
        """
        self.samples = list(samples)
        if not self.samples:
            raise ManifestEmptyError(
                "WindowBundleDataset built with zero samples -- did the gate filter "
                "drop everything? See fpgm.training.manifest.discover_windows's log output."
            )
        self.assembly_cfg = assembly_cfg or BundleAssemblyConfig()
        self.augment_cfg = augment_cfg
        if max_num_frames is not None and max_num_frames < 1:
            raise ValueError(f"max_num_frames must be >= 1, got {max_num_frames}")
        self.max_num_frames = max_num_frames

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, index: int) -> dict:
        sample = self.samples[index % len(self.samples)]
        arrays = load_window_arrays(sample, self.assembly_cfg)
        if self.max_num_frames is not None:
            n = _largest_4k_plus_1(min(self.max_num_frames, arrays["target"].shape[0]))
            arrays = {
                **arrays,
                "target": arrays["target"][:n],
                "control": arrays["control"][:n],
                "fg_mask": arrays["fg_mask"][:n],
            }
        control, fg_mask = arrays["control"], arrays["fg_mask"]

        if self.augment_cfg is not None:
            control, fg_mask = apply_augmentations(control, fg_mask, self.augment_cfg, index)

        target_frames = arrays["target"]
        target_h, target_w = target_frames.shape[1], target_frames.shape[2]
        ref_plate = letterbox_to_white(
            arrays["ref_plate"], target_w, target_h, self.assembly_cfg.letterbox_fill
        )

        data: dict = {
            "video": _to_pil_list(target_frames),
            "prompt": arrays["caption"],
            "vace_video": _to_pil_list(control),
            "vace_reference_image": [Image.fromarray(ref_plate)],
            # Not consumed by WanTrainingModule -- kept for logging/debugging
            # (e.g. which window a bad step came from) without a second lookup.
            "_window_name": sample.name,
        }
        if self.assembly_cfg.include_fg_mask:
            mask_rgb = np.repeat((fg_mask[..., None].astype(np.uint8)) * 255, 3, axis=-1)
            data["vace_video_mask"] = _to_pil_list(mask_rgb)
        return data


def build_dataset_from_manifest(
    root: Path,
    gate_filter: GateFilterConfig | None = None,
    assembly_cfg: BundleAssemblyConfig | None = None,
    augment_cfg: AugmentConfig | None = None,
) -> WindowBundleDataset:
    """Scan ``root`` and build a :class:`WindowBundleDataset` from the kept windows.

    Thin composition of :func:`fpgm.training.manifest.discover_windows` +
    :class:`WindowBundleDataset` -- the manifest scan's full drop-reason log
    (see that function's docstring) always runs before construction, so an
    empty/shrunken dataset is diagnosable from the log, not just a stack trace.
    """
    report = discover_windows(root, gate_filter)
    return WindowBundleDataset(report.kept, assembly_cfg=assembly_cfg, augment_cfg=augment_cfg)