From 18a205d69dd77e888ed195f86fde916bce95fa35 Mon Sep 17 00:00:00 2001
From: William Gill <william.gill@anomalous.dev>
Date: Sun, 19 Apr 2026 16:30:04 -0500
Subject: [PATCH] =?UTF-8?q?smart=5Fimage=5Fgen=20v0.7.5:=20fix=20black-ima?=
 =?UTF-8?q?ge=20bug=20=E2=80=94=20fetch=20from=20SaveImage=20explicitly?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

_submit_and_fetch was iterating history[prompt_id]['outputs'].values()
and grabbing the first image it saw. With the inpaint workflow that
includes nodes other than SaveImage that emit IMAGE outputs (the
GroundingDinoSAMSegment node returns an overlay/mask-applied image
in addition to the mask), and dict iteration order is undefined —
sometimes we'd return the overlay (which can render mostly black)
instead of the actual SaveImage result.

Fix: prefer outputs from the SaveImage node id ('9' in every workflow
the tool builds) explicitly. Fall back to scanning all outputs only
if SaveImage didn't appear (workflow drift, manual edit, etc).

User reported seeing the correct inpaint in ComfyUI's native UI but
black in chat — this is the gap.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 .../openwebui-tools/smart_image_gen.py        | 22 ++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/deployments/ai-stack/openwebui-tools/smart_image_gen.py b/deployments/ai-stack/openwebui-tools/smart_image_gen.py
index 681edf8..51a37b7 100644
--- a/deployments/ai-stack/openwebui-tools/smart_image_gen.py
+++ b/deployments/ai-stack/openwebui-tools/smart_image_gen.py
@@ -1,7 +1,7 @@
 """
 title: Smart Image Generator & Editor (ComfyUI)
 author: ai-stack
-version: 0.7.4
+version: 0.7.5
 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
@@ -648,6 +648,14 @@ async def _submit_and_fetch(
         f"CFG {settings['cfg']}, {settings['steps']} steps", False
     )
 
+    # The SaveImage node in every workflow this tool builds is id "9".
+    # We prefer it explicitly because intermediate nodes (e.g. the
+    # GroundingDinoSAMSegment IMAGE output in the inpaint workflow) can
+    # land in the outputs dict too, and dict iteration order is not
+    # stable across runs — without preferring "9" we sometimes returned
+    # an overlay or masked-only image that rendered mostly black.
+    SAVE_NODE_ID = "9"
+
     deadline = time.time() + timeout_seconds
     output_images: list = []
     while time.time() < deadline:
@@ -657,8 +665,16 @@ async def _submit_and_fetch(
                 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", []))
+            outputs = history[prompt_id].get("outputs", {}) or {}
+            # Prefer the canonical SaveImage output …
+            save_imgs = (outputs.get(SAVE_NODE_ID) or {}).get("images", [])
+            if save_imgs:
+                output_images.extend(save_imgs)
+            # … only fall back to other nodes if SaveImage didn't fire
+            # (workflow drift, manual override, etc.)
+            if not output_images:
+                for node_out in outputs.values():
+                    output_images.extend(node_out.get("images", []))
             if output_images:
                 break