| """Time the same workflow across several checkpoints, twice each (cold + warm), via the ComfyUI API. |
| |
| python3 time_checkpoints.py --host 127.0.0.1 --port 8188 --workflow smoke-prompt.json \\ |
| --checkpoints name1.safetensors name2.safetensors |
| |
| Prints, per run: wall time (submit -> finished) and the server's own "Prompt executed" figure is |
| captured separately from the container log. The first run of each checkpoint includes loading it |
| from disk; the second is the steady state. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| import time |
| import urllib.parse |
| import urllib.request |
| import uuid |
| from pathlib import Path |
|
|
|
|
| def post(url: str, payload: dict) -> dict: |
| req = urllib.request.Request(url, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"}) |
| with urllib.request.urlopen(req, timeout=120) as resp: |
| return json.loads(resp.read()) |
|
|
|
|
| def get(url: str) -> dict: |
| with urllib.request.urlopen(url, timeout=60) as resp: |
| return json.loads(resp.read()) |
|
|
|
|
| def run_once(base: str, graph: dict, checkpoint: str, prefix: str, timeout: float, seed: int | None = None) -> dict: |
| graph["13"]["inputs"]["unet_name"] = checkpoint |
| graph["21"]["inputs"]["filename_prefix"] = prefix |
| for node in graph.values(): |
| if node.get("class_type") == "UNETLoader": |
| node["inputs"]["unet_name"] = checkpoint |
| if node.get("class_type") == "VHS_VideoCombine": |
| node["inputs"]["filename_prefix"] = prefix |
| if node.get("class_type") == "RandomNoise" and seed is not None: |
| |
| |
| node["inputs"]["noise_seed"] = seed |
|
|
| t0 = time.time() |
| started = post(f"{base}/prompt", {"prompt": graph, "client_id": str(uuid.uuid4())})["prompt_id"] |
| first_seen = None |
| while time.time() - t0 < timeout: |
| history = get(f"{base}/history/{started}") |
| entry = history.get(started) |
| if entry and entry.get("outputs"): |
| status = entry.get("status", {}) |
| if status.get("status_str") != "success": |
| return {"ok": False, "error": json.dumps(status)[:200], "wall": time.time() - t0} |
| exec_ms = None |
| for name, payload in status.get("messages", []): |
| if name == "execution_start": |
| first_seen = payload.get("timestamp") |
| if name == "execution_success": |
| exec_ms = payload.get("timestamp", 0) - (first_seen or 0) |
| return {"ok": True, "wall": time.time() - t0, "exec": (exec_ms or 0) / 1000.0} |
| time.sleep(1) |
| return {"ok": False, "error": "timeout", "wall": time.time() - t0} |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--host", required=True) |
| ap.add_argument("--port", type=int, default=8188) |
| ap.add_argument("--workflow", required=True) |
| ap.add_argument("--checkpoints", nargs="+", required=True) |
| ap.add_argument("--repeats", type=int, default=2) |
| ap.add_argument("--seed-base", type=int, default=2347709, help="seed for run 1; run N uses base+N-1") |
| ap.add_argument("--timeout", type=float, default=900) |
| args = ap.parse_args(argv) |
|
|
| base = f"http://{args.host}:{args.port}" |
| graph = json.loads(Path(args.workflow).read_text()) |
| graph = graph.get("prompt", graph) |
|
|
| print(f"{'checkpoint':<48s} {'run':>3s} {'seed':>9s} {'wall s':>8s} {'service s':>10s}") |
| for ckpt in args.checkpoints: |
| for i in range(1, args.repeats + 1): |
| prefix = f"timing/{Path(ckpt).stem[:28]}_{i}" |
| seed = args.seed_base + i - 1 |
| result = run_once(base, json.loads(json.dumps(graph)), ckpt, prefix, args.timeout, seed=seed) |
| if result["ok"]: |
| print(f"{ckpt:<48s} {i:>3d} {seed:>9d} {result['wall']:>8.1f} {result['exec']:>10.1f}") |
| else: |
| print(f"{ckpt:<48s} {i:>3d} {seed:>9d} FAILED: {result['error']}") |
| sys.stdout.flush() |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
