From e2291f7148affcca44b626d62af3ed6bb7a82398 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Wed, 21 Aug 2024 01:39:30 +0200
Subject: [PATCH] refac

---
 backend/apps/images/main.py                   |  1 -
 backend/apps/images/utils/comfyui.py          | 30 +++++++++-----
 .../components/admin/Settings/Images.svelte   | 41 ++++++++++++++++++-
 3 files changed, 60 insertions(+), 12 deletions(-)

diff --git a/backend/apps/images/main.py b/backend/apps/images/main.py
index eea8bc2d5..35b84d75e 100644
--- a/backend/apps/images/main.py
+++ b/backend/apps/images/main.py
@@ -413,7 +413,6 @@ async def image_generations(
             return images
 
         elif app.state.config.ENGINE == "comfyui":
-
             data = {
                 "prompt": form_data.prompt,
                 "width": width,
diff --git a/backend/apps/images/utils/comfyui.py b/backend/apps/images/utils/comfyui.py
index 9117d7abf..7dff1c09a 100644
--- a/backend/apps/images/utils/comfyui.py
+++ b/backend/apps/images/utils/comfyui.py
@@ -74,7 +74,7 @@ def get_images(ws, prompt, client_id, base_url):
 
 class ComfyUINodeInput(BaseModel):
     field: Optional[str] = None
-    node_id: str
+    node_ids: list[str] = []
     key: Optional[str] = "text"
     value: Optional[str] = None
 
@@ -106,27 +106,37 @@ async def comfyui_generate_image(
     for node in payload.workflow.nodes:
         if node.field:
             if node.field == "model":
-                workflow[node.node_id]["inputs"][node.key] = model
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"][node.key] = model
             elif node.field == "prompt":
-                workflow[node.node_id]["inputs"]["text"] = payload.prompt
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"]["text"] = payload.prompt
             elif node.field == "negative_prompt":
-                workflow[node.node_id]["inputs"]["text"] = payload.negative_prompt
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"]["text"] = payload.negative_prompt
             elif node.field == "width":
-                workflow[node.node_id]["inputs"]["width"] = payload.width
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"]["width"] = payload.width
             elif node.field == "height":
-                workflow[node.node_id]["inputs"]["height"] = payload.height
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"]["height"] = payload.height
             elif node.field == "n":
-                workflow[node.node_id]["inputs"]["batch_size"] = payload.n
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"]["batch_size"] = payload.n
             elif node.field == "steps":
-                workflow[node.node_id]["inputs"]["steps"] = payload.steps
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"]["steps"] = payload.steps
             elif node.field == "seed":
-                workflow[node.node_id]["inputs"]["seed"] = (
+                seed = (
                     payload.seed
                     if payload.seed
                     else random.randint(0, 18446744073709551614)
                 )
+                for node_id in node.node_ids:
+                    workflow[node.node_id]["inputs"]["seed"] = seed
         else:
-            workflow[node.node_id]["inputs"][node.key] = node.value
+            for node_id in node.node_ids:
+                workflow[node_id]["inputs"][node.key] = node.value
 
     try:
         ws = websocket.WebSocket()
diff --git a/src/lib/components/admin/Settings/Images.svelte b/src/lib/components/admin/Settings/Images.svelte
index 2748b2b10..36260d48d 100644
--- a/src/lib/components/admin/Settings/Images.svelte
+++ b/src/lib/components/admin/Settings/Images.svelte
@@ -323,7 +323,46 @@
 						<div class="">
 							<div class=" mb-2 text-sm font-medium">{$i18n.t('ComfyUI Workflow Nodes')}</div>
 
-							<div>hi</div>
+							<div class="text-xs">
+								<div class="flex gap-2">
+									<div>Prompt Node</div>
+									<div>Node Ids</div>
+									<div>"text"</div>
+									<div>Default</div>
+								</div>
+
+								<div class="flex gap-2">
+									<div>Model Node</div>
+									<div>Node Ids</div>
+									<div>"text"</div>
+									<div>Default</div>
+								</div>
+
+								<div class="flex gap-2">
+									<div>Image Size Node</div>
+									<div>Node Ids</div>
+									<div>"text"</div>
+									<div>Default</div>
+								</div>
+
+								<div class="flex gap-2">
+									<div>Image Steps Node</div>
+									<div>Node Ids</div>
+									<div>"text"</div>
+									<div>Default</div>
+								</div>
+							</div>
+
+							<hr class=" dark:border-gray-850 my-2" />
+
+							<div class="text-xs">
+								<div class="flex gap-2">
+									<div>Custom Node</div>
+									<div>Node Ids</div>
+									<div>Key</div>
+									<div>Value</div>
+								</div>
+							</div>
 						</div>
 					{/if}
 				{:else if config?.engine === 'openai'}