hao-ai-lab-FastVideo/scripts/ltx2_sr_alignment.py

# SPDX-License-Identifier: Apache-2.0
"""LTX-2 SR (refine) numerical alignment harness — public vs internal.

Runs the LTX-2 stage-2 spatial-refine pipeline with a fixed prompt and
seed and writes the decoded video frames + intermediate latents to a
torch save file. Intended to be invoked twice with different
``PYTHONPATH`` prefixes — once against ``FastVideo-internal`` and once
against the public ``FastVideo`` package — so a third pass can diff the
two outputs tensor-by-tensor.

Usage::

    # 1) Internal reference run.
    PYTHONPATH=/home/william5lin/FastVideo-internal \
        python scripts/ltx2_sr_alignment.py --label internal \
            --output /tmp/ltx2_sr_internal.pt

    # 2) Public run against the new typed API.
    PYTHONPATH=/home/william5lin/FastVideo \
        python scripts/ltx2_sr_alignment.py --label public --use-typed-api \
            --output /tmp/ltx2_sr_public.pt

    # 3) Diff.
    python scripts/ltx2_sr_alignment.py --diff \
        --reference /tmp/ltx2_sr_internal.pt \
        --candidate /tmp/ltx2_sr_public.pt

The harness deliberately matches `examples/inference/basic/basic_ltx2_upscale.py`
defaults (LTX2-Distilled, 8 base steps, 3 refine steps, 1088x1920, seed=10)
but skips FP4 and Dreamverse-specific knobs to keep the bf16 path clean.
"""
from __future__ import annotations

import argparse
import os
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import torch

# Pinned alignment fixture — matches basic_ltx2_upscale.py upstream.
PROMPT = (
    "A warm sunny backyard. The camera starts in a tight cinematic close-up "
    "of a woman and a man in their 30s, facing each other with serious "
    "expressions. The woman, emotional and dramatic, says softly, \"That's "
    "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly "
    "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, "
    "then mutters defensively, \"He's just having fun.\" The camera slowly "
    "pans right, revealing the grandfather in the garden wearing enormous "
    "butterfly wings, waving his arms in the air like he's trying to take "
    "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. "
    "The woman covers her face, on the verge of tears. The tone is deadpan, "
    "absurd, and quietly tragic.")

DEFAULT_MODEL = "FastVideo/LTX2-Distilled-Diffusers"
DEFAULT_SEED = 10
DEFAULT_HEIGHT = 1088
DEFAULT_WIDTH = 1920
DEFAULT_FRAMES = 121
DEFAULT_FPS = 24
DEFAULT_BASE_STEPS = 8
DEFAULT_REFINE_STEPS = 3


@dataclass
class RunResult:
    label: str
    fastvideo_repo: str
    output_frames: torch.Tensor
    metadata: dict[str, Any]


def _detect_fastvideo_repo() -> str:
    import fastvideo  # noqa: PLC0415
    return os.path.realpath(os.path.dirname(fastvideo.__file__))


def _run_legacy(args: argparse.Namespace) -> RunResult:
    """Legacy ``VideoGenerator.from_pretrained(**flat_kwargs)`` path —
    used by the internal reference run because the internal API is
    still flat-kwarg-shaped."""
    from fastvideo import VideoGenerator  # noqa: PLC0415

    generator = VideoGenerator.from_pretrained(
        args.model,
        num_gpus=1,
        ltx2_refine_enabled=True,
        ltx2_refine_lora_path="",  # disable refine LoRA — distilled needs none
        ltx2_refine_num_inference_steps=args.refine_steps,
        ltx2_refine_guidance_scale=1.0,
        ltx2_refine_add_noise=True,
        dit_cpu_offload=False,
        vae_cpu_offload=False,
        text_encoder_cpu_offload=False,
        pin_cpu_memory=True,
        dit_layerwise_offload=False,
        enable_torch_compile=False,
        ltx2_vae_tiling=False,
    )

    result = generator.generate_video(
        prompt=PROMPT,
        output_path=str(args.output) + ".legacy.mp4",
        fps=args.fps,
        seed=args.seed,
        num_inference_steps=args.base_steps,
        guidance_scale=1.0,
        save_video=False,
        return_frames=True,
        height=args.height,
        width=args.width,
        num_frames=args.frames,
    )
    generator.shutdown()
    frames = _result_frames(result)
    return RunResult(
        label=args.label,
        fastvideo_repo=_detect_fastvideo_repo(),
        output_frames=frames,
        metadata={
            "api": "legacy_from_pretrained",
            "prompt": PROMPT,
            "seed": args.seed,
            "height": args.height,
            "width": args.width,
            "num_frames": args.frames,
            "fps": args.fps,
            "base_steps": args.base_steps,
            "refine_steps": args.refine_steps,
        },
    )


def _run_typed(args: argparse.Namespace) -> RunResult:
    """New typed public API: GeneratorConfig + GenerationRequest with
    pipeline.preset_overrides["refine"] driving the SR path."""
    from fastvideo import VideoGenerator  # noqa: PLC0415
    from fastvideo.api import (  # noqa: PLC0415
        ComponentConfig, EngineConfig, GenerationRequest, GeneratorConfig,
        InputConfig, OutputConfig, PipelineSelection, SamplingConfig)

    config = GeneratorConfig(
        model_path=args.model,
        engine=EngineConfig(num_gpus=1, ),
        pipeline=PipelineSelection(
            components=ComponentConfig(),
            preset_overrides={
                "refine": {
                    "enabled": True,
                    "add_noise": True,
                    "num_inference_steps": args.refine_steps,
                    "guidance_scale": 1.0,
                },
            },
        ),
    )

    generator = VideoGenerator.from_pretrained(config=config)

    request = GenerationRequest(
        prompt=PROMPT,
        sampling=SamplingConfig(
            num_frames=args.frames,
            height=args.height,
            width=args.width,
            fps=args.fps,
            num_inference_steps=args.base_steps,
            guidance_scale=1.0,
            seed=args.seed,
        ),
        inputs=InputConfig(),
        output=OutputConfig(
            save_video=False,
            return_frames=True,
        ),
    )

    result = generator.generate(request)
    generator.shutdown()
    frames = _result_frames(result)
    return RunResult(
        label=args.label,
        fastvideo_repo=_detect_fastvideo_repo(),
        output_frames=frames,
        metadata={
            "api": "typed",
            "prompt": PROMPT,
            "seed": args.seed,
            "height": args.height,
            "width": args.width,
            "num_frames": args.frames,
            "fps": args.fps,
            "base_steps": args.base_steps,
            "refine_steps": args.refine_steps,
        },
    )


def _result_frames(result: Any) -> torch.Tensor:
    """Coerce whatever generate() returned into ``[F, H, W, C]`` uint8."""
    frames = getattr(result, "frames", None)
    if frames is None and isinstance(result, dict):
        frames = result.get("frames")
    if frames is None:
        raise RuntimeError(
            "Generation returned no frames; ensure return_frames=True.")
    if isinstance(frames, list):
        frames = np.stack(frames, axis=0)
    if isinstance(frames, np.ndarray):
        return torch.from_numpy(frames)
    if torch.is_tensor(frames):
        return frames.detach().cpu()
    raise TypeError(f"Unsupported frames container: {type(frames)!r}")


def _run(args: argparse.Namespace) -> None:
    if args.use_typed_api:
        result = _run_typed(args)
    else:
        result = _run_legacy(args)

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "frames": result.output_frames,
            "metadata": result.metadata,
            "label": result.label,
            "fastvideo_repo": result.fastvideo_repo,
        },
        output_path,
    )
    print(
        f"[align] saved {result.output_frames.shape} frames "
        f"({result.fastvideo_repo}) -> {output_path}",
        flush=True,
    )


def _diff(args: argparse.Namespace) -> None:
    ref_path = Path(args.reference)
    cand_path = Path(args.candidate)
    if not ref_path.is_file() or not cand_path.is_file():
        raise FileNotFoundError(
            f"Need both --reference and --candidate to exist: {ref_path}, "
            f"{cand_path}")
    ref = torch.load(ref_path, map_location="cpu")
    cand = torch.load(cand_path, map_location="cpu")

    ref_frames = ref["frames"]
    cand_frames = cand["frames"]
    print(f"[align] reference: {ref['label']} {ref_frames.shape} from "
          f"{ref['fastvideo_repo']}")
    print(f"[align] candidate: {cand['label']} {cand_frames.shape} from "
          f"{cand['fastvideo_repo']}")

    if ref_frames.shape != cand_frames.shape:
        print(f"[align] shape MISMATCH: ref={tuple(ref_frames.shape)} "
              f"vs cand={tuple(cand_frames.shape)}")
        return

    diff = (ref_frames.float() - cand_frames.float()).abs()
    max_abs = float(diff.max().item())
    mean_abs = float(diff.mean().item())
    rms = float(torch.sqrt((diff**2).mean()).item())

    # PSNR over uint8 [0,255] data range.
    if rms == 0.0:
        psnr = float("inf")
    else:
        psnr = 20.0 * float(torch.log10(torch.tensor(255.0 / rms)).item())

    print("[align] frame-level uint8 metrics (per-pixel):")
    print(f"        max_abs_diff = {max_abs:.4f}")
    print(f"        mean_abs_diff = {mean_abs:.6f}")
    print(f"        rms_diff = {rms:.4f}")
    print(f"        psnr_db = {psnr:.2f}")

    # Bucket pixels by exactness so we have a sense of how close we are.
    eq = (diff == 0).float().mean().item()
    within_1 = (diff <= 1).float().mean().item()
    within_2 = (diff <= 2).float().mean().item()
    within_4 = (diff <= 4).float().mean().item()
    print("[align] pixel buckets:")
    print(f"        exact = {eq * 100:.2f}%")
    print(f"        within 1 = {within_1 * 100:.2f}%")
    print(f"        within 2 = {within_2 * 100:.2f}%")
    print(f"        within 4 = {within_4 * 100:.2f}%")


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="mode", required=False)

    run = parser
    run.add_argument("--label",
                     default="run",
                     help="Identifier saved alongside the output (e.g. internal/public).")
    run.add_argument("--output",
                     default="/tmp/ltx2_sr_alignment.pt",
                     help="Destination .pt file.")
    run.add_argument("--use-typed-api",
                     action="store_true",
                     help="Use the new typed GeneratorConfig API "
                          "(public side); otherwise use legacy from_pretrained.")
    run.add_argument("--model", default=DEFAULT_MODEL)
    run.add_argument("--seed", type=int, default=DEFAULT_SEED)
    run.add_argument("--height", type=int, default=DEFAULT_HEIGHT)
    run.add_argument("--width", type=int, default=DEFAULT_WIDTH)
    run.add_argument("--frames", type=int, default=DEFAULT_FRAMES)
    run.add_argument("--fps", type=int, default=DEFAULT_FPS)
    run.add_argument("--base-steps", type=int, default=DEFAULT_BASE_STEPS)
    run.add_argument("--refine-steps", type=int, default=DEFAULT_REFINE_STEPS)

    run.add_argument("--diff",
                     action="store_true",
                     help="Diff two .pt outputs instead of running.")
    run.add_argument("--reference")
    run.add_argument("--candidate")
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)
    if args.diff:
        if not (args.reference and args.candidate):
            parser.error("--diff requires --reference and --candidate")
        _diff(args)
        return 0
    _run(args)
    return 0


if __name__ == "__main__":
    sys.exit(main())