From 4461210f438fc7d18488be78b3bd7a2bdce3f13d Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Thu, 1 May 2025 16:07:42 -0400
Subject: [PATCH] Prune widgets on graph configure (#3716)

---
 src/scripts/app.ts       |  3 +++
 src/scripts/domWidget.ts | 15 +++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index a0237894b..3a033453f 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -62,6 +62,7 @@ import { deserialiseAndCreate } from '@/utils/vintageClipboard'
 
 import { type ComfyApi, PromptExecutionError, api } from './api'
 import { defaultGraph } from './defaultGraph'
+import { pruneWidgets } from './domWidget'
 import {
   getFlacMetadata,
   getLatentMetadata,
@@ -707,6 +708,8 @@ export class ComfyApp {
         node.onAfterGraphConfigured?.()
       }
 
+      pruneWidgets(this.nodes)
+
       return r
     }
   }
diff --git a/src/scripts/domWidget.ts b/src/scripts/domWidget.ts
index 981b0235d..933ef6f2b 100644
--- a/src/scripts/domWidget.ts
+++ b/src/scripts/domWidget.ts
@@ -314,3 +314,18 @@ LGraphNode.prototype.addDOMWidget = function <
 
   return widget
 }
+
+/**
+ * Prunes widgets that are no longer in the graph.
+ * @param nodes The nodes to prune widgets for.
+ */
+export const pruneWidgets = (nodes: LGraphNode[]) => {
+  const nodeSet = new Set(nodes)
+  const domWidgetStore = useDomWidgetStore()
+  for (const widgetState of domWidgetStore.widgetStates.values()) {
+    const widget = widgetState.widget
+    if (!nodeSet.has(widget.node as LGraphNode)) {
+      domWidgetStore.unregisterWidget(widget.id)
+    }
+  }
+}