143 lines
4.0 KiB
Python
143 lines
4.0 KiB
Python
"""
|
||
Speed and VRAM benchmark. Runs each config multiple times (1 cold + 10 warm).
|
||
"""
|
||
|
||
import argparse
|
||
import csv
|
||
import json
|
||
import os
|
||
import time
|
||
import copy
|
||
import threading
|
||
import subprocess
|
||
|
||
from config import BENCHMARK_CONFIGS, PROMPTS, STEPS
|
||
from utils import check_comfyui, queue_prompt, wait_for_completion, free_memory, \
|
||
build_uniform_workflow, build_dual_workflow, avg, std
|
||
|
||
|
||
RUNS_PER_CONFIG = 11
|
||
COLD_RUNS = 1
|
||
BENCH_PROMPT = PROMPTS["A01"]
|
||
BASE_SEED = 42
|
||
|
||
|
||
class VRAMMonitor:
|
||
def __init__(self, interval=0.3):
|
||
self.interval = interval
|
||
self.peak = 0
|
||
self._running = False
|
||
|
||
def _poll(self):
|
||
while self._running:
|
||
try:
|
||
r = subprocess.run(
|
||
["nvidia-smi", "--query-gpu=memory.used",
|
||
"--format=csv,noheader,nounits"],
|
||
capture_output=True, text=True, timeout=5)
|
||
if r.returncode == 0:
|
||
self.peak = max(self.peak, int(r.stdout.strip().split("\n")[0]))
|
||
except Exception:
|
||
pass
|
||
time.sleep(self.interval)
|
||
|
||
def start(self):
|
||
self.peak = 0
|
||
self._running = True
|
||
self._t = threading.Thread(target=self._poll, daemon=True)
|
||
self._t.start()
|
||
|
||
def stop(self):
|
||
self._running = False
|
||
self._t.join(timeout=2)
|
||
return self.peak
|
||
|
||
|
||
def make_workflow(cfg, seed):
|
||
if cfg["type"] == "uniform":
|
||
wf = build_uniform_workflow(cfg["model"], BENCH_PROMPT, seed, f"bench_{cfg['name']}")
|
||
else:
|
||
wf = build_dual_workflow(cfg["model_a"], cfg["model_b"], cfg["switch_step"],
|
||
BENCH_PROMPT, seed, f"bench_{cfg['name']}")
|
||
return wf
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--output-dir", required=True)
|
||
args = parser.parse_args()
|
||
|
||
os.makedirs(args.output_dir, exist_ok=True)
|
||
check_comfyui()
|
||
|
||
csv_path = os.path.join(args.output_dir, "benchmark_results.csv")
|
||
f = open(csv_path, "w", newline="", encoding="utf-8")
|
||
writer = csv.writer(f)
|
||
writer.writerow(["config", "type", "run_idx", "is_cold",
|
||
"time_sec", "peak_vram_mb", "status"])
|
||
|
||
vmon = VRAMMonitor()
|
||
|
||
print(f"{'='*60}")
|
||
print(f"Speed Benchmark: {len(BENCHMARK_CONFIGS)} configs × {RUNS_PER_CONFIG} runs")
|
||
print(f"{'='*60}")
|
||
|
||
all_stats = {}
|
||
|
||
for ci, cfg in enumerate(BENCHMARK_CONFIGS):
|
||
name = cfg["name"]
|
||
print(f"\n── [{ci+1}/{len(BENCHMARK_CONFIGS)}] {name} ──")
|
||
free_memory()
|
||
|
||
times_warm = []
|
||
|
||
for run in range(RUNS_PER_CONFIG):
|
||
is_cold = run < COLD_RUNS
|
||
label = "COLD" if is_cold else "WARM"
|
||
seed = BASE_SEED + run
|
||
|
||
print(f" Run {run+1:2d}/{RUNS_PER_CONFIG} ({label})...", end=" ", flush=True)
|
||
|
||
wf = make_workflow(cfg, seed)
|
||
vmon.start()
|
||
t0 = time.time()
|
||
|
||
try:
|
||
pid = queue_prompt(wf)
|
||
ok = wait_for_completion(pid)
|
||
dt = time.time() - t0
|
||
vram = vmon.stop()
|
||
status = "ok" if ok else "error"
|
||
print(f"{dt:.1f}s VRAM={vram}MB {status}")
|
||
if not is_cold and ok:
|
||
times_warm.append(dt)
|
||
except Exception as e:
|
||
dt = time.time() - t0
|
||
vram = vmon.stop()
|
||
status = f"exception"
|
||
print(f"FAIL: {e}")
|
||
|
||
writer.writerow([name, cfg["type"], run, is_cold,
|
||
f"{dt:.2f}", vram, status])
|
||
f.flush()
|
||
|
||
if times_warm:
|
||
all_stats[name] = times_warm
|
||
|
||
f.close()
|
||
|
||
# Summary
|
||
print(f"\n{'='*60}")
|
||
print("SUMMARY (warm runs)")
|
||
print(f"{'='*60}")
|
||
print(f"{'Config':<28} {'Avg':>8} {'Std':>6}")
|
||
print("─" * 44)
|
||
for name, times in all_stats.items():
|
||
print(f"{name:<28} {avg(times):>7.1f}s ±{std(times):.1f}s")
|
||
|
||
print(f"\nSaved: {csv_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|