From 47b0dabd8057e735f99d3ef7e4ed929a7b3e4c0c Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Tue, 23 May 2023 21:17:13 -0500 Subject: [PATCH] Reroute insert and splice --- litegraph | 2 +- src/lib/ComfyGraphCanvas.ts | 52 ++++++++++++++++++++++++++++----- src/lib/ComfyNodeDef.ts | 1 + src/lib/components/ComfyApp.ts | 3 ++ src/lib/nodes/ComfyGraphNode.ts | 5 ++-- src/lib/nodes/ComfyReroute.ts | 25 ++++++++++++++++ 6 files changed, 77 insertions(+), 11 deletions(-) diff --git a/litegraph b/litegraph index d8fa871..a57af62 160000 --- a/litegraph +++ b/litegraph @@ -1 +1 @@ -Subproject commit d8fa87185a258d362f159c04f9045c7930b884bd +Subproject commit a57af62d4648bc85b264329ff193c61d9ba84d6d diff --git a/src/lib/ComfyGraphCanvas.ts b/src/lib/ComfyGraphCanvas.ts index 4070323..ad71005 100644 --- a/src/lib/ComfyGraphCanvas.ts +++ b/src/lib/ComfyGraphCanvas.ts @@ -1,14 +1,11 @@ -import { BuiltInSlotShape, LGraph, LGraphCanvas, LGraphNode, LiteGraph, NodeMode, type MouseEventExt, type Vector2, type Vector4, TitleMode, type ContextMenuItem, type IContextMenuItem, Subgraph, LLink, type NodeID } from "@litegraph-ts/core"; -import type ComfyApp from "./components/ComfyApp"; -import queueState from "./stores/queueState"; +import { BuiltInSlotShape, LGraphCanvas, LGraphNode, LLink, LiteGraph, NodeMode, Subgraph, TitleMode, type ContextMenuItem, type IContextMenuItem, type NodeID, type Vector2, type Vector4, type MouseEventExt, ContextMenu } from "@litegraph-ts/core"; import { get, type Unsubscriber } from "svelte/store"; -import uiState from "./stores/uiState"; -import { Watch } from "@litegraph-ts/nodes-basic"; -import { ComfyReroute } from "./nodes"; -import type { Progress } from "./components/ComfyApp"; -import selectionState from "./stores/selectionState"; import type ComfyGraph from "./ComfyGraph"; +import type ComfyApp from "./components/ComfyApp"; +import { ComfyReroute } from "./nodes"; import layoutStates from "./stores/layoutStates"; +import queueState from "./stores/queueState"; +import selectionState from "./stores/selectionState"; export type SerializedGraphCanvasState = { offset: Vector2, @@ -403,6 +400,30 @@ export default class ComfyGraphCanvas extends LGraphCanvas { } } + private insertReroute(_value: IContextMenuItem, _options, mouseEvent: MouseEventExt, prevMenu?: ContextMenu, link?: LLink) { + if (link == null) + return + + const originNode = this.graph.getNodeById(link.origin_id); + const targetNode = this.graph.getNodeById(link.target_id); + + if (originNode == null || targetNode == null) + return; + + if (typeof prevMenu?.options?.event?.canvasX === "number") + mouseEvent = prevMenu.options.event as MouseEventExt; + + const reroute = LiteGraph.createNode(ComfyReroute); + const size = reroute.computeSize(); + const pos: Vector2 = [mouseEvent.canvasX - size[0] * 0.5, mouseEvent.canvasY - size[1] * 0.5]; + + this.graph.removeLink(link.id); + this.graph.add(reroute, { pos }); + + originNode.connect(link.origin_slot, reroute, 0) + reroute.connect(0, targetNode, link.target_slot); + } + private convertToSubgraph(_value: IContextMenuItem, _options, mouseEvent, prevMenu, callback?: (node: LGraphNode) => void) { if (Object.keys(this.selected_nodes).length === 0) return @@ -445,4 +466,19 @@ export default class ComfyGraphCanvas extends LGraphCanvas { return options } + + override getLinkMenuOptions(link: LLink): ContextMenuItem[] { + const options = super.getLinkMenuOptions(link); + + options.push( + { + content: "Insert Reroute", + has_submenu: false, + disabled: false, + callback: this.insertReroute.bind(this) + }, + ) + + return options + } } diff --git a/src/lib/ComfyNodeDef.ts b/src/lib/ComfyNodeDef.ts index 9f4ad13..70a726a 100644 --- a/src/lib/ComfyNodeDef.ts +++ b/src/lib/ComfyNodeDef.ts @@ -5,6 +5,7 @@ export type ComfyNodeDef = { name: string display_name?: string category: string + output_node?: boolean, input: ComfyNodeDefInputs /** Output type like "LATENT" or "IMAGE" */ output: string[] diff --git a/src/lib/components/ComfyApp.ts b/src/lib/components/ComfyApp.ts index 4fc0d62..6eeba23 100644 --- a/src/lib/components/ComfyApp.ts +++ b/src/lib/components/ComfyApp.ts @@ -375,6 +375,9 @@ export default class ComfyApp { reg("COMFYBOX_IMAGE") reg("COMFYBOX_IMAGES") reg("COMFYBOX_REGION") + + // hide base litegraph reroute from context menus since ComfyUI provides its own + LiteGraph.registered_node_types["basic/reroute"].hide_in_node_lists = true; } static registerDefaultSlotHandlers(nodeId: string, nodeDef: ComfyNodeDef) { diff --git a/src/lib/nodes/ComfyGraphNode.ts b/src/lib/nodes/ComfyGraphNode.ts index 71b5496..6ac2d45 100644 --- a/src/lib/nodes/ComfyGraphNode.ts +++ b/src/lib/nodes/ComfyGraphNode.ts @@ -196,8 +196,9 @@ export default class ComfyGraphNode extends LGraphNode { updateNodes.push(node); } else { // We've found an output - const nodeOutType = node.inputs && node.inputs[link?.target_slot] && node.inputs[link.target_slot].type ? node.inputs[link.target_slot].type : null; - if (inputType && nodeOutType !== inputType) { + let nodeOutType = node.inputs && node.inputs[link?.target_slot] != null ? node.inputs[link.target_slot].type : null; + nodeOutType ||= "*" + if (inputType && nodeOutType !== inputType && nodeOutType !== "*") { // The output doesnt match our input so disconnect it node.disconnectInput(link.target_slot); } else { diff --git a/src/lib/nodes/ComfyReroute.ts b/src/lib/nodes/ComfyReroute.ts index 96157dd..82567cc 100644 --- a/src/lib/nodes/ComfyReroute.ts +++ b/src/lib/nodes/ComfyReroute.ts @@ -76,6 +76,31 @@ export default class ComfyReroute extends ComfyGraphNode { }; override getExtraMenuOptions(_, options: ContextMenuItem[]): ContextMenuItem[] | null { + const canSplice = this.getInputLink(0) != null && this.getOutputLinks(0).length > 0; + + options.push({ + content: "Splice & Remove", + disabled: !canSplice, + callback: () => { + const inputLink = this.getInputLink(0); + const outputLinks = this.getOutputLinks(0); + if (!inputLink || !outputLinks) + return; + + const inputNode = this.graph.getNodeById(inputLink.origin_id) + this.graph.removeLink(inputLink.id); + + for (const outputLink of outputLinks) { + const outputNode = this.graph.getNodeById(outputLink.target_id) + this.graph.removeLink(outputLink.id); + + inputNode.connect(inputLink.origin_slot, outputNode, outputLink.target_slot); + } + + this.graph.remove(this); + } + }) + options.unshift( { content: (this.properties.showOutputText ? "Hide" : "Show") + " Type",