|
import {app} from "scripts/app.js";
|
|
import {
|
|
IoDirection,
|
|
followConnectionUntilType,
|
|
getConnectedInputInfosAndFilterPassThroughs,
|
|
} from "./utils.js";
|
|
import {rgthree} from "./rgthree.js";
|
|
import {
|
|
SERVICE as CONTEXT_SERVICE,
|
|
InputMutation,
|
|
InputMutationOperation,
|
|
} from "./services/context_service.js";
|
|
import {NodeTypesString} from "./constants.js";
|
|
import {removeUnusedInputsFromEnd} from "./utils_inputs_outputs.js";
|
|
import {INodeInputSlot, INodeOutputSlot, INodeSlot, LGraphNode, LLink} from "typings/litegraph.js";
|
|
import {ComfyNodeConstructor, ComfyObjectInfo} from "typings/comfy.js";
|
|
import {DynamicContextNodeBase} from "./dynamic_context_base.js";
|
|
import {SERVICE as CONFIG_SERVICE} from "./services/config_service.js";
|
|
|
|
const OWNED_PREFIX = "+";
|
|
const REGEX_OWNED_PREFIX = /^\+\s*/;
|
|
const REGEX_EMPTY_INPUT = /^\+\s*$/;
|
|
|
|
|
|
|
|
|
|
export class DynamicContextNode extends DynamicContextNodeBase {
|
|
static override title = NodeTypesString.DYNAMIC_CONTEXT;
|
|
static override type = NodeTypesString.DYNAMIC_CONTEXT;
|
|
static comfyClass = NodeTypesString.DYNAMIC_CONTEXT;
|
|
|
|
constructor(title = DynamicContextNode.title) {
|
|
super(title);
|
|
}
|
|
|
|
override onNodeCreated() {
|
|
this.addInput("base_ctx", "RGTHREE_DYNAMIC_CONTEXT");
|
|
this.ensureOneRemainingNewInputSlot();
|
|
super.onNodeCreated();
|
|
}
|
|
|
|
override onConnectionsChange(
|
|
type: number,
|
|
slotIndex: number,
|
|
isConnected: boolean,
|
|
link: LLink,
|
|
ioSlot: INodeSlot,
|
|
): void {
|
|
super.onConnectionsChange?.call(this, type, slotIndex, isConnected, link, ioSlot);
|
|
if (this.configuring) {
|
|
return;
|
|
}
|
|
if (type === LiteGraph.INPUT) {
|
|
if (isConnected) {
|
|
this.handleInputConnected(slotIndex);
|
|
} else {
|
|
this.handleInputDisconnected(slotIndex);
|
|
}
|
|
}
|
|
}
|
|
|
|
override onConnectInput(
|
|
inputIndex: number,
|
|
outputType: INodeOutputSlot["type"],
|
|
outputSlot: INodeOutputSlot,
|
|
outputNode: LGraphNode,
|
|
outputIndex: number,
|
|
): boolean {
|
|
let canConnect = true;
|
|
if (super.onConnectInput) {
|
|
canConnect = super.onConnectInput.apply(this, [...arguments] as any);
|
|
}
|
|
if (
|
|
canConnect &&
|
|
outputNode instanceof DynamicContextNode &&
|
|
outputIndex === 0 &&
|
|
inputIndex !== 0
|
|
) {
|
|
const [n, v] = rgthree.logger.warnParts(
|
|
"Currently, you can only connect a context node in the first slot.",
|
|
);
|
|
console[n]?.call(console, ...v);
|
|
canConnect = false;
|
|
}
|
|
return canConnect;
|
|
}
|
|
|
|
handleInputConnected(slotIndex: number) {
|
|
const ioSlot = this.inputs[slotIndex];
|
|
const connectedIndexes = [];
|
|
if (slotIndex === 0) {
|
|
let baseNodeInfos = getConnectedInputInfosAndFilterPassThroughs(this, this, 0);
|
|
const baseNodes = baseNodeInfos.map((n) => n.node)!;
|
|
const baseNodesDynamicCtx = baseNodes[0] as DynamicContextNodeBase;
|
|
if (baseNodesDynamicCtx?.provideInputsData) {
|
|
const inputsData = CONTEXT_SERVICE.getDynamicContextInputsData(baseNodesDynamicCtx);
|
|
console.log("inputsData", inputsData);
|
|
for (const input of baseNodesDynamicCtx.provideInputsData()) {
|
|
if (input.name === "base_ctx" || input.name === "+") {
|
|
continue;
|
|
}
|
|
this.addContextInput(input.name, input.type, input.index);
|
|
this.stabilizeNames();
|
|
}
|
|
}
|
|
} else if (this.isInputSlotForNewInput(slotIndex)) {
|
|
this.handleNewInputConnected(slotIndex);
|
|
}
|
|
}
|
|
|
|
isInputSlotForNewInput(slotIndex: number) {
|
|
const ioSlot = this.inputs[slotIndex];
|
|
return ioSlot && ioSlot.name === "+" && ioSlot.type === "*";
|
|
}
|
|
|
|
handleNewInputConnected(slotIndex: number) {
|
|
if (!this.isInputSlotForNewInput(slotIndex)) {
|
|
throw new Error('Expected the incoming slot index to be the "new input" input.');
|
|
}
|
|
const ioSlot = this.inputs[slotIndex]!;
|
|
let cxn = null;
|
|
if (ioSlot.link != null) {
|
|
cxn = followConnectionUntilType(this, IoDirection.INPUT, slotIndex, true);
|
|
}
|
|
if (cxn?.type && cxn?.name) {
|
|
let name = this.addOwnedPrefix(this.getNextUniqueNameForThisNode(cxn.name));
|
|
if (name.match(/^\+\s*[A-Z_]+(\.\d+)?$/)) {
|
|
name = name.toLowerCase();
|
|
}
|
|
ioSlot.name = name;
|
|
ioSlot.type = cxn.type as string;
|
|
ioSlot.removable = true;
|
|
while (!this.outputs[slotIndex]) {
|
|
this.addOutput("*", "*");
|
|
}
|
|
this.outputs[slotIndex]!.type = cxn.type as string;
|
|
this.outputs[slotIndex]!.name = this.stripOwnedPrefix(name).toLocaleUpperCase();
|
|
|
|
if (cxn.type === "COMBO" || cxn.type.includes(",") || Array.isArray(cxn.type)) {
|
|
(this.outputs[slotIndex] as any).widget = true;
|
|
}
|
|
this.inputsMutated({
|
|
operation: InputMutationOperation.ADDED,
|
|
node: this,
|
|
slotIndex,
|
|
slot: ioSlot,
|
|
});
|
|
this.stabilizeNames();
|
|
this.ensureOneRemainingNewInputSlot();
|
|
}
|
|
}
|
|
|
|
handleInputDisconnected(slotIndex: number) {
|
|
const inputs = this.getContextInputsList();
|
|
if (slotIndex === 0) {
|
|
for (let index = inputs.length - 1; index > 0; index--) {
|
|
if (index === 0 || index === inputs.length - 1) {
|
|
continue;
|
|
}
|
|
const input = inputs[index]!;
|
|
if (!this.isOwnedInput(input.name)) {
|
|
if (input.link || this.outputs[index]?.links?.length) {
|
|
this.renameContextInput(index, input.name, true);
|
|
} else {
|
|
this.removeContextInput(index);
|
|
}
|
|
}
|
|
}
|
|
this.setSize(this.computeSize());
|
|
this.setDirtyCanvas(true, true);
|
|
}
|
|
}
|
|
|
|
ensureOneRemainingNewInputSlot() {
|
|
removeUnusedInputsFromEnd(this, 1, REGEX_EMPTY_INPUT);
|
|
this.addInput(OWNED_PREFIX, "*");
|
|
}
|
|
|
|
getNextUniqueNameForThisNode(desiredName: string) {
|
|
const inputs = this.getContextInputsList();
|
|
const allExistingKeys = inputs.map((i) => this.stripOwnedPrefix(i.name).toLocaleUpperCase());
|
|
desiredName = this.stripOwnedPrefix(desiredName);
|
|
let newName = desiredName;
|
|
let n = 0;
|
|
while (allExistingKeys.includes(newName.toLocaleUpperCase())) {
|
|
newName = `${desiredName}.${++n}`;
|
|
}
|
|
return newName;
|
|
}
|
|
|
|
override removeInput(slotIndex: number) {
|
|
const slot = this.inputs[slotIndex]!;
|
|
super.removeInput(slotIndex);
|
|
if (this.outputs[slotIndex]) {
|
|
this.removeOutput(slotIndex);
|
|
}
|
|
this.inputsMutated({operation: InputMutationOperation.REMOVED, node: this, slotIndex, slot});
|
|
this.stabilizeNames();
|
|
}
|
|
|
|
stabilizeNames() {
|
|
const inputs = this.getContextInputsList();
|
|
const names: string[] = [];
|
|
for (const [index, input] of inputs.entries()) {
|
|
if (index === 0 || index === inputs.length - 1) {
|
|
continue;
|
|
}
|
|
input.label = undefined;
|
|
this.outputs[index]!.label = undefined;
|
|
let origName = this.stripOwnedPrefix(input.name).replace(/\.\d+$/, "");
|
|
let name = input.name;
|
|
if (!this.isOwnedInput(name)) {
|
|
names.push(name.toLocaleUpperCase());
|
|
} else {
|
|
let n = 0;
|
|
name = this.addOwnedPrefix(origName);
|
|
while (names.includes(this.stripOwnedPrefix(name).toLocaleUpperCase())) {
|
|
name = `${this.addOwnedPrefix(origName)}.${++n}`;
|
|
}
|
|
names.push(this.stripOwnedPrefix(name).toLocaleUpperCase());
|
|
if (input.name !== name) {
|
|
this.renameContextInput(index, name);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
override getSlotMenuOptions(slot: {
|
|
slot: number;
|
|
input?: INodeInputSlot | undefined;
|
|
output?: INodeOutputSlot | undefined;
|
|
}) {
|
|
const editable = this.isOwnedInput(slot.input!.name) && this.type !== "*";
|
|
return [
|
|
{
|
|
content: "✏️ Rename Input",
|
|
disabled: !editable,
|
|
callback: () => {
|
|
var dialog = app.canvas.createDialog(
|
|
"<span class='name'>Name</span><input autofocus type='text'/><button>OK</button>",
|
|
{},
|
|
);
|
|
var dialogInput = dialog.querySelector("input")!;
|
|
if (dialogInput) {
|
|
dialogInput.value = this.stripOwnedPrefix(slot.input!.name || "");
|
|
}
|
|
var inner = () => {
|
|
this.handleContextMenuRenameInputDialog(slot.slot, dialogInput.value);
|
|
dialog.close();
|
|
};
|
|
dialog.querySelector("button")!.addEventListener("click", inner);
|
|
dialogInput.addEventListener("keydown", (e) => {
|
|
dialog.is_modified = true;
|
|
if (e.keyCode == 27) {
|
|
dialog.close();
|
|
} else if (e.keyCode == 13) {
|
|
inner();
|
|
} else if (e.keyCode != 13 && (e.target as HTMLElement)?.localName != "textarea") {
|
|
return;
|
|
}
|
|
e.preventDefault();
|
|
e.stopPropagation();
|
|
});
|
|
dialogInput.focus();
|
|
},
|
|
},
|
|
{
|
|
content: "🗑️ Delete Input",
|
|
disabled: !editable,
|
|
callback: () => {
|
|
this.removeInput(slot.slot);
|
|
},
|
|
},
|
|
];
|
|
}
|
|
|
|
handleContextMenuRenameInputDialog(slotIndex: number, value: string) {
|
|
app.graph.beforeChange();
|
|
this.renameContextInput(slotIndex, value);
|
|
this.stabilizeNames();
|
|
this.setDirtyCanvas(true, true);
|
|
app.graph.afterChange();
|
|
}
|
|
}
|
|
|
|
const contextDynamicNodes = [DynamicContextNode];
|
|
app.registerExtension({
|
|
name: "rgthree.DynamicContext",
|
|
async beforeRegisterNodeDef(nodeType: ComfyNodeConstructor, nodeData: ComfyObjectInfo) {
|
|
if (!CONFIG_SERVICE.getConfigValue("unreleased.dynamic_context.enabled")) {
|
|
return;
|
|
}
|
|
if (nodeData.name === DynamicContextNode.type) {
|
|
DynamicContextNode.setUp(nodeType, nodeData);
|
|
}
|
|
},
|
|
});
|
|
|
