126 lines
5.1 KiB
Python
126 lines
5.1 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
End-to-end Test für ComfyUI Image Gen API.
|
||
Iteriert über model × quality × style, generiert Bilder, schreibt Report.
|
||
Usage: python3 scripts/comfy_e2e_test.py [--api-url URL] [--output-dir DIR]
|
||
"""
|
||
import os, time, json, base64, pathlib, argparse
|
||
|
||
|
||
def parse_args():
|
||
p = argparse.ArgumentParser(description="ComfyUI E2E image generation test")
|
||
p.add_argument("--api-url", default="http://localhost:8000/v1/images/generations")
|
||
p.add_argument("--output-dir", default=str(pathlib.Path.home() / ".hermes/cron/output/comfy_e2e"))
|
||
p.add_argument("--models", default="ideogram4,flux2", help="Comma-separated models to test")
|
||
p.add_argument("--qualities", default="low,medium,high", help="Comma-separated qualities")
|
||
p.add_argument("--styles", default=",photorealistic", help="Comma-separated styles (empty for none)")
|
||
p.add_argument("--size", default="1024x1024")
|
||
p.add_argument("--n", type=int, default=1)
|
||
p.add_argument("--timeout", type=int, default=1800, help="HTTP timeout per request in seconds")
|
||
return p.parse_args()
|
||
|
||
|
||
def generate_image(api_url: str, payload: dict, timeout: int):
|
||
import urllib.request
|
||
req = urllib.request.Request(
|
||
api_url,
|
||
data=json.dumps(payload).encode(),
|
||
headers={"Content-Type": "application/json"},
|
||
method="POST"
|
||
)
|
||
start = time.time()
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
data = json.loads(resp.read().decode())
|
||
elapsed = time.time() - start
|
||
b64 = data["data"][0]["b64_json"]
|
||
return base64.b64decode(b64), elapsed
|
||
except Exception as e:
|
||
return None, str(e)
|
||
|
||
|
||
def main():
|
||
args = parse_args()
|
||
out_dir = pathlib.Path(args.output_dir)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
models = [m.strip() for m in args.models.split(",")]
|
||
qualities = [q.strip() for q in args.qualities.split(",")]
|
||
styles = [s.strip() for s in args.styles.split(",")]
|
||
|
||
PROMPTS = {
|
||
"ideogram4": "A futuristic city skyline at sunset, neon lights, detailed architecture",
|
||
"flux2": "A serene mountain lake reflection, golden hour lighting, highly detailed",
|
||
"flux1-dev": "A cyberpunk street market at night, volumetric fog, neon signs",
|
||
"flux1-schnell": "A minimalist abstract painting, bold colors, geometric shapes",
|
||
}
|
||
|
||
RESULTS = []
|
||
total = len(models) * len(qualities) * len(styles)
|
||
print(f"Output: {out_dir}")
|
||
print(f"Total: {total} combinations")
|
||
print("=" * 60)
|
||
|
||
for model in models:
|
||
for quality in qualities:
|
||
for style in styles:
|
||
tag = f"{model}_{quality}_{style or 'default'}"
|
||
prompt = PROMPTS.get(model, "A beautiful landscape")
|
||
full_prompt = f"{style}, {prompt}".strip(", ") if style else prompt
|
||
payload = {
|
||
"model": model,
|
||
"prompt": full_prompt,
|
||
"n": args.n,
|
||
"size": args.size,
|
||
"quality": quality,
|
||
"response_format": "b64_json"
|
||
}
|
||
print(f"\n[{tag}] {full_prompt[:60]}...")
|
||
img_bytes, info = generate_image(args.api_url, payload, args.timeout)
|
||
if img_bytes:
|
||
fname = f"{tag}_{int(time.time())}.png"
|
||
fpath = out_dir / fname
|
||
fpath.write_bytes(img_bytes)
|
||
print(f" ✅ {info:.1f}s, {len(img_bytes)/1024:.0f} KB → {fname}")
|
||
RESULTS.append({
|
||
"model": model, "quality": quality, "style": style,
|
||
"status": "OK", "time_s": round(info, 1),
|
||
"size_kb": round(len(img_bytes)/1024, 1), "file": str(fpath)
|
||
})
|
||
else:
|
||
print(f" ❌ {info}")
|
||
RESULTS.append({
|
||
"model": model, "quality": quality, "style": style,
|
||
"status": "FAILED", "error": str(info)
|
||
})
|
||
|
||
# JSON report
|
||
report = out_dir / "report.json"
|
||
report.write_text(json.dumps(RESULTS, indent=2))
|
||
|
||
# Markdown report
|
||
md = out_dir / "report.md"
|
||
lines = [
|
||
"# ComfyUI E2E Test Report",
|
||
f"\n**Zeit:** {time.strftime('%Y-%m-%d %H:%M:%S')}",
|
||
]
|
||
ok = sum(1 for r in RESULTS if r["status"] == "OK")
|
||
fail = sum(1 for r in RESULTS if r["status"] == "FAILED")
|
||
lines.append(f"**Ergebnis:** {ok} OK / {fail} FAILED\n")
|
||
lines.append("| Modell | Quality | Style | Status | Zeit | Datei |")
|
||
lines.append("|---|---|---|---|---|---|")
|
||
for r in RESULTS:
|
||
style = r.get("style", "") or "—"
|
||
status = "✅" if r["status"] == "OK" else "❌"
|
||
t = f"{r.get('time_s', '—')}s" if "time_s" in r else "—"
|
||
f = r.get("file", "").split("/")[-1] if "file" in r else "—"
|
||
lines.append(f"| {r['model']} | {r['quality']} | {style} | {status} {r['status']} | {t} | {f} |")
|
||
md.write_text("\n".join(lines))
|
||
print(f"\n{'=' * 60}")
|
||
print(f"Reports: {report} , {md}")
|
||
print(f"Bilder: {out_dir}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|