diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 5a83cee9d..8dc1ab3d2 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -5,7 +5,7 @@ import {
   LGraphNode,
   LiteGraph
 } from '@comfyorg/litegraph'
-import type { Vector2 } from '@comfyorg/litegraph'
+import type { Subgraph, Vector2 } from '@comfyorg/litegraph'
 import type { IBaseWidget } from '@comfyorg/litegraph/dist/types/widgets'
 import _ from 'lodash'
 import type { ToastMessageOptions } from 'primevue/toast'
@@ -43,6 +43,7 @@ import { useSubgraphService } from '@/services/subgraphService'
 import { useWorkflowService } from '@/services/workflowService'
 import { useApiKeyAuthStore } from '@/stores/apiKeyAuthStore'
 import { useCommandStore } from '@/stores/commandStore'
+import { useDomWidgetStore } from '@/stores/domWidgetStore'
 import { useExecutionStore } from '@/stores/executionStore'
 import { useExtensionStore } from '@/stores/extensionStore'
 import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
@@ -73,6 +74,7 @@ import { deserialiseAndCreate } from '@/utils/vintageClipboard'
 
 import { type ComfyApi, PromptExecutionError, api } from './api'
 import { defaultGraph } from './defaultGraph'
+import type { BaseDOMWidget } from './domWidget'
 import { pruneWidgets } from './domWidget'
 import {
   getFlacMetadata,
@@ -716,7 +718,7 @@ export class ComfyApp {
   }
 
   #addAfterConfigureHandler() {
-    const { graph } = this
+    const { canvas, graph } = this
     const { onConfigure } = graph
     graph.onConfigure = function (...args) {
       fixLinkInputSlots(this)
@@ -733,7 +735,8 @@ export class ComfyApp {
         node.onAfterGraphConfigured?.()
       }
 
-      pruneWidgets(this.nodes)
+      const nodes = canvas.subgraph?.nodes ?? graph.nodes
+      pruneWidgets(nodes)
 
       return r
     }
@@ -792,6 +795,27 @@ export class ComfyApp {
     LiteGraph.alt_drag_do_clone_nodes = true
     LiteGraph.macGesturesRequireMac = false
 
+    this.canvas.canvas.addEventListener('litegraph:set-graph', (e) => {
+      type SetGraphType = CustomEvent<{
+        newGraph: LGraph | Subgraph
+        oldGraph: LGraph | Subgraph
+      }>
+
+      // Assertion: Not yet defined in litegraph.
+      const { newGraph } = (e as SetGraphType).detail
+
+      const nodeSet = new Set(newGraph.nodes)
+      const widgetStore = useDomWidgetStore()
+      for (const { widget } of widgetStore.widgetStates.values()) {
+        // Assertions: UnwrapRef
+        if (nodeSet.has(widget.node as LGraphNode)) {
+          widgetStore.registerWidget(widget as BaseDOMWidget)
+        } else {
+          widgetStore.unregisterWidget(widget.id)
+        }
+      }
+    })
+
     this.graph.start()
 
     // Ensure the canvas fills the window
diff --git a/src/scripts/domWidget.ts b/src/scripts/domWidget.ts
index 6acfcd95f..70f3bc31c 100644
--- a/src/scripts/domWidget.ts
+++ b/src/scripts/domWidget.ts
@@ -11,7 +11,7 @@ import type { InputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2'
 import { useDomWidgetStore } from '@/stores/domWidgetStore'
 import { generateUUID } from '@/utils/formatUtil'
 
-export interface BaseDOMWidget<V extends object | string>
+export interface BaseDOMWidget<V extends object | string = object | string>
   extends IBaseWidget<V, string, DOMWidgetOptions<V>> {
   // ICustomWidget properties
   type: string