From aaea9a5956f5f7526f4ba4acff110fc980122657 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Tue, 30 Dec 2025 12:24:48 +0400 Subject: [PATCH] refac/fix: comfyui filter output node type Co-Authored-By: Paul <239564541+mirrordna-reflection-protocol@users.noreply.github.com> --- backend/open_webui/utils/images/comfyui.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/backend/open_webui/utils/images/comfyui.py b/backend/open_webui/utils/images/comfyui.py index 0a097285a0..3c402cbc17 100644 --- a/backend/open_webui/utils/images/comfyui.py +++ b/backend/open_webui/utils/images/comfyui.py @@ -64,8 +64,8 @@ def get_history(prompt_id, base_url, api_key): return json.loads(response.read()) -def get_images(ws, prompt, client_id, base_url, api_key): - prompt_id = queue_prompt(prompt, client_id, base_url, api_key)["prompt_id"] +def get_images(ws, workflow, client_id, base_url, api_key): + prompt_id = queue_prompt(workflow, client_id, base_url, api_key)["prompt_id"] output_images = [] while True: out = ws.recv() @@ -81,12 +81,16 @@ def get_images(ws, prompt, client_id, base_url, api_key): history = get_history(prompt_id, base_url, api_key)[prompt_id] for node_id in history["outputs"]: node_output = history["outputs"][node_id] - if "images" in node_output: - for image in node_output["images"]: - url = get_image_url( - image["filename"], image["subfolder"], image["type"], base_url - ) - output_images.append({"url": url}) + if node_id in workflow and workflow[node_id].get("class_type") in [ + "SaveImage", + "PreviewImage", + ]: + if "images" in node_output: + for image in node_output["images"]: + url = get_image_url( + image["filename"], image["subfolder"], image["type"], base_url + ) + output_images.append({"url": url}) return {"data": output_images}