[backport 1.25] Implement subgraph unpacking (#4950)

Co-authored-by: AustinMroz <austin@comfy.org>

src/components/graph/selectionToolbox/ConvertToSubgraphButton.vue

@@ -1,6 +1,20 @@
 <template>
   <Button
-    v-show="isVisible"
+    v-if="isUnpackVisible"
+    v-tooltip.top="{
+      value: t('commands.Comfy_Graph_UnpackSubgraph.label'),
+      showDelay: 1000
+    }"
+    severity="secondary"
+    text
+    @click="() => commandStore.execute('Comfy.Graph.UnpackSubgraph')"
+  >
+    <template #icon>
+      <i-lucide:expand />
+    </template>
+  </Button>
+  <Button
+    v-else-if="isConvertVisible"
     v-tooltip.top="{
       value: t('commands.Comfy_Graph_ConvertToSubgraph.label'),
       showDelay: 1000
@@ -20,6 +34,7 @@ import Button from 'primevue/button'
 import { computed } from 'vue'
 import { useI18n } from 'vue-i18n'
 
+import { SubgraphNode } from '@/lib/litegraph/src/litegraph'
 import { useCommandStore } from '@/stores/commandStore'
 import { useCanvasStore } from '@/stores/graphStore'
 
@@ -27,7 +42,13 @@ const { t } = useI18n()
 const commandStore = useCommandStore()
 const canvasStore = useCanvasStore()
 
-const isVisible = computed(() => {
+const isUnpackVisible = computed(() => {
+  return (
+    canvasStore.selectedItems?.length === 1 &&
+    canvasStore.selectedItems[0] instanceof SubgraphNode
+  )
+})
+const isConvertVisible = computed(() => {
   return (
     canvasStore.groupSelected ||
     canvasStore.rerouteSelected ||

src/composables/useCoreCommands.ts

@@ -21,6 +21,7 @@ import { useWorkflowService } from '@/services/workflowService'
 import type { ComfyCommand } from '@/stores/commandStore'
 import { useExecutionStore } from '@/stores/executionStore'
 import { useCanvasStore, useTitleEditorStore } from '@/stores/graphStore'
+import { useNodeOutputStore } from '@/stores/imagePreviewStore'
 import { useQueueSettingsStore, useQueueStore } from '@/stores/queueStore'
 import { useSettingStore } from '@/stores/settingStore'
 import { useSubgraphNavigationStore } from '@/stores/subgraphNavigationStore'
@@ -797,6 +798,22 @@ export function useCoreCommands(): ComfyCommand[] {
         canvas.select(node)
       }
     },
+    {
+      id: 'Comfy.Graph.UnpackSubgraph',
+      icon: 'pi pi-sitemap',
+      label: 'Unpack the selected Subgraph',
+      versionAdded: '1.20.1',
+      category: 'essentials' as const,
+      function: () => {
+        const canvas = canvasStore.getCanvas()
+        const graph = canvas.subgraph ?? canvas.graph
+        if (!graph) throw new TypeError('Canvas has no graph or subgraph set.')
+
+        const subgraphNode = app.canvas.selectedItems.values().next().value
+        useNodeOutputStore().revokeSubgraphPreviews(subgraphNode)
+        graph.unpackSubgraph(subgraphNode)
+      }
+    },
     {
       id: 'Workspace.ToggleBottomPanel.Shortcuts',
       icon: 'pi pi-key',

src/lib/litegraph/src/LGraph.ts

@@ -1663,6 +1663,254 @@ export class LGraph
     return { subgraph, node: subgraphNode as SubgraphNode }
   }
 
+  unpackSubgraph(subgraphNode: SubgraphNode) {
+    if (!(subgraphNode instanceof SubgraphNode))
+      throw new Error('Can only unpack Subgraph Nodes')
+    this.beforeChange()
+    const center = [0, 0]
+    for (const node of subgraphNode.subgraph.nodes) {
+      center[0] += node.pos[0] + node.size[0] / 2
+      center[1] += node.pos[1] + node.size[1] / 2
+    }
+    center[0] /= subgraphNode.subgraph.nodes.length
+    center[1] /= subgraphNode.subgraph.nodes.length
+
+    const offsetX = subgraphNode.pos[0] - center[0] + subgraphNode.size[0] / 2
+    const offsetY = subgraphNode.pos[1] - center[1] + subgraphNode.size[1] / 2
+    const movedNodes = multiClone(subgraphNode.subgraph.nodes)
+    const nodeIdMap = new Map<NodeId, NodeId>()
+    for (const n_info of movedNodes) {
+      const node = LiteGraph.createNode(String(n_info.type), n_info.title)
+      if (!node) {
+        throw new Error('Node not found')
+      }
+
+      nodeIdMap.set(n_info.id, ++this.last_node_id)
+      node.id = this.last_node_id
+      n_info.id = this.last_node_id
+
+      this.add(node, true)
+      node.configure(n_info)
+      node.pos[0] += offsetX
+      node.pos[1] += offsetY
+      for (const input of node.inputs) {
+        input.link = null
+      }
+    }
+    //cleanup reoute.linkIds now, but leave link.parentIds dangling
+    for (const islot of subgraphNode.inputs) {
+      if (!islot.link) continue
+      const link = this.links.get(islot.link)
+      if (!link) {
+        console.warn('Broken link', islot, islot.link)
+        continue
+      }
+      for (const reroute of LLink.getReroutes(this, link)) {
+        reroute.linkIds.delete(link.id)
+      }
+    }
+    for (const oslot of subgraphNode.outputs) {
+      for (const linkId of oslot.links ?? []) {
+        const link = this.links.get(linkId)
+        if (!link) {
+          console.warn('Broken link', oslot, linkId)
+          continue
+        }
+        for (const reroute of LLink.getReroutes(this, link)) {
+          reroute.linkIds.delete(link.id)
+        }
+      }
+    }
+
+    const newLinks: [
+      NodeId,
+      number,
+      NodeId,
+      number,
+      LinkId,
+      RerouteId | undefined,
+      RerouteId | undefined,
+      boolean
+    ][] = []
+    for (const [, link] of subgraphNode.subgraph._links) {
+      let externalParentId: RerouteId | undefined
+      if (link.origin_id === SUBGRAPH_INPUT_ID) {
+        const outerLinkId = subgraphNode.inputs[link.origin_slot].link
+        if (!outerLinkId) {
+          console.error('Missing Link ID when unpacking')
+          continue
+        }
+        const outerLink = this.links[outerLinkId]
+        link.origin_id = outerLink.origin_id
+        link.origin_slot = outerLink.origin_slot
+        externalParentId = outerLink.parentId
+      } else {
+        const origin_id = nodeIdMap.get(link.origin_id)
+        if (!origin_id) {
+          console.error('Missing Link ID when unpacking')
+          continue
+        }
+        link.origin_id = origin_id
+      }
+      if (link.target_id === SUBGRAPH_OUTPUT_ID) {
+        for (const linkId of subgraphNode.outputs[link.target_slot].links ??
+          []) {
+          const sublink = this.links[linkId]
+          newLinks.push([
+            link.origin_id,
+            link.origin_slot,
+            sublink.target_id,
+            sublink.target_slot,
+            link.id,
+            link.parentId,
+            sublink.parentId,
+            true
+          ])
+          sublink.parentId = undefined
+        }
+        continue
+      } else {
+        const target_id = nodeIdMap.get(link.target_id)
+        if (!target_id) {
+          console.error('Missing Link ID when unpacking')
+          continue
+        }
+        link.target_id = target_id
+      }
+      newLinks.push([
+        link.origin_id,
+        link.origin_slot,
+        link.target_id,
+        link.target_slot,
+        link.id,
+        link.parentId,
+        externalParentId,
+        false
+      ])
+    }
+    this.remove(subgraphNode)
+    this.subgraphs.delete(subgraphNode.subgraph.id)
+    const linkIdMap = new Map<LinkId, LinkId[]>()
+    for (const newLink of newLinks) {
+      let created: LLink | null | undefined
+      if (newLink[0] == SUBGRAPH_INPUT_ID) {
+        if (!(this instanceof Subgraph)) {
+          console.error('Ignoring link to subgraph outside subgraph')
+          continue
+        }
+        const tnode = this._nodes_by_id[newLink[2]]
+        created = this.inputNode.slots[newLink[1]].connect(
+          tnode.inputs[newLink[3]],
+          tnode
+        )
+      } else if (newLink[2] == SUBGRAPH_OUTPUT_ID) {
+        if (!(this instanceof Subgraph)) {
+          console.error('Ignoring link to subgraph outside subgraph')
+          continue
+        }
+        const tnode = this._nodes_by_id[newLink[0]]
+        created = this.outputNode.slots[newLink[3]].connect(
+          tnode.outputs[newLink[1]],
+          tnode
+        )
+      } else {
+        created = this._nodes_by_id[newLink[0]].connect(
+          newLink[1],
+          this._nodes_by_id[newLink[2]],
+          newLink[3]
+        )
+      }
+      if (!created) {
+        console.error('Failed to create link')
+        continue
+      }
+      //This is a little unwieldy since Map.has isn't a type guard
+      const linkIds = linkIdMap.get(newLink[4]) ?? []
+      linkIds.push(created.id)
+      if (!linkIdMap.has(newLink[4])) {
+        linkIdMap.set(newLink[4], linkIds)
+      }
+      newLink[4] = created.id
+    }
+    const rerouteIdMap = new Map<RerouteId, RerouteId>()
+    for (const reroute of subgraphNode.subgraph.reroutes.values()) {
+      if (
+        reroute.parentId !== undefined &&
+        rerouteIdMap.get(reroute.parentId) === undefined
+      ) {
+        console.error('Missing Parent ID')
+      }
+      const migratedReroute = new Reroute(++this.state.lastRerouteId, this, [
+        reroute.pos[0] + offsetX,
+        reroute.pos[1] + offsetY
+      ])
+      rerouteIdMap.set(reroute.id, migratedReroute.id)
+      this.reroutes.set(migratedReroute.id, migratedReroute)
+    }
+    //iterate over newly created links to update reroute parentIds
+    for (const newLink of newLinks) {
+      const linkInstance = this.links.get(newLink[4])
+      if (!linkInstance) {
+        continue
+      }
+      let instance: Reroute | LLink | undefined = linkInstance
+      let parentId: RerouteId | undefined = newLink[6]
+      if (newLink[7]) {
+        parentId = newLink[6]
+        //TODO: recursion check/helper method? Probably exists, but wouldn't mesh with the reference tracking used by this implementation
+        while (parentId) {
+          instance.parentId = parentId
+          instance = this.reroutes.get(parentId)
+          if (!instance) throw new Error('Broken Id link when unpacking')
+          if (instance.linkIds.has(linkInstance.id))
+            throw new Error('Infinite parentId loop')
+          instance.linkIds.add(linkInstance.id)
+          parentId = instance.parentId
+        }
+      }
+      parentId = newLink[5]
+      while (parentId) {
+        const migratedId = rerouteIdMap.get(parentId)
+        if (!migratedId) throw new Error('Broken Id link when unpacking')
+        instance.parentId = migratedId
+        instance = this.reroutes.get(migratedId)
+        if (!instance) throw new Error('Broken Id link when unpacking')
+        if (instance.linkIds.has(linkInstance.id))
+          throw new Error('Infinite parentId loop')
+        instance.linkIds.add(linkInstance.id)
+        const oldReroute = subgraphNode.subgraph.reroutes.get(parentId)
+        if (!oldReroute) throw new Error('Broken Id link when unpacking')
+        parentId = oldReroute.parentId
+      }
+      if (!newLink[7]) {
+        parentId = newLink[6]
+        while (parentId) {
+          instance.parentId = parentId
+          instance = this.reroutes.get(parentId)
+          if (!instance) throw new Error('Broken Id link when unpacking')
+          if (instance.linkIds.has(linkInstance.id))
+            throw new Error('Infinite parentId loop')
+          instance.linkIds.add(linkInstance.id)
+          parentId = instance.parentId
+        }
+      }
+    }
+
+    const nodes: LGraphNode[] = []
+    for (const nodeId of nodeIdMap.values()) {
+      const node = this._nodes_by_id[nodeId]
+      nodes.push(node)
+      node._setConcreteSlots()
+      node.arrange()
+    }
+    const reroutes = [...rerouteIdMap.values()]
+      .map((i) => this.reroutes.get(i))
+      .filter((x): x is Reroute => !!x)
+
+    this.canvasAction((c) => c.selectItems([...nodes, ...reroutes]))
+    this.afterChange()
+  }
+
   /**
    * Resolve a path of subgraph node IDs into a list of subgraph nodes.
    * Not intended to be run from subgraphs.

src/lib/litegraph/test/subgraph/SubgraphConversion.test.ts

@@ -0,0 +1,197 @@
+import { assert, describe, expect, it } from 'vitest'
+
+import {
+  ISlotType,
+  LGraph,
+  LGraphNode,
+  LiteGraph
+} from '@/lib/litegraph/src/litegraph'
+
+import {
+  createTestSubgraph,
+  createTestSubgraphNode
+} from './fixtures/subgraphHelpers'
+
+function createNode(
+  graph: LGraph,
+  inputs: ISlotType[] = [],
+  outputs: ISlotType[] = [],
+  title?: string
+) {
+  const type = JSON.stringify({ inputs, outputs })
+  if (!LiteGraph.registered_node_types[type]) {
+    class testnode extends LGraphNode {
+      constructor(title: string) {
+        super(title)
+        let i_count = 0
+        for (const input of inputs) this.addInput('input_' + i_count++, input)
+        let o_count = 0
+        for (const output of outputs)
+          this.addOutput('output_' + o_count++, output)
+      }
+    }
+    LiteGraph.registered_node_types[type] = testnode
+  }
+  const node = LiteGraph.createNode(type, title)
+  if (!node) {
+    throw new Error('Failed to create node')
+  }
+  graph.add(node)
+  return node
+}
+describe('SubgraphConversion', () => {
+  describe('Subgraph Unpacking Functionality', () => {
+    it('Should keep interior nodes and links', () => {
+      const subgraph = createTestSubgraph()
+      const subgraphNode = createTestSubgraphNode(subgraph)
+      const graph = subgraphNode.graph
+      graph.add(subgraphNode)
+
+      const node1 = createNode(subgraph, [], ['number'])
+      const node2 = createNode(subgraph, ['number'])
+      node1.connect(0, node2, 0)
+
+      graph.unpackSubgraph(subgraphNode)
+
+      expect(graph.nodes.length).toBe(2)
+      expect(graph.links.size).toBe(1)
+    })
+    it('Should merge boundry links', () => {
+      const subgraph = createTestSubgraph({
+        inputs: [{ name: 'value', type: 'number' }],
+        outputs: [{ name: 'value', type: 'number' }]
+      })
+      const subgraphNode = createTestSubgraphNode(subgraph)
+      const graph = subgraphNode.graph
+      graph.add(subgraphNode)
+
+      const innerNode1 = createNode(subgraph, [], ['number'])
+      const innerNode2 = createNode(subgraph, ['number'], [])
+      subgraph.inputNode.slots[0].connect(innerNode2.inputs[0], innerNode2)
+      subgraph.outputNode.slots[0].connect(innerNode1.outputs[0], innerNode1)
+
+      const outerNode1 = createNode(graph, [], ['number'])
+      const outerNode2 = createNode(graph, ['number'])
+      outerNode1.connect(0, subgraphNode, 0)
+      subgraphNode.connect(0, outerNode2, 0)
+
+      graph.unpackSubgraph(subgraphNode)
+
+      expect(graph.nodes.length).toBe(4)
+      expect(graph.links.size).toBe(2)
+    })
+    it('Should keep reroutes', () => {
+      const subgraph = createTestSubgraph({
+        outputs: [{ name: 'value', type: 'number' }]
+      })
+      const subgraphNode = createTestSubgraphNode(subgraph)
+      const graph = subgraphNode.graph
+      graph.add(subgraphNode)
+
+      const inner = createNode(subgraph, [], ['number'])
+      const innerLink = subgraph.outputNode.slots[0].connect(
+        inner.outputs[0],
+        inner
+      )
+      assert(innerLink)
+
+      const outer = createNode(graph, ['number'])
+      const outerLink = subgraphNode.connect(0, outer, 0)
+      assert(outerLink)
+
+      subgraph.createReroute([10, 10], innerLink)
+      graph.createReroute([10, 10], outerLink)
+
+      graph.unpackSubgraph(subgraphNode)
+
+      expect(graph.reroutes.size).toBe(2)
+    })
+    it('Should map reroutes onto split outputs', () => {
+      const subgraph = createTestSubgraph({
+        outputs: [
+          { name: 'value1', type: 'number' },
+          { name: 'value2', type: 'number' }
+        ]
+      })
+      const subgraphNode = createTestSubgraphNode(subgraph)
+      const graph = subgraphNode.graph
+      graph.add(subgraphNode)
+
+      const inner = createNode(subgraph, [], ['number', 'number'])
+      const innerLink1 = subgraph.outputNode.slots[0].connect(
+        inner.outputs[0],
+        inner
+      )
+      const innerLink2 = subgraph.outputNode.slots[1].connect(
+        inner.outputs[1],
+        inner
+      )
+      const outer1 = createNode(graph, ['number'])
+      const outer2 = createNode(graph, ['number'])
+      const outer3 = createNode(graph, ['number'])
+      const outerLink1 = subgraphNode.connect(0, outer1, 0)
+      assert(innerLink1 && innerLink2 && outerLink1)
+      subgraphNode.connect(0, outer2, 0)
+      subgraphNode.connect(1, outer3, 0)
+
+      subgraph.createReroute([10, 10], innerLink1)
+      subgraph.createReroute([10, 20], innerLink2)
+      graph.createReroute([10, 10], outerLink1)
+
+      graph.unpackSubgraph(subgraphNode)
+
+      expect(graph.reroutes.size).toBe(3)
+      expect(graph.links.size).toBe(3)
+      let linkRefCount = 0
+      for (const reroute of graph.reroutes.values()) {
+        linkRefCount += reroute.linkIds.size
+      }
+      expect(linkRefCount).toBe(4)
+    })
+    it('Should map reroutes onto split inputs', () => {
+      const subgraph = createTestSubgraph({
+        inputs: [
+          { name: 'value1', type: 'number' },
+          { name: 'value2', type: 'number' }
+        ]
+      })
+      const subgraphNode = createTestSubgraphNode(subgraph)
+      const graph = subgraphNode.graph
+      graph.add(subgraphNode)
+
+      const inner1 = createNode(subgraph, ['number', 'number'])
+      const inner2 = createNode(subgraph, ['number'])
+      const innerLink1 = subgraph.inputNode.slots[0].connect(
+        inner1.inputs[0],
+        inner1
+      )
+      const innerLink2 = subgraph.inputNode.slots[1].connect(
+        inner1.inputs[1],
+        inner1
+      )
+      const innerLink3 = subgraph.inputNode.slots[1].connect(
+        inner2.inputs[0],
+        inner2
+      )
+      assert(innerLink1 && innerLink2 && innerLink3)
+      const outer = createNode(graph, [], ['number'])
+      const outerLink1 = outer.connect(0, subgraphNode, 0)
+      const outerLink2 = outer.connect(0, subgraphNode, 1)
+      assert(outerLink1 && outerLink2)
+
+      graph.createReroute([10, 10], outerLink1)
+      graph.createReroute([10, 20], outerLink2)
+      subgraph.createReroute([10, 10], innerLink1)
+
+      graph.unpackSubgraph(subgraphNode)
+
+      expect(graph.reroutes.size).toBe(3)
+      expect(graph.links.size).toBe(3)
+      let linkRefCount = 0
+      for (const reroute of graph.reroutes.values()) {
+        linkRefCount += reroute.linkIds.size
+      }
+      expect(linkRefCount).toBe(4)
+    })
+  })
+})

src/locales/en/commands.json

@@ -122,6 +122,9 @@
   "Comfy_Graph_ExitSubgraph": {
     "label": "Exit Subgraph"
   },
+  "Comfy_Graph_UnpackSubgraph": {
+    "label": "Unpack the selected Subgraph"
+  },
   "Comfy_Graph_FitGroupToContents": {
     "label": "Fit Group To Contents"
   },

src/services/litegraphService.ts

@@ -805,6 +805,15 @@ export const useLitegraphService = () => {
           })
         }
       }
+      if (this instanceof SubgraphNode) {
+        options.unshift({
+          content: 'Unpack Subgraph',
+          callback: () => {
+            useNodeOutputStore().revokeSubgraphPreviews(this)
+            this.graph.unpackSubgraph(this)
+          }
+        })
+      }
 
       return []
     }

src/stores/imagePreviewStore.ts

@@ -1,6 +1,6 @@
 import { defineStore } from 'pinia'
 
-import { LGraphNode } from '@/lib/litegraph/src/litegraph'
+import { LGraphNode, SubgraphNode } from '@/lib/litegraph/src/litegraph'
 import {
   ExecutedWsMessage,
   ResultItem,
@@ -268,6 +268,20 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
     app.nodePreviewImages = {}
   }
 
+  /**
+   * Revoke all preview of a subgraph node and the graph it contains.
+   * Does not recurse to contents of nested subgraphs.
+   */
+  function revokeSubgraphPreviews(subgraphNode: SubgraphNode) {
+    const graphId = subgraphNode.graph.isRootGraph
+      ? ''
+      : subgraphNode.graph.id + ':'
+    revokePreviewsByLocatorId(graphId + subgraphNode.id)
+    for (const node of subgraphNode.subgraph.nodes) {
+      revokePreviewsByLocatorId(subgraphNode.subgraph.id + node.id)
+    }
+  }
+
   return {
     getNodeOutputs,
     getNodeImageUrls,
@@ -279,6 +293,7 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
     setNodePreviewsByNodeId,
     revokePreviewsByExecutionId,
     revokeAllPreviews,
+    revokeSubgraphPreviews,
     getPreviewParam
   }
 })