From e9a98161ca429756c07e7dd4a070bf8c53fd919f Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Tue, 8 Apr 2025 20:20:19 -0400
Subject: [PATCH] [Bug] Fix converted widget compression on export (#3354)

---
 src/utils/executionUtil.ts                 |  9 +--
 src/utils/litegraphUtil.ts                 | 24 ++++++
 tests-ui/tests/utils/litegraphUtil.test.ts | 87 +++++++++++++++++++++-
 3 files changed, 113 insertions(+), 7 deletions(-)

diff --git a/src/utils/executionUtil.ts b/src/utils/executionUtil.ts
index 39be3a543..018e3b1c0 100644
--- a/src/utils/executionUtil.ts
+++ b/src/utils/executionUtil.ts
@@ -6,6 +6,8 @@ import type {
   ComfyWorkflowJSON
 } from '@/schemas/comfyWorkflowSchema'
 
+import { compressWidgetInputSlots } from './litegraphUtil'
+
 /**
  * Converts the current graph workflow for sending to the API.
  * Note: Node widgets are updated before serialization to prepare queueing.
@@ -38,12 +40,7 @@ export const graphToPrompt = async (
     }
   }
 
-  // Remove all unconnected widget input slots
-  for (const node of workflow.nodes) {
-    node.inputs = node.inputs?.filter(
-      (input) => !(input.widget && input.link === null)
-    )
-  }
+  compressWidgetInputSlots(workflow)
 
   const output: ComfyApiWorkflow = {}
   // Process nodes in order of execution
diff --git a/src/utils/litegraphUtil.ts b/src/utils/litegraphUtil.ts
index 384dcb03b..3c7cdcc67 100644
--- a/src/utils/litegraphUtil.ts
+++ b/src/utils/litegraphUtil.ts
@@ -1,5 +1,6 @@
 import type { ColorOption, LGraph } from '@comfyorg/litegraph'
 import { LGraphGroup, LGraphNode, isColorable } from '@comfyorg/litegraph'
+import type { ISerialisedGraph } from '@comfyorg/litegraph/dist/types/serialisation'
 import type {
   IComboWidget,
   IWidget
@@ -144,3 +145,26 @@ export function fixLinkInputSlots(graph: LGraph) {
     }
   }
 }
+
+/**
+ * Compress widget input slots by removing all unconnected widget input slots.
+ * This should match the serialization format of legacy widget conversion.
+ *
+ * @param graph - The graph to compress widget input slots for.
+ */
+export function compressWidgetInputSlots(graph: ISerialisedGraph) {
+  for (const node of graph.nodes) {
+    node.inputs = node.inputs?.filter(
+      (input) => !(input.widget && input.link === null)
+    )
+
+    for (const [inputIndex, input] of node.inputs?.entries() ?? []) {
+      if (input.link) {
+        const link = graph.links.find((link) => link[0] === input.link)
+        if (link) {
+          link[4] = inputIndex
+        }
+      }
+    }
+  }
+}
diff --git a/tests-ui/tests/utils/litegraphUtil.test.ts b/tests-ui/tests/utils/litegraphUtil.test.ts
index 237316efd..cc482588b 100644
--- a/tests-ui/tests/utils/litegraphUtil.test.ts
+++ b/tests-ui/tests/utils/litegraphUtil.test.ts
@@ -1,8 +1,12 @@
+import { ISerialisedGraph } from '@comfyorg/litegraph/dist/types/serialisation'
 import type { IWidget } from '@comfyorg/litegraph/dist/types/widgets'
 import { describe, expect, it } from 'vitest'
 
 import type { InputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2'
-import { migrateWidgetsValues } from '@/utils/litegraphUtil'
+import {
+  compressWidgetInputSlots,
+  migrateWidgetsValues
+} from '@/utils/litegraphUtil'
 
 describe('migrateWidgetsValues', () => {
   it('should remove widget values for forceInput inputs', () => {
@@ -86,3 +90,84 @@ describe('migrateWidgetsValues', () => {
     expect(result).toEqual(['first value', 'last value'])
   })
 })
+
+describe('compressWidgetInputSlots', () => {
+  it('should remove unconnected widget input slots', () => {
+    const graph: ISerialisedGraph = {
+      nodes: [
+        {
+          id: 1,
+          type: 'foo',
+          pos: [0, 0],
+          size: [100, 100],
+          flags: {},
+          order: 0,
+          mode: 0,
+          inputs: [
+            { widget: { name: 'foo' }, link: null, type: 'INT', name: 'foo' },
+            { widget: { name: 'bar' }, link: 2, type: 'INT', name: 'bar' },
+            { widget: { name: 'baz' }, link: null, type: 'INT', name: 'baz' }
+          ],
+          outputs: []
+        }
+      ],
+      links: [[2, 1, 0, 1, 0, 'INT']]
+    } as unknown as ISerialisedGraph
+
+    compressWidgetInputSlots(graph)
+
+    expect(graph.nodes[0].inputs).toEqual([
+      { widget: { name: 'bar' }, link: 2, type: 'INT', name: 'bar' }
+    ])
+  })
+
+  it('should update link target slots correctly', () => {
+    const graph: ISerialisedGraph = {
+      nodes: [
+        {
+          id: 1,
+          type: 'foo',
+          pos: [0, 0],
+          size: [100, 100],
+          flags: {},
+          order: 0,
+          mode: 0,
+          inputs: [
+            { widget: { name: 'foo' }, link: null, type: 'INT', name: 'foo' },
+            { widget: { name: 'bar' }, link: 2, type: 'INT', name: 'bar' },
+            { widget: { name: 'baz' }, link: 3, type: 'INT', name: 'baz' }
+          ],
+          outputs: []
+        }
+      ],
+      links: [
+        [2, 1, 0, 1, 1, 'INT'],
+        [3, 1, 0, 1, 2, 'INT']
+      ]
+    } as unknown as ISerialisedGraph
+
+    compressWidgetInputSlots(graph)
+
+    expect(graph.nodes[0].inputs).toEqual([
+      { widget: { name: 'bar' }, link: 2, type: 'INT', name: 'bar' },
+      { widget: { name: 'baz' }, link: 3, type: 'INT', name: 'baz' }
+    ])
+
+    expect(graph.links).toEqual([
+      [2, 1, 0, 1, 0, 'INT'],
+      [3, 1, 0, 1, 1, 'INT']
+    ])
+  })
+
+  it('should handle graphs with no nodes gracefully', () => {
+    const graph: ISerialisedGraph = {
+      nodes: [],
+      links: []
+    } as unknown as ISerialisedGraph
+
+    compressWidgetInputSlots(graph)
+
+    expect(graph.nodes).toEqual([])
+    expect(graph.links).toEqual([])
+  })
+})