Add smart_image_gen Tool for per-prompt checkpoint routing

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>

deployments/ai-stack/openwebui-tools/smart_image_gen.py

@@ -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})"