Add ComfyUI image generation (/v1/images/generations)

OpenAI-compatible image endpoint that proxies to a ComfyUI server: fills a
workflow template with per-request params (prompt/size/n/seed/steps/checkpoint),
submits to ComfyUI, polls history, fetches images, returns OpenAI image shape.

- router.py: _handle_image + fill_workflow + route; env ROUTER_COMFYUI /
  ROUTER_COMFY_CHECKPOINT / ROUTER_COMFY_WORKFLOW
- comfyui-workflow.json: default SD/SDXL text2img template with {{SENTINELS}}
- docs/IMAGES.md + README: config, usage, custom workflows
- verified: templating + clean error when ComfyUI is down; end-to-end pending
  a live ComfyUI

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Joseph Costa
2026-07-05 03:14:51 -05:00
co-authored by Claude Opus 4.8
parent 5bd8142d30
commit 2f95e6be96
4 changed files with 190 additions and 1 deletions
+83 -1
View File
@@ -9,7 +9,7 @@ forwards straight to Ollama's own OpenAI API. Point every client at
On upstream failure the router retries once with a mapped fallback model.
Response headers x-router-model / x-router-decided-by report the pick.
"""
import os, re, json, time, base64, urllib.request, urllib.error
import os, re, json, time, base64, urllib.request, urllib.error, urllib.parse
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
OLLAMA = os.environ.get("ROUTER_OLLAMA", "http://127.0.0.1:11434").rstrip("/") # client target; NOT OLLAMA_HOST (that's the server bind addr, e.g. 0.0.0.0)
@@ -18,6 +18,11 @@ HOST = os.environ.get("ROUTER_HOST", "0.0.0.0") # 0.0.0.0 = reachable on t
PORT = int(os.environ.get("ROUTER_PORT", "8080"))
API_KEY = os.environ.get("ROUTER_API_KEY", "").strip() # if set, require Bearer token
AUDIO_MODEL = os.environ.get("ROUTER_AUDIO_MODEL", "gemma4:e4b") # audio only works on official Gemma 4
# ComfyUI image generation (OpenAI /v1/images/generations -> ComfyUI submit/poll/fetch)
COMFYUI = os.environ.get("ROUTER_COMFYUI", "http://127.0.0.1:8188").rstrip("/")
COMFY_WORKFLOW = os.environ.get("ROUTER_COMFY_WORKFLOW",
os.path.join(os.path.dirname(os.path.abspath(__file__)), "comfyui-workflow.json"))
COMFY_CHECKPOINT = os.environ.get("ROUTER_COMFY_CHECKPOINT", "") # your ComfyUI checkpoint filename
M = {
"fast": "qwen3:8b",
@@ -153,6 +158,17 @@ def parse_multipart(body, boundary):
return fields
def fill_workflow(node, subs):
# recursively replace string sentinels ("{{PROMPT}}" etc.) with typed values
if isinstance(node, dict):
return {k: fill_workflow(v, subs) for k, v in node.items()}
if isinstance(node, list):
return [fill_workflow(x, subs) for x in node]
if isinstance(node, str) and node in subs:
return subs[node]
return node
def total_len(messages):
n = 0
for m in messages:
@@ -387,6 +403,65 @@ class Handler(BaseHTTPRequestHandler):
return self._text(200, text)
return self._json(200, {"text": text})
def _handle_image(self, body):
# OpenAI /v1/images/generations -> ComfyUI (submit workflow, poll, fetch images)
if not self._authed():
return self._json(401, {"error": "unauthorized"})
prompt = (body.get("prompt") or "").strip()
if not prompt:
return self._json(400, {"error": "missing 'prompt'"})
n = max(1, int(body.get("n", 1) or 1))
try:
w, h = (int(x) for x in str(body.get("size", "1024x1024")).lower().split("x"))
except Exception:
w, h = 1024, 1024
fmt = body.get("response_format", "b64_json")
try:
with open(COMFY_WORKFLOW) as f:
wf = json.load(f)
except Exception as e:
return self._json(500, {"error": f"workflow '{COMFY_WORKFLOW}' unreadable: {e}"})
subs = {"{{PROMPT}}": prompt, "{{NEGATIVE}}": body.get("negative_prompt", ""),
"{{WIDTH}}": w, "{{HEIGHT}}": h, "{{BATCH}}": n,
"{{SEED}}": int(time.time() * 1000) % 2_000_000_000,
"{{STEPS}}": int(body.get("steps", 25)),
"{{CHECKPOINT}}": COMFY_CHECKPOINT or "{{CHECKPOINT}}"}
wf = fill_workflow(wf, subs)
print(f"[router] image -> comfyui ({w}x{h} n={n})", flush=True)
try:
req = urllib.request.Request(COMFYUI + "/prompt",
data=json.dumps({"prompt": wf}).encode(),
headers={"content-type": "application/json"})
pid = json.loads(urllib.request.urlopen(req, timeout=30).read()).get("prompt_id")
if not pid:
return self._json(502, {"error": "comfyui returned no prompt_id"})
images, deadline = [], time.time() + 600
while time.time() < deadline:
time.sleep(1.5)
hist = json.loads(urllib.request.urlopen(COMFYUI + "/history/" + pid, timeout=15).read())
entry = hist.get(pid)
if entry and entry.get("outputs"):
for node in entry["outputs"].values():
images.extend(node.get("images", []))
break
if not images:
return self._json(504, {"error": "comfyui timed out producing images"})
out = []
for im in images[:n]:
q = urllib.parse.urlencode({"filename": im.get("filename", ""),
"subfolder": im.get("subfolder", ""),
"type": im.get("type", "output")})
data = urllib.request.urlopen(COMFYUI + "/view?" + q, timeout=60).read()
b64 = base64.b64encode(data).decode()
out.append({"b64_json": b64} if fmt == "b64_json"
else {"url": "data:image/png;base64," + b64})
return self._json(200, {"created": int(time.time()), "data": out},
extra={"x-router-model": "comfyui"})
except urllib.error.HTTPError as e:
return self._json(e.code, {"error": f"comfyui {e.code}: {e.read().decode()[:300]}"})
except Exception as e:
return self._json(502, {"error": f"comfyui unreachable at {COMFYUI}: {e}"})
def _handle_audio(self, body):
# OpenAI request carrying audio -> Ollama native /api/chat (images field), Gemma 4
model = AUDIO_MODEL
@@ -442,6 +517,13 @@ class Handler(BaseHTTPRequestHandler):
return self._handle_transcription()
if path == "/v1/audio/translations":
return self._handle_transcription(translate=True)
if path == "/v1/images/generations":
length = int(self.headers.get("content-length", 0) or 0)
try:
ibody = json.loads(self.rfile.read(length) or b"{}")
except Exception as e:
return self._json(400, {"error": f"bad request: {e}"})
return self._handle_image(ibody)
if path.startswith("/v1/messages"):
return self._proxy_anthropic(path)
if path == "/api/show": # ollama launch model probe (read-only)