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
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
5bd8142d30
commit
2f95e6be96
@@ -98,6 +98,8 @@ Everything is environment variables (read by `run.sh` / the router):
|
||||
| `ROUTER_AUDIO_MODEL` | `gemma4:e4b` | model used for audio/transcription |
|
||||
| `ROUTER_API_KEY` | *(unset)* | if set, require this Bearer token on /v1 |
|
||||
| `ROUTER_UNCENSORED_KEY` | *(unset)* | Bearer token that forces the uncensored fleet |
|
||||
| `ROUTER_COMFYUI` | `http://127.0.0.1:8188` | ComfyUI server for image generation |
|
||||
| `ROUTER_COMFY_CHECKPOINT` | *(unset)* | checkpoint filename for the default image workflow |
|
||||
|
||||
> ⚠️ `ROUTER_OLLAMA` is the router's **client** target — keep it `127.0.0.1`.
|
||||
> Do **not** confuse it with Ollama's own `OLLAMA_HOST` server-bind variable.
|
||||
@@ -117,6 +119,7 @@ main thing to customize for your machine. See [docs/ROUTING.md](docs/ROUTING.md)
|
||||
| POST | `/v1/messages` | Anthropic Messages (Claude Code) → LiteLLM |
|
||||
| POST | `/v1/audio/transcriptions` | Whisper-shaped speech-to-text |
|
||||
| POST | `/v1/audio/translations` | speech → English text |
|
||||
| POST | `/v1/images/generations` | OpenAI-shaped image gen → ComfyUI |
|
||||
| GET | `/v1/models` | lists route targets |
|
||||
| GET | `/`, `/api/version`, `/api/tags`, `/api/ps`, POST `/api/show` | Ollama-native probes (so `ollama launch` accepts the router) |
|
||||
| GET | `/healthz` | liveness |
|
||||
@@ -126,6 +129,7 @@ main thing to customize for your machine. See [docs/ROUTING.md](docs/ROUTING.md)
|
||||
- [docs/DEPLOY.md](docs/DEPLOY.md) — run as a service, expose it, tunnels, Caddy
|
||||
- [docs/ROUTING.md](docs/ROUTING.md) — how routing decides + customizing the map
|
||||
- [docs/AUDIO.md](docs/AUDIO.md) — voice input & transcription
|
||||
- [docs/IMAGES.md](docs/IMAGES.md) — image generation via ComfyUI
|
||||
- [docs/BENCHMARKS.md](docs/BENCHMARKS.md) — the benchmark harnesses & sample results
|
||||
|
||||
## Security
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
{
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"seed": "{{SEED}}",
|
||||
"steps": "{{STEPS}}",
|
||||
"cfg": 7,
|
||||
"sampler_name": "euler",
|
||||
"scheduler": "normal",
|
||||
"denoise": 1,
|
||||
"model": ["4", 0],
|
||||
"positive": ["6", 0],
|
||||
"negative": ["7", 0],
|
||||
"latent_image": ["5", 0]
|
||||
}
|
||||
},
|
||||
"4": {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"inputs": { "ckpt_name": "{{CHECKPOINT}}" }
|
||||
},
|
||||
"5": {
|
||||
"class_type": "EmptyLatentImage",
|
||||
"inputs": { "width": "{{WIDTH}}", "height": "{{HEIGHT}}", "batch_size": "{{BATCH}}" }
|
||||
},
|
||||
"6": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": { "text": "{{PROMPT}}", "clip": ["4", 1] }
|
||||
},
|
||||
"7": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": { "text": "{{NEGATIVE}}", "clip": ["4", 1] }
|
||||
},
|
||||
"8": {
|
||||
"class_type": "VAEDecode",
|
||||
"inputs": { "samples": ["3", 0], "vae": ["4", 2] }
|
||||
},
|
||||
"9": {
|
||||
"class_type": "SaveImage",
|
||||
"inputs": { "filename_prefix": "router", "images": ["8", 0] }
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
# Image generation (ComfyUI)
|
||||
|
||||
The router exposes an OpenAI-compatible **`POST /v1/images/generations`** endpoint
|
||||
backed by a [ComfyUI](https://github.com/comfyanonymous/ComfyUI) server. It submits
|
||||
a workflow to ComfyUI, polls until it finishes, fetches the images, and returns them
|
||||
in OpenAI's image-response shape — so any OpenAI image client works.
|
||||
|
||||
## Configure
|
||||
|
||||
| Env var | Default | Purpose |
|
||||
|---|---|---|
|
||||
| `ROUTER_COMFYUI` | `http://127.0.0.1:8188` | ComfyUI server URL |
|
||||
| `ROUTER_COMFY_CHECKPOINT` | *(unset)* | your checkpoint filename, e.g. `sd_xl_base_1.0.safetensors` |
|
||||
| `ROUTER_COMFY_WORKFLOW` | `comfyui-workflow.json` | the workflow template (see below) |
|
||||
|
||||
At minimum, **set `ROUTER_COMFY_CHECKPOINT`** to a model that exists in your
|
||||
ComfyUI (`ComfyUI/models/checkpoints/`), or hardcode it into the template's
|
||||
`ckpt_name`. Then restart the router.
|
||||
|
||||
## Use
|
||||
|
||||
```bash
|
||||
curl http://<host>:8080/v1/images/generations -H 'content-type: application/json' -d '{
|
||||
"prompt": "a red fox in a snowy forest, cinematic lighting",
|
||||
"n": 1, "size": "1024x1024"
|
||||
}'
|
||||
# -> {"created": ..., "data": [{"b64_json": "<png>"}]}
|
||||
```
|
||||
- `response_format`: `b64_json` (default) or `url` (returns a `data:` URL).
|
||||
- Extra (non-OpenAI) knobs the router honors: `negative_prompt`, `steps`.
|
||||
- Works with the OpenAI SDK: `images.generate(prompt=..., size=..., n=...)`.
|
||||
|
||||
## The workflow template
|
||||
|
||||
`comfyui-workflow.json` is a ComfyUI **API-format** workflow with string sentinels
|
||||
the router substitutes per request:
|
||||
|
||||
| Sentinel | Filled with |
|
||||
|---|---|
|
||||
| `{{PROMPT}}` | the request prompt |
|
||||
| `{{NEGATIVE}}` | `negative_prompt` |
|
||||
| `{{WIDTH}}` / `{{HEIGHT}}` | parsed from `size` |
|
||||
| `{{BATCH}}` | `n` |
|
||||
| `{{SEED}}` | a fresh seed |
|
||||
| `{{STEPS}}` | `steps` (default 25) |
|
||||
| `{{CHECKPOINT}}` | `ROUTER_COMFY_CHECKPOINT` |
|
||||
|
||||
The default is a standard SD/SDXL text-to-image graph. **To use your own** (Flux,
|
||||
a fancier pipeline, LoRAs, upscalers, etc.): build it in ComfyUI, export via
|
||||
**Save (API Format)**, then replace the values you want driven by requests with the
|
||||
sentinels above. Point `ROUTER_COMFY_WORKFLOW` at your file. Anything without a
|
||||
sentinel stays fixed.
|
||||
|
||||
## Notes
|
||||
|
||||
- **Not yet verified end-to-end** — it was written before ComfyUI was installed.
|
||||
Once your ComfyUI is up (`ROUTER_COMFYUI` reachable, checkpoint set), test with
|
||||
the curl above; if a node name/shape differs, adjust the template.
|
||||
- Generation can take a while; the router polls up to ~10 min per request.
|
||||
- Video: ComfyUI video workflows (AnimateDiff, SVD, etc.) also export to API format
|
||||
and return frames/files — the same endpoint can drive them, but the response
|
||||
mapping for multi-frame/video output may need a tweak. Ask when you get there.
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user