scripts/sample_appearance_lora.py

#!/usr/bin/env python
"""Generate from a control video with (and optionally without) the appearance LoRA.

This is the *only* honest read on whether the fine-tune is working. The training
loss is not one: Wan's diffusion objective samples a random timestep per step, so
its per-step value is dominated by that draw -- over the first 500 steps of this
run the per-100-step medians read 0.0299, 0.0313, 0.0273, 0.0339, 0.0255, which is
noise, not a plateau and not progress. What the LoRA is for is a specific,
visible failure: handed a control whose robot is already a clean, correctly-posed
white Franka + black Robotiq, zero-shot VACE-1.3B drew a **yellow toy arm** from
about frame 50. Either that stops happening or the run did nothing.

``--baseline`` runs the same control and seed with the LoRA cleared, so the pair
is comparable frame-for-frame rather than against a memory of an earlier run.

Inference goes through DiffSynth, not Wan's own ``generate.py``: the LoRA was
trained by DiffSynth's ``WanTrainingModule`` and upstream ``generate.py`` has no
flag to load one at all.

Env: ``wan-train``. Usage:
    CUDA_VISIBLE_DEVICES=2 PYTHONPATH=src \\
    /home/quang/miniconda3/envs/wan-train/bin/python scripts/sample_appearance_lora.py \\
        --lora outputs/lora_appearance_2k/step-500.safetensors \\
        --pair outputs/appearance_pairs/<uuid>__<serial> --out outputs/lora_samples
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import torch
from PIL import Image

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

_DEFAULT_CKPT = Path(
    "/home/quang/.cache/huggingface/hub/models--Wan-AI--Wan2.1-VACE-1.3B/"
    "snapshots/574e6a744642ce3bee319afc31496b88bde8aac4"
)
#: Wan's own default negative prompt, kept verbatim so a sample is comparable to
#: anything generated through the upstream launcher.
_NEGATIVE = (
    "Bright tones, overexposed, static, blurred details, subtitles, style, works, "
    "paintings, images, static, overall gray, worst quality, low quality, JPEG "
    "compression residue, ugly, incomplete, extra fingers, poorly drawn hands, "
    "poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, "
    "still picture, cluttered background, three legs, many people in the "
    "background, walking backwards"
)


def build_pipeline(ckpt_dir: Path, device: str, offload_text_encoder: bool = False):
    """The DiffSynth pipeline, optionally with the text encoder kept off the GPU.

    ``offload_text_encoder`` exists because of a measured allocation imbalance,
    not as a general-purpose knob. The three models this loads are wildly
    different sizes: ``umt5-xxl`` is ~11 GB in bf16 while the VACE-1.3B DiT it
    feeds is ~2.6 GB. On a shared host that difference decides whether the run
    starts at all -- this exact call OOM'd twice against ~16 GB and ~5 GB of
    free VRAM, each time inside the loader, before a single denoising step ran.

    The text encoder is also the one model whose cost does not scale with the
    job: the prompt is encoded **once** per run, then never touched again
    across 4 chunks x 30 denoising steps. Keeping 11 GB resident for one
    forward pass is the wrong trade whenever VRAM is contended, and the CPU
    round-trip costs a few seconds against a multi-minute generation.

    Off by default so an uncontended run keeps its current behaviour
    bit-for-bit: this changes *where a tensor lives*, and a silent default
    change to that is exactly the kind of thing that makes two runs
    incomparable for reasons nobody records.

    Implemented through DiffSynth's own ``vram_limit`` rather than by pinning
    the text encoder's ``ModelConfig`` to the CPU. Pinning was tried first and
    fails: ``WanVideoPipeline`` still moves the tokenised input ids to
    ``pipe.device``, so a CPU-resident encoder raises ``Expected all tensors to
    be on the same device ... cpu and cuda:0`` inside the embedding lookup.
    ``vram_limit`` is the path DiffSynth actually maintains for this, and it
    keeps the device plumbing consistent instead of half-overriding it.
    """
    from diffsynth.core.loader.config import ModelConfig
    from diffsynth.pipelines.wan_video import WanVideoPipeline

    kwargs = {}
    if offload_text_encoder:
        # Budget in GB for resident weights; below the ~11 GB text encoder, so it
        # is the model that gets streamed rather than the DiT doing the denoising.
        kwargs["vram_limit"] = 6.0

    return WanVideoPipeline.from_pretrained(
        torch_dtype=torch.bfloat16,
        device=device,
        model_configs=[
            ModelConfig(path=str(ckpt_dir / "diffusion_pytorch_model.safetensors")),
            ModelConfig(path=str(ckpt_dir / "models_t5_umt5-xxl-enc-bf16.pth")),
            ModelConfig(path=str(ckpt_dir / "Wan2.1_VAE.pth")),
        ],
        tokenizer_config=ModelConfig(path=str(ckpt_dir / "google" / "umt5-xxl")),
        **kwargs,
    )


def save_mp4(frames, out: Path, fps: float) -> None:
    import subprocess

    ff = "/home/quang/miniconda3/envs/ffmpeg_libs/bin/ffmpeg"
    arr = [np.asarray(f)[:, :, ::-1] for f in frames]  # RGB -> BGR
    h, w = arr[0].shape[:2]
    out.parent.mkdir(parents=True, exist_ok=True)
    p = subprocess.Popen(
        [ff, "-y", "-hide_banner", "-loglevel", "error", "-f", "rawvideo",
         "-pix_fmt", "bgr24", "-s", f"{w}x{h}", "-r", f"{fps:.4f}", "-i", "pipe:0",
         "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "16", str(out)],
        stdin=subprocess.PIPE)
    for f in arr:
        p.stdin.write(np.ascontiguousarray(f).tobytes())
    p.stdin.close()
    if p.wait() != 0:
        raise RuntimeError(f"ffmpeg failed writing {out}")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--pair", type=Path, required=True,
                    help="one outputs/appearance_pairs/<uuid>__<serial> directory")
    ap.add_argument("--lora", type=Path, default=None)
    ap.add_argument("--lora-alpha", type=float, default=1.0)
    ap.add_argument("--baseline", action="store_true",
                    help="also generate with the LoRA cleared, same seed")
    ap.add_argument("--out", type=Path, default=REPO_ROOT / "outputs/lora_samples")
    ap.add_argument("--ckpt-dir", type=Path, default=_DEFAULT_CKPT)
    ap.add_argument("--start", type=int, default=0, help="first frame of the window")
    ap.add_argument("--num-frames", type=int, default=81)
    ap.add_argument("--steps", type=int, default=30)
    ap.add_argument("--cfg-scale", type=float, default=5.0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--fps", type=float, default=15.0)
    ap.add_argument("--chunk-frames", type=int, default=None,
                    help="generate in overlapping chunks of this many frames instead of "
                         "one pass. Wan2.1-VACE-1.3B is trained at 81 frames and falls off "
                         "a cliff past it: over one 189-frame pass the robot region scored "
                         "18.4 dB on frames 0-27 and 11.6 dB on 162-189, the arm fading to "
                         "a transparent ghost. Chunking is done here in PIXEL space rather "
                         "than via the pipeline's own sliding_window_size, because that "
                         "path slices only [latents, y] and leaves vace_context at full "
                         "length -- with VACE it crashes on a negative concat dimension "
                         "before it can denoise anything.")
    ap.add_argument("--chunk-overlap", type=int, default=0,
                    help="frames shared between consecutive chunks, cross-faded. Only "
                         "meaningful with --no-chain: cross-fading blends two chunks that "
                         "were generated without knowledge of each other, which hides a "
                         "seam rather than preventing one.")
    ap.add_argument("--chain", action=argparse.BooleanOptionalAction, default=True,
                    help="autoregressive chunking: chunk N+1's reference image is chunk "
                         "N's own last generated frame, so each chunk continues from what "
                         "the previous one actually produced instead of restarting from "
                         "the episode's frame 0. The reference channel is where this model "
                         "gets appearance from -- measured: swapping it changes robot-region "
                         "PSNR by 8.8 dB while the control alone moves it 0.3 dB -- so it is "
                         "the right place to carry continuity. The cost is error "
                         "accumulation: a bad chunk is inherited, not averaged away.")
    ap.add_argument("--reference-image", type=Path, default=None,
                    help="override the reference with an arbitrary image, cover-cropped "
                         "to the same size. Use a real-Franka frame from ANOTHER episode "
                         "to test whether the LoRA copies robot appearance from any "
                         "reference containing one, or only from this clip's own frame 0 "
                         "-- the difference decides whether a canonical reference photo "
                         "is deployable in place of datagen's robot-free ref_plate.png.")
    ap.add_argument("--offload-text-encoder", action="store_true",
                    help="keep umt5-xxl (~11 GB bf16) on the CPU; the prompt is encoded "
                         "once per run, so this trades a few seconds for ~11 GB of VRAM "
                         "on a contended host. See build_pipeline.")
    ap.add_argument("--reference", choices=("target_frame0", "window_frame0"),
                    default="target_frame0",
                    help="what VACE gets as its reference. 'target_frame0' is what "
                         "training used and contains the real robot; 'window_frame0' "
                         "is the control's own first frame (CAD robot only), which is "
                         "the honest stand-in for datagen inference, where the "
                         "reference is a background plate with no robot in it.")
    args = ap.parse_args()

    from fpgm.training.appearance_dataset import AppearanceWindow, AppearancePairDataset

    meta = json.loads((args.pair / "meta.json").read_text())
    win = AppearanceWindow(pair_dir=args.pair, start=args.start,
                           n_frames=args.num_frames, caption=meta["caption"] or "",
                           uuid=meta["uuid"], camera_serial=str(meta["camera_serial"]))
    item = AppearancePairDataset([win], reference=args.reference)[0]
    if args.reference_image is not None:
        from fpgm.training.appearance_dataset import _cover_crop
        import cv2 as _cv2
        _bgr = _cv2.imread(str(args.reference_image))
        if _bgr is None:
            raise SystemExit(f"unreadable --reference-image {args.reference_image}")
        _w, _h = item["video"][0].size
        item["vace_reference_image"] = [
            Image.fromarray(_cover_crop(_bgr, _w, _h)[:, :, ::-1].copy())
        ]
    w, h = item["video"][0].size
    tag = (f"{meta['uuid']}__{meta['camera_serial']}_f{args.start:05d}"
           + ("" if args.reference == "target_frame0" else f"_ref-{args.reference}")
           + ("" if args.reference_image is None
              else f"_ref-{args.reference_image.stem}")
           + ("" if args.chunk_frames is None else
              f"_{'chain' if args.chain else 'chunk'}{args.chunk_frames}"))
    args.out.mkdir(parents=True, exist_ok=True)

    pipe = build_pipeline(args.ckpt_dir, "cuda",
                          offload_text_encoder=args.offload_text_encoder)

    def denoise(control, n_frames, seed, reference):
        return pipe(
            prompt=item["prompt"], negative_prompt=_NEGATIVE,
            vace_video=control, vace_reference_image=reference,
            height=h, width=w, num_frames=n_frames,
            num_inference_steps=args.steps, cfg_scale=args.cfg_scale, seed=seed,
        )

    def chunk_starts(total: int, size: int, overlap: int) -> list[int]:
        """Chunk starts covering ``total``, last one snapped flush to the end."""
        step = size - overlap
        starts = list(range(0, max(total - size, 0) + 1, step))
        if starts[-1] + size < total:
            starts.append(total - size)
        return starts

    def generate(label: str):
        # Same seed for both arms: the LoRA's effect has to show up against an
        # identical noise draw, otherwise a difference is just a different sample.
        ref0 = item["vace_reference_image"][0]
        if args.chunk_frames is None or args.chunk_frames >= args.num_frames:
            frames = denoise(item["vace_video"], args.num_frames, args.seed, ref0)
        elif args.chain:
            size = args.chunk_frames
            starts = list(range(0, args.num_frames, size))
            frames, ref = [], ref0
            for ci, st in enumerate(starts):
                n_c = min(size, args.num_frames - st)
                # Wan's VAE needs 4k+1 frames; a short tail is pulled back to end
                # flush with the clip rather than padded, so no frame is invented.
                if (n_c - 1) % 4:
                    st, n_c = max(0, args.num_frames - size), size
                out_c = denoise(item["vace_video"][st:st + n_c], n_c,
                                args.seed + ci, ref)
                keep = out_c[len(frames) - st:] if st < len(frames) else out_c
                frames.extend(keep)
                ref = out_c[-1]
                print(f"  {label}: chunk {ci+1}/{len(starts)} frames {st}-{st+n_c}"
                      f" -> total {len(frames)}", flush=True)
            frames = frames[:args.num_frames]
        else:
            size, ov = args.chunk_frames, args.chunk_overlap
            starts = chunk_starts(args.num_frames, size, ov)
            print(f"  {label}: {len(starts)} chunks of {size}f, overlap {ov}f, "
                  f"starts {starts}", flush=True)
            acc = np.zeros((args.num_frames, h, w, 3), np.float32)
            wsum = np.zeros((args.num_frames, 1, 1, 1), np.float32)
            for ci, st in enumerate(starts):
                out_c = denoise(item["vace_video"][st:st + size], size,
                                args.seed + ci, ref0)
                a = np.stack([np.asarray(f, np.float32) for f in out_c])
                # Ramp only where a neighbour actually overlaps, so the first and
                # last frames of the whole clip keep full weight instead of being
                # faded against nothing.
                ramp = np.ones(size, np.float32)
                if st > 0:
                    ramp[:ov] = np.linspace(0, 1, ov, endpoint=False)
                if st + size < args.num_frames:
                    ramp[-ov:] = np.linspace(1, 0, ov, endpoint=False)
                acc[st:st + size] += a * ramp[:, None, None, None]
                wsum[st:st + size] += ramp[:, None, None, None]
            frames = [Image.fromarray(f.astype(np.uint8))
                      for f in np.clip(acc / np.maximum(wsum, 1e-6), 0, 255)]
        out = args.out / f"{tag}__{label}.mp4"
        save_mp4(frames, out, args.fps)
        print(f"wrote {out}", flush=True)
        return out

    written = {}
    if args.baseline:
        pipe.clear_lora()
        written["baseline"] = str(generate("baseline"))
    if args.lora is not None:
        pipe.clear_lora()
        # Into pipe.vace, not pipe.dit: the trainer used lora_base_model="vace",
        # so every key is vace_blocks.*. Against the DiT, load_lora matches
        # nothing, says so only in a log line ("0 tensors are fused by LoRA"),
        # and then happily generates a plain baseline that looks like a result.
        # Hence the assertion rather than trust: a silently-unloaded LoRA is the
        # one failure that would make this whole measurement a lie.
        probe = pipe.vace.vace_blocks[0].cross_attn.q.weight
        before = probe.detach().float().clone()
        pipe.load_lora(pipe.vace, str(args.lora), alpha=args.lora_alpha)
        delta = (probe.detach().float() - before).abs().max().item()
        if delta == 0.0:
            raise SystemExit(
                f"{args.lora}: load_lora changed no weights -- wrong target module "
                f"or mismatched key names, refusing to report a baseline as a LoRA sample"
            )
        print(f"LoRA fused: max |delta| on vace_blocks.0.cross_attn.q = {delta:.3e}",
              flush=True)
        written["lora"] = str(generate(f"lora_{args.lora.stem}"))

    save_mp4(item["vace_video"], args.out / f"{tag}__control.mp4", args.fps)
    save_mp4(item["video"], args.out / f"{tag}__target.mp4", args.fps)
    (args.out / f"{tag}__meta.json").write_text(json.dumps({
        "pair": str(args.pair), "start": args.start, "n_frames": args.num_frames,
        "prompt": item["prompt"], "seed": args.seed, "steps": args.steps,
        "cfg_scale": args.cfg_scale, "lora": str(args.lora) if args.lora else None,
        "lora_alpha": args.lora_alpha, **written,
    }, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())