Support nodes that can inherit backend types into their outputs

src/lib/components/ComfyApp.ts

@@ -473,7 +473,7 @@ export default class ComfyApp {
 
             const inputs = {};
 
-            // Store all link values
+            // Store input values passed by frontend-only nodes
             if (node.inputs) {
                 for (let i = 0; i < node.inputs.length; i++) {
                     const inp = node.inputs[i];
@@ -520,14 +520,14 @@ export default class ComfyApp {
                 }
             }
 
-            // Store all links between nodes
+            // Store links between backend-only and hybrid nodes
             for (let i = 0; i < node.inputs.length; i++) {
                 let parent: ComfyGraphNode = node.getInputNode(i) as ComfyGraphNode;
                 if (parent) {
                     const seen = {}
                     let link = node.getInputLink(i);
 
-                    const isValidParent = (parent: ComfyGraphNode) => {
+                    const isFrontendParent = (parent: ComfyGraphNode) => {
                         if (!parent || parent.isBackendNode)
                             return false;
                         if (tag && !hasTag(parent, tag))
@@ -535,16 +535,31 @@ export default class ComfyApp {
                         return true;
                     }
 
-                    while (isValidParent(parent)) {
+                    // If there are frontend-only nodes between us and another
-                        link = parent.getInputLink(link.origin_slot);
+                    // backend node, we have to traverse them first. This
-                        if (link && !seen[link.id]) {
+                    // behavior is dependent on the type of node. Reroute nodes
-                            seen[link.id] = true
+                    // will simply follow their single input, while branching
-                            const inputNode = parent.getInputNode(link.origin_slot) as ComfyGraphNode;
+                    // nodes have conditional logic that determines which link
+                    // to follow backwards.
+                    while (isFrontendParent(parent)) {
+                        const nextLink = parent.getUpstreamLink()
+                        if (nextLink == null) {
+                            console.warn("[graphToPrompt] No upstream link found in frontend node", parent)
+                            break;
+                        }
+
+                        console.debug("[graphToPrompt] consider link", JSON.stringify(link), parent.id)
+
+                        if (nextLink && !seen[nextLink.id]) {
+                            seen[nextLink.id] = true
+                            const inputNode = parent.graph.getNodeById(nextLink.origin_id) as ComfyGraphNode;
                             if (inputNode && tag && !hasTag(inputNode, tag)) {
-                                console.debug("Skipping tagged parent node", tag, node.properties.tags)
+                                console.debug("[graphToPrompt] Skipping tagged intermediate frontend node", tag, node.properties.tags)
                                 parent = null;
                             }
                             else {
+                                console.debug("[graphToPrompt] Traverse upstream link", JSON.stringify(link), parent.id, inputNode?.id, inputNode?.isBackendNode)
+                                link = nextLink;
                                 parent = inputNode;
                             }
                         } else {
@@ -556,6 +571,7 @@ export default class ComfyApp {
                         if (tag && !hasTag(parent, tag))
                             continue;
 
+                        console.debug("[graphToPrompt] final link", JSON.stringify(link), parent.id, node.id)
                         const input = node.inputs[i]
                         // TODO can null be a legitimate value in some cases?
                         // Nodes like CLIPLoader will never have a value in the frontend, hence "null".
@@ -572,18 +588,19 @@ export default class ComfyApp {
         }
 
         // Remove inputs connected to removed nodes
-
+        console.debug("[graphToPrompt] before prune", JSON.stringify(output))
-        for (const o in output) {
+        for (const nodeId in output) {
-            for (const i in output[o].inputs) {
+            for (const inputName in output[nodeId].inputs) {
-                if (Array.isArray(output[o].inputs[i])
+                if (Array.isArray(output[nodeId].inputs[inputName])
-                    && output[o].inputs[i].length === 2
+                    && output[nodeId].inputs[inputName].length === 2
-                    && !output[output[o].inputs[i][0]]) {
+                    && !output[output[nodeId].inputs[inputName][0]]) {
-                    console.debug("Prune removed node link", o, i, output[o].inputs[i])
+                    console.debug("Prune removed node link", nodeId, inputName, output[nodeId].inputs[inputName])
-                    delete output[o].inputs[i];
+                    delete output[nodeId].inputs[inputName];
                 }
             }
         }
 
+        console.debug("[graphToPrompt] after prune", JSON.stringify(output))
         // console.debug({ workflow, output })
         // console.debug(promptToGraphVis({ workflow, output }))
 
@@ -615,6 +632,7 @@ export default class ComfyApp {
                     }
 
                     const p = await this.graphToPrompt(tag);
+                    console.debug(promptToGraphVis(p))
 
                     try {
                         await this.api.queuePrompt(num, p);

src/lib/nodes/ComfyGraphNode.ts

@@ -1,7 +1,7 @@
 import type { ComfyInputConfig } from "$lib/IComfyInputSlot";
 import type { SerializedPrompt } from "$lib/components/ComfyApp";
 import type ComfyWidget from "$lib/components/widgets/ComfyWidget";
-import { LGraph, LGraphNode, LiteGraph, NodeMode, type SerializedLGraphNode, type Vector2 } from "@litegraph-ts/core";
+import { LGraph, LGraphNode, LLink, LiteGraph, NodeMode, type INodeInputSlot, type SerializedLGraphNode, type Vector2, type INodeOutputSlot, LConnectionKind, type SlotType, LGraphCanvas, getStaticPropertyOnInstance, type PropertyLayout, type SlotLayout } from "@litegraph-ts/core";
 import type { SvelteComponentDev } from "svelte/internal";
 import type { ComfyWidgetNode } from "./ComfyWidgetNodes";
 import type IComfyInputSlot from "$lib/IComfyInputSlot";
@@ -26,6 +26,15 @@ export default class ComfyGraphNode extends LGraphNode {
 
     defaultWidgets?: DefaultWidgetLayout
 
+    /*
+     * If true, attempt to reconcile wildcard types in slots ("*")
+     * when a new input/output is connected
+     *
+     * Only set this to true if all output slots are wildcard typed in the
+     * static slotLayout property by default!
+     */
+    canInheritSlotTypes: boolean = false;
+
     /*
      * If false, don't serialize user-set properties into the workflow.
      * Useful for removing personal information from shared workflows.
@@ -39,6 +48,173 @@ export default class ComfyGraphNode extends LGraphNode {
         o.widgets_values = []
     }
 
+    /*
+     * Traverses this node backwards in the graph in order to reach a connecting
+     * backend node, if any. For example, reroute nodes will simply follow their
+     * single input, while branching nodes have conditional logic that
+     * determines which link to follow backwards.
+     */
+    getUpstreamLink(): LLink | null {
+        return null;
+    }
+
+    private inheritSlotTypes(type: LConnectionKind, isConnected: boolean) {
+        // Prevent multiple connections to different types when we have no input
+        if (isConnected && type === LConnectionKind.OUTPUT) {
+            // Ignore wildcard nodes as these will be updated to real types
+            const types = new Set(this.outputs.flatMap(o => o.links.map((l) => this.graph.links[l].type).filter((t) => t !== "*")));
+            if (types.size > 1) {
+                for (let j = 0; j < this.outputs.length; j++) {
+                    for (let i = 0; i < this.outputs[j].links.length - 1; i++) {
+                        const linkId = this.outputs[j].links[i];
+                        const link = this.graph.links[linkId];
+                        const node = this.graph.getNodeById(link.target_id);
+                        node.disconnectInput(link.target_slot);
+                    }
+                }
+            }
+        }
+
+        // Find root input
+        let currentNode: ComfyGraphNode = this;
+        let updateNodes: ComfyGraphNode[] = [];
+        let inputType: SlotType | null = null;
+        let inputNode = null;
+
+        while (currentNode) {
+            updateNodes.unshift(currentNode);
+            const link = currentNode.getUpstreamLink();
+            if (link !== null) {
+                const node = this.graph.getNodeById(link.origin_id) as ComfyGraphNode;
+                console.warn(node.type)
+                if (node.canInheritSlotTypes) {
+                    console.log("REROUTE2", node)
+                    if (node === this) {
+                        // We've found a circle
+                        currentNode.disconnectInput(link.target_slot);
+                        currentNode = null;
+                    }
+                    else {
+                        // Move the previous node
+                        currentNode = node;
+                    }
+                } else {
+                    // We've found the end
+                    inputNode = currentNode;
+                    inputType = node.outputs[link.origin_slot]?.type ?? null;
+                    break;
+                }
+            } else {
+                // This path has no input node
+                currentNode = null;
+                break;
+            }
+        }
+
+        // Find all outputs
+        const nodes: ComfyGraphNode[] = [this];
+        let outputType: SlotType | null = null;
+        while (nodes.length) {
+            currentNode = nodes.pop();
+            if (currentNode.outputs) {
+                for (let i = 0; i < currentNode.outputs.length; i++) {
+                    const outputs = currentNode.outputs[i].links || [];
+                    if (outputs.length) {
+                        for (const linkId of outputs) {
+                            const link = this.graph.links[linkId];
+
+                            // When disconnecting sometimes the link is still registered
+                            if (!link) continue;
+
+                            const node = this.graph.getNodeById(link.target_id) as ComfyGraphNode;
+
+                            if (node.canInheritSlotTypes) {
+                                console.log("REROUTE", node)
+                                // Follow reroute nodes
+                                nodes.push(node);
+                                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) {
+                                    // The output doesnt match our input so disconnect it
+                                    node.disconnectInput(link.target_slot);
+                                } else {
+                                    outputType = nodeOutType;
+                                }
+                            }
+                        }
+                    } else {
+                        // No more outputs for this path
+                    }
+                }
+            }
+        }
+
+        const displayType = inputType || outputType || "*";
+        const color = LGraphCanvas.DEFAULT_LINK_TYPE_COLORS[displayType];
+
+        // Update the types of each node
+        for (const node of updateNodes) {
+            // in lieu of static abstract properties
+            const slotLayout = getStaticPropertyOnInstance<SlotLayout>(node, "slotLayout");
+            if (!slotLayout)
+                continue
+
+            const layoutOutputs = slotLayout.outputs || []
+
+            for (let i = 0; i < node.outputs.length; i++) {
+                // Check if this output was defined as starting off as a
+                // wildcard. If for example it was something else like a string,
+                // it wouldn't make sense to change its type dynamically.
+                const isWildcardOutput = layoutOutputs.length > i && layoutOutputs[i].type === "*";
+                if (!isWildcardOutput) {
+                    console.error("not wildcard", node.outputs[i], layoutOutputs[i])
+                    continue;
+                }
+
+                // If we dont have an input type we are always wildcard but we'll show the output type
+                // This lets you change the output link to a different type and all nodes will update
+                node.outputs[i].type = inputType || "*";
+                (node as any).__outputType = displayType;
+                node.outputs[i].name = node.properties.showOutputText ? String(displayType) : "";
+                node.size = node.computeSize();
+
+                // TODO from ComfyReroute
+                if ("applyOrientation" in node && typeof node.applyOrientation === "function")
+                    node.applyOrientation();
+
+                for (const l of node.outputs[i].links || []) {
+                    const link = this.graph.links[l];
+                    if (link) {
+                        link.color = color;
+                    }
+                }
+            }
+        }
+
+        if (inputNode) {
+            for (let i = 0; i < inputNode.inputs.length; i++) {
+                const link = this.graph.links[inputNode.inputs[i].link];
+                if (link) {
+                    link.color = color;
+                }
+            }
+        }
+    }
+
+    override onConnectionsChange(
+        type: LConnectionKind,
+        slotIndex: number,
+        isConnected: boolean,
+        link: LLink,
+        ioSlot: (INodeInputSlot | INodeOutputSlot)
+    ) {
+        if (this.canInheritSlotTypes) {
+            this.inheritSlotTypes(type, isConnected);
+        }
+    }
+
     override onResize(size: Vector2) {
         if ((window as any)?.app?.shiftDown) {
             const w = LiteGraph.CANVAS_GRID_SIZE * Math.round(this.size[0] / LiteGraph.CANVAS_GRID_SIZE);

src/lib/nodes/ComfyPickFirstNode.ts

@@ -32,6 +32,8 @@ export default class ComfyPickFirstNode extends ComfyGraphNode {
         ],
     }
 
+    override canInheritSlotTypes = true;
+
     private selected: number = -1;
 
     displayWidget: ITextWidget;
@@ -63,6 +65,8 @@ export default class ComfyPickFirstNode extends ComfyGraphNode {
         link: LLink,
         ioSlot: (INodeInputSlot | INodeOutputSlot)
     ) {
+        super.onConnectionsChange(type, slotIndex, isConnected, link, ioSlot);
+
         if (type !== LConnectionKind.INPUT)
             return;
 
@@ -71,7 +75,7 @@ export default class ComfyPickFirstNode extends ComfyGraphNode {
                 // Add a new input
                 const lastInputName = this.inputs[this.inputs.length - 1].name
                 const inputName = nextLetter(lastInputName);
-                this.addInput(inputName, "*")
+                this.addInput(inputName, this.inputs[0].type)
             }
         }
         else {
@@ -97,6 +101,19 @@ export default class ComfyPickFirstNode extends ComfyGraphNode {
         }
     }
 
+    override getUpstreamLink(): LLink | null {
+        for (let index = 0; index < this.inputs.length; index++) {
+            const link = this.getInputLink(index);
+            if (link != null && (link.data != null || this.properties.acceptNullLinkData)) {
+                const node = this.getInputNode(index);
+                if (node != null && node.mode === NodeMode.ALWAYS) {
+                    return link;
+                }
+            }
+        }
+        return null;
+    }
+
     override onExecute() {
         for (let index = 0; index < this.inputs.length; index++) {
             const link = this.getInputLink(index);

src/lib/nodes/ComfyReroute.ts

@@ -1,4 +1,4 @@
-import { LiteGraph, type ContextMenuItem, type LGraphNode, type Vector2, LConnectionKind, LLink, LGraphCanvas, type SlotType, TitleMode, type SlotLayout } from "@litegraph-ts/core";
+import { LiteGraph, type ContextMenuItem, type LGraphNode, type Vector2, LConnectionKind, LLink, LGraphCanvas, type SlotType, TitleMode, type SlotLayout, NodeMode } from "@litegraph-ts/core";
 import ComfyGraphNode from "./ComfyGraphNode";
 
 export interface ComfyRerouteProperties extends Record<any, any> {
@@ -47,124 +47,20 @@ export default class ComfyReroute extends ComfyGraphNode {
         }
     }
 
-    override onConnectionsChange(type: LConnectionKind, slotIndex: number, isConnected: boolean, _link: LLink) {
+    override getUpstreamLink(): LLink | null {
+        const link = this.getInputLink(0)
+        const node = this.getInputNode(0)
+        if (link && node && node.mode === NodeMode.ALWAYS)
+            return link;
+        return null;
+    }
+
+    override canInheritSlotTypes = true;
+
+    override onConnectionsChange(type: LConnectionKind, slotIndex: number, isConnected: boolean, link: LLink, ioSlot: (INodeInputSlot | INodeOutputSlot)) {
         this.applyOrientation();
 
-        // Prevent multiple connections to different types when we have no input
+        super.onConnectionsChange(type, slotIndex, isConnected, link, ioSlot);
-        if (isConnected && type === LConnectionKind.OUTPUT) {
-            // Ignore wildcard nodes as these will be updated to real types
-            const types = new Set(this.outputs[0].links.map((l) => this.graph.links[l].type).filter((t) => t !== "*"));
-            if (types.size > 1) {
-                for (let i = 0; i < this.outputs[0].links.length - 1; i++) {
-                    const linkId = this.outputs[0].links[i];
-                    const link = this.graph.links[linkId];
-                    const node = this.graph.getNodeById(link.target_id);
-                    node.disconnectInput(link.target_slot);
-                }
-            }
-        }
-
-        // Find root input
-        let currentNode: ComfyReroute = this;
-        let updateNodes: ComfyReroute[] = [];
-        let inputType: SlotType | null = null;
-        let inputNode = null;
-        while (currentNode) {
-            updateNodes.unshift(currentNode);
-            const linkId = currentNode.inputs[0].link;
-            if (linkId !== null) {
-                const link = this.graph.links[linkId];
-                const node = this.graph.getNodeById(link.origin_id);
-                console.warn(node.type)
-                if (node.class === ComfyReroute) {
-                    console.log("REROUTE2")
-                    if (node === this) {
-                        // We've found a circle
-                        currentNode.disconnectInput(link.target_slot);
-                        currentNode = null;
-                    }
-                    else {
-                        // Move the previous node
-                        currentNode = node as ComfyReroute;
-                    }
-                } else {
-                    // We've found the end
-                    inputNode = currentNode;
-                    inputType = node.outputs[link.origin_slot]?.type ?? null;
-                    break;
-                }
-            } else {
-                // This path has no input node
-                currentNode = null;
-                break;
-            }
-        }
-
-        // Find all outputs
-        const nodes: ComfyReroute[] = [this];
-        let outputType: SlotType | null = null;
-        while (nodes.length) {
-            currentNode = nodes.pop();
-            const outputs = (currentNode.outputs ? currentNode.outputs[0].links : []) || [];
-            if (outputs.length) {
-                for (const linkId of outputs) {
-                    const link = this.graph.links[linkId];
-
-                    // When disconnecting sometimes the link is still registered
-                    if (!link) continue;
-
-                    const node = this.graph.getNodeById(link.target_id);
-
-                    if (node.class === ComfyReroute) {
-                        console.log("REROUTE")
-                        // Follow reroute nodes
-                        nodes.push(node as ComfyReroute);
-                        updateNodes.push(node as ComfyReroute);
-                    } 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) {
-                            // The output doesnt match our input so disconnect it
-                            node.disconnectInput(link.target_slot);
-                        } else {
-                            outputType = nodeOutType;
-                        }
-                    }
-                }
-            } else {
-                // No more outputs for this path
-            }
-        }
-
-        const displayType = inputType || outputType || "*";
-        const color = LGraphCanvas.DEFAULT_LINK_TYPE_COLORS[displayType];
-
-        // Update the types of each node
-        for (const node of updateNodes) {
-            // If we dont have an input type we are always wildcard but we'll show the output type
-            // This lets you change the output link to a different type and all nodes will update
-            node.outputs[0].type = inputType || "*";
-            (node as any).__outputType = displayType;
-            node.outputs[0].name = node.properties.showOutputText ? String(displayType) : "";
-            node.size = node.computeSize();
-
-            if ("applyOrientation" in node && typeof node.applyOrientation === "function")
-                node.applyOrientation();
-
-            for (const l of node.outputs[0].links || []) {
-                const link = this.graph.links[l];
-                if (link) {
-                    link.color = color;
-                }
-            }
-        }
-
-        if (inputNode) {
-            const link = this.graph.links[inputNode.inputs[0].link];
-            if (link) {
-                link.color = color;
-            }
-        }
     };
 
     override clone(): LGraphNode {

src/lib/nodes/ComfySelector.ts

@@ -1,4 +1,4 @@
-import { BuiltInSlotType, LiteGraph, type SlotLayout } from "@litegraph-ts/core";
+import { BuiltInSlotType, LConnectionKind, LLink, LiteGraph, NodeMode, type INodeInputSlot, type SlotLayout, type INodeOutputSlot } from "@litegraph-ts/core";
 import ComfyGraphNode from "./ComfyGraphNode";
 
 export interface ComfySelectorProperties extends Record<any, any> {
@@ -23,12 +23,29 @@ export default class ComfySelector extends ComfyGraphNode {
         ],
     }
 
+    override canInheritSlotTypes = true;
+
     private selected: number = 0;
 
     constructor(title?: string) {
         super(title);
     }
 
+    override getUpstreamLink(): LLink | null {
+        var sel = this.getInputData(0);
+        if (sel == null || sel.constructor !== Number)
+            sel = 0;
+
+        this.selected = sel = Math.round(sel) % (this.inputs.length - 1);
+
+        var link = this.getInputLink(sel + 1);
+        var node = this.getInputNode(sel + 1);
+        if (link != null && node != null && node.mode === NodeMode.ALWAYS)
+            return link;
+
+        return null
+    }
+
     override onDrawBackground(ctx: CanvasRenderingContext2D) {
         if (this.flags.collapsed) {
             return;
@@ -81,12 +98,41 @@ export class ComfySelectorTwo extends ComfyGraphNode {
         ],
     }
 
+    override canInheritSlotTypes = true;
+
     private selected: number = 0;
 
     constructor(title?: string) {
         super(title);
     }
 
+    override getUpstreamLink(): LLink | null {
+        var sel = this.getInputData(0);
+        if (sel == null || sel.constructor !== Boolean)
+            sel = 0;
+
+        this.selected = sel ? 0 : 1;
+
+        var link = this.getInputLink(this.selected + 1);
+        var node = this.getInputNode(this.selected + 1);
+        if (link != null && node != null && node.mode === NodeMode.ALWAYS)
+            return link
+
+        return null;
+    }
+
+    override onConnectionsChange(
+        type: LConnectionKind,
+        slotIndex: number,
+        isConnected: boolean,
+        link: LLink,
+        ioSlot: (INodeInputSlot | INodeOutputSlot)
+    ) {
+        if (type === LConnectionKind.INPUT) {
+
+        }
+    }
+
     override onDrawBackground(ctx: CanvasRenderingContext2D) {
         if (this.flags.collapsed) {
             return;
@@ -107,6 +153,8 @@ export class ComfySelectorTwo extends ComfyGraphNode {
         this.selected = sel ? 0 : 1;
         var v = this.getInputData(this.selected + 1);
         if (v !== undefined) {
+            const link = this.getInputLink(this.selected + 1);
+            const node = this.getInputNode(this.selected + 1);
             this.setOutputData(0, v);
         }
     }

src/lib/nodes/ComfyWidgetNodes.ts

@@ -189,13 +189,14 @@ export abstract class ComfyWidgetNode<T = any> extends ComfyGraphNode {
         this.propsChanged.set(get(this.propsChanged) + 1)
     }
 
-    onConnectionsChange(
+    override onConnectionsChange(
         type: LConnectionKind,
         slotIndex: number,
         isConnected: boolean,
         link: LLink,
         ioSlot: (INodeOutputSlot | INodeInputSlot)
     ): void {
+        super.onConnectionsChange(type, slotIndex, isConnected, link, ioSlot);
         this.clampConfig();
     }
 

src/lib/widgets/GalleryWidget.svelte

@@ -26,9 +26,10 @@
          nodeValue = node.value;
          propsChanged = node.propsChanged;
 
-         const len = $nodeValue.length
+         if ($nodeValue != null) {
-         if (node.properties.index < 0 || node.properties.index >= len) {
+             if (node.properties.index < 0 || node.properties.index >= $nodeValue.length) {
-             node.setProperty("index", clamp(node.properties.index, 0, len))
+                 node.setProperty("index", clamp(node.properties.index, 0, $nodeValue))
+             }
          }
      }
  };
@@ -123,7 +124,7 @@
 </script>
 
 <div class="wrapper comfy-gallery-widget gradio-gallery" bind:this={element}>
-    {#if widget && node && nodeValue}
+    {#if widget && node && nodeValue && $nodeValue != null}
         <Block variant="solid" padding={false}>
             <div class="padding">
                 <Gallery