refac

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.type:
-            if node.field == "model":
+            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

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="text-xs flex flex-col gap-1">
-								<div class="flex gap-2">
+								{#each workflowNodes as node}
-									<div>Prompt Node</div>
+									<div class="flex w-full items-center border border-gray-850 rounded-lg">
-									<div>Node Ids</div>
+										<div class="flex-shrink-0">
-									<div>"text"</div>
+											<div
-									<div>Default</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 class="w-full">
-									<div>Model Node</div>
+											<Tooltip
-									<div>Node Ids</div>
+												content="Comma separated Node Ids (e.g. 1,2,3)"
-									<div>"text"</div>
+												placement="top-start"
-									<div>Default</div>
+											>
-								</div>
+												<input
-
+													class="w-full py-1 px-4 rounded-r-lg text-xs bg-transparent outline-none"
-								<div class="flex gap-2">
+													placeholder="Node Ids"
-									<div>Image Size Node</div>
+													bind:value={node.node_ids}
-									<div>Node Ids</div>
+												/>
-									<div>"text"</div>
+											</Tooltip>
-									<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>
-
+								{/each}
-							<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}