From 428fd202c59f9faf121f2f8f5b1001981263f723 Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Fri, 11 Oct 2024 13:11:17 -0700 Subject: [PATCH] fix: custom comfyui prompt issue --- .../open_webui/apps/images/utils/comfyui.py | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/backend/open_webui/apps/images/utils/comfyui.py b/backend/open_webui/apps/images/utils/comfyui.py index 0a3e3a1d9..4c421d7c5 100644 --- a/backend/open_webui/apps/images/utils/comfyui.py +++ b/backend/open_webui/apps/images/utils/comfyui.py @@ -125,22 +125,34 @@ async def comfyui_generate_image( workflow[node_id]["inputs"][node.key] = model elif node.type == "prompt": for node_id in node.node_ids: - workflow[node_id]["inputs"]["text"] = payload.prompt + workflow[node_id]["inputs"][ + node.key if node.key else "text" + ] = payload.prompt elif node.type == "negative_prompt": for node_id in node.node_ids: - workflow[node_id]["inputs"]["text"] = payload.negative_prompt + workflow[node_id]["inputs"][ + node.key if node.key else "text" + ] = payload.negative_prompt elif node.type == "width": for node_id in node.node_ids: - workflow[node_id]["inputs"]["width"] = payload.width + workflow[node_id]["inputs"][ + node.key if node.key else "width" + ] = payload.width elif node.type == "height": for node_id in node.node_ids: - workflow[node_id]["inputs"]["height"] = payload.height + workflow[node_id]["inputs"][ + node.key if node.key else "height" + ] = payload.height elif node.type == "n": for node_id in node.node_ids: - workflow[node_id]["inputs"]["batch_size"] = payload.n + workflow[node_id]["inputs"][ + node.key if node.key else "batch_size" + ] = payload.n elif node.type == "steps": for node_id in node.node_ids: - workflow[node_id]["inputs"]["steps"] = payload.steps + workflow[node_id]["inputs"][ + node.key if node.key else "steps" + ] = payload.steps elif node.type == "seed": seed = ( payload.seed