Spaces:
Sleeping
Sleeping
| import { app } from "../../scripts/app.js"; | |
| import { IoDirection, addConnectionLayoutSupport, addMenuItem, matchLocalSlotsToServer, replaceNode, } from "./utils.js"; | |
| import { RgthreeBaseServerNode } from "./base_node.js"; | |
| import { SERVICE as KEY_EVENT_SERVICE } from "./services/key_events_services.js"; | |
| import { debounce, wait } from "../../rgthree/common/shared_utils.js"; | |
| import { removeUnusedInputsFromEnd } from "./utils_inputs_outputs.js"; | |
| import { NodeTypesString } from "./constants.js"; | |
| function findMatchingIndexByTypeOrName(otherNode, otherSlot, ctxSlots) { | |
| const otherNodeType = (otherNode.type || "").toUpperCase(); | |
| const otherNodeName = (otherNode.title || "").toUpperCase(); | |
| let otherSlotType = otherSlot.type; | |
| if (Array.isArray(otherSlotType) || otherSlotType.includes(",")) { | |
| otherSlotType = "COMBO"; | |
| } | |
| const otherSlotName = otherSlot.name.toUpperCase().replace("OPT_", "").replace("_NAME", ""); | |
| let ctxSlotIndex = -1; | |
| if (["CONDITIONING", "INT", "STRING", "FLOAT", "COMBO"].includes(otherSlotType)) { | |
| ctxSlotIndex = ctxSlots.findIndex((ctxSlot) => { | |
| const ctxSlotName = ctxSlot.name.toUpperCase().replace("OPT_", "").replace("_NAME", ""); | |
| let ctxSlotType = ctxSlot.type; | |
| if (Array.isArray(ctxSlotType) || ctxSlotType.includes(",")) { | |
| ctxSlotType = "COMBO"; | |
| } | |
| if (ctxSlotType !== otherSlotType) { | |
| return false; | |
| } | |
| if (ctxSlotName === otherSlotName || | |
| (ctxSlotName === "SEED" && otherSlotName.includes("SEED")) || | |
| (ctxSlotName === "STEP_REFINER" && otherSlotName.includes("AT_STEP")) || | |
| (ctxSlotName === "STEP_REFINER" && otherSlotName.includes("REFINER_STEP"))) { | |
| return true; | |
| } | |
| if ((otherNodeType.includes("POSITIVE") || otherNodeName.includes("POSITIVE")) && | |
| ((ctxSlotName === "POSITIVE" && otherSlotType === "CONDITIONING") || | |
| (ctxSlotName === "TEXT_POS_G" && otherSlotName.includes("TEXT_G")) || | |
| (ctxSlotName === "TEXT_POS_L" && otherSlotName.includes("TEXT_L")))) { | |
| return true; | |
| } | |
| if ((otherNodeType.includes("NEGATIVE") || otherNodeName.includes("NEGATIVE")) && | |
| ((ctxSlotName === "NEGATIVE" && otherSlotType === "CONDITIONING") || | |
| (ctxSlotName === "TEXT_NEG_G" && otherSlotName.includes("TEXT_G")) || | |
| (ctxSlotName === "TEXT_NEG_L" && otherSlotName.includes("TEXT_L")))) { | |
| return true; | |
| } | |
| return false; | |
| }); | |
| } | |
| else { | |
| ctxSlotIndex = ctxSlots.map((s) => s.type).indexOf(otherSlotType); | |
| } | |
| return ctxSlotIndex; | |
| } | |
| export class BaseContextNode extends RgthreeBaseServerNode { | |
| constructor(title) { | |
| super(title); | |
| this.___collapsed_width = 0; | |
| } | |
| get _collapsed_width() { | |
| return this.___collapsed_width; | |
| } | |
| set _collapsed_width(width) { | |
| const canvas = app.canvas; | |
| const ctx = canvas.canvas.getContext("2d"); | |
| const oldFont = ctx.font; | |
| ctx.font = canvas.title_text_font; | |
| let title = this.title.trim(); | |
| this.___collapsed_width = 30 + (title ? 10 + ctx.measureText(title).width : 0); | |
| ctx.font = oldFont; | |
| } | |
| connectByType(slot, sourceNode, sourceSlotType, optsIn) { | |
| let canConnect = super.connectByType && | |
| super.connectByType.call(this, slot, sourceNode, sourceSlotType, optsIn); | |
| if (!super.connectByType) { | |
| canConnect = LGraphNode.prototype.connectByType.call(this, slot, sourceNode, sourceSlotType, optsIn); | |
| } | |
| if (!canConnect && slot === 0) { | |
| const ctrlKey = KEY_EVENT_SERVICE.ctrlKey; | |
| for (const [index, input] of (sourceNode.inputs || []).entries()) { | |
| if (input.link && !ctrlKey) { | |
| continue; | |
| } | |
| const thisOutputSlot = findMatchingIndexByTypeOrName(sourceNode, input, this.outputs); | |
| if (thisOutputSlot > -1) { | |
| this.connect(thisOutputSlot, sourceNode, index); | |
| } | |
| } | |
| } | |
| return null; | |
| } | |
| connectByTypeOutput(slot, sourceNode, sourceSlotType, optsIn) { | |
| var _a; | |
| let canConnect = super.connectByTypeOutput && | |
| super.connectByTypeOutput.call(this, slot, sourceNode, sourceSlotType, optsIn); | |
| if (!super.connectByType) { | |
| canConnect = LGraphNode.prototype.connectByTypeOutput.call(this, slot, sourceNode, sourceSlotType, optsIn); | |
| } | |
| if (!canConnect && slot === 0) { | |
| const ctrlKey = KEY_EVENT_SERVICE.ctrlKey; | |
| for (const [index, output] of (sourceNode.outputs || []).entries()) { | |
| if (((_a = output.links) === null || _a === void 0 ? void 0 : _a.length) && !ctrlKey) { | |
| continue; | |
| } | |
| const thisInputSlot = findMatchingIndexByTypeOrName(sourceNode, output, this.inputs); | |
| if (thisInputSlot > -1) { | |
| sourceNode.connect(index, this, thisInputSlot); | |
| } | |
| } | |
| } | |
| return null; | |
| } | |
| static setUp(comfyClass, nodeData, ctxClass) { | |
| RgthreeBaseServerNode.registerForOverride(comfyClass, nodeData, ctxClass); | |
| wait(500).then(() => { | |
| LiteGraph.slot_types_default_out["RGTHREE_CONTEXT"] = | |
| LiteGraph.slot_types_default_out["RGTHREE_CONTEXT"] || []; | |
| LiteGraph.slot_types_default_out["RGTHREE_CONTEXT"].push(comfyClass.comfyClass); | |
| }); | |
| } | |
| static onRegisteredForOverride(comfyClass, ctxClass) { | |
| addConnectionLayoutSupport(ctxClass, app, [ | |
| ["Left", "Right"], | |
| ["Right", "Left"], | |
| ]); | |
| setTimeout(() => { | |
| ctxClass.category = comfyClass.category; | |
| }); | |
| } | |
| } | |
| class ContextNode extends BaseContextNode { | |
| constructor(title = ContextNode.title) { | |
| super(title); | |
| } | |
| static setUp(comfyClass, nodeData) { | |
| BaseContextNode.setUp(comfyClass, nodeData, ContextNode); | |
| } | |
| static onRegisteredForOverride(comfyClass, ctxClass) { | |
| BaseContextNode.onRegisteredForOverride(comfyClass, ctxClass); | |
| addMenuItem(ContextNode, app, { | |
| name: "Convert To Context Big", | |
| callback: (node) => { | |
| replaceNode(node, ContextBigNode.type); | |
| }, | |
| }); | |
| } | |
| } | |
| ContextNode.title = NodeTypesString.CONTEXT; | |
| ContextNode.type = NodeTypesString.CONTEXT; | |
| ContextNode.comfyClass = NodeTypesString.CONTEXT; | |
| class ContextBigNode extends BaseContextNode { | |
| constructor(title = ContextBigNode.title) { | |
| super(title); | |
| } | |
| static setUp(comfyClass, nodeData) { | |
| BaseContextNode.setUp(comfyClass, nodeData, ContextBigNode); | |
| } | |
| static onRegisteredForOverride(comfyClass, ctxClass) { | |
| BaseContextNode.onRegisteredForOverride(comfyClass, ctxClass); | |
| addMenuItem(ContextBigNode, app, { | |
| name: "Convert To Context (Original)", | |
| callback: (node) => { | |
| replaceNode(node, ContextNode.type); | |
| }, | |
| }); | |
| } | |
| } | |
| ContextBigNode.title = NodeTypesString.CONTEXT_BIG; | |
| ContextBigNode.type = NodeTypesString.CONTEXT_BIG; | |
| ContextBigNode.comfyClass = NodeTypesString.CONTEXT_BIG; | |
| class BaseContextMultiCtxInputNode extends BaseContextNode { | |
| constructor(title) { | |
| super(title); | |
| this.stabilizeBound = this.stabilize.bind(this); | |
| this.addContextInput(5); | |
| } | |
| addContextInput(num = 1) { | |
| for (let i = 0; i < num; i++) { | |
| this.addInput(`ctx_${String(this.inputs.length + 1).padStart(2, "0")}`, "RGTHREE_CONTEXT"); | |
| } | |
| } | |
| onConnectionsChange(type, slotIndex, isConnected, link, ioSlot) { | |
| var _a; | |
| (_a = super.onConnectionsChange) === null || _a === void 0 ? void 0 : _a.apply(this, [...arguments]); | |
| if (type === LiteGraph.INPUT) { | |
| this.scheduleStabilize(); | |
| } | |
| } | |
| scheduleStabilize(ms = 64) { | |
| return debounce(this.stabilizeBound, 64); | |
| } | |
| stabilize() { | |
| removeUnusedInputsFromEnd(this, 4); | |
| this.addContextInput(); | |
| } | |
| } | |
| class ContextSwitchNode extends BaseContextMultiCtxInputNode { | |
| constructor(title = ContextSwitchNode.title) { | |
| super(title); | |
| } | |
| static setUp(comfyClass, nodeData) { | |
| BaseContextNode.setUp(comfyClass, nodeData, ContextSwitchNode); | |
| } | |
| static onRegisteredForOverride(comfyClass, ctxClass) { | |
| BaseContextNode.onRegisteredForOverride(comfyClass, ctxClass); | |
| addMenuItem(ContextSwitchNode, app, { | |
| name: "Convert To Context Switch Big", | |
| callback: (node) => { | |
| replaceNode(node, ContextSwitchBigNode.type); | |
| }, | |
| }); | |
| } | |
| } | |
| ContextSwitchNode.title = NodeTypesString.CONTEXT_SWITCH; | |
| ContextSwitchNode.type = NodeTypesString.CONTEXT_SWITCH; | |
| ContextSwitchNode.comfyClass = NodeTypesString.CONTEXT_SWITCH; | |
| class ContextSwitchBigNode extends BaseContextMultiCtxInputNode { | |
| constructor(title = ContextSwitchBigNode.title) { | |
| super(title); | |
| } | |
| static setUp(comfyClass, nodeData) { | |
| BaseContextNode.setUp(comfyClass, nodeData, ContextSwitchBigNode); | |
| } | |
| static onRegisteredForOverride(comfyClass, ctxClass) { | |
| BaseContextNode.onRegisteredForOverride(comfyClass, ctxClass); | |
| addMenuItem(ContextSwitchBigNode, app, { | |
| name: "Convert To Context Switch", | |
| callback: (node) => { | |
| replaceNode(node, ContextSwitchNode.type); | |
| }, | |
| }); | |
| } | |
| } | |
| ContextSwitchBigNode.title = NodeTypesString.CONTEXT_SWITCH_BIG; | |
| ContextSwitchBigNode.type = NodeTypesString.CONTEXT_SWITCH_BIG; | |
| ContextSwitchBigNode.comfyClass = NodeTypesString.CONTEXT_SWITCH_BIG; | |
| class ContextMergeNode extends BaseContextMultiCtxInputNode { | |
| constructor(title = ContextMergeNode.title) { | |
| super(title); | |
| } | |
| static setUp(comfyClass, nodeData) { | |
| BaseContextNode.setUp(comfyClass, nodeData, ContextMergeNode); | |
| } | |
| static onRegisteredForOverride(comfyClass, ctxClass) { | |
| BaseContextNode.onRegisteredForOverride(comfyClass, ctxClass); | |
| addMenuItem(ContextMergeNode, app, { | |
| name: "Convert To Context Merge Big", | |
| callback: (node) => { | |
| replaceNode(node, ContextMergeBigNode.type); | |
| }, | |
| }); | |
| } | |
| } | |
| ContextMergeNode.title = NodeTypesString.CONTEXT_MERGE; | |
| ContextMergeNode.type = NodeTypesString.CONTEXT_MERGE; | |
| ContextMergeNode.comfyClass = NodeTypesString.CONTEXT_MERGE; | |
| class ContextMergeBigNode extends BaseContextMultiCtxInputNode { | |
| constructor(title = ContextMergeBigNode.title) { | |
| super(title); | |
| } | |
| static setUp(comfyClass, nodeData) { | |
| BaseContextNode.setUp(comfyClass, nodeData, ContextMergeBigNode); | |
| } | |
| static onRegisteredForOverride(comfyClass, ctxClass) { | |
| BaseContextNode.onRegisteredForOverride(comfyClass, ctxClass); | |
| addMenuItem(ContextMergeBigNode, app, { | |
| name: "Convert To Context Switch", | |
| callback: (node) => { | |
| replaceNode(node, ContextMergeNode.type); | |
| }, | |
| }); | |
| } | |
| } | |
| ContextMergeBigNode.title = NodeTypesString.CONTEXT_MERGE_BIG; | |
| ContextMergeBigNode.type = NodeTypesString.CONTEXT_MERGE_BIG; | |
| ContextMergeBigNode.comfyClass = NodeTypesString.CONTEXT_MERGE_BIG; | |
| const contextNodes = [ | |
| ContextNode, | |
| ContextBigNode, | |
| ContextSwitchNode, | |
| ContextSwitchBigNode, | |
| ContextMergeNode, | |
| ContextMergeBigNode, | |
| ]; | |
| const contextTypeToServerDef = {}; | |
| function fixBadConfigs(node) { | |
| const wrongName = node.outputs.find((o, i) => o.name === "CLIP_HEIGTH"); | |
| if (wrongName) { | |
| wrongName.name = "CLIP_HEIGHT"; | |
| } | |
| } | |
| app.registerExtension({ | |
| name: "rgthree.Context", | |
| async beforeRegisterNodeDef(nodeType, nodeData) { | |
| for (const ctxClass of contextNodes) { | |
| if (nodeData.name === ctxClass.type) { | |
| contextTypeToServerDef[ctxClass.type] = nodeData; | |
| ctxClass.setUp(nodeType, nodeData); | |
| break; | |
| } | |
| } | |
| }, | |
| async nodeCreated(node) { | |
| const type = node.type || node.constructor.type; | |
| const serverDef = type && contextTypeToServerDef[type]; | |
| if (serverDef) { | |
| fixBadConfigs(node); | |
| matchLocalSlotsToServer(node, IoDirection.OUTPUT, serverDef); | |
| if (!type.includes("Switch") && !type.includes("Merge")) { | |
| matchLocalSlotsToServer(node, IoDirection.INPUT, serverDef); | |
| } | |
| } | |
| }, | |
| async loadedGraphNode(node) { | |
| const type = node.type || node.constructor.type; | |
| const serverDef = type && contextTypeToServerDef[type]; | |
| if (serverDef) { | |
| fixBadConfigs(node); | |
| matchLocalSlotsToServer(node, IoDirection.OUTPUT, serverDef); | |
| if (!type.includes("Switch") && !type.includes("Merge")) { | |
| matchLocalSlotsToServer(node, IoDirection.INPUT, serverDef); | |
| } | |
| } | |
| }, | |
| }); | |