From 3070ed36763772ce775e5a901380ad0d6b6e62e3 Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Sun, 14 May 2023 12:29:49 -0500 Subject: [PATCH] Refactor prompt serializer --- src/lib/components/ComfyApp.ts | 176 +---------------- src/lib/components/ComfyPromptSerializer.ts | 197 ++++++++++++++++++++ src/lib/nodes/ComfyBackendNode.ts | 2 +- 3 files changed, 202 insertions(+), 173 deletions(-) create mode 100644 src/lib/components/ComfyPromptSerializer.ts diff --git a/src/lib/components/ComfyApp.ts b/src/lib/components/ComfyApp.ts index acb560f..e5b21e4 100644 --- a/src/lib/components/ComfyApp.ts +++ b/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 { @@ -559,157 +541,7 @@ export default class ComfyApp { * @returns The workflow and node links */ graphToPrompt(tag: string | null = null): SerializedPrompt { - // Run frontend-only logic - this.lGraph.runStep(1) - - const workflow = this.lGraph.serialize(); - - const output = {}; - // Process nodes in order of execution - for (const node_ of this.lGraph.computeExecutionOrder(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 }; + return this.promptSerializer.serializePrompt(this.lGraph, tag) } async queuePrompt(num: number, batchCount: number = 1, tag: string | null = null) { diff --git a/src/lib/components/ComfyPromptSerializer.ts b/src/lib/components/ComfyPromptSerializer.ts new file mode 100644 index 0000000..46fe181 --- /dev/null +++ b/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 { + // 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 { + 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(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 }; + } +} diff --git a/src/lib/nodes/ComfyBackendNode.ts b/src/lib/nodes/ComfyBackendNode.ts index a75c968..fd17f00 100644 --- a/src/lib/nodes/ComfyBackendNode.ts +++ b/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" }); } }