Add Subgraphs (#3905)

Co-authored-by: github-actions <github-actions@github.com>

src/scripts/app.ts

@@ -39,9 +39,11 @@ import { getSvgMetadata } from '@/scripts/metadata/svg'
 import { useDialogService } from '@/services/dialogService'
 import { useExtensionService } from '@/services/extensionService'
 import { useLitegraphService } from '@/services/litegraphService'
+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'
@@ -72,7 +74,6 @@ import { deserialiseAndCreate } from '@/utils/vintageClipboard'
 
 import { type ComfyApi, PromptExecutionError, api } from './api'
 import { defaultGraph } from './defaultGraph'
-import { pruneWidgets } from './domWidget'
 import {
   getFlacMetadata,
   getLatentMetadata,
@@ -717,25 +718,23 @@ export class ComfyApp {
   }
 
   #addAfterConfigureHandler() {
-    const app = this
-    const onConfigure = app.graph.onConfigure
-    app.graph.onConfigure = function (this: LGraph, ...args) {
+    const { graph } = this
+    const { onConfigure } = graph
+    graph.onConfigure = function (...args) {
       fixLinkInputSlots(this)
 
       // Fire callbacks before the onConfigure, this is used by widget inputs to setup the config
-      for (const node of app.graph.nodes) {
+      for (const node of graph.nodes) {
         node.onGraphConfigured?.()
       }
 
       const r = onConfigure?.apply(this, args)
 
       // Fire after onConfigure, used by primitives to generate widget using input nodes config
-      for (const node of app.graph.nodes) {
+      for (const node of graph.nodes) {
         node.onAfterGraphConfigured?.()
       }
 
-      pruneWidgets(this.nodes)
-
       return r
     }
   }
@@ -767,6 +766,21 @@ export class ComfyApp {
 
     this.#graph = new LGraph()
 
+    // Register the subgraph - adds type wrapper for Litegraph's `createNode` factory
+    this.graph.events.addEventListener('subgraph-created', (e) => {
+      try {
+        const { subgraph, data } = e.detail
+        useSubgraphService().registerNewSubgraph(subgraph, data)
+      } catch (err) {
+        console.error('Failed to register subgraph', err)
+        useToastStore().add({
+          severity: 'error',
+          summary: 'Failed to register subgraph',
+          detail: err instanceof Error ? err.message : String(err)
+        })
+      }
+    })
+
     this.#addAfterConfigureHandler()
 
     this.canvas = new LGraphCanvas(canvasEl, this.graph)
@@ -779,6 +793,30 @@ export class ComfyApp {
     LiteGraph.alt_drag_do_clone_nodes = true
     LiteGraph.macGesturesRequireMac = false
 
+    this.canvas.canvas.addEventListener<'litegraph:set-graph'>(
+      'litegraph:set-graph',
+      (e) => {
+        // Assertion: Not yet defined in litegraph.
+        const { newGraph } = e.detail
+
+        const nodeSet = new Set(newGraph.nodes)
+        const widgetStore = useDomWidgetStore()
+
+        // Assertions: UnwrapRef
+        for (const { widget } of widgetStore.activeWidgetStates) {
+          if (!nodeSet.has(widget.node)) {
+            widgetStore.deactivateWidget(widget.id)
+          }
+        }
+
+        for (const { widget } of widgetStore.inactiveWidgetStates) {
+          if (nodeSet.has(widget.node)) {
+            widgetStore.activateWidget(widget.id)
+          }
+        }
+      }
+    )
+
     this.graph.start()
 
     // Ensure the canvas fills the window
@@ -1015,6 +1053,7 @@ export class ComfyApp {
       })
     }
     useWorkflowService().beforeLoadNewGraph()
+    useSubgraphService().loadSubgraphs(graphData)
 
     const missingNodeTypes: MissingNodeType[] = []
     const missingModels: ModelFile[] = []
@@ -1212,6 +1251,9 @@ export class ComfyApp {
           // Allow widgets to run callbacks before a prompt has been queued
           // e.g. random seed before every gen
           executeWidgetsCallback(this.graph.nodes, 'beforeQueued')
+          for (const subgraph of this.graph.subgraphs.values()) {
+            executeWidgetsCallback(subgraph.nodes, 'beforeQueued')
+          }
 
           const p = await this.graphToPrompt(this.graph, { queueNodeIds })
           try {
@@ -1254,9 +1296,13 @@ export class ComfyApp {
           executeWidgetsCallback(
             p.workflow.nodes
               .map((n) => this.graph.getNodeById(n.id))
-              .filter((n) => !!n) as LGraphNode[],
+              .filter((n) => !!n),
             'afterQueued'
           )
+          for (const subgraph of this.graph.subgraphs.values()) {
+            executeWidgetsCallback(subgraph.nodes, 'afterQueued')
+          }
+
           this.canvas.draw(true, true)
           await this.ui.queue.update()
         }
@@ -1652,6 +1698,8 @@ export class ComfyApp {
     const executionStore = useExecutionStore()
     executionStore.lastNodeErrors = null
     executionStore.lastExecutionError = null
+
+    useDomWidgetStore().clear()
   }
 
   clientPosToCanvasPos(pos: Vector2): Vector2 {

src/scripts/changeTracker.ts

@@ -6,6 +6,7 @@ import log from 'loglevel'
 import type { ExecutedWsMessage } from '@/schemas/apiSchema'
 import type { ComfyWorkflowJSON } from '@/schemas/comfyWorkflowSchema'
 import { useExecutionStore } from '@/stores/executionStore'
+import { useSubgraphNavigationStore } from '@/stores/subgraphNavigationStore'
 import { ComfyWorkflow, useWorkflowStore } from '@/stores/workflowStore'
 
 import { api } from './api'
@@ -37,6 +38,10 @@ export class ChangeTracker {
   ds?: { scale: number; offset: [number, number] }
   nodeOutputs?: Record<string, any>
 
+  private subgraphState?: {
+    navigation: string[]
+  }
+
   constructor(
     /**
      * The workflow that this change tracker is tracking
@@ -67,6 +72,8 @@ export class ChangeTracker {
       scale: app.canvas.ds.scale,
       offset: [app.canvas.ds.offset[0], app.canvas.ds.offset[1]]
     }
+    const navigation = useSubgraphNavigationStore().exportState()
+    this.subgraphState = navigation.length ? { navigation } : undefined
   }
 
   restore() {
@@ -77,6 +84,16 @@ export class ChangeTracker {
     if (this.nodeOutputs) {
       app.nodeOutputs = this.nodeOutputs
     }
+    if (this.subgraphState) {
+      const { navigation } = this.subgraphState
+      useSubgraphNavigationStore().restoreState(navigation)
+
+      const activeId = navigation.at(-1)
+      if (activeId) {
+        const subgraph = app.graph.subgraphs.get(activeId)
+        if (subgraph) app.canvas.setGraph(subgraph)
+      }
+    }
   }
 
   updateModified() {
@@ -376,7 +393,14 @@ export class ChangeTracker {
         return false
 
       // Compare other properties normally
-      for (const key of ['links', 'floatingLinks', 'reroutes', 'groups']) {
+      for (const key of [
+        'links',
+        'floatingLinks',
+        'reroutes',
+        'groups',
+        'definitions',
+        'subgraphs'
+      ]) {
         if (!_.isEqual(a[key], b[key])) {
           return false
         }
@@ -392,7 +416,12 @@ export class ChangeTracker {
     function sortGraphNodes(graph: ComfyWorkflowJSON) {
       return {
         links: graph.links,
+        floatingLinks: graph.floatingLinks,
+        reroutes: graph.reroutes,
         groups: graph.groups,
+        extra: graph.extra,
+        definitions: graph.definitions,
+        subgraphs: graph.subgraphs,
         nodes: graph.nodes.sort((a, b) => {
           if (typeof a.id === 'number' && typeof b.id === 'number') {
             return a.id - b.id

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
@@ -330,9 +330,8 @@ LGraphNode.prototype.addDOMWidget = function <
 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)) {
+  for (const { widget } of domWidgetStore.widgetStates.values()) {
+    if (!nodeSet.has(widget.node)) {
       domWidgetStore.unregisterWidget(widget.id)
     }
   }