src/fpgm/training/train.py

"""LoRA-finetune Wan2.1-VACE-1.3B on our ``vace/window_*`` bundles via DiffSynth-Studio.

**Trainer verification (read from code, not the README).** DiffSynth-Studio
(``third_party/diffsynth``, pinned commit -- see ``third_party/PINNED_COMMITS.txt``)
genuinely drives the VACE control branches during training:
``examples/wanvideo/model_training/lora/Wan2.1-VACE-1.3B.sh`` sets
``--lora_base_model vace --extra_inputs vace_video,vace_reference_image``;
``WanTrainingModule.__init__`` (this same directory's ``train.py``) resolves
``lora_base_model="vace"`` against ``pipe.vace``
(``diffsynth/pipelines/wan_video.py:159-163``, a real
``VaceWanModel`` instance -- ``diffsynth/models/wan_video_vace.py``), and the
forward pass's ``WanVideoUnit_VACE.process`` (``wan_video.py:652-707``) builds
a real ``vace_context`` tensor from ``vace_video``/``vace_video_mask``/
``vace_reference_image`` that is fed through ``VaceWanModel.forward``'s own
``vace_blocks``/``vace_patch_embedding`` and injected into the base DiT at
every layer in ``VaceWanModel``'s ``vace_layers`` (``wan_video.py:1525-1572``).
This is not the base T2V/I2V path with the VACE weights merely present but
unused.

**flash_attn is optional here** (unlike ``third_party/wan2.1``'s native
inference CLI, which hard-asserts without it --
``scripts/setup_wan_env.sh``): ``diffsynth/core/attention/attention.py``'s
``initialize_attention_priority`` genuinely falls back to
``torch.nn.functional.scaled_dot_product_attention`` when
``flash_attn`` isn't importable (``ATTENTION_IMPLEMENTATION = "torch"``),
so training still runs, just slower, if the wheel install in
``scripts/setup_wan_train_env.sh`` fails on some future host.

This module imports ``diffsynth``/``accelerate``/``peft`` **lazily**, inside
:func:`main`/:func:`build_training_module`, not at module scope -- so
``import fpgm.training.train`` succeeds in the plain ``fpgm`` conda env
(no GPU, no ``diffsynth`` installed), matching this package's "lazy heavy
imports" convention. Only :mod:`fpgm.training.dataset`/``manifest``/``augment``
are meant to be imported from that env; this module is meant to run inside
the ``wan-train`` env created by ``scripts/setup_wan_train_env.sh``.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

from fpgm.training.dataset import build_dataset_from_manifest
from fpgm.training.types import AugmentConfig, BundleAssemblyConfig, GateFilterConfig
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

#: The DiffSynth-Studio checkout this module was written and pinned against
#: -- see third_party/PINNED_COMMITS.txt. Only used to locate the example
#: script we reuse WanTrainingModule from (that script is not part of the
#: installed `diffsynth` package, see _load_diffsynth_train_module).
_REPO_ROOT = Path(__file__).resolve().parents[3]
_DIFFSYNTH_DIR = _REPO_ROOT / "third_party" / "diffsynth"
_DIFFSYNTH_TRAIN_SCRIPT_DIR = _DIFFSYNTH_DIR / "examples" / "wanvideo" / "model_training"


def _load_diffsynth_train_module():
    """Import ``examples/wanvideo/model_training/train.py`` for its
    ``WanTrainingModule`` class.

    That file is an example script, not part of the installable ``diffsynth``
    package (``pip install -e third_party/diffsynth`` does not put it on
    ``sys.path``), so it's loaded by path. Only the module-level class/function
    definitions are used -- its own ``if __name__ == "__main__":`` launcher
    block never runs because we import it, not execute it as a script.
    """
    if str(_DIFFSYNTH_TRAIN_SCRIPT_DIR) not in sys.path:
        sys.path.insert(0, str(_DIFFSYNTH_TRAIN_SCRIPT_DIR))
    import train as diffsynth_train  # noqa: PLC0415 (intentionally lazy, see module docstring)

    return diffsynth_train


def _default_model_paths(ckpt_dir: Path) -> list[str]:
    """The 3 local checkpoint files DiffSynth's ``ModelConfig(path=...)`` needs.

    Points directly at the already-downloaded snapshot
    (see the owning plan's "Already established" section) -- never triggers
    a modelscope/HuggingFace download, unlike ``--model_id_with_origin_paths``.
    """
    return [
        str(ckpt_dir / "diffusion_pytorch_model.safetensors"),
        str(ckpt_dir / "models_t5_umt5-xxl-enc-bf16.pth"),
        str(ckpt_dir / "Wan2.1_VAE.pth"),
    ]


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--dataset-kind", choices=("bundle", "appearance"), default="bundle",
                    help="'bundle' = the datagen vace/window_* exports (geometry-buffer "
                         "controls, gate-filtered). 'appearance' = the (control, target) "
                         "pairs from scripts/appearance_control.py, whose only job is to "
                         "teach the model what a Franka looks like. The two carry different "
                         "supervision and are selected, not merged, so a run's log says "
                         "which one it trained on.")
    p.add_argument("--dataset-root", type=Path, default=_REPO_ROOT / "outputs" / "datagen",
                    help="bundle: root scanned for */*/vace/window_*/sample.json. "
                         "appearance: root holding <uuid>__<serial>/{control,target}.mp4.")
    p.add_argument("--window-stride", type=int, default=None,
                    help="appearance only: window start step in frames (default: no overlap). "
                         "A smaller stride multiplies sample count from the same clips, at the "
                         "cost of the extra windows being highly correlated.")
    p.add_argument("--reference", choices=("target_frame0", "window_frame0"),
                    default="target_frame0",
                    help="appearance only: what VACE gets as its reference image. See "
                         "AppearancePairDataset.__init__ for why neither is silently correct.")
    p.add_argument("--ckpt-dir", type=Path,
                    default=Path("/home/quang/.cache/huggingface/hub/models--Wan-AI--Wan2.1-VACE-1.3B/"
                                 "snapshots/574e6a744642ce3bee319afc31496b88bde8aac4"),
                    help="Local Wan2.1-VACE-1.3B snapshot dir (never re-downloaded).")
    p.add_argument("--output-path", type=Path, required=True,
                    help="Where to save LoRA checkpoints.")
    p.add_argument("--learning-rate", type=float, default=1e-4)
    p.add_argument("--lora-rank", type=int, default=32)
    p.add_argument("--lora-target-modules", type=str, default="q,k,v,o,ffn.0,ffn.2")
    p.add_argument("--num-epochs", type=int, default=5)
    p.add_argument("--max-steps", type=int, default=None,
                    help="Stop after this many optimizer steps regardless of --num-epochs "
                         "(the smoke train's mechanism for '30 steps', not a full run).")
    p.add_argument("--save-steps", type=int, default=None)
    p.add_argument("--gradient-accumulation-steps", type=int, default=1)
    p.add_argument("--use-gradient-checkpointing", action="store_true", default=True)
    p.add_argument("--no-gradient-checkpointing", dest="use_gradient_checkpointing",
                    action="store_false")
    p.add_argument("--use-gradient-checkpointing-offload", action="store_true", default=False)
    p.add_argument("--max-num-frames", type=int, default=None,
                    help="Truncate every window to its largest 4k+1 <= this (VRAM-vs-frame-count "
                         "measurement knob, see fpgm.training.dataset.WindowBundleDataset).")

    # Gate filter (fpgm.training.types.GateFilterConfig)
    p.add_argument("--allow-pose-gap-overlap", action="store_true", default=False,
                    help="Include windows with sample.json overlaps_pose_gap=True. Only intended "
                         "for the mechanics-only smoke train -- see GateFilterConfig's docstring.")
    p.add_argument("--ignored-gate-keys", type=str, default="",
                    help="Comma-separated gate keys to tolerate failing, e.g. for a smoke train.")

    # Bundle assembly (fpgm.training.types.BundleAssemblyConfig)
    p.add_argument("--include-fg-mask", action="store_true", default=False,
                    help="Also pass fg_mask.mp4 as vace_video_mask. Off by default -- see "
                         "BundleAssemblyConfig's docstring (measured worse in the project's own "
                         "Wan-VACE ablation, outputs/wan_ablation/comparison.md).")

    # Augmentation (fpgm.training.types.AugmentConfig)
    p.add_argument("--no-augment", action="store_true", default=False,
                    help="Disable every augmentation (pose noise, mask dilate/erode, channel "
                         "dropout) -- e.g. for the VRAM measurement, where they're irrelevant.")
    p.add_argument("--no-pose-noise", action="store_true", default=False)
    p.add_argument("--no-mask-dilate-erode", action="store_true", default=False)
    p.add_argument("--no-channel-dropout", action="store_true", default=False)

    p.add_argument("--measure-vram-only", action="store_true", default=False,
                    help="Run exactly 1 forward+backward step, report peak VRAM to "
                         "--report-json, and exit before any optimizer step or checkpoint save.")
    p.add_argument("--report-json", type=Path, default=None,
                    help="Where to write the {steps, steps_per_sec, peak_vram_gb, ...} report.")
    return p


def _build_configs(
    args: argparse.Namespace,
) -> tuple[GateFilterConfig, BundleAssemblyConfig, AugmentConfig | None]:
    gate_filter = GateFilterConfig(
        allow_pose_gap_overlap=args.allow_pose_gap_overlap,
        ignored_gate_keys=tuple(k for k in args.ignored_gate_keys.split(",") if k),
    )
    assembly_cfg = BundleAssemblyConfig(include_fg_mask=args.include_fg_mask)
    if args.no_augment:
        augment_cfg = None
    else:
        augment_cfg = AugmentConfig(
            enable_pose_noise=not args.no_pose_noise,
            enable_mask_dilate_erode=not args.no_mask_dilate_erode,
            enable_channel_dropout=not args.no_channel_dropout,
        )
    return gate_filter, assembly_cfg, augment_cfg


def build_training_module(args: argparse.Namespace, device: str):
    """Construct DiffSynth's ``WanTrainingModule`` pointed at the local checkpoint.

    Kept as its own function (rather than inlined in :func:`main`) so
    ``scripts/train_wan_vace_lora.sh``'s VRAM-measurement mode can build the
    module once and reuse it across multiple frame-count trials without
    reloading the ~1.3B-parameter DiT + T5 + VAE from disk each time.
    """
    diffsynth_train = _load_diffsynth_train_module()
    return diffsynth_train.WanTrainingModule(
        model_paths=json.dumps(_default_model_paths(args.ckpt_dir)),
        tokenizer_path=str(args.ckpt_dir / "google" / "umt5-xxl"),
        lora_base_model="vace",
        lora_target_modules=args.lora_target_modules,
        lora_rank=args.lora_rank,
        use_gradient_checkpointing=args.use_gradient_checkpointing,
        use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload,
        extra_inputs="vace_video,vace_reference_image"
        + (",vace_video_mask" if args.include_fg_mask else ""),
        remove_prefix_in_ckpt="pipe.vace.",
        device=device,
    )


def main(argv: list[str] | None = None) -> dict:
    args = build_parser().parse_args(argv)

    import accelerate
    import torch

    if args.dataset_kind == "appearance":
        from fpgm.training.appearance_dataset import (
            AppearancePairDataset, discover_windows,
        )

        windows = discover_windows(args.dataset_root, stride=args.window_stride)
        dataset = AppearancePairDataset(windows, reference=args.reference)
        logger.info(
            "appearance dataset: %d window(s) from %d clip(s), reference=%s",
            len(dataset), len({w.pair_dir for w in windows}), args.reference,
        )
        if args.max_num_frames is not None:
            # Not silently ignored: the appearance windows are a fixed 81 frames
            # by construction, and truncating them is a bundle-path knob.
            raise SystemExit("--max-num-frames applies to --dataset-kind bundle only")
    else:
        gate_filter, assembly_cfg, augment_cfg = _build_configs(args)
        dataset = build_dataset_from_manifest(
            args.dataset_root, gate_filter, assembly_cfg, augment_cfg,
        )
        # Frame-count truncation is a dataset-level knob applied after the
        # gate-filtered manifest is built (see WindowBundleDataset.__init__).
        dataset.max_num_frames = args.max_num_frames
        logger.info("training dataset: %d window(s) after gate filter", len(dataset))

    accelerator = accelerate.Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps
    )
    # Load the ~1.3B-param DiT/T5/VAE directly onto the training device (no
    # CPU-offload path in this minimal driver -- see the VRAM measurement
    # job for whether that's needed on this contended host).
    model = build_training_module(args, device=accelerator.device)

    from diffsynth.diffusion.logger import ModelLogger

    args.output_path.mkdir(parents=True, exist_ok=True)
    model_logger = ModelLogger(str(args.output_path), remove_prefix_in_ckpt="pipe.vace.")

    optimizer = torch.optim.AdamW(model.trainable_modules(), lr=args.learning_rate)
    scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer)
    dataloader = torch.utils.data.DataLoader(
        dataset, shuffle=True, collate_fn=lambda x: x[0], num_workers=0
    )

    model.to(device=accelerator.device)
    model, optimizer, dataloader, scheduler = accelerator.prepare(
        model, optimizer, dataloader, scheduler
    )

    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats(accelerator.device)

    # DiffSynth's own WanTrainingModule.__init__ forcibly re-enables gradient
    # checkpointing (with a UserWarning) if asked to disable it -- "to
    # prevent out-of-memory errors" is its own stated reason. So the
    # *requested* value and the *effective* one can differ; record both
    # rather than let the report silently claim gradient checkpointing was
    # off when DiffSynth overrode that. accelerator.prepare() may wrap
    # `model` (e.g. DDP) but does not change this attribute's value.
    unwrapped = accelerator.unwrap_model(model)
    effective_gc = getattr(unwrapped, "use_gradient_checkpointing", args.use_gradient_checkpointing)
    report = {
        "dataset_kind": args.dataset_kind,
        "max_num_frames": args.max_num_frames,
        "requested_gradient_checkpointing": args.use_gradient_checkpointing,
        "effective_gradient_checkpointing": effective_gc,
        "n_windows": len(dataset),
        "steps": 0,
    }
    step_times: list[float] = []
    global_step = 0
    stop = False
    t_prev = time.monotonic()
    for _epoch in range(args.num_epochs):
        if stop:
            break
        for data in dataloader:
            with accelerator.accumulate(model):
                loss = model(data)
                accelerator.backward(loss)
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()
            t_now = time.monotonic()
            step_times.append(t_now - t_prev)
            t_prev = t_now
            global_step += 1
            logger.info(
                "step %d window=%s loss=%.4f dt=%.2fs",
                global_step, data.get("_window_name"), float(loss.detach().item()), step_times[-1],
            )
            model_logger.on_step_end(accelerator, model, args.save_steps, loss=loss)
            if args.measure_vram_only:
                stop = True
                break
            if args.max_steps is not None and global_step >= args.max_steps:
                stop = True
                break
        if not args.measure_vram_only and not stop and args.save_steps is None:
            model_logger.on_epoch_end(accelerator, model, _epoch)
    if not args.measure_vram_only:
        model_logger.on_training_end(accelerator, model, args.save_steps)

    report["steps"] = global_step
    mean_step_seconds = sum(step_times) / len(step_times) if step_times else None
    report["mean_step_seconds"] = mean_step_seconds
    report["steps_per_sec"] = (1.0 / mean_step_seconds) if mean_step_seconds else None
    if torch.cuda.is_available():
        report["peak_vram_allocated_gb"] = torch.cuda.max_memory_allocated(accelerator.device) / 1e9
        report["peak_vram_reserved_gb"] = torch.cuda.max_memory_reserved(accelerator.device) / 1e9

    if args.report_json is not None:
        args.report_json.parent.mkdir(parents=True, exist_ok=True)
        args.report_json.write_text(json.dumps(report, indent=2))
    logger.info("training report: %s", json.dumps(report, indent=2))
    return report


if __name__ == "__main__":
    main()