diff --git a/litegraph b/litegraph index 90204e2..4c88137 160000 --- a/litegraph +++ b/litegraph @@ -1 +1 @@ -Subproject commit 90204e22e74feae2154242f68afc6ad178e9dff3 +Subproject commit 4c8813722e42388a8f3e20f2b9c2290641912ef4 diff --git a/src/lib/ComfyGraph.ts b/src/lib/ComfyGraph.ts index d1cdd91..d6cab43 100644 --- a/src/lib/ComfyGraph.ts +++ b/src/lib/ComfyGraph.ts @@ -1,4 +1,4 @@ -import { LConnectionKind, LGraph, LGraphNode, type INodeSlot, type SlotIndex, LiteGraph, getStaticProperty, type LGraphAddNodeOptions, LGraphCanvas, type LGraphRemoveNodeOptions } from "@litegraph-ts/core"; +import { LConnectionKind, LGraph, LGraphNode, type INodeSlot, type SlotIndex, LiteGraph, getStaticProperty, type LGraphAddNodeOptions, LGraphCanvas, type LGraphRemoveNodeOptions, Subgraph, type LGraphAddNodeMode } from "@litegraph-ts/core"; import GraphSync from "./GraphSync"; import EventEmitter from "events"; import type TypedEmitter from "typed-emitter"; @@ -41,8 +41,9 @@ export default class ComfyGraph extends LGraph { } override onNodeAdded(node: LGraphNode, options: LGraphAddNodeOptions) { - if (options.subgraphs && options.subgraphs.length > 0) - return + // Don't add detached subgraphs + if (node.getRootGraph() == null || this._is_subgraph) + return; layoutState.nodeAdded(node, options) @@ -112,7 +113,19 @@ export default class ComfyGraph extends LGraph { const dragItemIDs = widgetNodesAdded.map(wn => get(layoutState).allItemsByNode[wn.id]?.dragItem?.id).filter(Boolean) console.debug("[ComfyGraph] Group new widgets", dragItemIDs) - layoutState.groupItems(dragItemIDs, { title: node.title }) + // Use the default node title instead of custom node title, in + // case node was cloned + const reg = LiteGraph.registered_node_types[node.type] + + layoutState.groupItems(dragItemIDs, { title: reg.title }) + } + } + + // Handle subgraphs being attached + if (node.is(Subgraph)) { + console.error("ISSUBGRAPH") + for (const child of node.subgraph.iterateNodesInOrder()) { + this.onNodeAdded(child, options) } } @@ -123,6 +136,13 @@ export default class ComfyGraph extends LGraph { override onNodeRemoved(node: LGraphNode, options: LGraphRemoveNodeOptions) { layoutState.nodeRemoved(node, options); + // Handle subgraphs being removed + if (node.is(Subgraph)) { + for (const child of node.subgraph.iterateNodesInOrder()) { + this.onNodeRemoved(child, options) + } + } + // console.debug("Removed", node); this.eventBus.emit("nodeRemoved", node); } diff --git a/src/lib/stores/layoutState.ts b/src/lib/stores/layoutState.ts index d86eecf..525d263 100644 --- a/src/lib/stores/layoutState.ts +++ b/src/lib/stores/layoutState.ts @@ -291,7 +291,7 @@ const setNodeTitle = (arg: IDragItem, value: any) => { if (reg == null) return - if (value) + if (value && value !== reg.title) widget.node.title = `${reg.title} (${value})` else widget.node.title = reg.title @@ -699,6 +699,20 @@ const store: Writable = writable({ } }) +function clear() { + store.set({ + root: null, + allItems: {}, + allItemsByNode: {}, + isMenuOpen: false, + isConfiguring: true, + refreshPropsPanel: writable(0), + attrs: { + ...defaultWorkflowAttributes + } + }) +} + function findDefaultContainerForInsertion(): ContainerLayout | null { const state = get(store); @@ -741,6 +755,7 @@ function addContainer(parent: ContainerLayout | null, attrs: Partial ...attrs } } + const entry: DragItemEntry = { dragItem, children: [], parent: null }; if (state.allItemsByNode[dragItem.id] != null) @@ -750,6 +765,7 @@ function addContainer(parent: ContainerLayout | null, attrs: Partial if (parent) { moveItem(dragItem, parent, index) } + console.debug("[layoutState] addContainer", state) store.set(state) runOnChangedForWidgetDefaults(dragItem) @@ -758,7 +774,7 @@ function addContainer(parent: ContainerLayout | null, attrs: Partial function addWidget(parent: ContainerLayout, node: ComfyWidgetNode, attrs: Partial = {}, index?: number): WidgetLayout { const state = get(store); - const widgetName = "Widget" + const widgetName = node.title || "Widget" const dragItem: WidgetLayout = { type: "widget", id: uuidv4(), @@ -771,7 +787,7 @@ function addWidget(parent: ContainerLayout, node: ComfyWidgetNode, attrs: Partia ...attrs } } - const parentEntry = state.allItems[parent.id] + const entry: DragItemEntry = { dragItem, children: [], parent: null }; if (state.allItems[dragItem.id] != null) @@ -821,6 +837,7 @@ function removeEntry(state: LayoutState, id: DragItemID) { function nodeAdded(node: LGraphNode, options: LGraphAddNodeOptions) { const state = get(store) + console.error("NODEADDED", node.type, (node as any).svelteComponentType) if (state.isConfiguring) return; @@ -834,7 +851,7 @@ function nodeAdded(node: LGraphNode, options: LGraphAddNodeOptions) { const parent = findDefaultContainerForInsertion(); - console.debug("[layoutState] nodeAdded", node) + console.debug("[layoutState] nodeAdded", node.id) if ("svelteComponentType" in node) { addWidget(parent, node as ComfyWidgetNode); } diff --git a/src/lib/widgets/utils.ts b/src/lib/widgets/utils.ts index 96e1dd5..9580ccc 100644 --- a/src/lib/widgets/utils.ts +++ b/src/lib/widgets/utils.ts @@ -8,6 +8,9 @@ export function isNodeDisabled(node: LGraphNode): boolean { if (node.mode !== NodeMode.ALWAYS) { return true; } + if (node.graph == null) { + return true + } node = node.graph._subgraph_node; } return false; diff --git a/src/tests/ComfyGraphTests.ts b/src/tests/ComfyGraphTests.ts new file mode 100644 index 0000000..91882b8 --- /dev/null +++ b/src/tests/ComfyGraphTests.ts @@ -0,0 +1,94 @@ +import { LGraph, LiteGraph, Subgraph, type SlotLayout } from "@litegraph-ts/core" +import { Watch } from "@litegraph-ts/nodes-basic" +import { expect } from 'vitest' +import UnitTest from "./UnitTest" +import ComfyGraph from "$lib/ComfyGraph"; +import ComfyPromptSerializer from "$lib/components/ComfyPromptSerializer"; +import { ComfyBackendNode } from "$lib/nodes/ComfyBackendNode"; +import ComfyGraphNode from "$lib/nodes/ComfyGraphNode"; +import { graphToGraphVis } from "$lib/utils"; +import layoutState from "$lib/stores/layoutState"; +import { ComfyNumberNode } from "$lib/nodes/widgets"; +import { get } from "svelte/store"; + +export default class ComfyGraphTests extends UnitTest { + test__onNodeAdded__updatesLayoutState() { + const graph = new ComfyGraph(); + layoutState.initDefaultLayout() // adds 3 containers + + const state = get(layoutState) + expect(Object.keys(state.allItems)).toHaveLength(3) + expect(Object.keys(state.allItemsByNode)).toHaveLength(0) + + const widget = LiteGraph.createNode(ComfyNumberNode); + graph.add(widget) + + expect(Object.keys(state.allItems)).toHaveLength(4) + expect(Object.keys(state.allItemsByNode)).toHaveLength(1) + expect(state.allItemsByNode[widget.id]).toBeTruthy(); + + graph.add(widget) + } + + test__correctSubgraphFactory() { + const graph = new ComfyGraph(); + const subgraph = LiteGraph.createNode(Subgraph); + graph.add(subgraph) + expect(subgraph.graph).toBeInstanceOf(ComfyGraph) + } + + test__onNodeAdded__handlesNodesAddedInSubgraphs() { + const graph = new ComfyGraph(); + layoutState.initDefaultLayout() + + const subgraph = LiteGraph.createNode(Subgraph); + graph.add(subgraph) + + const state = get(layoutState) + expect(Object.keys(state.allItems)).toHaveLength(3) + expect(Object.keys(state.allItemsByNode)).toHaveLength(0) + + const widget = LiteGraph.createNode(ComfyNumberNode); + subgraph.subgraph.add(widget) + + expect(Object.keys(state.allItems)).toHaveLength(4) + expect(Object.keys(state.allItemsByNode)).toHaveLength(1) + expect(state.allItemsByNode[widget.id]).toBeTruthy(); + } + + test__onNodeAdded__handlesSubgraphsWithNodes() { + const graph = new ComfyGraph(); + layoutState.initDefaultLayout() + + const state = get(layoutState) + expect(Object.keys(state.allItems)).toHaveLength(3) + expect(Object.keys(state.allItemsByNode)).toHaveLength(0) + + const subgraph = LiteGraph.createNode(Subgraph); + const widget = LiteGraph.createNode(ComfyNumberNode); + subgraph.subgraph.add(widget) + graph.add(subgraph) + + expect(Object.keys(state.allItems)).toHaveLength(4) + expect(Object.keys(state.allItemsByNode)).toHaveLength(1) + expect(state.allItemsByNode[widget.id]).toBeTruthy(); + } + + test__onNodeRemoved__updatesLayoutState() { + const graph = new ComfyGraph(); + layoutState.initDefaultLayout() + + const widget = LiteGraph.createNode(ComfyNumberNode); + graph.add(widget) + + const state = get(layoutState) + expect(Object.keys(state.allItems)).toHaveLength(4) + expect(Object.keys(state.allItemsByNode)).toHaveLength(1) + expect(state.allItemsByNode[widget.id]).toBeTruthy(); + + graph.remove(widget) + + expect(Object.keys(state.allItems)).toHaveLength(3) + expect(Object.keys(state.allItemsByNode)).toHaveLength(0) + } +} diff --git a/src/tests/main.ts b/src/tests/main.ts index b603b26..204b9f2 100644 --- a/src/tests/main.ts +++ b/src/tests/main.ts @@ -1,3 +1,6 @@ +console.debug = (...msg) => { +} + import { vi, describe, it } from "vitest" import UnitTest from "./UnitTest" import * as testSuite from "./testSuite" diff --git a/src/tests/testSuite.ts b/src/tests/testSuite.ts index f11df41..b0dd6be 100644 --- a/src/tests/testSuite.ts +++ b/src/tests/testSuite.ts @@ -1 +1,2 @@ export { default as ComfyPromptSerializerTests } from "./ComfyPromptSerializerTests" +export { default as ComfyGraphTests } from "./ComfyGraphTests"