import { app } from "../../scripts/app.js"; import { IoDirection, followConnectionUntilType, getConnectedInputInfosAndFilterPassThroughs, } from "./utils.js"; import { rgthree } from "./rgthree.js"; import { SERVICE as CONTEXT_SERVICE, InputMutationOperation, } from "./services/context_service.js"; import { NodeTypesString } from "./constants.js"; import { removeUnusedInputsFromEnd } from "./utils_inputs_outputs.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 { constructor(title = DynamicContextNode.title) { super(title); } onNodeCreated() { this.addInput("base_ctx", "RGTHREE_DYNAMIC_CONTEXT"); this.ensureOneRemainingNewInputSlot(); super.onNodeCreated(); } onConnectionsChange(type, slotIndex, isConnected, link, ioSlot) { var _a; (_a = super.onConnectionsChange) === null || _a === void 0 ? void 0 : _a.call(this, type, slotIndex, isConnected, link, ioSlot); if (this.configuring) { return; } if (type === LiteGraph.INPUT) { if (isConnected) { this.handleInputConnected(slotIndex); } else { this.handleInputDisconnected(slotIndex); } } } onConnectInput(inputIndex, outputType, outputSlot, outputNode, outputIndex) { var _a; let canConnect = true; if (super.onConnectInput) { canConnect = super.onConnectInput.apply(this, [...arguments]); } 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."); (_a = console[n]) === null || _a === void 0 ? void 0 : _a.call(console, ...v); canConnect = false; } return canConnect; } handleInputConnected(slotIndex) { 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]; if (baseNodesDynamicCtx === null || baseNodesDynamicCtx === void 0 ? void 0 : 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) { const ioSlot = this.inputs[slotIndex]; return ioSlot && ioSlot.name === "+" && ioSlot.type === "*"; } handleNewInputConnected(slotIndex) { 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 === null || cxn === void 0 ? void 0 : cxn.type) && (cxn === null || cxn === void 0 ? void 0 : 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; ioSlot.removable = true; while (!this.outputs[slotIndex]) { this.addOutput("*", "*"); } this.outputs[slotIndex].type = cxn.type; this.outputs[slotIndex].name = this.stripOwnedPrefix(name).toLocaleUpperCase(); if (cxn.type === "COMBO" || cxn.type.includes(",") || Array.isArray(cxn.type)) { this.outputs[slotIndex].widget = true; } this.inputsMutated({ operation: InputMutationOperation.ADDED, node: this, slotIndex, slot: ioSlot, }); this.stabilizeNames(); this.ensureOneRemainingNewInputSlot(); } } handleInputDisconnected(slotIndex) { var _a, _b; 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 || ((_b = (_a = this.outputs[index]) === null || _a === void 0 ? void 0 : _a.links) === null || _b === void 0 ? void 0 : _b.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) { 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; } removeInput(slotIndex) { 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 = []; 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); } } } } getSlotMenuOptions(slot) { const editable = this.isOwnedInput(slot.input.name) && this.type !== "*"; return [ { content: "✏️ Rename Input", disabled: !editable, callback: () => { var dialog = app.canvas.createDialog("Name", {}); 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) => { var _a; dialog.is_modified = true; if (e.keyCode == 27) { dialog.close(); } else if (e.keyCode == 13) { inner(); } else if (e.keyCode != 13 && ((_a = e.target) === null || _a === void 0 ? void 0 : _a.localName) != "textarea") { return; } e.preventDefault(); e.stopPropagation(); }); dialogInput.focus(); }, }, { content: "🗑️ Delete Input", disabled: !editable, callback: () => { this.removeInput(slot.slot); }, }, ]; } handleContextMenuRenameInputDialog(slotIndex, value) { app.graph.beforeChange(); this.renameContextInput(slotIndex, value); this.stabilizeNames(); this.setDirtyCanvas(true, true); app.graph.afterChange(); } } DynamicContextNode.title = NodeTypesString.DYNAMIC_CONTEXT; DynamicContextNode.type = NodeTypesString.DYNAMIC_CONTEXT; DynamicContextNode.comfyClass = NodeTypesString.DYNAMIC_CONTEXT; const contextDynamicNodes = [DynamicContextNode]; app.registerExtension({ name: "rgthree.DynamicContext", async beforeRegisterNodeDef(nodeType, nodeData) { if (!CONFIG_SERVICE.getConfigValue("unreleased.dynamic_context.enabled")) { return; } if (nodeData.name === DynamicContextNode.type) { DynamicContextNode.setUp(nodeType, nodeData); } }, });