src/fpgm/training/types.py

"""Data contracts for the training package: manifest entries + configs.

Mirrors the discipline of :mod:`fpgm.types` and :mod:`fpgm.datagen.types`: every
stage (manifest discovery, bundle decode, augmentation, dataset assembly)
consumes/produces a dataclass defined here, exceptions are named and rooted in
:class:`fpgm.types.FpgmError`, and shape/contract violations fail loudly at
construction rather than three stages downstream.

No heavy imports here (no torch/cv2/DiffSynth) -- this module must import
cleanly on a machine with no GPU and no conda env active, since the manifest
filter (which only reads ``sample.json``) is exercised by CI-style unit tests.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path

from fpgm.types import FpgmError

# --------------------------------------------------------------------------- #
# Exceptions
# --------------------------------------------------------------------------- #


class TrainingError(FpgmError):
    """Base class for every error raised by :mod:`fpgm.training`."""


class BundleReadError(TrainingError):
    """A window bundle's files are missing, unreadable, or violate the shape
    contract (wrong resolution/frame count/dtype).

    Distinct from a gate *filter* decision (which is a policy choice recorded
    in ``dropped_by_reason``): this is raised for bundles that claim to have
    passed export but are actually broken on disk -- a bug, not a quality flag.
    """


class ManifestEmptyError(TrainingError):
    """The gate filter kept zero windows.

    Raised rather than silently constructing an empty dataset (which would
    make ``accelerate launch`` hang or crash confusingly many steps later) --
    see the module docstring's "a silently shrunken training set is a bug we
    would not notice" concern in the owning plan.
    """


# --------------------------------------------------------------------------- #
# Manifest
# --------------------------------------------------------------------------- #


#: Every gate key ``sample.json["gates_passed"]`` is currently known to carry
#: (see ``fpgm.datagen.export_vace``'s ``sample.json`` writer). Kept as an
#: explicit tuple rather than "whatever keys happen to be in the dict" so a
#: gate silently renamed/dropped upstream fails the manifest scan loudly
#: instead of just quietly no longer being checked.
KNOWN_GATE_KEYS: tuple[str, ...] = (
    "s2_robot_alignment",
    "s7_drawer_prismatic_gated",
    "s7_support_plane_fit",
    "resolution_contract_hw16",
    "window_frame_count_4n1",
)


@dataclass(frozen=True)
class WindowSample:
    """One ``vace/window_*`` export directory, resolved to absolute paths.

    Every path is validated to exist by :func:`fpgm.training.manifest.discover_windows`
    at scan time (a manifest entry pointing at a file that vanished between
    S8 export and training time must fail at manifest-build time, not mid-epoch).
    """

    episode_uuid: str
    camera_serial: str
    window_dir: Path
    target_video: Path
    control_depth: Path
    control_seg: Path
    control_normal: Path
    fg_mask: Path
    ref_plate: Path
    caption_path: Path
    sample_json: Path
    resolution: tuple[int, int]  # (width, height), from sample.json
    frame_range: tuple[int, int]  # (video_frame_start, video_frame_end)
    gates_passed: dict[str, bool]
    overlaps_pose_gap: bool

    @property
    def name(self) -> str:
        return f"{self.episode_uuid}/{self.camera_serial}/{self.window_dir.name}"


@dataclass(frozen=True)
class GateFilterConfig:
    """Policy for which ``vace/window_*`` directories are trainable.

    Defaults are the strict, "don't teach the model a wrong association"
    reading of the owning plan: a window is kept only if every known gate in
    ``sample.json["gates_passed"]`` is ``True`` *and* it does not overlap a
    pose gap (``overlaps_pose_gap`` -- 2 of the 3 windows that exist as of
    this build have it, per the S6 brick-GAP defect recorded in
    ``fpgm.datagen.export_vace``'s own module docstring).

    ``allow_pose_gap_overlap`` and ``ignore_gate_failures`` exist for exactly
    one legitimate use in this codebase: the Job 4 smoke train, which needs
    to exercise the training loop's mechanics on all 3 currently-exported
    windows regardless of data quality (see ``scripts/train_wan_vace_lora.sh
    --smoke``). Every caller that flips them must say so out loud in an
    invocation log -- :func:`fpgm.training.manifest.discover_windows` does,
    unconditionally, so a permissive run can never look like a strict one in
    the logs.
    """

    allow_pose_gap_overlap: bool = False
    #: Gate keys whose failure is tolerated. Empty by default: every gate in
    #: KNOWN_GATE_KEYS must pass. Not "ignore_gate_failures: bool" -- a named
    #: allowlist makes a partial relaxation ("skip a broken gate we know is
    #: too strict") visible in the config, instead of an unreviewable blanket
    #: bypass.
    ignored_gate_keys: tuple[str, ...] = ()
    #: Unknown gate keys are only a problem if this is True -- during rapid
    #: iteration on the S8 exporter (owned by a sibling agent) a new gate key
    #: appearing is expected; a *missing* one of KNOWN_GATE_KEYS is not.
    require_known_gate_keys_only: bool = False


@dataclass(frozen=True)
class ManifestFilterReport:
    """What :func:`fpgm.training.manifest.discover_windows` dropped, and why.

    Logged in full at manifest-build time (never just a count) -- see the
    owning plan: "log how many samples were dropped and why -- a silently
    shrunken training set is a bug we would not notice."
    """

    total_found: int
    kept: tuple[WindowSample, ...]
    dropped: tuple[tuple[str, str], ...] = field(default_factory=tuple)  # (window name, reason)

    @property
    def n_kept(self) -> int:
        return len(self.kept)

    @property
    def n_dropped(self) -> int:
        return len(self.dropped)

    def reason_counts(self) -> dict[str, int]:
        counts: dict[str, int] = {}
        for _, reason in self.dropped:
            counts[reason] = counts.get(reason, 0) + 1
        return counts


# --------------------------------------------------------------------------- #
# Bundle assembly (control-video compositing, reference letterboxing)
# --------------------------------------------------------------------------- #


@dataclass(frozen=True)
class BundleAssemblyConfig:
    """How a :class:`WindowSample` becomes the tensor format
    ``diffsynth.pipelines.wan_video.WanVideoPipeline``'s VACE branch consumes.

    **Why a composited RGB control video, not three separate control
    inputs.** Read from the trainer's own code
    (``diffsynth/pipelines/wan_video.py``'s ``WanVideoUnit_VACE.process``,
    ``diffsynth/models/wan_video_vace.py``'s ``VaceWanModel``): the VACE
    branch's patch embedding (``vace_patch_embedding``, ``in_channels=96``)
    is fed exactly one ``vace_context`` tensor built from exactly one
    ``vace_video`` + one ``vace_video_mask`` + one ``vace_reference_image``
    -- there is no multi-control-video collate path to plug three separate
    channels into. So ``control_depth``/``control_seg``/``control_normal``
    must be composited into a single RGB video before it reaches the
    trainer, as the owning plan's fallback suggestion says: depth in R,
    normalised seg-id in G, normal-Z in B.

    **Why no ``vace_video_mask`` by default.** ``fpgm.datagen.export_vace``'s
    own module docstring records a measured, not assumed, finding: "this
    exporter never emits a training ``src_mask``" and "The real background
    must come from the reference image ..., never from a mask carve-out of
    ``src_video``" -- backed by the ablation in
    ``outputs/wan_ablation/comparison.md`` (variant D, no mask + ref image,
    LPIPS 0.428 vs variant C, mask + no ref image, LPIPS 0.714). Following
    that precedent here keeps train-time input construction consistent with
    the export design it is built on. ``include_fg_mask`` exists to make
    this a policy switch, not a silent omission, and to support a future
    ablation that re-tests it against *this* project's now-complete
    depth/seg/normal control (the cited numbers are from an earlier control
    version with no brick/drawer wired in -- see that file's own caveat).
    """

    #: Composite channel order: R=depth, G=seg id (normalised), B=normal-Z.
    #: ``max_seg_id`` sets the normalisation divisor -- fpgm.config's
    #: SegmentationConfig.max_num_objects default (16) is reused here so the
    #: composite has headroom for every id this pipeline can ever emit,
    #: without recomputing a per-episode max (which would make the same
    #: physical id map to a different R-channel value in different windows).
    max_seg_id: int = 16
    #: If True, also emit ``vace_video_mask`` from ``fg_mask.mp4``. Off by
    #: default -- see class docstring.
    include_fg_mask: bool = False
    #: Reference-image letterbox background (white, matching the established
    #: Wan/VACE inference-time convention -- see fpgm.datagen.export_vace's
    #: "Reference image: letterbox on a white canvas" contract clause).
    letterbox_fill: int = 255


@dataclass(frozen=True)
class AugmentConfig:
    """Train<->inference-gap augmentations. Every knob individually toggleable.

    These exist because training conditioning is exporter-perfect (poses
    solved by PnP against the *same* video the model must reproduce, masks
    from the same render) while inference conditioning comes from a MuJoCo
    simulation whose pose estimate, mask, and per-channel availability will
    all be slightly different. Skipping these would let the LoRA overfit to
    exporter-perfect conditioning it will never see again at inference.
    """

    enable_pose_noise: bool = True
    #: Foreground-region rigid jitter applied to the *composited control*
    #: only (never to ``target.mp4``, which is real ground truth and must
    #: stay untouched). See fpgm.training.augment.pose_noise's docstring for
    #: why this is a 2D affine approximation of a 3D pose perturbation, not a
    #: mesh re-render.
    pose_noise_translate_mm_std: float = 2.0
    pose_noise_rotate_deg_std: float = 1.0
    #: Nominal px-per-mm at the image plane used to convert the physical
    #: pose-noise spec above into a pixel-space affine jitter, since the
    #: Dataset only has the exported raster bundle, not per-sample camera
    #: intrinsics (those live in ``master/meta.json``, outside the ``vace/``
    #: tier this package reads -- see module docstring on why we don't reach
    #: into ``master/``). Derived from a 60 deg horizontal FOV pinhole at
    #: this export's 832 px width and a ~0.6 m nominal DROID desk working
    #: distance: focal_px = 832 / (2*tan(30 deg)) ~= 720 px;
    #: px_per_mm = focal_px / 600 mm ~= 1.2. Documented as an approximation:
    #: the augmentation's job is to blur exact-pixel-alignment reliance, not
    #: to simulate a physically exact reprojection.
    pose_noise_px_per_mm: float = 1.2

    enable_mask_dilate_erode: bool = True
    #: Max structuring-element half-width (px) for the random dilate/erode of
    #: the foreground region (fg_mask.mp4), applied identically to every
    #: frame in a window (one random draw per __getitem__, not per frame --
    #: a per-frame-random kernel would flicker in a way no real mask-quality
    #: variation looks like).
    mask_dilate_erode_max_px: int = 5

    enable_channel_dropout: bool = True
    #: Each of the 3 composited channels (depth/seg/normal) is independently
    #: zeroed with this probability, for the whole window (not per frame).
    #: 0.10-0.20 per the owning plan; 0.15 is the midpoint.
    channel_dropout_prob: float = 0.15

    #: Base seed for the per-sample RNG (mixed with the dataset index so
    #: different windows/epochs draw independent augmentations while a given
    #: (seed, index) pair stays reproducible for debugging).
    seed: int = 0