From 156013aa24704e7b5f4d9998c88db112bfc16b1c Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Tue, 25 Feb 2025 10:54:21 -0500
Subject: [PATCH] [Reland] Restrict applyToGraph to PrimitiveNode (#2724)

---
 src/extensions/core/widgetInputs.ts   | 9 ++-------
 src/types/litegraph-augmentation.d.ts | 2 --
 src/utils/executionUtil.ts            | 4 ++--
 src/utils/typeGuardUtil.ts            | 9 +++++++++
 4 files changed, 13 insertions(+), 11 deletions(-)
 create mode 100644 src/utils/typeGuardUtil.ts

diff --git a/src/extensions/core/widgetInputs.ts b/src/extensions/core/widgetInputs.ts
index 4827e8294..5056b2ce3 100644
--- a/src/extensions/core/widgetInputs.ts
+++ b/src/extensions/core/widgetInputs.ts
@@ -11,6 +11,7 @@ import { INodeSlot } from '@comfyorg/litegraph'
 import { useNodeDefStore } from '@/stores/nodeDefStore'
 import { useSettingStore } from '@/stores/settingStore'
 import type { InputSpec } from '@/types/apiTypes'
+import { isPrimitiveNode } from '@/utils/typeGuardUtil'
 
 import { app } from '../../scripts/app'
 import { applyTextReplacements } from '../../scripts/utils'
@@ -30,10 +31,8 @@ const CONFIG = Symbol()
 const GET_CONFIG = Symbol()
 const TARGET = Symbol() // Used for reroutes to specify the real target widget
 
-interface PrimitiveNode extends LGraphNode {}
-
 const replacePropertyName = 'Run widget replace on values'
-class PrimitiveNode extends LGraphNode {
+export class PrimitiveNode extends LGraphNode {
   controlValues: any[]
   lastType: string
   static category: string
@@ -577,10 +576,6 @@ function isValidCombo(combo: string[], obj: unknown) {
   return true
 }
 
-function isPrimitiveNode(node: LGraphNode): node is PrimitiveNode {
-  return node.type === 'PrimitiveNode'
-}
-
 export function setWidgetConfig(slot, config, target?: IWidget) {
   if (!slot.widget) return
   if (config) {
diff --git a/src/types/litegraph-augmentation.d.ts b/src/types/litegraph-augmentation.d.ts
index 33be37010..f2cc7c5f2 100644
--- a/src/types/litegraph-augmentation.d.ts
+++ b/src/types/litegraph-augmentation.d.ts
@@ -82,8 +82,6 @@ declare module '@comfyorg/litegraph' {
     convertToNodes?(): LGraphNode[]
     recreate?(): Promise<LGraphNode>
     refreshComboInNode?(defs: Record<string, ComfyNodeDef>)
-    /** Used by virtual nodes (primitives) to insert their values into the graph prior to queueing. */
-    applyToGraph?(extraLinks?: LLink[]): void
     /** @deprecated groupNode */
     updateLink?(link: LLink): LLink | null
     onExecutionStart?(): unknown
diff --git a/src/utils/executionUtil.ts b/src/utils/executionUtil.ts
index 9229f5e21..c537848c0 100644
--- a/src/utils/executionUtil.ts
+++ b/src/utils/executionUtil.ts
@@ -2,6 +2,7 @@ import type { LGraph } from '@comfyorg/litegraph'
 import { LGraphEventMode } from '@comfyorg/litegraph'
 
 import type { ComfyApiWorkflow, ComfyWorkflowJSON } from '@/types/comfyWorkflow'
+import { isPrimitiveNode } from '@/utils/typeGuardUtil'
 
 /**
  * Converts the current graph workflow for sending to the API.
@@ -17,8 +18,7 @@ export const graphToPrompt = async (
   for (const node of graph.computeExecutionOrder(false)) {
     const innerNodes = node.getInnerNodes ? node.getInnerNodes() : [node]
     for (const innerNode of innerNodes) {
-      // Don't serialize frontend only nodes but let them make changes
-      if (innerNode.isVirtualNode) {
+      if (isPrimitiveNode(innerNode)) {
         innerNode.applyToGraph?.()
       }
     }
diff --git a/src/utils/typeGuardUtil.ts b/src/utils/typeGuardUtil.ts
new file mode 100644
index 000000000..a8ac480d1
--- /dev/null
+++ b/src/utils/typeGuardUtil.ts
@@ -0,0 +1,9 @@
+import { LGraphNode } from '@comfyorg/litegraph'
+
+import type { PrimitiveNode } from '@/extensions/core/widgetInputs'
+
+export function isPrimitiveNode(
+  node: LGraphNode
+): node is PrimitiveNode & LGraphNode {
+  return node.type === 'PrimitiveNode'
+}