unsloth/scripts/diffusion_bench.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""Standalone GPU benchmark + regression harness for the Unsloth diffusion backend.

Drives ``DiffusionBackend`` directly (no HTTP server) to measure load time, peak VRAM and generation latency for a single GGUF image model, plus an accuracy guard: a fixed-seed image is rendered and compared (PSNR) against a stored reference, so a precision/dtype/guard regression that silently changes output is caught, not just speed and memory.

Two modes: --write-baseline PATH runs once and saves the metrics JSON plus reference.png next to it, and --compare PATH runs again, diffs against the baseline and exits nonzero if a latency, VRAM or PSNR threshold is exceeded.

torch and diffusers are imported lazily (only after argument parsing and only inside functions) so ``--help`` works on a host without them. Not part of CPU CI; this needs a real GPU and a downloadable model.

Example: `python scripts/diffusion_bench.py --write-baseline outputs/diffusion_bench/baseline.json --model unsloth/Z-Image-Turbo-GGUF --gguf z-image-turbo-Q4_K_M.gguf`, then the same command with --compare after making changes.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import platform
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

# Backend root on sys.path so `core.inference.diffusion` imports as the server does (deferred into main() so --help never triggers torch).
_BACKEND_ROOT = Path(__file__).resolve().parent.parent / "studio" / "backend"
if str(_BACKEND_ROOT) not in sys.path:
    sys.path.insert(0, str(_BACKEND_ROOT))


# ── small helpers ──────────────────────────────────────────────────────────


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _git_commit() -> Optional[str]:
    try:
        out = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd = str(Path(__file__).resolve().parent),
            capture_output = True,
            text = True,
            timeout = 10,
        )
        return out.stdout.strip() or None if out.returncode == 0 else None
    except Exception:
        return None


def _percentile(values: list[float], pct: float) -> float:
    """Nearest-rank percentile over a small sample (no numpy)."""
    if not values:
        return 0.0
    ordered = sorted(values)
    rank = int(math.ceil(pct / 100.0 * len(ordered))) - 1
    rank = max(0, min(rank, len(ordered) - 1))
    return ordered[rank]


def _is_cuda(device: Optional[str]) -> bool:
    return bool(device) and device.split(":", 1)[0] == "cuda"


def _cuda_reset_peak() -> None:
    import torch
    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats()


def _cuda_sync() -> None:
    import torch
    if torch.cuda.is_available():
        torch.cuda.synchronize()


def _cuda_peak_alloc() -> Optional[int]:
    import torch
    return int(torch.cuda.max_memory_allocated()) if torch.cuda.is_available() else None


def _cuda_peak_reserved() -> Optional[int]:
    import torch
    return int(torch.cuda.max_memory_reserved()) if torch.cuda.is_available() else None


def _cuda_alloc() -> Optional[int]:
    import torch
    return int(torch.cuda.memory_allocated()) if torch.cuda.is_available() else None


def _gpu_name() -> Optional[str]:
    try:
        import torch
        if torch.cuda.is_available():
            return torch.cuda.get_device_name(0)
    except Exception:
        pass
    return None


def _process_rss_bytes() -> Optional[int]:
    """Best-effort current-process RSS, without making the benchmark depend on psutil."""
    try:
        import psutil
        return int(psutil.Process().memory_info().rss)
    except Exception:
        pass
    try:
        if sys.platform.startswith("linux"):
            resident_pages = int(Path("/proc/self/statm").read_text().split()[1])
            return resident_pages * int(os.sysconf("SC_PAGE_SIZE"))
    except Exception:
        pass
    return None


def _finite_json_number(value: float) -> Optional[float]:
    """Keep benchmark JSON RFC-compliant when a metric is infinite or unavailable."""
    return value if math.isfinite(value) else None


def _versions() -> dict[str, Optional[str]]:
    out: dict[str, Optional[str]] = {"torch": None, "diffusers": None}
    try:
        import torch
        out["torch"] = torch.__version__
    except Exception:
        pass
    try:
        import diffusers
        out["diffusers"] = diffusers.__version__
    except Exception:
        pass
    return out


def _psnr(ref_png: Path, cand_png: Path) -> float:
    """PSNR (dB) between two PNGs; inf when identical."""
    import numpy as np
    from PIL import Image

    with Image.open(ref_png) as im_a:
        a = np.asarray(im_a.convert("RGB"), dtype = np.float64)
    with Image.open(cand_png) as im_b:
        b = np.asarray(im_b.convert("RGB"), dtype = np.float64)
    if a.shape != b.shape:
        # Different geometry means the comparison is meaningless; report worst case.
        return 0.0
    mse = float(((a - b) ** 2).mean())
    if mse == 0.0:
        return math.inf
    return 20.0 * math.log10(255.0) - 10.0 * math.log10(mse)


# ── load + generate ────────────────────────────────────────────────────────


def _wait_for_load(backend: Any, timeout_s: int = 2400) -> None:
    deadline = time.time() + timeout_s
    last = None
    while time.time() < deadline:
        p = backend.load_progress()
        phase = p.get("phase")
        if phase != last:
            last = phase
            frac = p.get("fraction") or 0.0
            bt = (p.get("bytes_total") or 0) / 1e9
            print(f"  load phase={phase} frac={frac:.3f} total={bt:.2f}GB", flush = True)
        if phase == "ready":
            return
        if phase == "error":
            raise RuntimeError(f"load error: {p.get('error')}")
        time.sleep(2)
    raise TimeoutError(f"model load did not reach ready within {timeout_s}s")


def _generate_once(backend: Any, args: argparse.Namespace) -> Any:
    """One generation at the fixed seed; returns the first PIL image."""
    result = backend.generate(
        prompt = args.prompt,
        width = args.width,
        height = args.height,
        steps = args.steps,
        guidance = args.guidance,
        seed = args.seed,
        batch_size = args.batch_size,
    )
    images = result["images"]
    return images[0]


def _run(args: argparse.Namespace) -> dict[str, Any]:
    """Load the model, measure load and generation, render the fixed-seed image. Returns the metrics dict; writes the rendered image to ``args._image_out``."""
    from core.inference.diffusion import get_diffusion_backend

    backend = get_diffusion_backend()
    status: dict[str, Any] = {}
    load_metrics: dict[str, Any] = {}
    gen_metrics: dict[str, Any] = {}

    try:
        _cuda_reset_peak()
        t0 = time.time()
        backend.begin_load(
            args.model,
            gguf_filename = args.gguf,
            base_repo = args.base_repo,
            family_override = args.family_override,
            hf_token = os.environ.get("HF_TOKEN"),
            cpu_offload = args.cpu_offload,
            memory_mode = args.memory_mode,
            speed_mode = args.speed_mode,
            text_encoder_quant = args.text_encoder_quant,
            transformer_quant = args.transformer_quant,
            transformer_quant_fast_accum = {"auto": None, "on": True, "off": False}[
                args.fp8_fast_accum
            ],
        )
        _wait_for_load(backend)
        _cuda_sync()
        load_metrics = {
            "wall_seconds": round(time.time() - t0, 2),
            "peak_vram_bytes": _cuda_peak_alloc(),
            "peak_reserved_bytes": _cuda_peak_reserved(),
            "final_vram_bytes": _cuda_alloc(),
        }
        status = backend.status()
        print(f"  loaded: {status}", flush = True)

        rss_after_load = _process_rss_bytes()

        # ── warmup (discarded) ──
        for _ in range(max(0, args.warmup)):
            _generate_once(backend, args)

        rss_after_warmup = _process_rss_bytes()

        # ── measured generations (fixed seed -> deterministic) ──
        _cuda_reset_peak()
        latencies: list[float] = []
        rss_after_generations: list[Optional[int]] = []
        first_image = None
        for i in range(max(1, args.iters)):
            _cuda_sync()
            g0 = time.time()
            image = _generate_once(backend, args)
            _cuda_sync()
            latencies.append(time.time() - g0)
            rss_after_generations.append(_process_rss_bytes())
            if first_image is None:
                first_image = image
            print(f"  gen[{i}] {latencies[-1]:.3f}s", flush = True)

        total = sum(latencies)
        measured_rss = [value for value in rss_after_generations if value is not None]
        post_warmup_rss_growth = (
            max(0, max(measured_rss) - rss_after_warmup)
            if rss_after_warmup is not None and measured_rss
            else None
        )

        gen_metrics = {
            "iters": len(latencies),
            "warmup": max(0, args.warmup),
            "latencies_s": [round(x, 4) for x in latencies],
            "median_latency_s": round(_percentile(latencies, 50), 4),
            "p90_latency_s": round(_percentile(latencies, 90), 4),
            "images_per_sec": round((args.batch_size * len(latencies)) / total, 4)
            if total > 0
            else None,
            "peak_vram_bytes": _cuda_peak_alloc(),
            "host_rss": {
                "after_load_bytes": rss_after_load,
                "after_warmup_bytes": rss_after_warmup,
                "after_each_generation_bytes": rss_after_generations,
                "post_warmup_growth_bytes": post_warmup_rss_growth,
            },
        }

        # The fixed-seed image is the accuracy anchor.
        args._image_out.parent.mkdir(parents = True, exist_ok = True)
        first_image.save(args._image_out)
        print(f"  saved image -> {args._image_out}", flush = True)
    finally:
        try:
            backend.unload()
        except Exception as exc:  # noqa: BLE001 — best-effort cleanup
            print(f"  warn: unload failed: {exc}", flush = True)

    return {
        "env": {
            "timestamp": _now_iso(),
            "git_commit": _git_commit(),
            "python": platform.python_version(),
            "platform": platform.platform(),
            "versions": _versions(),
            "gpu_name": _gpu_name(),
            "status": status,
        },
        "load": load_metrics,
        "generate": gen_metrics,
        "config": {
            "model": args.model,
            "gguf": args.gguf,
            "base_repo": args.base_repo,
            "family_override": args.family_override,
            "prompt": args.prompt,
            "width": args.width,
            "height": args.height,
            "steps": args.steps,
            "guidance": args.guidance,
            "seed": args.seed,
            "batch_size": args.batch_size,
            "memory_mode": args.memory_mode,
            "speed_mode": args.speed_mode,
            "cpu_offload": args.cpu_offload,
            "text_encoder_quant": args.text_encoder_quant,
            "transformer_quant": args.transformer_quant,
            "fp8_fast_accum": args.fp8_fast_accum,
        },
    }


# ── modes ──────────────────────────────────────────────────────────────────


def _write_baseline(args: argparse.Namespace) -> int:
    baseline_path = Path(args.write_baseline).resolve()
    ref_png = baseline_path.parent / "reference.png"
    args._image_out = ref_png

    metrics = _run(args)
    metrics["accuracy"] = {
        "reference_png": str(ref_png),
        "width": args.width,
        "height": args.height,
        "steps": args.steps,
        "seed": args.seed,
        "dtype": (metrics["env"]["status"] or {}).get("dtype"),
    }

    baseline_path.parent.mkdir(parents = True, exist_ok = True)
    baseline_path.write_text(json.dumps(metrics, indent = 2))
    print("\n=== BASELINE WRITTEN ===", flush = True)
    print(f"  json:      {baseline_path}", flush = True)
    print(f"  reference: {ref_png}", flush = True)
    print(f"  load: {metrics['load']}", flush = True)
    print(
        f"  generate: median={metrics['generate'].get('median_latency_s')}s "
        f"p90={metrics['generate'].get('p90_latency_s')}s "
        f"img/s={metrics['generate'].get('images_per_sec')} "
        f"peak_vram={metrics['generate'].get('peak_vram_bytes')}",
        flush = True,
    )

    rss_growth = (metrics["generate"].get("host_rss") or {}).get("post_warmup_growth_bytes")
    if rss_growth is not None:
        print(f"  host RSS growth after warmup: {rss_growth / 2**20:.1f} MiB", flush = True)
    return 0


def _compare(args: argparse.Namespace) -> int:
    baseline_path = Path(args.compare).resolve()
    baseline = json.loads(baseline_path.read_text())
    out_dir = Path(args.out_dir).resolve()
    args._image_out = out_dir / "compare.png"
    # --write-baseline takes any path, so a baseline can be sitting on a name this run writes. Refuse before the generation is paid for, not after.
    for written in (out_dir / "compare.json", args._image_out):
        if baseline_path == written:
            print(
                f"error: baseline {baseline_path} is the file this run writes; "
                f"pass a different --out-dir or rename the baseline",
                file = sys.stderr,
            )
            return 2

    # Refuse a noisy cross-hardware / cross-dtype comparison unless forced.
    base_env = baseline.get("env", {})
    base_status = base_env.get("status") or {}
    cur_gpu = _gpu_name()
    base_gpu = base_env.get("gpu_name")
    metrics = _run(args)
    cur_status = metrics["env"]["status"] or {}

    mismatch = []
    if base_gpu != cur_gpu:
        mismatch.append(f"gpu {base_gpu!r} -> {cur_gpu!r}")
    if base_status.get("device") != cur_status.get("device"):
        mismatch.append(f"device {base_status.get('device')!r} -> {cur_status.get('device')!r}")
    if base_status.get("dtype") != cur_status.get("dtype"):
        mismatch.append(f"dtype {base_status.get('dtype')!r} -> {cur_status.get('dtype')!r}")
    if mismatch:
        print("\n!! environment mismatch vs baseline: " + "; ".join(mismatch), flush = True)
        if not args.force_compare:
            print("   refusing noisy comparison (pass --force-compare to override).", flush = True)
            return 2

    # PSNR against the stored reference; reference_png is absolute, so fall back to reference.png beside the baseline.
    ref_png = Path(baseline.get("accuracy", {}).get("reference_png", ""))
    if not ref_png.is_file():
        ref_png = baseline_path.parent / "reference.png"
    psnr = _psnr(ref_png, args._image_out) if ref_png.is_file() else float("nan")

    base_gen = baseline.get("generate", {})
    cur_gen = metrics["generate"]
    base_median = base_gen.get("median_latency_s") or 0.0
    cur_median = cur_gen.get("median_latency_s") or 0.0
    latency_reg = (cur_median - base_median) / base_median if base_median > 0 else 0.0

    base_peak = base_gen.get("peak_vram_bytes")
    cur_peak = cur_gen.get("peak_vram_bytes")
    vram_reg = ((cur_peak - base_peak) / base_peak) if (base_peak and cur_peak) else 0.0

    base_rss_growth = (base_gen.get("host_rss") or {}).get("post_warmup_growth_bytes")
    cur_rss_growth = (cur_gen.get("host_rss") or {}).get("post_warmup_growth_bytes")

    print("\n=== REGRESSION REPORT ===", flush = True)
    print(f"  {'metric':<22}{'baseline':>16}{'current':>16}{'delta':>12}", flush = True)
    print(
        f"  {'median_latency_s':<22}{base_median:>16.4f}{cur_median:>16.4f}{latency_reg * 100:>11.1f}%",
        flush = True,
    )
    if base_peak and cur_peak:
        print(
            f"  {'peak_vram_MB':<22}{base_peak / 1e6:>16.1f}{cur_peak / 1e6:>16.1f}{vram_reg * 100:>11.1f}%",
            flush = True,
        )

    if base_rss_growth is not None or cur_rss_growth is not None:
        base_rss_label = f"{base_rss_growth / 2**20:.1f}" if base_rss_growth is not None else "-"
        cur_rss_label = f"{cur_rss_growth / 2**20:.1f}" if cur_rss_growth is not None else "-"
        rss_delta_label = (
            f"{(cur_rss_growth - base_rss_growth) / 2**20:+.1f}"
            if base_rss_growth is not None and cur_rss_growth is not None
            else "-"
        )
        print(
            f"  {'host_rss_growth_MiB':<22}{base_rss_label:>16}{cur_rss_label:>16}{rss_delta_label:>12}",
            flush = True,
        )
    print(f"  {'psnr_dB(vs ref)':<22}{'-':>16}{psnr:>16.2f}{'':>12}", flush = True)

    failures = []
    if latency_reg > args.max_latency_regression:
        failures.append(
            f"latency +{latency_reg * 100:.1f}% > {args.max_latency_regression * 100:.0f}%"
        )
    if base_peak and cur_peak and vram_reg > args.max_vram_regression:
        failures.append(f"peak VRAM +{vram_reg * 100:.1f}% > {args.max_vram_regression * 100:.0f}%")

    if args.max_host_rss_growth_mib is not None:
        if cur_rss_growth is None:
            failures.append("host RSS unavailable; cannot verify the configured growth limit")
        elif cur_rss_growth > args.max_host_rss_growth_mib * 2**20:
            failures.append(
                f"host RSS growth {cur_rss_growth / 2**20:.1f} MiB > "
                f"{args.max_host_rss_growth_mib:.1f} MiB"
            )
    if math.isnan(psnr):
        failures.append("PSNR reference image missing; cannot verify output quality")
    elif psnr < args.min_psnr:
        failures.append(f"PSNR {psnr:.2f}dB < {args.min_psnr:.1f}dB (output changed)")

    metrics["comparison"] = {
        "baseline_json": str(baseline_path),
        "psnr_db": _finite_json_number(psnr),
        "latency_regression": latency_reg,
        "vram_regression": vram_reg,
        "failures": list(failures),
    }
    out_dir.mkdir(parents = True, exist_ok = True)
    (out_dir / "compare.json").write_text(json.dumps(metrics, indent = 2, allow_nan = False))

    if failures:
        print("\n  FAIL: " + "; ".join(failures), flush = True)
        return 1
    print("\n  PASS: no regression beyond thresholds.", flush = True)
    return 0


# ── cli ────────────────────────────────────────────────────────────────────


def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description = "Benchmark + regression guard for the Unsloth diffusion backend.",
        formatter_class = argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument(
        "--model", default = "unsloth/Z-Image-Turbo-GGUF", help = "GGUF repo id or local path"
    )
    p.add_argument(
        "--gguf",
        default = "z-image-turbo-Q4_K_M.gguf",
        help = "transformer GGUF filename inside --model",
    )
    p.add_argument("--base-repo", default = None, help = "override the diffusers base repo")
    p.add_argument("--family-override", default = None, help = "force a diffusion family")
    p.add_argument(
        "--prompt",
        default = "A cozy reading nook by a rain-streaked window, warm lamplight, "
        "a cat asleep on a stack of books, highly detailed",
    )
    p.add_argument("--width", type = int, default = 1024)
    p.add_argument("--height", type = int, default = 1024)
    p.add_argument("--steps", type = int, default = 9)
    p.add_argument("--guidance", type = float, default = 0.0)
    p.add_argument("--seed", type = int, default = 12345, help = "fixed seed -> deterministic image")
    p.add_argument("--batch-size", type = int, default = 1)
    p.add_argument("--warmup", type = int, default = 1, help = "discarded warmup generations")
    p.add_argument("--iters", type = int, default = 3, help = "measured generations")
    p.add_argument(
        "--memory-mode",
        default = None,
        choices = ["auto", "fast", "balanced", "low_vram"],
        help = "memory policy (default: backend auto)",
    )
    p.add_argument(
        "--speed-mode",
        default = None,
        choices = ["off", "default", "max"],
        help = "speed profile: off is bit-identical; default adds compile + "
        "cudnn.benchmark (near-lossless); max also adds TF32 + fused QKV",
    )
    p.add_argument(
        "--text-encoder-quant",
        default = None,
        choices = ["fp8", "nvfp4"],
        help = "quantise the companion text encoder (fp8 or nvfp4)",
    )
    p.add_argument(
        "--transformer-quant",
        default = None,
        choices = ["auto", "int8", "fp8", "nvfp4", "mxfp8"],
        help = "opt-in fast transformer: load the DENSE bf16 transformer and torchao-"
        "quantise it onto the low-precision tensor cores (faster than GGUF, higher "
        "VRAM). auto picks per GPU; falls back to GGUF if unsupported / no VRAM",
    )
    p.add_argument(
        "--fp8-fast-accum",
        default = "auto",
        choices = ["auto", "on", "off"],
        help = "fp8 accumulate: auto picks by GPU class (fast on consumer, precise on "
        "data-center); on/off force it",
    )
    p.add_argument(
        "--cpu-offload", action = "store_true", help = "legacy: force whole-module CPU offload"
    )
    p.add_argument(
        "--write-baseline",
        metavar = "PATH",
        default = None,
        help = "run once and save metrics JSON + reference.png",
    )
    p.add_argument(
        "--compare", metavar = "PATH", default = None, help = "run again and diff against a baseline JSON"
    )
    p.add_argument(
        "--max-latency-regression",
        type = float,
        default = 0.10,
        help = "fail if median latency rises by more than this fraction",
    )
    p.add_argument(
        "--max-vram-regression",
        type = float,
        default = 0.10,
        help = "fail if peak generation VRAM rises by more than this fraction",
    )
    p.add_argument(
        "--max-host-rss-growth-mib",
        type = float,
        default = None,
        help = "fail comparison if peak post-warmup process RSS growth exceeds this many MiB",
    )
    p.add_argument(
        "--min-psnr",
        type = float,
        default = 35.0,
        help = "fail if the fixed-seed image PSNR vs reference drops below this",
    )
    p.add_argument(
        "--force-compare",
        action = "store_true",
        help = "compare even when GPU/device/dtype differ from the baseline",
    )
    p.add_argument(
        "--out-dir", default = "outputs/diffusion_bench", help = "where compare.png is written"
    )
    return p


def main(argv: Optional[list[str]] = None) -> int:
    args = _build_parser().parse_args(argv)
    if "nvfp4" in (args.transformer_quant, args.text_encoder_quant):
        os.environ["UNSLOTH_NVFP4_DIFFUSION"] = "1"
    if isinstance(args.gguf, str):
        args.gguf = args.gguf.strip() or None
    if bool(args.write_baseline) == bool(args.compare):
        print("error: pass exactly one of --write-baseline / --compare", file = sys.stderr)
        return 2
    if args.write_baseline:
        return _write_baseline(args)
    return _compare(args)


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