Convert to subgraph command

package.json

@@ -8,7 +8,7 @@
     "preview": "turbo run preview",
     "check": "svelte-kit sync && svelte-check --tsconfig ./jsconfig.json",
     "check:watch": "svelte-kit sync && svelte-check --tsconfig ./jsconfig.json --watch",
-    "test:unit": "vitest",
+    "test": "vitest run --root .",
     "lint": "prettier --plugin-search-dir . --check . && eslint .",
     "format": "prettier --plugin-search-dir . --write .",
     "svelte-check": "svelte-check",

src/lib/ComfyGraph.ts

@@ -1,4 +1,4 @@
-import { LConnectionKind, LGraph, LGraphNode, type INodeSlot, type SlotIndex, LiteGraph, getStaticProperty, type LGraphAddNodeOptions, LGraphCanvas } from "@litegraph-ts/core";
+import { LConnectionKind, LGraph, LGraphNode, type INodeSlot, type SlotIndex, LiteGraph, getStaticProperty, type LGraphAddNodeOptions, LGraphCanvas, type LGraphRemoveNodeOptions } from "@litegraph-ts/core";
 import GraphSync from "./GraphSync";
 import EventEmitter from "events";
 import type TypedEmitter from "typed-emitter";
@@ -22,12 +22,10 @@ type ComfyGraphEvents = {
 }
 
 export default class ComfyGraph extends LGraph {
-    graphSync: GraphSync;
     eventBus: TypedEmitter<ComfyGraphEvents> = new EventEmitter() as TypedEmitter<ComfyGraphEvents>;
 
     constructor() {
         super();
-        this.graphSync = new GraphSync(this)
     }
 
     override onConfigure() {
@@ -50,8 +48,7 @@ export default class ComfyGraph extends LGraph {
     }
 
     override onNodeAdded(node: LGraphNode, options: LGraphAddNodeOptions) {
-        layoutState.nodeAdded(node)
+        layoutState.nodeAdded(node, options)
-        this.graphSync.onNodeAdded(node);
 
         // All nodes whether they come from base litegraph or ComfyBox should
         // have tags added to them. Can't override serialization for existing
@@ -92,7 +89,7 @@ export default class ComfyGraph extends LGraph {
 
         if (get(uiState).autoAddUI) {
             console.warn("ADD", node.type, options)
-            if (!("svelteComponentType" in node) && options.addedByDeserialize == null) {
+            if (!("svelteComponentType" in node) && options.addedBy == null) {
                 console.debug("[ComfyGraph] AutoAdd UI")
                 const comfyNode = node as ComfyGraphNode;
                 const widgetNodesAdded = []
@@ -127,9 +124,8 @@ export default class ComfyGraph extends LGraph {
         this.eventBus.emit("nodeAdded", node);
     }
 
-    override onNodeRemoved(node: LGraphNode) {
+    override onNodeRemoved(node: LGraphNode, options: LGraphRemoveNodeOptions) {
-        layoutState.nodeRemoved(node);
+        layoutState.nodeRemoved(node, options);
-        this.graphSync.onNodeRemoved(node);
 
         console.debug("Removed", node);
         this.eventBus.emit("nodeRemoved", node);

src/lib/ComfyGraphCanvas.ts

@@ -1,4 +1,4 @@
-import { BuiltInSlotShape, LGraph, LGraphCanvas, LGraphNode, LiteGraph, NodeMode, type MouseEventExt, type Vector2, type Vector4, TitleMode, type ContextMenuItem, type IContextMenuItem } from "@litegraph-ts/core";
+import { BuiltInSlotShape, LGraph, LGraphCanvas, LGraphNode, LiteGraph, NodeMode, type MouseEventExt, type Vector2, type Vector4, TitleMode, type ContextMenuItem, type IContextMenuItem, Subgraph } from "@litegraph-ts/core";
 import type ComfyApp from "./components/ComfyApp";
 import queueState from "./stores/queueState";
 import { get } from "svelte/store";
@@ -322,6 +322,34 @@ export default class ComfyGraphCanvas extends LGraphCanvas {
         }
     }
 
+    private convertToSubgraph(_value: IContextMenuItem, _options, mouseEvent, prevMenu, callback?: (node: LGraphNode) => void) {
+        if (Object.keys(this.selected_nodes).length === 0)
+            return
+
+        const selected = Object.values(this.selected_nodes).filter(n => n != null);
+        this.selected_nodes = {}
+
+        const subgraph = LiteGraph.createNode(Subgraph);
+        subgraph.buildFromNodes(selected)
+
+        this.graph.add(subgraph)
+    }
+
+    override getCanvasMenuOptions(): ContextMenuItem[] {
+        const options = super.getCanvasMenuOptions();
+
+        options.push(
+            {
+                content: "Convert to Subgraph",
+                has_submenu: false,
+                disabled: Object.keys(this.selected_nodes).length === 0,
+                callback: this.convertToSubgraph.bind(this)
+            },
+        )
+
+        return options
+    }
+
     override getNodeMenuOptions(node: LGraphNode): ContextMenuItem[] {
         const options = super.getNodeMenuOptions(node);
 

src/lib/GraphSync.ts

@@ -1,77 +0,0 @@
-import type { LGraph, LGraphNode } from "@litegraph-ts/core";
-import type ComfyApp from "./components/ComfyApp";
-import type { Unsubscriber, Writable } from "svelte/store";
-import type { ComfyWidgetNode } from "./nodes";
-import type ComfyGraph from "./ComfyGraph";
-
-type WidgetSubStore = {
-    store: WidgetUIStateStore,
-    unsubscribe: Unsubscriber
-}
-
-/*
- * Responsible for watching for and synchronizing state changes from the
- * frontend to the litegraph instance.
- *
- * The other way around is unnecessary since the nodes in ComfyBox can't be
- * interacted with. If that were true the implementation would be way more
- * complex since litegraph doesn't (currently) expose a global
- * event-emitter-like thing for when nodes/widgets are changed.
- *
- * Assumptions:
- * - Widgets can't be added to a node after they're created (messes up the indices in WidgetSubStore[])
- * - Widgets can't be interacted with from the graph, only from the frontend
- * - Only one workflow/graph can ever be loaded into the program
- */
-export default class GraphSync {
-    graph: LGraph;
-
-    // nodeId -> widgetSubStore
-    private stores: Record<string, WidgetSubStore> = {}
-
-    constructor(graph: ComfyGraph) {
-        this.graph = graph;
-    }
-
-    onNodeAdded(node: LGraphNode) {
-        // TODO assumes only a single graph's widget state.
-
-        if ("svelteComponentType" in node) {
-            this.addStore(node as ComfyWidgetNode);
-        }
-
-        this.graph.setDirtyCanvas(true, true);
-    }
-
-    onNodeRemoved(node: LGraphNode) {
-        if ("svelteComponentType" in node) {
-            this.removeStore(node as ComfyWidgetNode);
-        }
-
-        this.graph.setDirtyCanvas(true, true);
-    }
-
-    private addStore(node: ComfyWidgetNode) {
-        if (this.stores[node.id]) {
-            console.warn("[GraphSync] Stores already exist!", node.id, this.stores[node.id])
-        }
-
-        const unsub = node.value.subscribe((v) => this.onWidgetStateChanged(node, v))
-        this.stores[node.id] = ({ store: node.value, unsubscribe: unsub });
-
-        console.debug("[GraphSync] NEWSTORE", this.stores[node.id])
-    }
-
-    private removeStore(node: ComfyWidgetNode) {
-        console.debug("[GraphSync] DELSTORE", this.stores[node.id])
-        this.stores[node.id].unsubscribe()
-        delete this.stores[node.id]
-    }
-
-    /*
-     * Fired when a single widget's value changes.
-     */
-    private onWidgetStateChanged(node: ComfyWidgetNode, value: any) {
-        this.graph.setDirtyCanvas(true, true);
-    }
-}

src/lib/components/ComfyApp.ts

@@ -37,6 +37,7 @@ export const COMFYBOX_SERIAL_VERSION = 1;
 
 LiteGraph.catch_exceptions = false;
 LiteGraph.CANVAS_GRID_SIZE = 32;
+LiteGraph.default_subgraph_lgraph_factory = () => new ComfyGraph();
 
 if (typeof window !== "undefined") {
     // Load default visibility

src/lib/stores/layoutState.ts

@@ -1,7 +1,7 @@
 import { get, writable } from 'svelte/store';
 import type { Readable, Writable } from 'svelte/store';
 import type ComfyApp from "$lib/components/ComfyApp"
-import { type LGraphNode, type IWidget, type LGraph, NodeMode } from "@litegraph-ts/core"
+import { type LGraphNode, type IWidget, type LGraph, NodeMode, type LGraphRemoveNodeOptions, type LGraphAddNodeOptions } from "@litegraph-ts/core"
 import { dndzone, SHADOW_PLACEHOLDER_ITEM_ID } from 'svelte-dnd-action';
 import type { ComfyWidgetNode } from '$lib/nodes';
 
@@ -79,7 +79,7 @@ export type LayoutState = {
      * If true, a saved workflow is being deserialized, so ignore any
      * nodeAdded/nodeRemoved events.
      *
-     * TODO: instead use LGraphAddNodeOptions.addedByDeserialize
+     * TODO: instead use LGraphAddNodeOptions.addedBy
      */
     isConfiguring: boolean,
 
@@ -649,8 +649,8 @@ type LayoutStateOps = {
     addWidget: (parent: ContainerLayout, node: ComfyWidgetNode, attrs: Partial<Attributes>, index?: number) => WidgetLayout,
     findDefaultContainerForInsertion: () => ContainerLayout | null,
     updateChildren: (parent: IDragItem, children: IDragItem[]) => IDragItem[],
-    nodeAdded: (node: LGraphNode) => void,
+    nodeAdded: (node: LGraphNode, options: LGraphAddNodeOptions) => void,
-    nodeRemoved: (node: LGraphNode) => void,
+    nodeRemoved: (node: LGraphNode, options: LGraphRemoveNodeOptions) => void,
     groupItems: (dragItems: IDragItem[], attrs?: Partial<Attributes>) => ContainerLayout,
     ungroup: (container: ContainerLayout) => void,
     getCurrentSelection: () => IDragItem[],
@@ -757,19 +757,6 @@ function updateChildren(parent: IDragItem, newChildren?: IDragItem[]): IDragItem
     return state.allItems[parent.id].children
 }
 
-function nodeAdded(node: LGraphNode) {
-    const state = get(store)
-    if (state.isConfiguring)
-        return;
-
-    const parent = findDefaultContainerForInsertion();
-
-    console.debug("[layoutState] nodeAdded", node)
-    if ("svelteComponentType" in node) {
-        addWidget(parent, node as ComfyWidgetNode);
-    }
-}
-
 function removeEntry(state: LayoutState, id: DragItemID) {
     const entry = state.allItems[id]
     if (entry.children && entry.children.length > 0) {
@@ -788,7 +775,30 @@ function removeEntry(state: LayoutState, id: DragItemID) {
     delete state.allItems[id]
 }
 
-function nodeRemoved(node: LGraphNode) {
+function nodeAdded(node: LGraphNode, options: LGraphAddNodeOptions) {
+    const state = get(store)
+    if (state.isConfiguring)
+        return;
+
+    if (options.addedBy === "moveIntoSubgraph" || options.addedBy === "moveOutOfSubgraph") {
+        // All we need to do is update the nodeID linked to this node.
+    }
+
+    const parent = findDefaultContainerForInsertion();
+
+    console.debug("[layoutState] nodeAdded", node)
+    if ("svelteComponentType" in node) {
+        addWidget(parent, node as ComfyWidgetNode);
+    }
+}
+
+function nodeRemoved(node: LGraphNode, options: LGraphRemoveNodeOptions) {
+    if (options.removedBy === "moveIntoSubgraph" || options.removedBy === "moveOutOfSubgraph") {
+        // This node is being moved into a subgraph, so it will be readded under
+        // a new node ID shortly.
+        return
+    }
+
     const state = get(store)
 
     console.debug("[layoutState] nodeRemoved", node)

src/main-mobile.ts

@@ -5,6 +5,7 @@ import { f7 } from 'framework7-svelte';
 import ComfyApp from '$lib/components/ComfyApp';
 import uiState from '$lib/stores/uiState';
 import { LiteGraph } from '@litegraph-ts/core';
+import ComfyGraph from '$lib/ComfyGraph';
 
 Framework7.use(Framework7Svelte);