neuro_llc/ComfyBox
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases 1 Wiki Activity

Support nodes that can inherit backend types into their outputs

Browse Source
This commit is contained in:
space-nuko
2023-05-07 23:16:06 -05:00
parent c3af1244ae
commit fbef7cd744
7 changed files with 299 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files

130
src/lib/nodes/ComfyReroute.ts
View File

@@ -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
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;
}
}
super.onConnectionsChange(type, slotIndex, isConnected, link, ioSlot);
};
override clone(): LGraphNode {