benchmarks/gate_g1.py

"""Gate G1: FLUX.2 Klein latency benchmark.

Parameter contribution: 0B. This script loads runtime model weights only when
executed in benchmark mode. Local dry-runs are non-authoritative and must never
be used as Space benchmark numbers.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import statistics
import sys
import time
from typing import Any

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

from lightloom.core.config import MODEL_REFS  # noqa: E402

RESULTS_DIR = ROOT / "benchmarks" / "results"
DEFAULT_JSON = RESULTS_DIR / "g1.json"
DEFAULT_MD = RESULTS_DIR / "g1.md"
PROMPT = (
    "cinematic illustrated film still, an old lighthouse at the edge of the world, "
    "stormy dusk, warm lantern glow, painterly texture, volumetric light, 2.39:1 frame"
)
RESOLUTIONS = ((768, 432), (1024, 576))
DTYPES = ("fp8", "bf16")
AOT_FLAGS = (False, True)


def _percentile(values: list[float], q: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    idx = min(len(ordered) - 1, max(0, round((len(ordered) - 1) * q)))
    return ordered[idx]


def _torch_dtype(name: str) -> Any:
    import torch

    if name == "bf16":
        return torch.bfloat16
    if name == "fp8":
        return torch.float8_e4m3fn
    raise ValueError(f"unsupported dtype: {name}")


def _hardware_profile() -> dict[str, Any]:
    import os
    import platform

    profile: dict[str, Any] = {
        "lightloom_profile": os.getenv("LIGHTLOOM_PROFILE", "local"),
        "space_id": os.getenv("SPACE_ID") or os.getenv("LIGHTLOOM_DEV_SPACE_ID"),
        "platform": platform.platform(),
    }
    try:
        import torch

        profile.update(
            {
                "torch": torch.__version__,
                "cuda": getattr(torch.version, "cuda", None),
                "cuda_available": bool(torch.cuda.is_available()),
                "device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
                "arch_list": list(torch.cuda.get_arch_list()) if torch.cuda.is_available() else [],
            }
        )
    except Exception as exc:  # noqa: BLE001
        profile["torch_error"] = f"{type(exc).__name__}: {exc}"
    return profile


def _dry_run_combination(width: int, height: int, dtype: str, aot: bool, reps: int) -> dict[str, Any]:
    timings = []
    for i in range(reps):
        start = time.perf_counter()
        time.sleep(0.002 + (i * 0.0001))
        timings.append((time.perf_counter() - start) * 1000)
    return {
        "width": width,
        "height": height,
        "dtype": dtype,
        "aot": aot,
        "status": "dry_run",
        "authoritative": False,
        "timings_ms": [round(v, 3) for v in timings],
        "p50_ms": round(statistics.median(timings), 3),
        "p95_ms": round(_percentile(timings, 0.95), 3),
        "note": "plumbing validation only; not a benchmark",
    }


def _load_pipeline(dtype: str) -> Any:
    import torch
    from diffusers import Flux2KleinPipeline

    ref = MODEL_REFS["painter"]
    return Flux2KleinPipeline.from_pretrained(
        ref.repo_id,
        revision=ref.revision,
        torch_dtype=_torch_dtype(dtype),
    ).to("cuda")


def _maybe_compile(pipe: Any) -> str:
    import torch

    if not hasattr(pipe, "transformer"):
        return "no_transformer_attr"
    pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead", fullgraph=False)
    return "torch.compile(reduce-overhead)"


def _benchmark_combination(
    pipe: Any,
    width: int,
    height: int,
    dtype: str,
    aot: bool,
    reps: int,
) -> dict[str, Any]:
    import torch

    compile_mode = "off"
    if aot:
        compile_mode = _maybe_compile(pipe)
    generator = torch.Generator(device="cuda").manual_seed(1901)
    timings: list[float] = []
    try:
        with torch.inference_mode():
            pipe(
                prompt=PROMPT,
                height=height,
                width=width,
                num_inference_steps=4,
                guidance_scale=1.0,
                generator=generator,
                output_type="pil",
            )
            torch.cuda.synchronize()
            for rep in range(reps):
                generator = torch.Generator(device="cuda").manual_seed(1901 + rep)
                start = time.perf_counter()
                pipe(
                    prompt=PROMPT,
                    height=height,
                    width=width,
                    num_inference_steps=4,
                    guidance_scale=1.0,
                    generator=generator,
                    output_type="pil",
                )
                torch.cuda.synchronize()
                timings.append((time.perf_counter() - start) * 1000)
        return {
            "width": width,
            "height": height,
            "dtype": dtype,
            "aot": aot,
            "status": "ok",
            "authoritative": True,
            "compile_mode": compile_mode,
            "timings_ms": [round(v, 3) for v in timings],
            "p50_ms": round(statistics.median(timings), 3),
            "p95_ms": round(_percentile(timings, 0.95), 3),
        }
    except Exception as exc:  # noqa: BLE001
        return {
            "width": width,
            "height": height,
            "dtype": dtype,
            "aot": aot,
            "status": "failed",
            "authoritative": True,
            "compile_mode": compile_mode,
            "error": f"{type(exc).__name__}: {exc}",
        }


def _decide(results: list[dict[str, Any]]) -> dict[str, Any]:
    successful = [item for item in results if item.get("status") == "ok" and item.get("authoritative")]
    if not successful:
        return {"status": "blocked", "reason": "no authoritative successful G1 run"}
    preferred = [item for item in successful if item["dtype"] == "fp8"] or successful
    by_resolution: dict[tuple[int, int], dict[str, Any]] = {}
    for item in sorted(preferred, key=lambda x: (x["p50_ms"], x["dtype"] != "fp8", x["aot"] is False)):
        by_resolution.setdefault((item["width"], item["height"]), item)
    high = by_resolution.get((1024, 576))
    low = by_resolution.get((768, 432))
    reason = "fp8" if any(item["dtype"] == "fp8" for item in preferred) else "fp8 unavailable; selected fastest successful dtype"
    if high and high["p50_ms"] <= 1800:
        return {"status": "pass", "resolution": [1024, 576], "selected": high, "reason": reason}
    if low and low["p50_ms"] <= 3000:
        return {"status": "fallback", "resolution": [768, 432], "selected": low, "reason": reason}
    return {"status": "fallback_polaroid", "resolution": [768, 432], "selected": low or high, "reason": reason}


def _write_markdown(path: Path, data: dict[str, Any]) -> None:
    rows = [
        "# Gate G1 Results",
        "",
        f"Authoritative: `{data['authoritative']}`",
        f"Hardware profile: `{data['hardware_profile'].get('lightloom_profile')}`",
        "",
        "| Resolution | dtype | AoT | status | p50 ms | p95 ms |",
        "|---|---|---:|---|---:|---:|",
    ]
    for item in data["results"]:
        rows.append(
            "| {w}x{h} | {dtype} | {aot} | {status} | {p50} | {p95} |".format(
                w=item["width"],
                h=item["height"],
                dtype=item["dtype"],
                aot="yes" if item["aot"] else "no",
                status=item["status"],
                p50=item.get("p50_ms", ""),
                p95=item.get("p95_ms", ""),
            )
        )
    rows.extend(["", f"Decision: `{data['decision']['status']}`"])
    path.write_text("\n".join(rows) + "\n", encoding="utf-8")


def run(dry_run: bool, reps: int, allow_local: bool) -> dict[str, Any]:
    import os

    profile = os.getenv("LIGHTLOOM_PROFILE", "local")
    if not dry_run and profile != "space" and not allow_local:
        raise SystemExit(
            "Refusing to run G1 benchmark outside LIGHTLOOM_PROFILE=space. "
            "Use --dry-run for plumbing or --allow-local for private debugging; "
            "local timings are non-authoritative."
        )
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    results: list[dict[str, Any]] = []
    if dry_run:
        for width, height in RESOLUTIONS:
            for dtype in DTYPES:
                for aot in AOT_FLAGS:
                    results.append(_dry_run_combination(width, height, dtype, aot, reps))
    else:
        for dtype in DTYPES:
            try:
                pipe = _load_pipeline(dtype)
            except Exception as exc:  # noqa: BLE001
                for width, height in RESOLUTIONS:
                    for aot in AOT_FLAGS:
                        results.append(
                            {
                                "width": width,
                                "height": height,
                                "dtype": dtype,
                                "aot": aot,
                                "status": "failed",
                                "authoritative": profile == "space",
                                "error": f"load failed: {type(exc).__name__}: {exc}",
                            }
                        )
                continue
            for width, height in RESOLUTIONS:
                for aot in AOT_FLAGS:
                    results.append(_benchmark_combination(pipe, width, height, dtype, aot, reps))
            del pipe
    data = {
        "schema_version": 1,
        "gate": "G1",
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "authoritative": (not dry_run and profile == "space"),
        "hardware_profile": _hardware_profile(),
        "reps": reps,
        "prompt": PROMPT,
        "results": results,
        "decision": _decide(results),
    }
    DEFAULT_JSON.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    _write_markdown(DEFAULT_MD, data)
    return data


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--allow-local", action="store_true")
    parser.add_argument("--reps", type=int, default=5)
    args = parser.parse_args()
    data = run(dry_run=args.dry_run, reps=args.reps, allow_local=args.allow_local)
    print(f"G1 wrote {DEFAULT_JSON}")
    print(f"G1 decision: {data['decision']['status']}")
    if args.dry_run:
        return 0
    return 0 if data["decision"]["status"] in {"pass", "fallback", "fallback_polaroid"} else 1


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