diff --git a/deployments/ai-stack/README.md b/deployments/ai-stack/README.md index 5cf088f..47959af 100644 --- a/deployments/ai-stack/README.md +++ b/deployments/ai-stack/README.md @@ -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 diff --git a/deployments/ai-stack/openwebui-tools/smart_image_gen.py b/deployments/ai-stack/openwebui-tools/smart_image_gen.py new file mode 100644 index 0000000..f3711ee --- /dev/null +++ b/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})"