alphacentri/comfyui-nvidia
1
0
Fork 0
Code Issues Pull Requests Actions Packages Releases Activity

Add smart_image_gen Tool for per-prompt checkpoint routing

Browse Source 
Open WebUI Tool the LLM invokes instead of the built-in image action.
Auto-routes among the seven SDXL checkpoints (photo / juggernaut /
pony / general / furry-{nai,noob,il}) based on either an explicit
`style` arg or first-match-wins regex over the prompt. Constructs the
ComfyUI workflow inline, submits via /prompt, polls /history, returns
the result as a base64 data-URI markdown image so no extra hosting is
needed. Per-style default negatives. ComfyUI URL / steps / CFG /
timeout are admin-tunable Valves.

Filters can't see image-gen requests in Open WebUI (the routers skip
the filter chain), so the LLM-driven Tool is the only path that
gives intent-aware routing without changing the chat UX.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
57_Wolve
2026-04-19 12:17:02 -05:00
co-authored by Claude Opus 4.7
parent 704bcfdf13
commit 392b26167f
2 changed files with 258 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
deployments/ai-stack/README.md
+41 -7
View File
@@ -10,13 +10,14 @@ production `srvno.de` deployment.
## Files
| File | Purpose |
| -------------------------- | -------------------------------------------------------- |
| `docker-compose.yml` | Service definitions, volumes, GPU reservations |
| `Caddyfile` | TLS + reverse proxy config (one site block per hostname) |
| `init-models.sh` | LLMs to preseed into Ollama on first boot |
| `comfyui-init-models.sh` | Checkpoints/VAEs/LoRAs to preseed into ComfyUI on first boot |
| `.env.example` | Secrets and image-tag pins. Copy to `.env` |
| File | Purpose |
| --------------------------------------- | -------------------------------------------------------- |
| `docker-compose.yml` | Service definitions, volumes, GPU reservations |
| `Caddyfile` | TLS + reverse proxy config (one site block per hostname) |
| `init-models.sh` | LLMs to preseed into Ollama on first boot |
| `comfyui-init-models.sh` | Checkpoints/VAEs/LoRAs to preseed into ComfyUI on first boot |
| `openwebui-tools/smart_image_gen.py` | Tool that auto-routes image generation to the right SDXL checkpoint |
| `.env.example` | Secrets and image-tag pins. Copy to `.env` |
## 1. Host prerequisites
@@ -146,6 +147,39 @@ Open WebUI submits the workflow to ComfyUI; the result drops back into
the chat when KSampler finishes. To test img2img, attach an image and
use the edit action.
## 8. (Optional) Install the smart-routing Tool
The image-button path always uses the admin's **Default Model**. To get
per-prompt checkpoint routing — e.g. "draw me a cyberpunk city" picks
CyberRealistic, "anthro fox warrior" picks one of the furry checkpoints —
install the `smart_image_gen.py` Tool. The LLM calls it instead of the
built-in image action and chooses the right SDXL checkpoint per request.
1. **Workspace -> Tools -> +** (top-right).
2. Paste the contents of
[`openwebui-tools/smart_image_gen.py`](openwebui-tools/smart_image_gen.py).
3. Save. Optionally adjust the Valves (ComfyUI URL, default steps, CFG,
timeout) via the gear icon.
4. **Workspace -> Models** (or pick an existing chat model) -> edit ->
under **Tools**, enable `smart_image_gen` -> save.
5. Make sure the model has **native function calling** enabled
(Workspace -> Models -> the model -> Advanced Params -> Function
Calling: Native). Mistral, Qwen, and Llama 3.1+ all support this.
In a chat with that model, ask for an image — "make me a photoreal
portrait of a cyberpunk samurai" — the LLM should call
`generate_image(prompt=..., style="photo")`. The status bar shows
"Routing to photo (CyberRealisticXLPlay…)" while it generates.
To extend (new checkpoint, new style):
- Add the filename to `comfyui-init-models.sh` so it gets pulled.
- Add a key to the `CHECKPOINTS` dict in `smart_image_gen.py`.
- Optionally add style-specific negatives to `NEGATIVES`.
- Optionally add keyword routing rules to `ROUTING_RULES` for the
auto-detect path.
- Re-paste the Tool source in Workspace -> Tools.
## Enabling Anubis (later)
The `anubis-owui` service is defined in compose but no Caddy site block
deployments/ai-stack/openwebui-tools/smart_image_gen.py
+217
View File
@@ -0,0 +1,217 @@
"""
title: Smart Image Generator (ComfyUI)
author: ai-stack
version: 0.1.0
description: Generate images via ComfyUI with automatic SDXL checkpoint
routing. The LLM picks (or auto-detects) the right model — photoreal,
anime/score-tag, furry-IL, etc. — based on the user's request.
required_open_webui_version: 0.5.0
"""
import asyncio
import base64
import re
import time
import uuid
from typing import Awaitable, Callable, Optional
import aiohttp
from pydantic import BaseModel, Field
# Filename → use case. Edit alongside `comfyui-init-models.sh` so the
# files actually exist in /opt/comfyui/models/checkpoints/.
CHECKPOINTS = {
"photo": "CyberRealisticXLPlay_V8.0_FP16.safetensors",
"juggernaut": "Juggernaut-XL_v9_RunDiffusionPhoto_v2.safetensors",
"pony": "ponyDiffusionV6XL_v6StartWithThisOne.safetensors",
"general": "talmendoxlSDXL_v11Beta.safetensors",
"furry-nai": "reedFURRYMixSDXL_v23nai.safetensors",
"furry-noob": "indigoVoidFurryFusedXL_noobaiV32.safetensors",
"furry-il": "novaFurryXL_ilV170.safetensors",
}
# Style-specific negative prompts. Appended to whatever the caller supplies.
NEGATIVES = {
"photo": "cartoon, drawing, illustration, anime, painting, sketch, lowres, blurry",
"juggernaut": "cartoon, drawing, illustration, anime, painting, sketch, lowres, blurry",
"pony": "score_6, score_5, score_4, lowres, blurry, worst quality, bad anatomy",
"general": "lowres, blurry, jpeg artifacts, watermark, text, signature",
"furry-nai": "human, lowres, blurry, worst quality, bad anatomy",
"furry-noob": "human, lowres, blurry, worst quality, bad anatomy",
"furry-il": "human, lowres, blurry, worst quality, bad anatomy",
}
# First-match-wins keyword router used when the caller didn't pick a style.
# Order matters — narrower patterns above broader ones.
ROUTING_RULES = [
(re.compile(r"\bscore_\d", re.I), "pony"),
(re.compile(r"\bpony\b", re.I), "pony"),
(re.compile(r"\b(noobai|noob)\b", re.I), "furry-noob"),
(re.compile(r"\b(illustrious|ilxl)\b", re.I), "furry-il"),
(re.compile(r"\b(furry|anthro|feral|kemono|fursona|species)\b", re.I), "furry-nai"),
(re.compile(r"\b(juggernaut)\b", re.I), "juggernaut"),
(re.compile(r"\b(photo|photograph|realistic|portrait|selfie|cinematic)\b", re.I), "photo"),
(re.compile(r"\b(anime|manga|2d|illustration)\b", re.I), "pony"),
]
DEFAULT_STYLE = "general"
def _route_style(prompt: str) -> str:
for pattern, style in ROUTING_RULES:
if pattern.search(prompt):
return style
return DEFAULT_STYLE
def _build_workflow(prompt, negative, ckpt, width, height, steps, cfg, seed):
return {
"3": {"class_type": "KSampler", "inputs": {
"seed": seed if seed > 0 else int(time.time() * 1000) % (2**31),
"steps": steps, "cfg": cfg,
"sampler_name": "dpmpp_2m", "scheduler": "karras",
"denoise": 1.0,
"model": ["4", 0], "positive": ["6", 0],
"negative": ["7", 0], "latent_image": ["5", 0],
}},
"4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": ckpt}},
"5": {"class_type": "EmptyLatentImage",
"inputs": {"width": width, "height": height, "batch_size": 1}},
"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": "smartgen", "images": ["8", 0]}},
}
class Tools:
class Valves(BaseModel):
COMFYUI_BASE_URL: str = Field(
default="http://comfyui:8188",
description="ComfyUI server URL reachable from the open-webui container.",
)
DEFAULT_STEPS: int = Field(default=25, description="KSampler steps.")
DEFAULT_CFG: float = Field(default=7.0, description="CFG scale.")
TIMEOUT_SECONDS: int = Field(
default=180,
description="Maximum wait for a single generation to complete.",
)
def __init__(self):
self.valves = self.Valves()
async def generate_image(
self,
prompt: str,
style: Optional[str] = None,
negative_prompt: Optional[str] = None,
width: int = 1024,
height: int = 1024,
seed: int = 0,
__event_emitter__: Optional[Callable[[dict], Awaitable[None]]] = None,
) -> str:
"""
Generate an image with the right SDXL checkpoint for the request.
Pick `style` based on what the user wants:
- "photo": photorealistic photographs, portraits, cinematic shots.
- "juggernaut": versatile photoreal alternative — sharper, more saturated.
- "pony": anime / illustration with score tags (score_9, score_8_up, ...).
- "general": catch-all SDXL when no specific style applies.
- "furry-nai": anthropomorphic characters, NAI-trained mix.
- "furry-noob": anthropomorphic characters, NoobAI base.
- "furry-il": anthropomorphic characters, Illustrious base.
If `style` is omitted, the tool auto-detects from `prompt` keywords.
:param prompt: The image description.
:param style: One of the keys above. Omit to auto-route.
:param negative_prompt: Extra negatives appended to the style default.
:param width: Output width in pixels (default 1024, SDXL native).
:param height: Output height in pixels (default 1024, SDXL native).
:param seed: Specific seed, or 0 to randomize.
:return: Markdown embedding the generated image.
"""
chosen = style or _route_style(prompt)
ckpt = CHECKPOINTS.get(chosen)
if not ckpt:
return (
f"Unknown style '{chosen}'. "
f"Available: {', '.join(CHECKPOINTS.keys())}"
)
async def emit(msg: str, done: bool = False):
if __event_emitter__:
await __event_emitter__({
"type": "status",
"data": {"description": msg, "done": done},
})
await emit(f"Routing to {chosen} ({ckpt})")
negative = NEGATIVES.get(chosen, "")
if negative_prompt:
negative = f"{negative}, {negative_prompt}" if negative else negative_prompt
workflow = _build_workflow(
prompt=prompt,
negative=negative,
ckpt=ckpt,
width=width,
height=height,
steps=self.valves.DEFAULT_STEPS,
cfg=self.valves.DEFAULT_CFG,
seed=seed,
)
client_id = str(uuid.uuid4())
base = self.valves.COMFYUI_BASE_URL.rstrip("/")
async with aiohttp.ClientSession() as session:
async with session.post(
f"{base}/prompt",
json={"prompt": workflow, "client_id": client_id},
) as resp:
if resp.status != 200:
return f"ComfyUI rejected the prompt: {resp.status} {await resp.text()}"
submit = await resp.json()
prompt_id = submit.get("prompt_id")
if not prompt_id:
return f"ComfyUI didn't return a prompt_id: {submit}"
await emit("Queued, sampling…")
deadline = time.time() + self.valves.TIMEOUT_SECONDS
output_images = []
while time.time() < deadline:
await asyncio.sleep(1.5)
async with session.get(f"{base}/history/{prompt_id}") as resp:
if resp.status != 200:
continue
history = await resp.json()
if prompt_id in history:
for node_out in history[prompt_id].get("outputs", {}).values():
output_images.extend(node_out.get("images", []))
if output_images:
break
if not output_images:
return f"Timed out after {self.valves.TIMEOUT_SECONDS}s waiting for image."
await emit("Fetching result…")
img = output_images[0]
params = {
"filename": img["filename"],
"subfolder": img.get("subfolder", ""),
"type": img.get("type", "output"),
}
async with session.get(f"{base}/view", params=params) as resp:
if resp.status != 200:
return f"Failed to fetch image: {resp.status}"
raw = await resp.read()
b64 = base64.b64encode(raw).decode("ascii")
await emit(f"Done — {chosen}", done=True)
return f"![{chosen}](data:image/png;base64,{b64})"