UndefinedUo-ComfyUi-DuoNodes/web/multi_lora_stack_model_only_ortho.js

import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";

async function getAvailableLoras() {
    try {
        const response = await api.fetchApi("/object_info");
        const objectInfo = await response.json();
        for (const [, nodeInfo] of Object.entries(objectInfo)) {
            if (nodeInfo.input?.required?.lora_name) {
                return nodeInfo.input.required.lora_name[0] || ["None"];
            }
        }
        return ["None"];
    } catch (error) {
        console.error("Error fetching LoRAs:", error);
        return ["None"];
    }
}

app.registerExtension({
    name: "MultiLoRAStackModelOnlyOrtho",

    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        if (nodeData.name !== "MultiLoRAStackModelOnlyOrtho") return;

        const onNodeCreated = nodeType.prototype.onNodeCreated;
        nodeType.prototype.onNodeCreated = async function () {
            const result = onNodeCreated?.apply(this, arguments);

            // Hide the raw lora_stack JSON widget synchronously before the first
            // render fires. Must happen before the await below.
            const stackWidget = this.widgets?.find(w => w.name === "lora_stack");
            if (stackWidget) {
                stackWidget.computeSize = () => [0, -4];
                stackWidget.draw = () => {};
            }

            this.size = [450, 340];
            this.loraData = [];
            this.availableLoras = await getAvailableLoras();
            // onConfigure fires during the await above for workflow loads,
            // setting _configureWasRun = true before we reach this point.

            // For fresh nodes: strip the strength_1…10 sockets that Python's
            // forceInput:True declarations auto-create. Workflow loads skip this
            // because configure() has already restored the correct saved sockets.
            if (!this._configureWasRun) {
                for (let i = (this.inputs?.length ?? 0) - 1; i >= 0; i--) {
                    if (/^strength_\d+$/.test(this.inputs[i].name)) this.removeInput(i);
                }
            }

            this.addWidget("button", "➕ Add LoRA", null, () => this.addLoRA());
            this.addWidget("button", "Toggle All LoRAs", null, () => this.toggleAllLoRAs());

            this.addLoRA();
            return result;
        };

        nodeType.prototype.addLoRA = function (initialData = null) {
            const index    = this.loraData.length;
            const on       = initialData?.on       ?? true;
            const lora     = initialData?.lora     ?? "None";
            const strength = initialData?.strength ?? 1.0;

            this.loraData.push({ on, lora, strength });

            const enableWidget = this.addWidget("toggle", `Enable LoRA ${index + 1}`, on, (v) => {
                this.loraData[index].on = v;
                this.updateLoraStack();
            });

            const loraWidget = this.addWidget("combo", `LoRA ${index + 1}`, lora, (v) => {
                this.loraData[index].lora = v;
                this.updateLoraStack();
            }, { values: this.availableLoras });

            const strengthWidget = this.addWidget("number", `Strength ${index + 1}`, strength, (v) => {
                this.loraData[index].strength = v;
                this.updateLoraStack();
            }, { min: -2.0, max: 2.0, step: 0.1 });

            const removeWidget = this.addWidget("button", `Remove LoRA ${index + 1}`, null, () => {
                this.removeLoRA(index);
            });

            this.loraData[index]._widgets = {
                enable: enableWidget, lora: loraWidget,
                strength: strengthWidget, remove: removeWidget,
            };

            // Add the strength input socket inline with the strength widget row.
            // The {widget} extra_info tells LiteGraph/ComfyUI to render the socket
            // next to that widget rather than in the top socket list.
            const slotName = `strength_${index + 1}`;
            const existing = this.inputs?.find(inp => inp.name === slotName);
            if (existing) {
                // Patch widget info onto a socket restored from a saved workflow.
                if (!existing.widget) existing.widget = { name: strengthWidget.name };
            } else {
                this.addInput(slotName, "FLOAT", { widget: { name: strengthWidget.name } });
            }

            this.updateLoraStack();
            this.setSize([this.size[0], this.computeSize()[1]]);
        };

        const onConfigure = nodeType.prototype.onConfigure;
        nodeType.prototype.onConfigure = function (info) {
            // Mark that configure has run so onNodeCreated knows not to strip inputs.
            this._configureWasRun = true;

            const result = onConfigure?.apply(this, arguments);

            // Capture the saved JSON before onNodeCreated's addLoRA() can overwrite it.
            const stackWidget = this.widgets?.find(w => w.name === "lora_stack");
            const savedJson = stackWidget?.value;

            const tryRestore = () => {
                if (!this.availableLoras) { setTimeout(tryRestore, 50); return; }
                if (!savedJson) return;
                let saved;
                try { saved = JSON.parse(savedJson); } catch (e) { return; }
                if (!Array.isArray(saved) || saved.length === 0) return;

                for (const item of this.loraData) {
                    if (item._widgets) {
                        Object.values(item._widgets).forEach(w => {
                            const i = this.widgets.indexOf(w);
                            if (i !== -1) this.widgets.splice(i, 1);
                        });
                    }
                }
                this.loraData = [];

                for (const item of saved) {
                    this.addLoRA(item);
                }
                this.setSize([this.size[0], this.computeSize()[1]]);
            };
            setTimeout(tryRestore, 0);
            return result;
        };

        nodeType.prototype.removeLoRA = function (index) {
            const loraItem = this.loraData[index];
            if (loraItem?._widgets) {
                Object.values(loraItem._widgets).forEach(w => {
                    const i = this.widgets.indexOf(w);
                    if (i !== -1) this.widgets.splice(i, 1);
                });
            }

            const slotName = `strength_${index + 1}`;
            const inputIdx = this.inputs?.findIndex(inp => inp.name === slotName) ?? -1;
            if (inputIdx !== -1) this.removeInput(inputIdx);

            this.loraData.splice(index, 1);
            this._refreshWidgetNames();
            this.updateLoraStack();
            this.setSize([this.size[0], this.computeSize()[1]]);
        };

        nodeType.prototype._refreshWidgetNames = function () {
            this.loraData.forEach((item, i) => {
                if (!item._widgets) return;
                item._widgets.enable.name   = `Enable LoRA ${i + 1}`;
                item._widgets.lora.name     = `LoRA ${i + 1}`;
                item._widgets.strength.name = `Strength ${i + 1}`;
                item._widgets.remove.name   = `Remove LoRA ${i + 1}`;
            });
            let n = 1;
            for (const inp of (this.inputs || [])) {
                if (/^strength_\d+$/.test(inp.name)) {
                    inp.name   = `strength_${n}`;
                    inp.widget = { name: `Strength ${n}` };
                    n++;
                }
            }
        };

        nodeType.prototype.toggleAllLoRAs = function () {
            const allOn = this.loraData.every(l => l.on);
            const next = !allOn;
            this.loraData.forEach((item) => {
                item.on = next;
                if (item._widgets?.enable) item._widgets.enable.value = next;
            });
            this.updateLoraStack();
        };

        nodeType.prototype.updateLoraStack = function () {
            const widget = this.widgets?.find(w => w.name === "lora_stack");
            if (widget) {
                widget.value = JSON.stringify(
                    this.loraData.map(({ on, lora, strength }) => ({ on, lora, strength }))
                );
            }
        };

        const originalComputeSize = nodeType.prototype.computeSize;
        nodeType.prototype.computeSize = function (out) {
            const size = originalComputeSize ? originalComputeSize.call(this, out) : [300, 200];
            size[0] = Math.max(size[0], 450);
            return size;
        };
    },
});