85 lines
2.5 KiB
JavaScript
85 lines
2.5 KiB
JavaScript
import { app } from "/scripts/app.js";
|
||
|
||
const DEFAULT_SLOTS = 2;
|
||
|
||
function serializeSlots(node) {
|
||
const entries = (node._loraSlots ?? []).map(s => ({
|
||
name: s.nameW.value,
|
||
strength_model: s.smW.value,
|
||
enabled: s.enW.value,
|
||
}));
|
||
const jw = node.widgets?.find(w => w.name === "loras_json");
|
||
if (jw) jw.value = JSON.stringify(entries);
|
||
}
|
||
|
||
function addSlot(node, def = {}) {
|
||
const loras = node._loraNames ?? ["None"];
|
||
|
||
const nameW = node.addWidget("combo", "lora", def.name ?? "None",
|
||
() => serializeSlots(node), { values: loras });
|
||
|
||
const smW = node.addWidget("number", "strength", def.strength_model ?? 1.0,
|
||
() => serializeSlots(node), { min: 0.0, max: 2.0, step: 0.1, precision: 2 });
|
||
|
||
const enW = node.addWidget("toggle", "enabled", def.enabled ?? true,
|
||
() => serializeSlots(node));
|
||
|
||
node._loraSlots = node._loraSlots ?? [];
|
||
node._loraSlots.push({ nameW, smW, enW });
|
||
|
||
serializeSlots(node);
|
||
node.setSize([node.size[0], node.computeSize()[1]]);
|
||
node.graph?.setDirtyCanvas(true, true);
|
||
}
|
||
|
||
function removeSlot(node) {
|
||
if (!node._loraSlots?.length) return;
|
||
const s = node._loraSlots.pop();
|
||
for (const w of [s.nameW, s.smW, s.enW]) {
|
||
const i = node.widgets.indexOf(w);
|
||
if (i !== -1) node.widgets.splice(i, 1);
|
||
}
|
||
serializeSlots(node);
|
||
node.setSize([node.size[0], node.computeSize()[1]]);
|
||
node.graph?.setDirtyCanvas(true, true);
|
||
}
|
||
|
||
function init(node) {
|
||
if (node._atomSimpleInit) return;
|
||
node._atomSimpleInit = true;
|
||
node._loraSlots = [];
|
||
|
||
const jw = node.widgets?.find(w => w.name === "loras_json");
|
||
if (jw) {
|
||
jw.hidden = true;
|
||
jw.computeSize = () => [0, -4];
|
||
}
|
||
|
||
node.addWidget("button", "+ Add LoRA", null, () => addSlot(node));
|
||
node.addWidget("button", "− Remove LoRA", null, () => removeSlot(node));
|
||
|
||
let restored = false;
|
||
if (jw?.value && jw.value !== "[]") {
|
||
try {
|
||
const entries = JSON.parse(jw.value);
|
||
if (entries.length > 0) {
|
||
entries.forEach(e => addSlot(node, e));
|
||
restored = true;
|
||
}
|
||
} catch {}
|
||
}
|
||
if (!restored) {
|
||
for (let i = 0; i < DEFAULT_SLOTS; i++) addSlot(node);
|
||
}
|
||
}
|
||
|
||
app.registerExtension({
|
||
name: "Atom_AI.SimpleLoRANode",
|
||
|
||
nodeCreated(node) {
|
||
if (node.comfyClass !== "AtomSimpleLoRANode") return;
|
||
node._loraNames = app._atomLoraNames ?? ["None"];
|
||
init(node);
|
||
},
|
||
});
|