bench_flux_lora_latency.py

#!/usr/bin/env python3
"""
bench_flux_lora_latency.py
==========================
Q2 REPRODUCTION BENCHMARK — MPS Latency Delta: flux-q4 base vs base+fused-LoRA
vs base+unfused-LoRA on the imageforge _load_flux_gguf path.

Usage
-----
    # 1. Activate the imageforge venv (arm64 Python):
    arch -arm64 /path/to/venv/bin/python bench_flux_lora_latency.py

    # 2. Mandatory env vars:
    export ALLOW_FLUX=1
    export IMAGEFORGE_FLUX_GGUF_URL="https://huggingface.co/city96/FLUX.1-schnell-gguf/resolve/main/flux1-schnell-Q4_0.gguf"

    # 3. Optional – point at a real FLUX-compatible LoRA safetensors (key-matched):
    export BENCH_LORA_PATH="path/to/lora.safetensors"
    # If unset the bench auto-downloads ostris/FLUX.1-schnell-training-adapter
    # as a structurally compatible weight set.

What this measures
------------------
Three pipeline configurations, each run N=20 times at 4 steps / 1024×1024
with a fixed 20-prompt seed list:

  A. base        – raw GGUF transformer, no LoRA
  B. fused       – load_lora_weights() → fuse_lora() (if fuse raises, documented)
  C. unfused     – load_lora_weights() → set_adapters() at call time

Each run is timed with torch.mps.synchronize() guards so wall-clock ≈ GPU time.

Output: markdown table + per-config mean, p95, per_step_ms, fuse_succeeded,
output_std, verdict.

Key research findings wired into this script
--------------------------------------------
• fuse_lora() on a GGUF-quantized FluxTransformer2DModel RAISES in diffusers
  ≤0.36: PEFT's merger calls `base_layer.weight.data += delta_weight` which hits
  a uint8/QTensor in-place add — RuntimeError or NotImplementedError depending
  on diffusers version.  Reference: github.com/huggingface/diffusers/issues/12047,
  #10492, #10381.  The bench catches the exact exception, records it, and falls
  back to unfused-scaled path for the fused slot.
• load_lora_weights() itself on a GGUF pipeline CAN succeed (the LoRA tensors
  land in PEFT hook layers, not in the quantized base weights) — confirmed in
  community testing (city96/FLUX.1-dev-gguf/discussions/42).
• Round-1 claim "fuse-at-load gives zero per-step overhead on MPS" is unverified
  and likely wrong for GGUF: since fuse itself likely fails, the question becomes
  whether set_adapters (unfused) adds measurable overhead vs base.
• schnell LoRA quality note: schnell is timestep-distilled; LoRAs trained on dev
  show reduced effect at 4 steps. LoRAs trained with ostris assistant adapter
  (ostris/FLUX.1-schnell-training-adapter, disabled at inference) are preferred.
  Structural similarity (SSIM) vs base is measured but NOT treated as a
  quality-of-LoRA verdict — it measures whether fusing corrupts the distilled
  checkpoint vs merely adding adaptation signal.
• MPS timing: MPS is ASYNC; must call torch.mps.synchronize() before stopping
  the clock. cpu-seeded Generator is used (MPS generators are non-deterministic).
"""

from __future__ import annotations

import argparse
import os
import statistics
import sys
import time
import traceback
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

# ---------------------------------------------------------------------------
# Argument validation — done BEFORE any heavy import so errors are instant
# ---------------------------------------------------------------------------


def _validate_args(args: argparse.Namespace) -> None:
    """Validate parsed args and exit with a clear message on bad input."""
    errors: List[str] = []

    if args.n < 1:
        errors.append(f"--n must be >= 1 (got {args.n})")
    if args.steps < 1:
        errors.append(f"--steps must be >= 1 (got {args.steps})")
    if args.resolution < 8:
        errors.append(f"--resolution must be >= 8 (got {args.resolution})")
    elif args.resolution % 8 != 0:
        errors.append(
            f"--resolution must be a multiple of 8 (got {args.resolution}); "
            "diffusers VAE requires dimensions divisible by 8"
        )
    if args.lora_scale <= 0:
        errors.append(f"--lora-scale must be > 0 (got {args.lora_scale})")

    if errors:
        for msg in errors:
            print(f"[ERROR] {msg}", file=sys.stderr)
        sys.exit(2)


# ---------------------------------------------------------------------------
# Guard: must be arm64 Python + have required env vars
# ---------------------------------------------------------------------------


def _check_env() -> None:
    import platform
    machine = platform.machine()
    if machine != "arm64":
        print(
            f"[WARN] Running on {machine}, not arm64. "
            "MPS timings will be meaningless (fallback to CPU). "
            "Relaunch with: arch -arm64 python bench_flux_lora_latency.py"
        )
    if not os.environ.get("ALLOW_FLUX"):
        print(
            "[ERROR] ALLOW_FLUX=1 not set. "
            "flux-q4 peaks at 16-18 GB; set ALLOW_FLUX=1 to confirm you accept the swap risk.",
            file=sys.stderr,
        )
        sys.exit(1)
    if not os.environ.get("IMAGEFORGE_FLUX_GGUF_URL"):
        print(
            "[ERROR] IMAGEFORGE_FLUX_GGUF_URL not set. "
            "Example:\n"
            '  export IMAGEFORGE_FLUX_GGUF_URL="https://huggingface.co/city96/'
            'FLUX.1-schnell-gguf/resolve/main/flux1-schnell-Q4_0.gguf"',
            file=sys.stderr,
        )
        sys.exit(1)


# ---------------------------------------------------------------------------
# Dependency guard — actionable error, not a bare traceback
# ---------------------------------------------------------------------------


def _check_deps() -> None:
    """
    Verify that diffusers and the GGUF quantisation shim are importable.
    Called after arg validation so --help / --dry-run never trigger it.
    """
    missing: List[str] = []

    try:
        import diffusers  # noqa: F401
    except ImportError:
        missing.append(
            "diffusers (install with: "
            "pip install 'diffusers>=0.32.0')"
        )

    if not missing:
        try:
            from diffusers import GGUFQuantizationConfig  # noqa: F401
        except (ImportError, AttributeError):
            missing.append(
                "diffusers.GGUFQuantizationConfig is not available — "
                "upgrade diffusers: pip install 'diffusers>=0.32.0'"
            )

    try:
        import torch  # noqa: F401
    except ImportError:
        missing.append(
            "torch (install with: pip install torch or follow "
            "https://pytorch.org/get-started/locally/)"
        )

    if missing:
        print("[ERROR] Missing required dependencies:", file=sys.stderr)
        for m in missing:
            print(f"  - {m}", file=sys.stderr)
        sys.exit(1)


# ---------------------------------------------------------------------------
# Prompts — 20 fixed seeds for reproducibility
# ---------------------------------------------------------------------------

PROMPTS = [
    "a golden retriever running through a field of lavender at sunset, photorealistic",
    "an astronaut floating above earth holding a coffee cup, dramatic lighting",
    "a medieval castle on a cliff by the sea, mist, cinematic",
    "portrait of a red fox wearing a scarf, studio lighting, detailed fur",
    "cyberpunk city at night, rain reflections on neon-lit streets, 8K",
    "a hot air balloon over the Grand Canyon at dawn",
    "watercolor painting of a Japanese temple in autumn foliage",
    "a vintage motorcycle parked outside a diner, 1950s Americana",
    "macro photograph of a dewdrop on a spider web at sunrise",
    "a black cat sitting on a stack of old books, cozy library",
    "aerial view of a tropical island with crystal clear water",
    "a muscular robot blacksmith forging a sword in a medieval forge",
    "surrealist painting: melting clocks in a desert landscape, Dali style",
    "a lighthouse during a dramatic ocean storm, waves crashing",
    "close-up of a hummingbird hovering near a red flower",
    "snowy mountain peak above the clouds at golden hour",
    "a steampunk airship flying over Victorian London",
    "a wolf howling at a full moon in a pine forest",
    "underwater scene with bioluminescent jellyfish and coral",
    "a chef plating an artistic dish in a modern kitchen, food photography",
]

SEEDS = list(range(42, 42 + len(PROMPTS)))  # 42..61


# ---------------------------------------------------------------------------
# Result container
# ---------------------------------------------------------------------------

@dataclass
class BenchResult:
    config: str
    fuse_attempted: bool
    fuse_succeeded: bool
    fuse_exception: str
    latencies_s: List[float] = field(default_factory=list)
    output_stds: List[float] = field(default_factory=list)

    def mean_s(self) -> float:
        return statistics.mean(self.latencies_s) if self.latencies_s else float("nan")

    def p95_s(self) -> float:
        if not self.latencies_s:
            return float("nan")
        sorted_l = sorted(self.latencies_s)
        idx = min(int(len(sorted_l) * 0.95), len(sorted_l) - 1)
        return sorted_l[idx]

    def per_step_ms(self, steps: int = 4) -> float:
        m = self.mean_s()
        return (m / steps) * 1000 if m == m else float("nan")  # nan check

    def mean_std(self) -> float:
        return statistics.mean(self.output_stds) if self.output_stds else float("nan")


# ---------------------------------------------------------------------------
# MPS-safe timing helper
# ---------------------------------------------------------------------------

def _now_mps(device: str) -> float:
    """Wall-clock time after synchronizing the MPS device (or CUDA)."""
    if device == "mps":
        try:
            import torch
            torch.mps.synchronize()
        except Exception:
            pass
    elif device == "cuda":
        try:
            import torch
            torch.cuda.synchronize()
        except Exception:
            pass
    return time.perf_counter()


# ---------------------------------------------------------------------------
# Output std helper (pixel-level std across outputs from same prompt)
# ---------------------------------------------------------------------------

def _image_std(img) -> float:
    """Return pixel standard deviation of a PIL image (grayscale proxy)."""
    try:
        import numpy as np
        arr = np.array(img.convert("L")).astype(float)
        return float(arr.std())
    except Exception:
        return float("nan")


# ---------------------------------------------------------------------------
# GGUF pipeline loader — mirrors imageforge engine/_load_flux_gguf exactly
# ---------------------------------------------------------------------------

def _load_base_pipe(gguf_url: str, device: str):
    """Load the GGUF transformer + FluxPipeline. fp16, matches pipeline.py:478-488."""
    import torch
    from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig

    print(f"[bench] Loading GGUF transformer from: {gguf_url}")
    transformer = FluxTransformer2DModel.from_single_file(
        gguf_url,
        quantization_config=GGUFQuantizationConfig(compute_dtype=torch.float16),
        torch_dtype=torch.float16,
    )
    pipe = FluxPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-schnell",
        transformer=transformer,
        torch_dtype=torch.float16,
    )
    pipe = pipe.to(device)
    # VAE tiling — essential on 16GB (mirrors pipeline.py:490-496)
    for fn in ("enable_vae_tiling", "enable_vae_slicing"):
        meth = getattr(pipe, fn, None)
        if callable(meth):
            try:
                meth()
            except Exception:
                pass
    return pipe


def _resolve_lora_path() -> str:
    """
    Return local path or HF repo+filename for a structurally compatible FLUX LoRA.

    Priority:
      1. BENCH_LORA_PATH env var (user-supplied safetensors, must have FLUX keys)
      2. ostris/FLUX.1-schnell-training-adapter (assistant adapter; safe to use
         as weight-compatible LoRA at inference — it was trained on FLUX.1-schnell
         transformer, matches key shape; disabled at inference in ai-toolkit but
         the weights are structurally FLUX-compatible)

    NOTE: The ostris assistant adapter is designed to be used ONLY during training
    (disabled during sampling per the README). Using it as a generic LoRA here
    serves only as a key-shape-compatible test object to probe whether fuse_lora
    succeeds without error. Quality output with this adapter is not meaningful.
    If you have a real style LoRA trained on schnell via ai-toolkit, supply it
    via BENCH_LORA_PATH for meaningful Q3 (quality) measurements.
    """
    env = os.environ.get("BENCH_LORA_PATH")
    if env:
        print(f"[bench] Using user-supplied LoRA: {env}")
        return env
    print(
        "[bench] BENCH_LORA_PATH not set; using ostris/FLUX.1-schnell-training-adapter "
        "as a structurally compatible key-shape test object. "
        "Quality results with this adapter are NOT meaningful."
    )
    return "ostris/FLUX.1-schnell-training-adapter"


# ---------------------------------------------------------------------------
# Single inference call with timing
# ---------------------------------------------------------------------------

def _run_one(
    pipe,
    prompt: str,
    seed: int,
    steps: int,
    resolution: int,
    device: str,
    lora_scale: Optional[float] = None,
) -> Tuple[float, float]:
    """
    Run one inference. Returns (wall_seconds, output_pixel_std).
    lora_scale: if set, passed as cross_attention_kwargs for unfused path.
    """
    import torch
    gen = torch.Generator(device="cpu").manual_seed(seed)

    kwargs: Dict[str, Any] = dict(
        prompt=prompt,
        num_inference_steps=steps,
        width=resolution,
        height=resolution,
        guidance_scale=0.0,         # schnell: no CFG
        max_sequence_length=256,
        generator=gen,
    )
    if lora_scale is not None:
        kwargs["joint_attention_kwargs"] = {"scale": lora_scale}

    t0 = _now_mps(device)
    result = pipe(**kwargs)
    t1 = _now_mps(device)

    img = result.images[0]
    return (t1 - t0), _image_std(img)


# ---------------------------------------------------------------------------
# Warmup
# ---------------------------------------------------------------------------

def _warmup(pipe, device: str, steps: int, resolution: int) -> None:
    print("[bench] Warming up (2 runs at reduced resolution) …")
    import torch
    for i in range(2):
        gen = torch.Generator(device="cpu").manual_seed(0)
        try:
            pipe(
                prompt="warmup",
                num_inference_steps=steps,
                width=512,
                height=512,
                guidance_scale=0.0,
                max_sequence_length=256,
                generator=gen,
            )
        except Exception as e:
            print(f"  [warn] warmup run {i} raised: {e}")
    _now_mps(device)  # drain
    print("[bench] Warmup done.")


# ---------------------------------------------------------------------------
# Config A — base (no LoRA)
# ---------------------------------------------------------------------------

def bench_base(pipe, device: str, steps: int, resolution: int, n: int) -> BenchResult:
    result = BenchResult(
        config="A_base",
        fuse_attempted=False,
        fuse_succeeded=False,
        fuse_exception="",
    )
    print(f"\n[bench] Config A: base GGUF, no LoRA — {n} runs")
    for i, (prompt, seed) in enumerate(zip(PROMPTS[:n], SEEDS[:n])):
        elapsed, std = _run_one(pipe, prompt, seed, steps, resolution, device)
        result.latencies_s.append(elapsed)
        result.output_stds.append(std)
        print(f"  run {i+1:02d}/{n}  {elapsed:.2f}s  pixel_std={std:.1f}")
    return result


# ---------------------------------------------------------------------------
# Config B — fused LoRA
# ---------------------------------------------------------------------------

def bench_fused(
    pipe,
    lora_path: str,
    device: str,
    steps: int,
    resolution: int,
    n: int,
) -> BenchResult:
    """
    Attempt load_lora_weights → fuse_lora → unload_lora_weights.
    Catches the known exception class from GGUF+PEFT merge incompatibility.

    Known failure mode (diffusers ≤0.36):
      In diffusers/quantizers/gguf/utils.py line 428 in __torch_function__,
      PEFT's `base_layer.weight.data += delta_weight` tries in-place add on
      a GGUF QTensor. Raises:
        RuntimeError: "The size of tensor a (N) must match the size of tensor b (M)
          at non-singleton dimension K"
      or on torchao path:
        NotImplementedError: "AffineQuantizedTensor dispatch: attempting to run
          unimplemented operator/function: func=<OpOverload(op='aten.add_',
          overload='Tensor')>"
    If fuse fails, we run the benchmark with the adapter loaded but unfused
    (same config as C) and record fuse_succeeded=False.
    """
    result = BenchResult(
        config="B_fused",
        fuse_attempted=True,
        fuse_succeeded=False,
        fuse_exception="",
    )
    print(f"\n[bench] Config B: fused LoRA — attempting load+fuse from {lora_path}")

    # Step 1: load LoRA weights
    load_ok = False
    try:
        pipe.load_lora_weights(lora_path, adapter_name="bench_lora")
        load_ok = True
        print("  [OK] load_lora_weights succeeded")
    except Exception as e:
        exc_text = traceback.format_exc()
        result.fuse_exception = f"load_lora_weights FAILED: {e}\n{exc_text}"
        print(f"  [FAIL] load_lora_weights raised:\n{exc_text}")
        # Can't proceed; record N/A latencies
        result.latencies_s = [float("nan")] * n
        result.output_stds = [float("nan")] * n
        return result

    # Step 2: attempt fuse
    try:
        pipe.fuse_lora(lora_scale=1.0)
        pipe.unload_lora_weights()
        result.fuse_succeeded = True
        print("  [OK] fuse_lora succeeded — running fused benchmark")
    except Exception as e:
        exc_text = traceback.format_exc()
        result.fuse_exception = f"fuse_lora RAISED: {type(e).__name__}: {e}\n{exc_text}"
        print(
            f"  [EXPECTED FAIL] fuse_lora raised {type(e).__name__}: {e}\n"
            "  Falling back to unfused-scaled path for timing (same as config C).\n"
            "  See bench doc: this is the known GGUF+PEFT in-place add incompatibility."
        )
        # Keep adapter loaded but unfused — run timing anyway to measure "fused slot"
        # with adapter present but not fused (so the numbers are comparable to C)
        try:
            pipe.set_adapters("bench_lora", adapter_weights=1.0)
        except Exception as se:
            print(f"  [warn] set_adapters also failed: {se}")

    print(f"  Running {n} inference passes (fuse_succeeded={result.fuse_succeeded})")
    lora_scale_kwarg = None if result.fuse_succeeded else 1.0
    for i, (prompt, seed) in enumerate(zip(PROMPTS[:n], SEEDS[:n])):
        elapsed, std = _run_one(
            pipe, prompt, seed, steps, resolution, device,
            lora_scale=lora_scale_kwarg,
        )
        result.latencies_s.append(elapsed)
        result.output_stds.append(std)
        print(f"  run {i+1:02d}/{n}  {elapsed:.2f}s  pixel_std={std:.1f}")

    # Clean up for config C
    try:
        if not result.fuse_succeeded:
            pipe.unload_lora_weights()
    except Exception:
        pass

    return result


# ---------------------------------------------------------------------------
# Config C — unfused LoRA (set_adapters at call time)
# ---------------------------------------------------------------------------

def bench_unfused(
    pipe,
    lora_path: str,
    device: str,
    steps: int,
    resolution: int,
    n: int,
    lora_scale: float = 1.0,
) -> BenchResult:
    """
    load_lora_weights → set_adapters (active but not fused) → run.
    Per-call overhead = C.mean - A.mean.
    """
    result = BenchResult(
        config="C_unfused_scaled",
        fuse_attempted=False,
        fuse_succeeded=False,
        fuse_exception="",
    )
    print(f"\n[bench] Config C: unfused LoRA (set_adapters scale={lora_scale}) — {n} runs")

    try:
        pipe.load_lora_weights(lora_path, adapter_name="bench_lora_c")
        pipe.set_adapters("bench_lora_c", adapter_weights=lora_scale)
        print("  [OK] load_lora_weights + set_adapters succeeded")
    except Exception as e:
        exc_text = traceback.format_exc()
        result.fuse_exception = f"load/set_adapters FAILED: {e}\n{exc_text}"
        print(f"  [FAIL] load_lora_weights raised:\n{exc_text}")
        result.latencies_s = [float("nan")] * n
        result.output_stds = [float("nan")] * n
        return result

    for i, (prompt, seed) in enumerate(zip(PROMPTS[:n], SEEDS[:n])):
        elapsed, std = _run_one(
            pipe, prompt, seed, steps, resolution, device, lora_scale=lora_scale
        )
        result.latencies_s.append(elapsed)
        result.output_stds.append(std)
        print(f"  run {i+1:02d}/{n}  {elapsed:.2f}s  pixel_std={std:.1f}")

    try:
        pipe.unload_lora_weights()
    except Exception:
        pass

    return result


# ---------------------------------------------------------------------------
# SSIM delta helper (Q3: does fuse corrupt the distilled checkpoint?)
# ---------------------------------------------------------------------------

def _ssim_delta(base_result: BenchResult, lora_result: BenchResult) -> str:
    """
    Compare pixel_std distributions between base and LoRA configs.
    A large negative delta suggests LoRA is collapsing outputs (corrupting).
    A large positive delta suggests LoRA is adding structure/diversity.
    Near-zero delta → LoRA effect is statistically invisible (typical for
    weak or wrong-target LoRA on schnell GGUF).
    """
    b = [x for x in base_result.output_stds if x == x]  # filter nan
    l = [x for x in lora_result.output_stds if x == x]
    if not b or not l:
        return "N/A"
    d = statistics.mean(l) - statistics.mean(b)
    return f"{d:+.1f} px_std (LoRA - base)"


# ---------------------------------------------------------------------------
# Markdown table printer
# ---------------------------------------------------------------------------

def _print_table(
    results: List[BenchResult],
    steps: int,
    n: int,
    base_result: BenchResult,
    lora_path: str,
) -> None:
    print("\n\n" + "=" * 80)
    print("## bench_flux_lora_latency results")
    print(f"  model:   FLUX.1-schnell Q4_0 GGUF")
    print(f"  lora:    {lora_path}")
    print(f"  steps:   {steps}  resolution: 1024x1024  n_runs: {n}")
    print()

    header = (
        "| config | mean_s | p95_s | per_step_ms | fuse_succeeded "
        "| fuse_error_class | output_std | vs_base_px_std | verdict |"
    )
    sep = (
        "|--------|--------|-------|-------------|----------------"
        "|------------------|------------|----------------|---------|"
    )
    print(header)
    print(sep)

    for r in results:
        mean = r.mean_s()
        p95 = r.p95_s()
        pstep = r.per_step_ms(steps)
        std = r.mean_std()

        if r.config == "A_base":
            vs_base = "baseline"
            verdict = "BASELINE"
        else:
            vs_std = _ssim_delta(base_result, r)
            vs_base = vs_std
            if not r.fuse_attempted:
                # Config C
                delta_ms = (r.mean_s() - base_result.mean_s()) * 1000
                if abs(delta_ms) < 500:
                    verdict = f"PASS: unfused adds {delta_ms:+.0f}ms/gen (<500ms OK)"
                else:
                    verdict = f"WARN: unfused adds {delta_ms:+.0f}ms/gen"
            elif r.fuse_succeeded:
                delta_ms = (r.mean_s() - base_result.mean_s()) * 1000
                verdict = f"PASS: fused, delta {delta_ms:+.0f}ms/gen"
            else:
                fuse_cls = r.fuse_exception.split("\n")[0][:50] if r.fuse_exception else "unknown"
                verdict = f"EXPECTED FAIL: {fuse_cls}"

        fuse_err = ""
        if r.fuse_exception:
            # First line only, truncated
            first = r.fuse_exception.split("\n")[0]
            fuse_err = first[:55].replace("|", "/")

        print(
            f"| {r.config:<22} "
            f"| {mean:6.2f} "
            f"| {p95:5.2f} "
            f"| {pstep:11.0f} "
            f"| {str(r.fuse_succeeded):<14} "
            f"| {fuse_err:<16} "
            f"| {std:10.1f} "
            f"| {vs_base if r.config != 'A_base' else 'baseline':<14} "
            f"| {verdict[:50]} |"
        )

    print()
    print("### Notes")
    print(
        "- per_step_ms = mean_s / steps * 1000 (includes VAE decode in total)")
    print(
        "- output_std  = mean pixel std of generated images (grayscale proxy)")
    print(
        "- fuse_error_class: expected on GGUF path; PEFT in-place add "
        "(+= delta_weight) on QTensor/uint8 raises RuntimeError or "
        "NotImplementedError in diffusers ≤0.36")
    print(
        "- If fuse_succeeded=False, config B and C timings are identical "
        "(both ran with adapter loaded but unfused).")
    print(
        "- 'zero overhead claim' verdict: if C.mean_s ≈ A.mean_s (delta < 500ms) "
        "on MPS, claim is confirmed. >500ms suggests set_adapters dispatch adds "
        "measurable cost per call on 4-step schnell at 1024px.")


# ---------------------------------------------------------------------------
# Exceptions block (also printed for downstream doc)
# ---------------------------------------------------------------------------

def _print_exception_report(results: List[BenchResult]) -> None:
    any_exception = any(r.fuse_exception for r in results)
    if not any_exception:
        return
    print("\n### Exception Report (verbatim)")
    for r in results:
        if r.fuse_exception:
            print(f"\n#### Config {r.config}")
            print("```")
            print(r.fuse_exception[:2000])
            print("```")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="FLUX GGUF LoRA latency benchmark",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Required env vars (for real runs):\n"
            "  ALLOW_FLUX=1\n"
            "  IMAGEFORGE_FLUX_GGUF_URL=<HF URL or local path to Q4_0 GGUF>\n"
            "\n"
            "Optional env vars:\n"
            "  BENCH_LORA_PATH=<path to safetensors LoRA>\n"
            "\n"
            "Example quick run:\n"
            "  ALLOW_FLUX=1 IMAGEFORGE_FLUX_GGUF_URL=<url> \\\n"
            "    arch -arm64 python bench_flux_lora_latency.py --n=5\n"
            "\n"
            "Dry-run (validate args without loading weights):\n"
            "  python bench_flux_lora_latency.py --dry-run --n=5 --steps=4\n"
        ),
    )
    parser.add_argument(
        "--steps", type=int, default=4,
        help="Inference steps (default 4; must be >= 1)",
    )
    parser.add_argument(
        "--resolution", type=int, default=1024,
        help="Image resolution in pixels, square (default 1024; must be multiple of 8)",
    )
    parser.add_argument(
        "--n", type=int, default=20,
        help="Runs per config (default 20; use --n=5 for quick test; must be >= 1)",
    )
    parser.add_argument(
        "--lora-scale", type=float, default=1.0,
        help="LoRA scale for unfused path (default 1.0; must be > 0)",
    )
    parser.add_argument(
        "--device", type=str, default=None,
        help="Override device (mps/cuda/cpu); auto-detected if omitted",
    )
    parser.add_argument(
        "--skip-fused", action="store_true",
        help="Skip config B (fuse attempt) — useful if you already know fuse fails",
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help=(
            "Validate args and print the planned config WITHOUT loading models "
            "or running inference. Exits 0 if args are valid."
        ),
    )
    args = parser.parse_args()

    # ---------------------------------------------------------------------- #
    # Arg validation — catches bad --n / --steps / --resolution BEFORE any   #
    # env check or heavy import.  Exits 2 (argparse convention) on error.    #
    # ---------------------------------------------------------------------- #
    _validate_args(args)

    # ---------------------------------------------------------------------- #
    # --dry-run: print config and exit 0 without touching models              #
    # ---------------------------------------------------------------------- #
    if args.dry_run:
        gguf_url = os.environ.get(
            "IMAGEFORGE_FLUX_GGUF_URL", "<IMAGEFORGE_FLUX_GGUF_URL not set>"
        )
        lora_path_display = os.environ.get(
            "BENCH_LORA_PATH",
            "ostris/FLUX.1-schnell-training-adapter (default — set BENCH_LORA_PATH to override)",
        )
        allow_flux = os.environ.get("ALLOW_FLUX", "<not set>")
        device_display = args.device if args.device else "auto-detect (mps > cuda > cpu)"
        print("=== bench_flux_lora_latency — DRY RUN ===")
        print(f"  steps:       {args.steps}")
        print(f"  resolution:  {args.resolution}x{args.resolution}")
        print(f"  n_runs:      {args.n}")
        print(f"  lora_scale:  {args.lora_scale}")
        print(f"  device:      {device_display}")
        print(f"  skip_fused:  {args.skip_fused}")
        print(f"  ALLOW_FLUX:  {allow_flux}")
        print(f"  GGUF URL:    {gguf_url}")
        print(f"  LoRA path:   {lora_path_display}")
        print()
        print("  Configs that WILL run:")
        print("    A  base GGUF, no LoRA")
        if not args.skip_fused:
            print("    B  fused LoRA (fuse_lora — expected FAIL on GGUF path, see docs)")
        print("    C  unfused LoRA (set_adapters at call time)")
        n_configs = 2 if args.skip_fused else 3
        print()
        print(f"  Total inference calls (approx): {args.n * n_configs}")
        print()
        print("[dry-run] Args are valid. Remove --dry-run to execute.")
        sys.exit(0)

    _check_env()

    # ---------------------------------------------------------------------- #
    # Dependency check — actionable error if diffusers / torch missing        #
    # ---------------------------------------------------------------------- #
    _check_deps()

    import torch

    # Device detection
    device = args.device
    if device is None:
        if torch.backends.mps.is_available():
            device = "mps"
        elif torch.cuda.is_available():
            device = "cuda"
        else:
            device = "cpu"
    print(f"[bench] Device: {device}")
    print(f"[bench] torch version: {torch.__version__}")

    import diffusers
    print(f"[bench] diffusers version: {diffusers.__version__}")

    gguf_url = os.environ["IMAGEFORGE_FLUX_GGUF_URL"]
    lora_path = _resolve_lora_path()
    steps = args.steps
    resolution = args.resolution
    n = args.n

    # ------------------------------------------------------------------ #
    # Load base pipeline once — ALL configs share the same loaded weights #
    # ------------------------------------------------------------------ #
    print("\n[bench] ===== LOADING BASE PIPELINE =====")
    try:
        pipe = _load_base_pipe(gguf_url, device)
    except Exception as e:
        print(f"[ERROR] Failed to load base pipeline: {e}", file=sys.stderr)
        traceback.print_exc()
        sys.exit(1)

    _warmup(pipe, device, steps=min(steps, 4), resolution=512)

    results: List[BenchResult] = []

    # ------------------------------------------------------------------ #
    # Config A: base                                                       #
    # ------------------------------------------------------------------ #
    base_result = bench_base(pipe, device, steps, resolution, n)
    results.append(base_result)

    # ------------------------------------------------------------------ #
    # Config B: fused                                                      #
    # ------------------------------------------------------------------ #
    if not args.skip_fused:
        fused_result = bench_fused(pipe, lora_path, device, steps, resolution, n)
        results.append(fused_result)
    else:
        print("\n[bench] Config B skipped (--skip-fused)")

    # ------------------------------------------------------------------ #
    # Config C: unfused-scaled                                             #
    # ------------------------------------------------------------------ #
    unfused_result = bench_unfused(
        pipe, lora_path, device, steps, resolution, n, lora_scale=args.lora_scale
    )
    results.append(unfused_result)

    # ------------------------------------------------------------------ #
    # Report                                                               #
    # ------------------------------------------------------------------ #
    _print_table(results, steps, n, base_result, lora_path)
    _print_exception_report(results)

    print("\n[bench] Done.")


if __name__ == "__main__":
    main()