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

smart_image_gen v0.5: surface images via files event (canonical path)

Browse Source 
The data-URI message-event approach didn't render — Open WebUI's chat
frontend ignores data URIs from tool-emitted message events because
the markdown-base64 rewriter (utils/files.py convert_markdown_base64
_images) only runs on assistant streaming content, not on tool emits.

Switched to the path Open WebUI's own image-generation flow uses
(backend/open_webui/utils/middleware.py ~1325):

  1. Upload image bytes via open_webui.routers.files.upload_file_handler
     (gets back a file_item with id)
  2. Resolve the served URL via request.app.url_path_for(
     "get_file_content_by_id", id=file_item.id) → /api/v1/files/{id}/content
  3. Emit a `files` event:
        {"type": "files", "data": {"files": [{"type": "image", "url": ...}]}}

Tools now take __request__, __user__, __metadata__ params for the
upload (Open WebUI auto-injects these). Falls back to data-URI
message event if the runtime imports aren't available (e.g. running
the file standalone for tests). The internal upload bypasses
get_verified_user via the user= kwarg, so no token plumbing.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
57_Wolve
2026-04-19 13:21:48 -05:00
co-authored by Claude Opus 4.7
parent 4d996e1205
commit b604e3f509
1 changed files with 109 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
deployments/ai-stack/openwebui-tools/smart_image_gen.py
+109 -15
View File
@@ -1,30 +1,45 @@
"""
title: Smart Image Generator & Editor (ComfyUI)
author: ai-stack
version: 0.4.0
version: 0.5.0
description: Generate or edit images via ComfyUI with automatic SDXL
checkpoint routing. Two methods — generate_image (txt2img) and
edit_image (img2img on the user's most recently attached image). The
LLM picks (or auto-detects) the right model — photoreal, Pony
score-tag, NoobAI/Illustrious furry, etc. — and each style ships
with the creator-recommended sampler, scheduler, CFG, steps, CLIP
skip, prompt-prefix dialect, and negatives. The actual image is
pushed into the chat via the event emitter; the function return is
just a short confirmation so the LLM doesn't try to describe or
re-emit the data URI.
skip, prompt-prefix dialect, and negatives. The image is uploaded
to Open WebUI's file store and surfaced via a `files` event (the
canonical pattern used by Open WebUI's own image-gen path); the
function return is a short confirmation so the LLM doesn't try to
describe or re-emit the image.
required_open_webui_version: 0.5.0
"""
import asyncio
import base64
import inspect
import io
import re
import time
import uuid
from typing import Awaitable, Callable, Optional
from typing import Awaitable, Callable, Literal, Optional
import aiohttp
from pydantic import BaseModel, Field
from typing import Literal
# Open WebUI's runtime — only available when the tool is loaded inside the
# Open WebUI process. Guarded so the module still imports for standalone
# linting/testing; if the imports fail at runtime, _push_image_to_chat
# falls back to emitting a markdown data-URI message.
try:
from fastapi import UploadFile
from open_webui.models.users import Users
from open_webui.routers.files import upload_file_handler
_OPENWEBUI_RUNTIME = True
except ImportError:
_OPENWEBUI_RUNTIME = False
StyleName = Literal[
"photo", "juggernaut", "pony", "general",
@@ -344,6 +359,70 @@ async def _upload_to_comfyui(
return (await resp.json()).get("name", name)
async def _push_image_to_chat(
raw: bytes,
filename_prefix: str,
request,
user_dict: Optional[dict],
metadata: Optional[dict],
event_emitter: Optional[Callable[[dict], Awaitable[None]]],
) -> bool:
"""
Surface a generated image in the chat using Open WebUI's canonical
pattern: upload the bytes via the internal file store, then emit a
`files` event referencing the served URL. This is the same path Open
WebUI's own image-generation code uses (utils/middleware.py ~1325).
Returns True if the image was uploaded and emitted via the files
event. Returns False if anything is missing — caller should fall
back to a data-URI markdown message in that case.
"""
if not (_OPENWEBUI_RUNTIME and request and user_dict and event_emitter):
return False
try:
user = Users.get_user_by_id(user_dict.get("id"))
if not user:
return False
upload = UploadFile(
file=io.BytesIO(raw),
filename=f"{filename_prefix}_{uuid.uuid4().hex[:8]}.png",
headers={"content-type": "image/png"},
)
meta = metadata or {}
result = upload_file_handler(
request=request,
file=upload,
metadata={
"chat_id": meta.get("chat_id"),
"message_id": meta.get("message_id"),
},
process=False,
user=user,
)
# upload_file_handler may be sync or async depending on the Open
# WebUI version — handle either.
if inspect.iscoroutine(result):
file_item = await result
else:
file_item = result
url = request.app.url_path_for(
"get_file_content_by_id", id=file_item.id
)
await event_emitter({
"type": "files",
"data": {"files": [{"type": "image", "url": url}]},
})
return True
except Exception:
# Any failure (signature drift, missing route, etc.) falls back
# to the data-URI path in the caller.
return False
async def _submit_and_fetch(
session: aiohttp.ClientSession,
base: str,
@@ -421,6 +500,9 @@ class Tools:
width: int = 1024,
height: int = 1024,
seed: int = 0,
__request__=None,
__user__: Optional[dict] = None,
__metadata__: Optional[dict] = None,
__event_emitter__: Optional[Callable[[dict], Awaitable[None]]] = None,
) -> str:
"""
@@ -485,16 +567,21 @@ class Tools:
if err:
return err
b64 = base64.b64encode(raw).decode("ascii")
# Push the image straight into the chat. The function return is just
# a confirmation for the LLM — if we returned the markdown, the LLM
# would either echo the base64 to the user as text, or hallucinate
# a description of what it thinks the image looks like.
if __event_emitter__:
# Surface the image in the chat. Preferred path uploads to Open
# WebUI's file store and emits a `files` event (matches the built-
# in image-gen flow). Fallback inlines a data-URI markdown via a
# `message` event for environments where the file API isn't
# reachable from the tool process.
pushed = await _push_image_to_chat(
raw, "smartgen", __request__, __user__, __metadata__, __event_emitter__,
)
if not pushed and __event_emitter__:
b64 = base64.b64encode(raw).decode("ascii")
await __event_emitter__({
"type": "message",
"data": {"content": f"![{chosen}](data:image/png;base64,{b64})"},
})
await emit(f"Done — {chosen}", done=True)
return (
f"Image generated and shown to the user above (style: {chosen}, "
@@ -511,6 +598,9 @@ class Tools:
denoise: float = 0.7,
negative_prompt: Optional[str] = None,
seed: int = 0,
__request__=None,
__user__: Optional[dict] = None,
__metadata__: Optional[dict] = None,
__files__: Optional[list] = None,
__messages__: Optional[list] = None,
__event_emitter__: Optional[Callable[[dict], Awaitable[None]]] = None,
@@ -599,12 +689,16 @@ class Tools:
if err:
return err
b64 = base64.b64encode(raw_out).decode("ascii")
if __event_emitter__:
pushed = await _push_image_to_chat(
raw_out, "smartedit", __request__, __user__, __metadata__, __event_emitter__,
)
if not pushed and __event_emitter__:
b64 = base64.b64encode(raw_out).decode("ascii")
await __event_emitter__({
"type": "message",
"data": {"content": f"![edit:{chosen}](data:image/png;base64,{b64})"},
})
await emit(f"Done — {chosen} (denoise {denoise:.2f})", done=True)
return (
f"Edited image shown to the user above (style: {chosen}, "