From 29bef261be9771eea755d4d1d1ed52be439a70bf Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Wed, 21 Aug 2024 14:16:33 +0200
Subject: [PATCH] refac

---
 backend/apps/images/utils/comfyui.py          |  20 +--
 .../components/admin/Settings/Images.svelte   | 120 ++++++++++++------
 2 files changed, 90 insertions(+), 50 deletions(-)

diff --git a/backend/apps/images/utils/comfyui.py b/backend/apps/images/utils/comfyui.py
index 7dff1c09a..cece3d737 100644
--- a/backend/apps/images/utils/comfyui.py
+++ b/backend/apps/images/utils/comfyui.py
@@ -73,7 +73,7 @@ def get_images(ws, prompt, client_id, base_url):
 
 
 class ComfyUINodeInput(BaseModel):
-    field: Optional[str] = None
+    type: Optional[str] = None
     node_ids: list[str] = []
     key: Optional[str] = "text"
     value: Optional[str] = None
@@ -104,29 +104,29 @@ async def comfyui_generate_image(
     workflow = json.loads(payload.workflow.workflow)
 
     for node in payload.workflow.nodes:
-        if node.field:
-            if node.field == "model":
+        if node.type:
+            if node.type == "model":
                 for node_id in node.node_ids:
                     workflow[node_id]["inputs"][node.key] = model
-            elif node.field == "prompt":
+            elif node.type == "prompt":
                 for node_id in node.node_ids:
                     workflow[node_id]["inputs"]["text"] = payload.prompt
-            elif node.field == "negative_prompt":
+            elif node.type == "negative_prompt":
                 for node_id in node.node_ids:
                     workflow[node_id]["inputs"]["text"] = payload.negative_prompt
-            elif node.field == "width":
+            elif node.type == "width":
                 for node_id in node.node_ids:
                     workflow[node_id]["inputs"]["width"] = payload.width
-            elif node.field == "height":
+            elif node.type == "height":
                 for node_id in node.node_ids:
                     workflow[node_id]["inputs"]["height"] = payload.height
-            elif node.field == "n":
+            elif node.type == "n":
                 for node_id in node.node_ids:
                     workflow[node_id]["inputs"]["batch_size"] = payload.n
-            elif node.field == "steps":
+            elif node.type == "steps":
                 for node_id in node.node_ids:
                     workflow[node_id]["inputs"]["steps"] = payload.steps
-            elif node.field == "seed":
+            elif node.type == "seed":
                 seed = (
                     payload.seed
                     if payload.seed
diff --git a/src/lib/components/admin/Settings/Images.svelte b/src/lib/components/admin/Settings/Images.svelte
index 36260d48d..a0ffe4e71 100644
--- a/src/lib/components/admin/Settings/Images.svelte
+++ b/src/lib/components/admin/Settings/Images.svelte
@@ -14,6 +14,7 @@
 	} from '$lib/apis/images';
 	import SensitiveInput from '$lib/components/common/SensitiveInput.svelte';
 	import Switch from '$lib/components/common/Switch.svelte';
+	import Tooltip from '$lib/components/common/Tooltip.svelte';
 	const dispatch = createEventDispatcher();
 
 	const i18n = getContext('i18n');
@@ -25,7 +26,33 @@
 
 	let models = null;
 
-	let inputFiles = null;
+	let workflowNodes = [
+		{
+			type: 'prompt',
+			key: 'text',
+			node_ids: []
+		},
+		{
+			type: 'model',
+			key: 'ckpt_name',
+			node_ids: []
+		},
+		{
+			type: 'width',
+			key: 'width',
+			node_ids: []
+		},
+		{
+			type: 'height',
+			key: 'height',
+			node_ids: []
+		},
+		{
+			type: 'steps',
+			key: 'steps',
+			node_ids: []
+		}
+	];
 
 	const getModels = async () => {
 		models = await getImageGenerationModels(localStorage.token).catch((error) => {
@@ -50,9 +77,28 @@
 		}
 	};
 
+	const validateJSON = (json) => {
+		try {
+			const obj = JSON.parse(json);
+
+			if (obj && typeof obj === 'object') {
+				return true;
+			}
+		} catch (e) {}
+		return false;
+	};
+
 	const saveHandler = async () => {
 		loading = true;
 
+		if (config?.comfyui?.COMFYUI_WORKFLOW) {
+			if (!validateJSON(config.comfyui.COMFYUI_WORKFLOW)) {
+				toast.error('Invalid JSON format for ComfyUI Workflow.');
+				loading = false;
+				return;
+			}
+		}
+
 		await updateConfig(localStorage.token, config).catch((error) => {
 			toast.error(error);
 			loading = false;
@@ -276,7 +322,6 @@
 								class="w-full rounded-lg mb-1 py-2 px-4 text-xs bg-gray-50 dark:text-gray-300 dark:bg-gray-850 outline-none disabled:text-gray-600 resize-none"
 								rows="10"
 								value={JSON.stringify(JSON.parse(config.comfyui.COMFYUI_WORKFLOW), null, 2)}
-								disabled
 							/>
 						{/if}
 
@@ -323,45 +368,40 @@
 						<div class="">
 							<div class=" mb-2 text-sm font-medium">{$i18n.t('ComfyUI Workflow Nodes')}</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="text-xs flex flex-col gap-1">
+								{#each workflowNodes as node}
+									<div class="flex w-full items-center border border-gray-850 rounded-lg">
+										<div class="flex-shrink-0">
+											<div
+												class=" capitalize line-clamp-1 font-medium px-3 py-1 w-20 text-center rounded-l-lg bg-green-500/20 text-green-700 dark:text-green-200"
+											>
+												{node.type}
+											</div>
+										</div>
+										<div class="">
+											<Tooltip content="Input Key (e.g. text, unet_name, steps)">
+												<input
+													class="py-1 px-3 w-24 text-xs text-center bg-transparent outline-none border-r border-gray-850"
+													placeholder="Key"
+													bind:value={node.key}
+												/>
+											</Tooltip>
+										</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 class="w-full">
+											<Tooltip
+												content="Comma separated Node Ids (e.g. 1,2,3)"
+												placement="top-start"
+											>
+												<input
+													class="w-full py-1 px-4 rounded-r-lg text-xs bg-transparent outline-none"
+													placeholder="Node Ids"
+													bind:value={node.node_ids}
+												/>
+											</Tooltip>
+										</div>
+									</div>
+								{/each}
 							</div>
 						</div>
 					{/if}