Merge pull request #63 from space-nuko/subgraph-templates

Reroute insert and splice actions

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 { 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 type ComfyApp from "./components/ComfyApp";
-import queueState from "./stores/queueState";
 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
+    }
 }

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[]

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) {

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;
+                                let nodeOutType = node.inputs && node.inputs[link?.target_slot] != null ? node.inputs[link.target_slot].type : null;
-                                if (inputType && nodeOutType !== inputType) {
+                                nodeOutType ||= "*"
+                                if (inputType && nodeOutType !== inputType && nodeOutType !== "*") {
                                     // The output doesnt match our input so disconnect it
                                     node.disconnectInput(link.target_slot);
                                 } else {

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",