diff --git a/README.md b/README.md index 3108cfd..6b4980b 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/comfyui-workflow.json b/comfyui-workflow.json new file mode 100644 index 0000000..b6a388e --- /dev/null +++ b/comfyui-workflow.json @@ -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] } + } +} diff --git a/docs/IMAGES.md b/docs/IMAGES.md new file mode 100644 index 0000000..e6b21b9 --- /dev/null +++ b/docs/IMAGES.md @@ -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://: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": ""}]} +``` +- `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. diff --git a/router.py b/router.py index b9d98a0..2f71289 100644 --- a/router.py +++ b/router.py @@ -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)