"""Sampler-only speed benchmark: this repo's W4A4 against int8 and bf16, with and without compile. python tools/speed_bench.py --server http://127.0.0.1:8188 python tools/speed_bench.py --arms int8 r256aa --compile off on --out speed.json Why this exists. Every speed number this repo has published is either *end to end* (which BENCHMARKS.md measures at 10.4 s for int8 against 10.1 s for W4A4 -- a 3% gap that says almost nothing about the kernels) or *per layer* (`svdquant_diag.py --mode bench`, which times one Linear in isolation against bf16 and never sees int8 at all). Neither answers the question users actually ask, which is "how much faster is a step". Two things make that hard to measure through ComfyUI, and this script handles both: * **Fixed cost swamps the signal.** A 1024x1024 8-step turbo render spends several seconds on CLIP text-encode, model staging and VAE decode. So instead of timing one render, every arm is run at *two* step counts and the per-step cost is the **slope**, `(t_hi - t_lo) / (steps_hi - steps_lo)`. Everything that does not scale with step count cancels out. The intercept is reported too, because a suspicious intercept is how you catch a run where the model was silently reloaded mid-measurement. * **Cold runs are a different measurement.** The first render after switching checkpoints pays disk-to-VRAM load. Each arm therefore does one discarded warmup at each step count before anything is timed. `--compile on` inserts a **TorchCompileModel** node (backend `inductor`) between the loader and the sampler, which is what the report this script was written for was running. That arm is the one that matters: on `int8-fast`'s own numbers compile is worth 1.58x to them (1.64 -> 1.04 s/it), while this repo shields all 224 quantized linears from Dynamo (`svdquant_w4a4.py:_shield_from_dynamo`) and gets far less. Measuring both sides with the same harness is the point. Needs a running ComfyUI on `--server`, and for the `int8` arm the `int8-fast` custom nodes (node `OTUNetLoaderW8A8`) -- the checkpoint it loads is per-row int8 + convrot in that pack's own format, not one ComfyUI's `UNETLoader` can read. """ from __future__ import annotations import argparse import json import os import statistics import sys import time import urllib.error import urllib.request HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) # The graph is the same one fidelity_bench builds, minus the parts that would add variance: # one fixed prompt, one fixed seed, no LoRA. Sharing the module keeps the sampler settings, # CLIP and VAE names from drifting apart between the two benchmarks. import fidelity_bench as fb # noqa: E402 # One prompt, deliberately mid-length: text-encode cost is constant across arms and cancels # in the slope anyway, but a pathologically short prompt would make the intercept unrealistic # and hide a staging regression. PROMPT = fb.PROMPTS_OBJECTS[9][1] # Every render gets its own seed. ComfyUI caches node outputs by input hash, so submitting # the same graph twice returns in ~20 ms without sampling anything -- the first version of # this script measured `2.01 0.02 0.03` and would have reported a per-step cost of zero. # Bumping the seed invalidates the KSampler and everything downstream while leaving the # CLIP encode cached, which is what we want: text-encode is constant across arms anyway, # and keeping it out of the timing makes the intercept easier to read. SEED_BASE = 987654321 _seed_counter = 0 # Two step counts far enough apart that the difference is many times the timing noise, and # low enough that the long arm still finishes in a sane time on a 3090. STEPS_LO = 8 STEPS_HI = 24 # An arm is a checkpoint plus the loader node that can read it. `int8` is not a ComfyUI-native # format: per-tensor `comfy_quant` blobs written by int8-fast, per-row scales, convrot baked # into the weight. Its loader also insists on a `model_type`, which for Krea 2 is only used # for on-the-fly quantization -- this file is already quantized, so the value is inert. ARMS = { "bf16": { "node": "UNETLoader", "inputs": {"unet_name": "turbo.safetensors", "weight_dtype": "default"}, }, "int8": { "node": "OTUNetLoaderW8A8", "inputs": {"unet_name": "krea2turboint8convrot.safetensors", "weight_dtype": "default", "model_type": "qwen", "on_the_fly_quantization": False, "enable_convrot": True, "lora_mode": "None"}, }, "nolowrank": { "node": "UNETLoader", "inputs": {"unet_name": "Krea2-Turbo-W4A4-noLowRank.safetensors", "weight_dtype": "default"}, }, "r256aa": { "node": "Krea2SVDQuantW4A4Loader", "inputs": {"model_name": "Krea2-Turbo-SVDQuant-W4A4-rank256-actaware.safetensors"}, }, # The rank sweep exists to price the low-rank branch. `--act-stats` cannot change any of # these numbers -- same shapes, same kernels, only different values inside the factors -- # so the plain builds stand in for their act-aware equivalents. "r16": { "node": "Krea2SVDQuantW4A4Loader", "inputs": {"model_name": "Krea2-Turbo-SVDQuant-W4A4-rank16.safetensors"}, }, "r64": { "node": "Krea2SVDQuantW4A4Loader", "inputs": {"model_name": "Krea2-Turbo-SVDQuant-W4A4-rank64.safetensors"}, }, "r128": { "node": "Krea2SVDQuantW4A4Loader", "inputs": {"model_name": "Krea2-Turbo-SVDQuant-W4A4-rank128.safetensors"}, }, } # bf16 is 24.5 GB and does not fit a 3090 alongside the text encoder, so it spends the run # being shuttled over PCIe. It stays out of the default set: it is a reference point for the # README table, not something to pay for on every run. DEFAULT_ARMS = ["int8", "r256aa", "nolowrank"] def build_graph(arm: str, steps: int, compile_backend: str | None, seed: int) -> dict: spec = ARMS[arm] g = { "1": {"class_type": spec["node"], "inputs": dict(spec["inputs"])}, "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": fb.CLIP_NAME, "type": "krea2", "device": "default"}}, "3": {"class_type": "VAELoader", "inputs": {"vae_name": fb.VAE_NAME}}, "4": {"class_type": "CLIPTextEncode", "inputs": {"text": PROMPT, "clip": ["2", 0]}}, "5": {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["4", 0]}}, "6": {"class_type": "EmptySD3LatentImage", "inputs": {"width": fb.WIDTH, "height": fb.HEIGHT, "batch_size": 1}}, "7": {"class_type": "KSampler", "inputs": {"seed": seed, "steps": steps, "cfg": 1.0, "sampler_name": fb.SAMPLER, "scheduler": fb.SCHEDULER, "denoise": 1.0, "model": ["1", 0], "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["6", 0]}}, "8": {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["3", 0]}}, "9": {"class_type": "PreviewImage", "inputs": {"images": ["8", 0]}}, } if compile_backend: # Note for whoever reads a FAILED row: a branchless w4a4 checkpoint arrives through # the stock UNETLoader, which never registers the kitchen kernel as a custom op, and # compiling one *raises* rather than running slowly. That is a real limitation of the # project, not of this script -- see TROUBLESHOOTING.md. g["11"] = {"class_type": "TorchCompileModel", "inputs": {"model": ["1", 0], "backend": compile_backend}} g["7"]["inputs"]["model"] = ["11", 0] return g def check_nodes(server: str, arms: list[str], want_compile: bool) -> list[str]: """Names of graph nodes this run needs that the server does not have. Checked up front rather than discovered on the first submit: a missing custom node comes back as a validation error against node id "1", which reads like a bad checkpoint path. """ needed = {ARMS[a]["node"] for a in arms} if want_compile: needed.add("TorchCompileModel") try: info = fb._get(server, "/object_info") except Exception as exc: raise SystemExit("could not reach ComfyUI at {}: {}".format(server, exc)) return sorted(n for n in needed if n not in info) def _server_seconds(entry: dict) -> float | None: """Execution time as ComfyUI itself measured it, from the history entry's timestamps. Polling `/history` on a 2 s timer put up to 2 s of quantization noise on every sample -- an 8.27 s render came back as 10.11 s. That bias mostly cancels in the slope, but its *variance* does not, and at a 15 s difference between the two step counts a couple of seconds of jitter is several percent on the answer. ComfyUI already records exactly what we want: `execution_start` and `execution_success` in milliseconds. """ stamps = {} for name, payload in entry.get("status", {}).get("messages", []): if name in ("execution_start", "execution_success"): stamps[name] = payload.get("timestamp") start, end = stamps.get("execution_start"), stamps.get("execution_success") if start is None or end is None: return None return (end - start) / 1000.0 def time_run(server: str, arm: str, steps: int, compile_backend: str | None, timeout: int) -> float: global _seed_counter _seed_counter += 1 graph = build_graph(arm, steps, compile_backend, SEED_BASE + _seed_counter) try: result = fb._post(server, "/prompt", {"prompt": graph}) except urllib.error.HTTPError as exc: raise RuntimeError("{} @ {} steps: submit failed: {}".format( arm, steps, exc.read().decode("utf-8", "replace")[:400])) if "error" in result: raise RuntimeError("{} @ {} steps: {}".format( arm, steps, json.dumps(result["error"])[:400])) pid = result["prompt_id"] deadline = time.time() + timeout while time.time() < deadline: entry = fb._get(server, "/history/" + pid).get(pid) if entry: status = entry.get("status", {}) if status.get("status_str") == "error": raise RuntimeError("{} @ {} steps: {}".format( arm, steps, json.dumps(status)[:400])) if status.get("completed") or status.get("status_str") == "success": seconds = _server_seconds(entry) if seconds is None: raise RuntimeError( "{} @ {} steps: history entry carries no execution timestamps; this " "ComfyUI is too old for a per-step measurement".format(arm, steps)) return seconds time.sleep(1) raise RuntimeError("{} @ {} steps: timeout".format(arm, steps)) def measure(server: str, arm: str, compile_backend: str | None, reps: int, timeout: int) -> dict: """Median wall time at two step counts, and the per-step slope between them. The warmup at each step count is what makes the numbers "warm": it pays the checkpoint load, the CLIP encode that ComfyUI then caches, and -- when compiling -- inductor's first-run compilation, which is tens of seconds and would otherwise land entirely in the first timed rep. """ samples: dict[int, list[float]] = {} for steps in (STEPS_LO, STEPS_HI): time_run(server, arm, steps, compile_backend, timeout) # warmup, discarded runs = [time_run(server, arm, steps, compile_backend, timeout) for _ in range(reps)] samples[steps] = runs print(" {:>2} steps: {} median {:.2f}s".format( steps, " ".join("{:.2f}".format(r) for r in runs), statistics.median(runs)), flush=True) lo = statistics.median(samples[STEPS_LO]) hi = statistics.median(samples[STEPS_HI]) per_step = (hi - lo) / (STEPS_HI - STEPS_LO) return { "arm": arm, "compile": compile_backend or "off", "steps_lo": STEPS_LO, "steps_hi": STEPS_HI, "median_lo": lo, "median_hi": hi, "runs_lo": samples[STEPS_LO], "runs_hi": samples[STEPS_HI], "s_per_step": per_step, # Everything that does not scale with steps: text-encode, staging, VAE decode, queue. # It should be near-identical across arms of the same checkpoint size; a big outlier # means that arm was reloading the model between runs and its slope is not trustworthy. "fixed_overhead": lo - per_step * STEPS_LO, } def main() -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--server", default="http://127.0.0.1:8188") ap.add_argument("--arms", nargs="+", default=DEFAULT_ARMS, choices=sorted(ARMS)) ap.add_argument("--compile", nargs="+", default=["off", "on"], choices=["off", "on", "cudagraphs"], help="'on' is TorchCompileModel with backend inductor") ap.add_argument("--reps", type=int, default=3, help="timed runs per (arm, step count)") ap.add_argument("--timeout", type=int, default=1800, help="per render; the first compiled run pays inductor's compilation") ap.add_argument("--out", default=None, help="write the raw results here as JSON") args = ap.parse_args() backends = {"off": None, "on": "inductor", "cudagraphs": "cudagraphs"} missing = check_nodes(args.server, args.arms, any(c != "off" for c in args.compile)) if missing: raise SystemExit("ComfyUI at {} has no node(s): {}. The int8 arm needs the int8-fast " "custom nodes; drop it with --arms.".format(args.server, ", ".join(missing))) results = [] for arm in args.arms: for mode in args.compile: print("{} / compile={}".format(arm, mode), flush=True) t0 = time.time() try: results.append(measure(args.server, arm, backends[mode], args.reps, args.timeout)) except RuntimeError as exc: # One arm failing (a checkpoint that will not load, compile blowing up on a # kernel) must not throw away the arms that already ran. print(" SKIPPED: {}".format(exc), flush=True) results.append({"arm": arm, "compile": mode, "error": str(exc)}) continue print(" -> {:.3f} s/step (fixed {:.2f}s, {:.0f}s elapsed)".format( results[-1]["s_per_step"], results[-1]["fixed_overhead"], time.time() - t0), flush=True) ok = [r for r in results if "error" not in r] print("\n{:<12} {:>8} {:>10} {:>10} {:>9}".format( "arm", "compile", "s/step", "fixed s", "vs best")) best = min((r["s_per_step"] for r in ok), default=None) for r in results: if "error" in r: print("{:<12} {:>8} {:>10}".format(r["arm"], r["compile"], "FAILED")) continue print("{:<12} {:>8} {:>10.3f} {:>10.2f} {:>8.2f}x".format( r["arm"], r["compile"], r["s_per_step"], r["fixed_overhead"], r["s_per_step"] / best)) if args.out: with open(args.out, "w", encoding="utf-8") as fh: json.dump({"resolution": [fb.WIDTH, fb.HEIGHT], "sampler": fb.SAMPLER, "scheduler": fb.SCHEDULER, "prompt": PROMPT, "seed_base": SEED_BASE, "results": results}, fh, indent=2) print("\nwrote {}".format(args.out)) return 1 if any("error" in r for r in results) else 0 if __name__ == "__main__": raise SystemExit(main())