Refactor prompt serializer

src/lib/components/ComfyApp.ts

@@ -33,6 +33,7 @@ import notify from "$lib/notify";
 import configState from "$lib/stores/configState";
 import { blankGraph } from "$lib/defaultGraph";
 import type { ComfyExecutionResult } from "$lib/nodes/ComfyWidgetNodes";
+import ComfyPromptSerializer from "./ComfyPromptSerializer";
 
 export const COMFYBOX_SERIAL_VERSION = 1;
 
@@ -84,27 +85,6 @@ type BackendComboNode = {
     backendNode: ComfyBackendNode
 }
 
-function isActiveBackendNode(node: ComfyGraphNode, tag: string | null): boolean {
-    if (!node.isBackendNode)
-        return false;
-
-    if (tag && !hasTag(node, tag)) {
-        console.debug("Skipping tagged node", tag, node.properties.tags, node)
-        return false;
-    }
-
-    if (node.mode === NodeMode.NEVER) {
-        // Don't serialize muted nodes
-        return false;
-    }
-
-    return true;
-}
-
-function hasTag(node: LGraphNode, tag: string): boolean {
-    return "tags" in node.properties && node.properties.tags.indexOf(tag) !== -1
-}
-
 export default class ComfyApp {
     api: ComfyAPI;
     rootEl: HTMLDivElement | null = null;
@@ -122,9 +102,11 @@ export default class ComfyApp {
     private queueItems: QueueItem[] = [];
     private processingQueue: boolean = false;
     private alreadySetup = false;
+    private promptSerializer: ComfyPromptSerializer;
 
     constructor() {
         this.api = new ComfyAPI();
+        this.promptSerializer = new ComfyPromptSerializer();
     }
 
     async setup(): Promise<void> {
@@ -559,157 +541,7 @@ export default class ComfyApp {
      * @returns The workflow and node links
      */
     graphToPrompt(tag: string | null = null): SerializedPrompt {
-        // Run frontend-only logic
+        return this.promptSerializer.serializePrompt(this.lGraph, tag)
-        this.lGraph.runStep(1)
-
-        const workflow = this.lGraph.serialize();
-
-        const output = {};
-        // Process nodes in order of execution
-        for (const node_ of this.lGraph.computeExecutionOrder<ComfyGraphNode>(false, null)) {
-            const n = workflow.nodes.find((n) => n.id === node_.id);
-
-            if (!isActiveBackendNode(node_, tag)) {
-                continue;
-            }
-
-            const node = node_ as ComfyBackendNode;
-
-            const inputs = {};
-
-            // Store input values passed by frontend-only nodes
-            if (node.inputs) {
-                for (let i = 0; i < node.inputs.length; i++) {
-                    const inp = node.inputs[i];
-                    const inputLink = node.getInputLink(i)
-                    const inputNode = node.getInputNode(i)
-
-                    // We don't check tags for non-backend nodes.
-                    // Just check for node inactivity (so you can toggle groups of
-                    // tagged frontend nodes on/off)
-                    if (inputNode && inputNode.mode === NodeMode.NEVER) {
-                        console.debug("Skipping inactive node", inputNode)
-                        continue;
-                    }
-
-                    if (!inputLink || !inputNode) {
-                        if ("config" in inp) {
-                            const defaultValue = (inp as IComfyInputSlot).config?.defaultValue
-                            if (defaultValue !== null && defaultValue !== undefined)
-                                inputs[inp.name] = defaultValue
-                        }
-                        continue;
-                    }
-
-                    let serialize = true;
-                    if ("config" in inp)
-                        serialize = (inp as IComfyInputSlot).serialize
-
-                    let isBackendNode = node.isBackendNode;
-                    let isInputBackendNode = false;
-                    if ("isBackendNode" in inputNode)
-                        isInputBackendNode = (inputNode as ComfyGraphNode).isBackendNode;
-
-                    // The reasoning behind this check:
-                    // We only want to serialize inputs to nodes with backend equivalents.
-                    // And in ComfyBox, the backend nodes in litegraph *never* have widgets, instead they're all inputs.
-                    // All values are passed by separate frontend-only nodes,
-                    // either UI-bound or something like ConstantInteger.
-                    // So we know that any value passed into a backend node *must* come from
-                    // a frontend node.
-                    // The rest (links between backend nodes) will be serialized after this bit runs.
-                    if (serialize && isBackendNode && !isInputBackendNode) {
-                        inputs[inp.name] = inputLink.data
-                    }
-                }
-            }
-
-            // Store links between backend-only and hybrid nodes
-            for (let i = 0; i < node.inputs.length; i++) {
-                let parent: ComfyGraphNode = node.getInputNode(i) as ComfyGraphNode;
-                if (parent) {
-                    const seen = {}
-                    let link = node.getInputLink(i);
-
-                    const isFrontendParent = (parent: ComfyGraphNode) => {
-                        if (!parent || parent.isBackendNode)
-                            return false;
-                        if (tag && !hasTag(parent, tag))
-                            return false;
-                        return true;
-                    }
-
-                    // If there are frontend-only nodes between us and another
-                    // backend node, we have to traverse them first. This
-                    // behavior is dependent on the type of node. Reroute nodes
-                    // will simply follow their single input, while branching
-                    // nodes have conditional logic that determines which link
-                    // to follow backwards.
-                    while (isFrontendParent(parent)) {
-                        if (!("getUpstreamLink" in parent)) {
-                            console.warn("[graphToPrompt] Node does not support getUpstreamLink", parent.type)
-                            break;
-                        }
-
-                        const nextLink = parent.getUpstreamLink()
-                        if (nextLink == null) {
-                            console.warn("[graphToPrompt] No upstream link found in frontend node", parent)
-                            break;
-                        }
-
-                        if (nextLink && !seen[nextLink.id]) {
-                            seen[nextLink.id] = true
-                            const inputNode = parent.graph.getNodeById(nextLink.origin_id) as ComfyGraphNode;
-                            if (inputNode && tag && !hasTag(inputNode, tag)) {
-                                console.debug("[graphToPrompt] Skipping tagged intermediate frontend node", tag, node.properties.tags)
-                                parent = null;
-                            }
-                            else {
-                                console.debug("[graphToPrompt] Traverse upstream link", parent.id, inputNode?.id, inputNode?.isBackendNode)
-                                link = nextLink;
-                                parent = inputNode;
-                            }
-                        } else {
-                            parent = null;
-                        }
-                    }
-
-                    if (link && parent && parent.isBackendNode) {
-                        if (tag && !hasTag(parent, tag))
-                            continue;
-
-                        console.debug("[graphToPrompt] final link", parent.id, node.id)
-                        const input = node.inputs[i]
-                        // TODO can null be a legitimate value in some cases?
-                        // Nodes like CLIPLoader will never have a value in the frontend, hence "null".
-                        if (!(input.name in inputs))
-                            inputs[input.name] = [String(link.origin_id), link.origin_slot];
-                    }
-                }
-            }
-
-            output[String(node.id)] = {
-                inputs,
-                class_type: node.comfyClass,
-            };
-        }
-
-        // Remove inputs connected to removed nodes
-        for (const nodeId in output) {
-            for (const inputName in output[nodeId].inputs) {
-                if (Array.isArray(output[nodeId].inputs[inputName])
-                    && output[nodeId].inputs[inputName].length === 2
-                    && !output[output[nodeId].inputs[inputName][0]]) {
-                    console.debug("Prune removed node link", nodeId, inputName, output[nodeId].inputs[inputName])
-                    delete output[nodeId].inputs[inputName];
-                }
-            }
-        }
-
-        // console.debug({ workflow, output })
-        // console.debug(promptToGraphVis({ workflow, output }))
-
-        return { workflow, output };
     }
 
     async queuePrompt(num: number, batchCount: number = 1, tag: string | null = null) {

src/lib/components/ComfyPromptSerializer.ts

@@ -0,0 +1,197 @@
+import type ComfyGraph from "$lib/ComfyGraph";
+import type { ComfyBackendNode } from "$lib/nodes/ComfyBackendNode";
+import type ComfyGraphNode from "$lib/nodes/ComfyGraphNode";
+import { LGraphNode, NodeMode } from "@litegraph-ts/core";
+import type { SerializedPrompt, SerializedPromptInput, SerializedPromptInputs, SerializedPromptInputsAll } from "./ComfyApp";
+import type IComfyInputSlot from "$lib/IComfyInputSlot";
+
+function hasTag(node: LGraphNode, tag: string): boolean {
+    return "tags" in node.properties && node.properties.tags.indexOf(tag) !== -1
+}
+
+function isActiveBackendNode(node: ComfyGraphNode, tag: string | null): node is ComfyBackendNode {
+    if (!node.isBackendNode)
+        return false;
+
+    if (tag && !hasTag(node, tag)) {
+        console.debug("Skipping tagged node", tag, node.properties.tags, node)
+        return false;
+    }
+
+    if (node.mode === NodeMode.NEVER) {
+        // Don't serialize muted nodes
+        return false;
+    }
+
+    return true;
+}
+
+export default class ComfyPromptSerializer {
+    serializeInputValues(node: ComfyBackendNode): Record<string, SerializedPromptInput> {
+        // Store input values passed by frontend-only nodes
+        if (!node.inputs) {
+            return {}
+        }
+
+        const inputs = {}
+
+        for (let i = 0; i < node.inputs.length; i++) {
+            const inp = node.inputs[i];
+            const inputLink = node.getInputLink(i)
+            const inputNode = node.getInputNode(i)
+
+            // We don't check tags for non-backend nodes.
+            // Just check for node inactivity (so you can toggle groups of
+            // tagged frontend nodes on/off)
+            if (inputNode && inputNode.mode === NodeMode.NEVER) {
+                console.debug("Skipping inactive node", inputNode)
+                continue;
+            }
+
+            if (!inputLink || !inputNode) {
+                if ("config" in inp) {
+                    const defaultValue = (inp as IComfyInputSlot).config?.defaultValue
+                    if (defaultValue !== null && defaultValue !== undefined)
+                        inputs[inp.name] = defaultValue
+                }
+                continue;
+            }
+
+            let serialize = true;
+            if ("config" in inp)
+                serialize = (inp as IComfyInputSlot).serialize
+
+            let isBackendNode = node.isBackendNode;
+            let isInputBackendNode = false;
+            if ("isBackendNode" in inputNode)
+                isInputBackendNode = (inputNode as ComfyGraphNode).isBackendNode;
+
+            // The reasoning behind this check:
+            // We only want to serialize inputs to nodes with backend equivalents.
+            // And in ComfyBox, the backend nodes in litegraph *never* have widgets, instead they're all inputs.
+            // All values are passed by separate frontend-only nodes,
+            // either UI-bound or something like ConstantInteger.
+            // So we know that any value passed into a backend node *must* come from
+            // a frontend node.
+            // The rest (links between backend nodes) will be serialized after this bit runs.
+            if (serialize && isBackendNode && !isInputBackendNode) {
+                inputs[inp.name] = inputLink.data
+            }
+        }
+
+        return inputs
+    }
+
+    serializeBackendLinks(node: ComfyBackendNode, tag: string | null): Record<string, SerializedPromptInput> {
+        const inputs = {}
+
+        // Store links between backend-only and hybrid nodes
+        for (let i = 0; i < node.inputs.length; i++) {
+            let parent: ComfyGraphNode = node.getInputNode(i) as ComfyGraphNode;
+            if (parent) {
+                const seen = {}
+                let link = node.getInputLink(i);
+
+                const isFrontendParent = (parent: ComfyGraphNode) => {
+                    if (!parent || parent.isBackendNode)
+                        return false;
+                    if (tag && !hasTag(parent, tag))
+                        return false;
+                    return true;
+                }
+
+                // If there are frontend-only nodes between us and another
+                // backend node, we have to traverse them first. This
+                // behavior is dependent on the type of node. Reroute nodes
+                // will simply follow their single input, while branching
+                // nodes have conditional logic that determines which link
+                // to follow backwards.
+                while (isFrontendParent(parent)) {
+                    if (!("getUpstreamLink" in parent)) {
+                        console.warn("[graphToPrompt] Node does not support getUpstreamLink", parent.type)
+                        break;
+                    }
+
+                    const nextLink = parent.getUpstreamLink()
+                    if (nextLink == null) {
+                        console.warn("[graphToPrompt] No upstream link found in frontend node", parent)
+                        break;
+                    }
+
+                    if (nextLink && !seen[nextLink.id]) {
+                        seen[nextLink.id] = true
+                        const inputNode = parent.graph.getNodeById(nextLink.origin_id) as ComfyGraphNode;
+                        if (inputNode && tag && !hasTag(inputNode, tag)) {
+                            console.debug("[graphToPrompt] Skipping tagged intermediate frontend node", tag, node.properties.tags)
+                            parent = null;
+                        }
+                        else {
+                            console.debug("[graphToPrompt] Traverse upstream link", parent.id, inputNode?.id, inputNode?.isBackendNode)
+                            link = nextLink;
+                            parent = inputNode;
+                        }
+                    } else {
+                        parent = null;
+                    }
+                }
+
+                if (link && parent && parent.isBackendNode) {
+                    if (tag && !hasTag(parent, tag))
+                        continue;
+
+                    console.debug("[graphToPrompt] final link", parent.id, node.id)
+                    const input = node.inputs[i]
+                    // TODO can null be a legitimate value in some cases?
+                    // Nodes like CLIPLoader will never have a value in the frontend, hence "null".
+                    if (!(input.name in inputs))
+                        inputs[input.name] = [String(link.origin_id), link.origin_slot];
+                }
+            }
+        }
+
+        return inputs
+    }
+
+    serializePrompt(graph: ComfyGraph, tag: string | null = null): SerializedPrompt {
+        // Run frontend-only logic
+        graph.runStep(1)
+
+        const workflow = graph.serialize();
+
+        const output: SerializedPromptInputsAll = {};
+
+        // Process nodes in order of execution
+        for (const node of graph.computeExecutionOrder<ComfyGraphNode>(false, null)) {
+            const n = workflow.nodes.find((n) => n.id === node.id);
+
+            if (!isActiveBackendNode(node, tag)) {
+                continue;
+            }
+
+            const inputs = this.serializeInputValues(node);
+            const links = this.serializeBackendLinks(node, tag);
+
+            output[String(node.id)] = {
+                inputs: { ...inputs, ...links },
+                class_type: node.comfyClass,
+            };
+        }
+
+        // Remove inputs connected to removed nodes
+        for (const nodeId in output) {
+            for (const inputName in output[nodeId].inputs) {
+                if (Array.isArray(output[nodeId].inputs[inputName])
+                    && output[nodeId].inputs[inputName].length === 2
+                    && !output[output[nodeId].inputs[inputName][0]]) {
+                    console.debug("Prune removed node link", nodeId, inputName, output[nodeId].inputs[inputName])
+                    delete output[nodeId].inputs[inputName];
+                }
+            }
+        }
+
+        // console.debug({ workflow, output })
+        // console.debug(promptToGraphVis({ workflow, output }))
+
+        return { workflow, output };
+    }
+}

src/lib/nodes/ComfyBackendNode.ts

@@ -30,7 +30,7 @@ export class ComfyBackendNode extends ComfyGraphNode {
         // It just returns a hash like { "ui": { "images": results } } internally.
         // So this will need to be hardcoded for now.
         if (["PreviewImage", "SaveImage"].indexOf(comfyClass) !== -1) {
-            this.addOutput("onExecuted", BuiltInSlotType.EVENT, { color_off: "rebeccapurple", color_on: "rebeccapurple" });
+            this.addOutput("OUTPUT", BuiltInSlotType.EVENT, { color_off: "rebeccapurple", color_on: "rebeccapurple" });
         }
     }