From 6cb33d9431890c614af74dfec66d82e007cd6d06 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Fri, 21 Feb 2025 15:53:33 -0500
Subject: [PATCH] Restrict applyToGraph to PrimitiveNode (#2669)

---
 src/nodes/PrimitiveNode.ts            | 2 +-
 src/types/litegraph-augmentation.d.ts | 1 -
 src/utils/executionUtil.ts            | 8 +++-----
 3 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/src/nodes/PrimitiveNode.ts b/src/nodes/PrimitiveNode.ts
index e40dd7ad2..40e4daaa7 100644
--- a/src/nodes/PrimitiveNode.ts
+++ b/src/nodes/PrimitiveNode.ts
@@ -569,7 +569,7 @@ export function isValidCombo(combo: string[], obj: unknown) {
   return true
 }
 
-function isPrimitiveNode(node: LGraphNode): node is PrimitiveNode {
+export function isPrimitiveNode(node: LGraphNode): node is PrimitiveNode {
   return node.type === 'PrimitiveNode'
 }
 
diff --git a/src/types/litegraph-augmentation.d.ts b/src/types/litegraph-augmentation.d.ts
index 4ee151dbe..6e6fe22f6 100644
--- a/src/types/litegraph-augmentation.d.ts
+++ b/src/types/litegraph-augmentation.d.ts
@@ -73,7 +73,6 @@ declare module '@comfyorg/litegraph' {
     convertToNodes?(): LGraphNode[]
     recreate?(): Promise<LGraphNode>
     refreshComboInNode?(defs: Record<string, ComfyNodeDef>)
-    applyToGraph?(extraLinks?: LLink[]): void
     updateLink?(link: LLink): LLink | null
     onExecutionStart?(): unknown
     /**
diff --git a/src/utils/executionUtil.ts b/src/utils/executionUtil.ts
index a759623fe..884d4e971 100644
--- a/src/utils/executionUtil.ts
+++ b/src/utils/executionUtil.ts
@@ -1,6 +1,7 @@
 import type { LGraph } from '@comfyorg/litegraph'
 import { LGraphEventMode } from '@comfyorg/litegraph'
 
+import { isPrimitiveNode } from '@/nodes/PrimitiveNode'
 import type { ComfyApiWorkflow, ComfyWorkflowJSON } from '@/types/comfyWorkflow'
 
 /**
@@ -19,11 +20,8 @@ export const graphToPrompt = async (
       ? 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()
-        }
+      if (isPrimitiveNode(node)) {
+        node.applyToGraph()
       }
     }
   }