[Refactor] Move app.graphToPrompt to executionUtil (#2664)

src/utils/executionUtil.ts

@@ -0,0 +1,184 @@
+import type { LGraph } from '@comfyorg/litegraph'
+import { LGraphEventMode } from '@comfyorg/litegraph'
+
+import type { ComfyApiWorkflow, ComfyWorkflowJSON } from '@/types/comfyWorkflow'
+
+/**
+ * Converts the current graph workflow for sending to the API.
+ * Note: Node widgets are updated before serialization to prepare queueing.
+ * @returns The workflow and node links
+ */
+export const graphToPrompt = async (
+  graph: LGraph,
+  options: { clean?: boolean; sortNodes?: boolean } = {}
+): Promise<{ workflow: ComfyWorkflowJSON; output: ComfyApiWorkflow }> => {
+  const { clean = true, sortNodes = false } = options
+
+  for (const outerNode of graph.computeExecutionOrder(false)) {
+    if (outerNode.widgets) {
+      for (const widget of outerNode.widgets) {
+        // Allow widgets to run callbacks before a prompt has been queued
+        // e.g. random seed before every gen
+        widget.beforeQueued?.()
+      }
+    }
+
+    const innerNodes = outerNode.getInnerNodes
+      ? outerNode.getInnerNodes()
+      : [outerNode]
+    for (const node of innerNodes) {
+      if (node.isVirtualNode) {
+        // Don't serialize frontend only nodes but let them make changes
+        if (node.applyToGraph) {
+          node.applyToGraph()
+        }
+      }
+    }
+  }
+
+  const workflow = graph.serialize({ sortNodes })
+
+  // Remove localized_name from the workflow
+  for (const node of workflow.nodes) {
+    for (const slot of node.inputs ?? []) {
+      delete slot.localized_name
+    }
+    for (const slot of node.outputs ?? []) {
+      delete slot.localized_name
+    }
+  }
+
+  const output: ComfyApiWorkflow = {}
+  // Process nodes in order of execution
+  for (const outerNode of graph.computeExecutionOrder(false)) {
+    const skipNode =
+      outerNode.mode === LGraphEventMode.NEVER ||
+      outerNode.mode === LGraphEventMode.BYPASS
+    const innerNodes =
+      !skipNode && outerNode.getInnerNodes
+        ? outerNode.getInnerNodes()
+        : [outerNode]
+    for (const node of innerNodes) {
+      if (node.isVirtualNode) {
+        continue
+      }
+
+      if (
+        node.mode === LGraphEventMode.NEVER ||
+        node.mode === LGraphEventMode.BYPASS
+      ) {
+        // Don't serialize muted nodes
+        continue
+      }
+
+      const inputs: ComfyApiWorkflow[string]['inputs'] = {}
+      const widgets = node.widgets
+
+      // Store all widget values
+      if (widgets) {
+        for (let i = 0; i < widgets.length; i++) {
+          const widget = widgets[i]
+          if (
+            widget.name &&
+            (!widget.options || widget.options.serialize !== false)
+          ) {
+            inputs[widget.name] = widget.serializeValue
+              ? await widget.serializeValue(node, i)
+              : widget.value
+          }
+        }
+      }
+
+      // Store all node links
+      for (let i = 0; i < node.inputs.length; i++) {
+        let parent = node.getInputNode(i)
+        if (parent) {
+          let link = node.getInputLink(i)
+          while (
+            parent.mode === LGraphEventMode.BYPASS ||
+            parent.isVirtualNode
+          ) {
+            let found = false
+            if (parent.isVirtualNode) {
+              link = link ? parent.getInputLink(link.origin_slot) : null
+              if (link) {
+                parent = parent.getInputNode(link.target_slot)
+                if (parent) {
+                  found = true
+                }
+              }
+            } else if (link && parent.mode === LGraphEventMode.BYPASS) {
+              let all_inputs = [link.origin_slot]
+              if (parent.inputs) {
+                // @ts-expect-error convert list of strings to list of numbers
+                all_inputs = all_inputs.concat(Object.keys(parent.inputs))
+                for (let parent_input in all_inputs) {
+                  // @ts-expect-error assign string to number
+                  parent_input = all_inputs[parent_input]
+                  if (
+                    parent.inputs[parent_input]?.type === node.inputs[i].type
+                  ) {
+                    // @ts-expect-error convert string to number
+                    link = parent.getInputLink(parent_input)
+                    if (link) {
+                      // @ts-expect-error convert string to number
+                      parent = parent.getInputNode(parent_input)
+                    }
+                    found = true
+                    break
+                  }
+                }
+              }
+            }
+
+            if (!found) {
+              break
+            }
+          }
+
+          if (link) {
+            if (parent?.updateLink) {
+              link = parent.updateLink(link)
+            }
+            if (link) {
+              inputs[node.inputs[i].name] = [
+                String(link.origin_id),
+                // @ts-expect-error link.origin_slot is already number.
+                parseInt(link.origin_slot)
+              ]
+            }
+          }
+        }
+      }
+
+      output[String(node.id)] = {
+        inputs,
+        // TODO(huchenlei): Filter out all nodes that cannot be mapped to a
+        // comfyClass.
+        class_type: node.comfyClass!,
+        // Ignored by the backend.
+        _meta: {
+          title: node.title
+        }
+      }
+    }
+  }
+
+  // Remove inputs connected to removed nodes
+  if (clean) {
+    for (const o in output) {
+      for (const i in output[o].inputs) {
+        if (
+          Array.isArray(output[o].inputs[i]) &&
+          output[o].inputs[i].length === 2 &&
+          !output[output[o].inputs[i][0]]
+        ) {
+          delete output[o].inputs[i]
+        }
+      }
+    }
+  }
+
+  // @ts-expect-error Convert ISerializedGraph to ComfyWorkflowJSON
+  return { workflow: workflow as ComfyWorkflowJSON, output }
+}