245 lines
9.3 KiB
TypeScript
245 lines
9.3 KiB
TypeScript
import { LiteGraph, type ContextMenuItem, type LGraphNode, type Vector2, LConnectionKind, LLink, LGraphCanvas, type SlotType, TitleMode, type SlotLayout } from "@litegraph-ts/core";
|
|
import ComfyGraphNode from "./ComfyGraphNode";
|
|
|
|
export interface ComfyRerouteProperties extends Record<any, any> {
|
|
showOutputText: boolean;
|
|
horizontal: boolean;
|
|
}
|
|
|
|
export default class ComfyReroute extends ComfyGraphNode {
|
|
static defaultVisibility: boolean = true;
|
|
|
|
static setDefaultTextVisibility(visible: boolean) {
|
|
ComfyReroute.defaultVisibility = visible;
|
|
if (visible) {
|
|
localStorage["Comfy.ComfyReroute.DefaultVisibility"] = "true";
|
|
} else {
|
|
delete localStorage["Comfy.ComfyReroute.DefaultVisibility"];
|
|
}
|
|
}
|
|
|
|
override titleMode: TitleMode = TitleMode.NO_TITLE;
|
|
override collapsable: boolean = false;
|
|
|
|
override properties: ComfyRerouteProperties = {
|
|
showOutputText: ComfyReroute.defaultVisibility,
|
|
horizontal: false
|
|
}
|
|
|
|
static slotLayout: SlotLayout = {
|
|
inputs: [
|
|
{ name: "", type: "*" }
|
|
],
|
|
outputs: [
|
|
{ name: "", type: "*" }
|
|
],
|
|
}
|
|
|
|
constructor(title?: string) {
|
|
super(title);
|
|
|
|
this.properties ||= {} as any;
|
|
this.properties.showOutputText = ComfyReroute.defaultVisibility;
|
|
this.properties.horizontal = false;
|
|
|
|
if (this.properties.showOutputText) {
|
|
this.outputs[0].name = "*"
|
|
}
|
|
}
|
|
|
|
override onConnectionsChange(type: LConnectionKind, slotIndex: number, isConnected: boolean, _link: LLink) {
|
|
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;
|
|
}
|
|
}
|
|
};
|
|
|
|
override clone(): LGraphNode {
|
|
const cloned = super.clone.apply(this) as ComfyReroute;
|
|
cloned.removeOutput(0);
|
|
cloned.addOutput(this.properties.showOutputText ? "*" : "", "*");
|
|
cloned.size = cloned.computeSize();
|
|
return cloned;
|
|
};
|
|
|
|
override getExtraMenuOptions(_, options: ContextMenuItem[]): ContextMenuItem[] | null {
|
|
options.unshift(
|
|
{
|
|
content: (this.properties.showOutputText ? "Hide" : "Show") + " Type",
|
|
callback: () => {
|
|
this.properties.showOutputText = !this.properties.showOutputText;
|
|
if (this.properties.showOutputText) {
|
|
this.outputs[0].name = (this as any).__outputType || this.outputs[0].type;
|
|
} else {
|
|
this.outputs[0].name = "";
|
|
}
|
|
this.size = this.computeSize();
|
|
this.applyOrientation();
|
|
this.graph.setDirtyCanvas(true, true);
|
|
},
|
|
},
|
|
{
|
|
content: (ComfyReroute.defaultVisibility ? "Hide" : "Show") + " Type By Default",
|
|
callback: () => {
|
|
ComfyReroute.setDefaultTextVisibility(!ComfyReroute.defaultVisibility);
|
|
},
|
|
},
|
|
{
|
|
// naming is inverted with respect to LiteGraphNode.horizontal
|
|
// LiteGraphNode.horizontal == true means that
|
|
// each slot in the inputs and outputs are layed out horizontally,
|
|
// which is the opposite of the visual orientation of the inputs and outputs as a node
|
|
content: "Set " + (this.properties.horizontal ? "Horizontal" : "Vertical"),
|
|
callback: () => {
|
|
this.properties.horizontal = !this.properties.horizontal;
|
|
this.applyOrientation();
|
|
},
|
|
}
|
|
);
|
|
|
|
return null
|
|
}
|
|
|
|
applyOrientation() {
|
|
this.horizontal = this.properties.horizontal;
|
|
if (this.horizontal) {
|
|
// we correct the input position, because LiteGraphNode.horizontal
|
|
// doesn't account for title presence
|
|
// which reroute nodes don't have
|
|
this.inputs[0].pos = [this.size[0] / 2, 0];
|
|
} else {
|
|
delete this.inputs[0].pos;
|
|
}
|
|
this.graph.setDirtyCanvas(true, true);
|
|
}
|
|
|
|
override computeSize(): Vector2 {
|
|
return [
|
|
this.properties.showOutputText && this.outputs && this.outputs.length
|
|
? Math.max(75, LiteGraph.NODE_TEXT_SIZE * this.outputs[0].name.length * 0.6 + 40)
|
|
: 75,
|
|
26,
|
|
];
|
|
}
|
|
}
|
|
|
|
LiteGraph.registerNodeType({
|
|
class: ComfyReroute,
|
|
title: "Comfy.Reroute",
|
|
desc: "Reroutes nodes preserving input/output types",
|
|
type: "utils/reroute"
|
|
})
|