139 lines
4.4 KiB
JavaScript
139 lines
4.4 KiB
JavaScript
import { app } from "/scripts/app.js";
|
||
|
||
const DEFAULT_SLOTS = 1;
|
||
|
||
function serializeSlots(node) {
|
||
const entries = (node._loraSlots ?? []).map(s => ({
|
||
name: s.nameW.value,
|
||
strength_model: s.smW.value,
|
||
strength_clip: s.scW.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", "model str.", def.strength_model ?? 1.0,
|
||
() => serializeSlots(node), { min: -10, max: 10, step: 0.1, precision: 2 });
|
||
|
||
const scW = node.addWidget("number", "clip str.", def.strength_clip ?? 1.0,
|
||
() => serializeSlots(node), { min: -10, max: 10, 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, scW, 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.scW, 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._atomDynInit) return;
|
||
node._atomDynInit = true;
|
||
node._loraSlots = [];
|
||
|
||
// Hide loras_json
|
||
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));
|
||
|
||
// Restore from saved JSON or add defaults
|
||
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);
|
||
}
|
||
}
|
||
|
||
// Fetch lora list directly from ComfyUI's /object_info endpoint
|
||
async function fetchLoraNames() {
|
||
try {
|
||
const r = await fetch("/object_info");
|
||
const all = await r.json();
|
||
|
||
// Try every known lora loader node name
|
||
const candidates = [
|
||
"LoraLoader",
|
||
"LoRALoader",
|
||
"Load LoRA",
|
||
"LoraLoaderModelOnly",
|
||
];
|
||
|
||
for (const name of candidates) {
|
||
const def = all[name];
|
||
if (!def) continue;
|
||
// Search all required inputs for a list that looks like lora filenames
|
||
for (const [, val] of Object.entries(def.input?.required ?? {})) {
|
||
if (Array.isArray(val[0]) && val[0].some(v =>
|
||
typeof v === "string" && (v.endsWith(".safetensors") || v.endsWith(".pt") || v.endsWith(".ckpt"))
|
||
)) {
|
||
return ["None", ...val[0]];
|
||
}
|
||
}
|
||
}
|
||
|
||
// Last resort: scan ALL node defs for a safetensors list
|
||
for (const def of Object.values(all)) {
|
||
for (const [, val] of Object.entries(def.input?.required ?? {})) {
|
||
if (Array.isArray(val[0]) && val[0].length > 0 &&
|
||
val[0].some(v => typeof v === "string" && v.includes("lora"))) {
|
||
return ["None", ...val[0]];
|
||
}
|
||
}
|
||
}
|
||
} catch (e) {
|
||
console.warn("[Atom_AI] Could not fetch lora list:", e);
|
||
}
|
||
return ["None"];
|
||
}
|
||
|
||
app.registerExtension({
|
||
name: "Atom_AI.DynamicLoRANode",
|
||
|
||
async setup() {
|
||
// Fetch and cache lora list at startup
|
||
app._atomLoraNames = await fetchLoraNames();
|
||
console.log("[Atom_AI] LoRA list loaded:", app._atomLoraNames.length, "entries");
|
||
},
|
||
|
||
nodeCreated(node) {
|
||
if (node.comfyClass !== "AtomDynamicLoRANode") return;
|
||
node._loraNames = app._atomLoraNames ?? ["None"];
|
||
init(node);
|
||
},
|
||
});
|