Keep the alibaba-pai file connection and note that the stack row UI matured from rgthree-comfy's Power Lora Loader.
781 lines
27 KiB
JavaScript
781 lines
27 KiB
JavaScript
// Stack row UI matured from rgthree-comfy's Power Lora Loader.
|
||
import { app } from "../../scripts/app.js";
|
||
import { api } from "../../scripts/api.js";
|
||
|
||
const NODE_TYPE = "H3PowerLoraStack";
|
||
const ROW_HEIGHT = 22;
|
||
const MARGIN = 12;
|
||
|
||
const BALANCE_ROUTE = "/h3_power_lora_stack/balance";
|
||
const BALANCE_LABEL = "⚖ Auto-balance strengths";
|
||
const RESTORE_LABEL = "↺ Restore manual strengths";
|
||
|
||
// Canonical row list. ComfyUI's configure() runs migrateWidgetsValues before
|
||
// onConfigure and drops widgets_values slots that line up with serialize:false
|
||
// buttons, so the stack is snapshotted here and rebuilt after that migrator.
|
||
const ROWS_KEY = "h3Rows";
|
||
|
||
/**
|
||
* The stack node declares no ``lora_name`` input of its own -- rows are added
|
||
* in the browser -- so the current folder listing is borrowed from a node that
|
||
* does each time the picker opens.
|
||
*/
|
||
function fetchLoraList() {
|
||
return api
|
||
.fetchApi("/object_info/LoraLoaderModelOnly")
|
||
.then((res) => res.json())
|
||
.then(
|
||
(info) => info?.LoraLoaderModelOnly?.input?.required?.lora_name?.[0] ?? []
|
||
)
|
||
.catch((err) => {
|
||
console.warn("[H3PowerLoraStack] could not fetch lora list", err);
|
||
return [];
|
||
});
|
||
}
|
||
|
||
/**
|
||
* Ask the server what each LoRA actually does to the weights.
|
||
*
|
||
* "Strength 1.0" is not a unit -- across the local H3 collection the
|
||
* perturbation it produces spans 65x -- so the server measures each file and
|
||
* returns the multiplier that puts it on the same scale as the rest. Results
|
||
* are cached server-side on (path, mtime, size), so pressing the button again
|
||
* costs one round trip and no disk.
|
||
*/
|
||
async function fetchBalance(names) {
|
||
const res = await api.fetchApi(BALANCE_ROUTE, {
|
||
method: "POST",
|
||
headers: { "Content-Type": "application/json" },
|
||
body: JSON.stringify({ loras: names }),
|
||
});
|
||
if (!res.ok) throw new Error(`balance request failed (${res.status})`);
|
||
return res.json();
|
||
}
|
||
|
||
function drawToggle(ctx, x, y, size, on) {
|
||
ctx.save();
|
||
ctx.beginPath();
|
||
ctx.roundRect(x, y + 3, size * 1.7, size - 6, (size - 6) / 2);
|
||
ctx.fillStyle = on ? "#3d8c5a" : "#3a3a3a";
|
||
ctx.fill();
|
||
ctx.beginPath();
|
||
const knob = (size - 8) / 2;
|
||
ctx.arc(x + (on ? size * 1.7 - knob - 3 : knob + 3), y + size / 2, knob, 0, Math.PI * 2);
|
||
ctx.fillStyle = on ? "#d8f5e4" : "#8a8a8a";
|
||
ctx.fill();
|
||
ctx.restore();
|
||
}
|
||
|
||
function shortName(name) {
|
||
if (!name || name === "None") return "click to choose";
|
||
const base = String(name).replace(/\\/g, "/").split("/").pop();
|
||
return base.replace(/\.(safetensors|ckpt|pt|bin)$/i, "");
|
||
}
|
||
|
||
/**
|
||
* One stack row: enable toggle, lora name, strength, remove button.
|
||
* Values serialize as {on, lora, strength}, which is what the Python side's
|
||
* flexible `lora_N` inputs expect.
|
||
*/
|
||
function makeLoraWidget(node, name, value) {
|
||
const widget = {
|
||
name,
|
||
type: "H3_LORA",
|
||
value: Object.assign({ on: true, lora: "None", strength: 1.0 }, value || {}),
|
||
// Widget-level flag, not only options.serialize. The frontend skips a row
|
||
// when `serialize === false`; undefined is kept, but be explicit.
|
||
serialize: true,
|
||
options: { serialize: true },
|
||
hitAreas: {},
|
||
|
||
computeSize() {
|
||
return [node.size[0], ROW_HEIGHT];
|
||
},
|
||
|
||
draw(ctx, node_, widgetWidth, y, height) {
|
||
const margin = MARGIN;
|
||
const left = margin;
|
||
const right = widgetWidth - margin;
|
||
const midY = y + height / 2;
|
||
ctx.save();
|
||
|
||
ctx.beginPath();
|
||
ctx.roundRect(left, y + 1, right - left, height - 2, 6);
|
||
ctx.fillStyle = this.value.on ? "#2b2b2b" : "#232323";
|
||
ctx.fill();
|
||
|
||
let cursor = left + 8;
|
||
drawToggle(ctx, cursor, y, height, this.value.on);
|
||
this.hitAreas.toggle = [cursor, cursor + height * 1.7];
|
||
cursor += height * 1.7 + 10;
|
||
|
||
// remove button, laid out from the right edge inward
|
||
const removeX = right - 20;
|
||
ctx.fillStyle = "#8a8a8a";
|
||
ctx.font = "14px sans-serif";
|
||
ctx.textAlign = "center";
|
||
ctx.textBaseline = "middle";
|
||
ctx.fillText("✕", removeX + 6, midY);
|
||
this.hitAreas.remove = [removeX, removeX + 16];
|
||
|
||
const strengthRight = removeX - 10;
|
||
const strengthLeft = strengthRight - 74;
|
||
ctx.fillStyle = "#1e1e1e";
|
||
ctx.beginPath();
|
||
ctx.roundRect(strengthLeft, y + 3, strengthRight - strengthLeft, height - 6, 4);
|
||
ctx.fill();
|
||
ctx.fillStyle = "#b8b8b8";
|
||
ctx.font = "11px sans-serif";
|
||
ctx.fillText("◀", strengthLeft + 9, midY);
|
||
ctx.fillText("▶", strengthRight - 9, midY);
|
||
ctx.fillStyle = this.value.on ? "#e8e8e8" : "#777";
|
||
ctx.font = "12px sans-serif";
|
||
ctx.fillText(
|
||
Number(this.value.strength).toFixed(2),
|
||
(strengthLeft + strengthRight) / 2,
|
||
midY
|
||
);
|
||
this.hitAreas.strengthDown = [strengthLeft, strengthLeft + 18];
|
||
this.hitAreas.strengthUp = [strengthRight - 18, strengthRight];
|
||
this.hitAreas.strengthValue = [strengthLeft + 18, strengthRight - 18];
|
||
|
||
// auto-balance badge: what the measurement did to this row. Amber flags
|
||
// a note from the server (a duplicate adapter, or an unreadable file).
|
||
let badgeLeft = strengthLeft - 10;
|
||
if (this.value.autoApplied && this.value.factor !== undefined) {
|
||
const badge = `⚖×${Number(this.value.factor).toFixed(2)}`;
|
||
ctx.font = "10px sans-serif";
|
||
ctx.textAlign = "right";
|
||
ctx.fillStyle = this.value.note
|
||
? "#d99a4e"
|
||
: this.value.factor < 0.995
|
||
? "#7fd6a0"
|
||
: "#6a6a6a";
|
||
ctx.fillText(badge, strengthLeft - 8, midY);
|
||
badgeLeft = strengthLeft - 14 - ctx.measureText(badge).width;
|
||
}
|
||
|
||
const nameLeft = cursor;
|
||
const nameRight = badgeLeft;
|
||
ctx.textAlign = "left";
|
||
ctx.fillStyle = this.value.on ? "#e8e8e8" : "#777";
|
||
ctx.font = "12px sans-serif";
|
||
const label = shortName(this.value.lora);
|
||
let text = label;
|
||
const maxWidth = nameRight - nameLeft;
|
||
if (ctx.measureText(text).width > maxWidth) {
|
||
while (text.length > 4 && ctx.measureText(text + "…").width > maxWidth) {
|
||
text = text.slice(0, -1);
|
||
}
|
||
text += "…";
|
||
}
|
||
ctx.fillText(text, nameLeft, midY);
|
||
this.hitAreas.name = [nameLeft, nameRight];
|
||
|
||
ctx.restore();
|
||
},
|
||
|
||
mouse(event, pos, node_) {
|
||
if (event.type !== "pointerdown") return false;
|
||
const x = pos[0];
|
||
const hit = (area) => area && x >= area[0] && x <= area[1];
|
||
|
||
if (hit(this.hitAreas.toggle)) {
|
||
this.value = { ...this.value, on: !this.value.on };
|
||
node_.setDirtyCanvas(true, true);
|
||
markChanged(node_);
|
||
// a row switched on inside a balanced stack has not been measured yet
|
||
if (this.value.on && isBalanced(node_)) applyBalance(node_);
|
||
return true;
|
||
}
|
||
if (hit(this.hitAreas.remove)) {
|
||
removeLoraWidget(node_, this);
|
||
return true;
|
||
}
|
||
if (hit(this.hitAreas.strengthDown)) {
|
||
this.setStrength(node_, round2(this.value.strength - 0.05));
|
||
return true;
|
||
}
|
||
if (hit(this.hitAreas.strengthUp)) {
|
||
this.setStrength(node_, round2(this.value.strength + 0.05));
|
||
return true;
|
||
}
|
||
if (hit(this.hitAreas.strengthValue)) {
|
||
app.canvas.prompt(
|
||
"Strength",
|
||
this.value.strength,
|
||
(v) => {
|
||
const parsed = parseFloat(v);
|
||
if (!Number.isNaN(parsed)) this.setStrength(node_, parsed);
|
||
},
|
||
event
|
||
);
|
||
return true;
|
||
}
|
||
if (hit(this.hitAreas.name)) {
|
||
showLoraMenu(node_, this, event);
|
||
return true;
|
||
}
|
||
return false;
|
||
},
|
||
|
||
/**
|
||
* Editing a strength by hand overrides the balance for this row: the badge
|
||
* clears and a later recompute leaves it alone, but the stashed manual
|
||
* value stays put so Restore still returns to what was there before.
|
||
*/
|
||
setStrength(node_, value) {
|
||
this.value = { ...this.value, strength: value, autoApplied: false };
|
||
node_.setDirtyCanvas(true, true);
|
||
markChanged(node_);
|
||
},
|
||
|
||
serializeValue() {
|
||
return { ...this.value };
|
||
},
|
||
};
|
||
return widget;
|
||
}
|
||
|
||
function round2(v) {
|
||
return Math.round(v * 100) / 100;
|
||
}
|
||
|
||
/* ---------------------------------------------------------------- picker -- */
|
||
|
||
const PICKER_CSS = `
|
||
.h3pls-picker{position:fixed;z-index:10000;display:flex;flex-direction:column;
|
||
width:460px;max-width:92vw;background:#1e1e1e;border:1px solid #4a4a4a;
|
||
border-radius:8px;box-shadow:0 12px 32px rgba(0,0,0,.6);overflow:hidden;
|
||
font-family:system-ui,-apple-system,sans-serif;}
|
||
.h3pls-picker input{margin:8px;padding:7px 9px;background:#111;color:#eee;
|
||
border:1px solid #555;border-radius:5px;font-size:13px;outline:none;}
|
||
.h3pls-picker input:focus{border-color:#3d8c5a;}
|
||
.h3pls-count{padding:0 12px 6px;font-size:11px;color:#777;}
|
||
.h3pls-list{overflow-y:auto;max-height:46vh;}
|
||
.h3pls-item{padding:6px 12px;font-size:12.5px;color:#ddd;cursor:pointer;
|
||
white-space:nowrap;overflow:hidden;text-overflow:ellipsis;}
|
||
.h3pls-item:hover{background:#2c2c2c;}
|
||
.h3pls-item.sel{background:#31543f;}
|
||
.h3pls-item .dir{color:#7a7a7a;}
|
||
.h3pls-item .hit{color:#7fd6a0;font-weight:600;}
|
||
.h3pls-item.cur::after{content:"current";float:right;color:#666;font-size:10px;}
|
||
.h3pls-empty{padding:14px 12px;font-size:12px;color:#888;}
|
||
`;
|
||
|
||
function ensurePickerCss() {
|
||
if (document.getElementById("h3pls-css")) return;
|
||
const style = document.createElement("style");
|
||
style.id = "h3pls-css";
|
||
style.textContent = PICKER_CSS;
|
||
document.head.appendChild(style);
|
||
}
|
||
|
||
/** Case-insensitive AND-match of every whitespace-separated term. */
|
||
function matchTerms(name, terms) {
|
||
const hay = name.toLowerCase();
|
||
const spans = [];
|
||
for (const term of terms) {
|
||
const at = hay.indexOf(term);
|
||
if (at < 0) return null;
|
||
spans.push([at, at + term.length]);
|
||
}
|
||
return spans;
|
||
}
|
||
|
||
/** Basename hits and early hits rank first; then shortest name. */
|
||
function score(name, terms, spans) {
|
||
if (!terms.length) return 0;
|
||
const slash = name.lastIndexOf("/") + 1;
|
||
let s = 0;
|
||
for (const [start] of spans) {
|
||
if (start >= slash) s -= 1000;
|
||
s += start;
|
||
}
|
||
return s + name.length * 0.01;
|
||
}
|
||
|
||
function highlight(name, spans) {
|
||
const slash = name.lastIndexOf("/") + 1;
|
||
const merged = [];
|
||
for (const [a, b] of [...spans].sort((p, q) => p[0] - q[0])) {
|
||
const last = merged[merged.length - 1];
|
||
if (last && a <= last[1]) last[1] = Math.max(last[1], b);
|
||
else merged.push([a, b]);
|
||
}
|
||
const esc = (t) =>
|
||
t.replace(/&/g, "&").replace(/</g, "<").replace(/>/g, ">");
|
||
let html = "";
|
||
let at = 0;
|
||
for (const [a, b] of merged) {
|
||
html += wrapDir(esc(name.slice(at, a)), at, slash);
|
||
html += `<span class="hit">${esc(name.slice(a, b))}</span>`;
|
||
at = b;
|
||
}
|
||
html += wrapDir(esc(name.slice(at)), at, slash);
|
||
return html;
|
||
}
|
||
|
||
function wrapDir(text, offset, slash) {
|
||
if (!text) return "";
|
||
if (offset + text.length <= slash) return `<span class="dir">${text}</span>`;
|
||
if (offset >= slash) return text;
|
||
const cut = slash - offset;
|
||
return `<span class="dir">${text.slice(0, cut)}</span>${text.slice(cut)}`;
|
||
}
|
||
|
||
/**
|
||
* Filterable LoRA picker. LiteGraph's ContextMenu has no search, which is
|
||
* unusable once the loras folder holds more than a screenful, so this is a
|
||
* self-contained overlay: type to narrow, arrows to move, Enter to take it.
|
||
*/
|
||
async function showLoraMenu(node, widget, event, { removeOnCancel = false } = {}) {
|
||
const names = await fetchLoraList();
|
||
if (!names.length) {
|
||
if (removeOnCancel) removeLoraWidget(node, widget);
|
||
alert("No LoRAs found in models/loras.");
|
||
return;
|
||
}
|
||
ensurePickerCss();
|
||
|
||
const current = widget.value.lora;
|
||
const root = document.createElement("div");
|
||
root.className = "h3pls-picker";
|
||
const input = document.createElement("input");
|
||
input.type = "text";
|
||
input.placeholder = "Type to filter… (space-separated terms all must match)";
|
||
input.spellcheck = false;
|
||
const count = document.createElement("div");
|
||
count.className = "h3pls-count";
|
||
const list = document.createElement("div");
|
||
list.className = "h3pls-list";
|
||
root.append(input, count, list);
|
||
document.body.appendChild(root);
|
||
|
||
// place near the click, clamped into the viewport
|
||
const px = event?.clientX ?? window.innerWidth / 2;
|
||
const py = event?.clientY ?? window.innerHeight / 2;
|
||
root.style.left = `${Math.min(px, window.innerWidth - 470)}px`;
|
||
root.style.top = `${Math.min(py, window.innerHeight - 320)}px`;
|
||
|
||
let rows = [];
|
||
let sel = 0;
|
||
|
||
let committed = false;
|
||
const close = () => {
|
||
document.removeEventListener("pointerdown", onOutside, true);
|
||
root.remove();
|
||
// a row opened straight from "Add LoRA" and then dismissed was never
|
||
// wanted, so take it back out rather than leaving an empty slot behind
|
||
if (removeOnCancel && !committed) removeLoraWidget(node, widget);
|
||
};
|
||
const commit = (value) => {
|
||
const wasBalanced = isBalanced(node);
|
||
committed = true;
|
||
const next = { ...widget.value, lora: value };
|
||
// the balance factor belonged to the file that was here before, so drop the
|
||
// whole measurement and hand the row back its manual strength; clearing the
|
||
// stash too is what marks it for re-measuring
|
||
if (next.manual !== undefined || next.autoApplied) {
|
||
next.strength = next.manual !== undefined ? next.manual : next.strength;
|
||
delete next.manual;
|
||
delete next.factor;
|
||
delete next.rel;
|
||
delete next.note;
|
||
delete next.autoApplied;
|
||
}
|
||
widget.value = next;
|
||
node.setDirtyCanvas(true, true);
|
||
markChanged(node);
|
||
close();
|
||
if (wasBalanced) applyBalance(node);
|
||
};
|
||
const onOutside = (e) => {
|
||
if (!root.contains(e.target)) close();
|
||
};
|
||
|
||
const render = () => {
|
||
const query = input.value.trim().toLowerCase();
|
||
const terms = query ? query.split(/\s+/) : [];
|
||
rows = [];
|
||
for (const name of names) {
|
||
const spans = matchTerms(name, terms);
|
||
if (spans === null) continue;
|
||
rows.push({ name, spans, s: score(name, terms, spans) });
|
||
}
|
||
if (terms.length) rows.sort((a, b) => a.s - b.s);
|
||
sel = 0;
|
||
if (!terms.length) {
|
||
const at = rows.findIndex((r) => r.name === current);
|
||
if (at > 0) sel = at;
|
||
}
|
||
|
||
count.textContent = `${rows.length} of ${names.length}`;
|
||
list.replaceChildren();
|
||
if (!rows.length) {
|
||
const empty = document.createElement("div");
|
||
empty.className = "h3pls-empty";
|
||
empty.textContent = "No LoRA matches that filter.";
|
||
list.appendChild(empty);
|
||
return;
|
||
}
|
||
rows.forEach((row, i) => {
|
||
const item = document.createElement("div");
|
||
item.className =
|
||
"h3pls-item" + (i === sel ? " sel" : "") + (row.name === current ? " cur" : "");
|
||
item.innerHTML = highlight(row.name, row.spans);
|
||
item.title = row.name;
|
||
item.addEventListener("pointerdown", (e) => {
|
||
e.preventDefault();
|
||
e.stopPropagation();
|
||
commit(row.name);
|
||
});
|
||
list.appendChild(item);
|
||
});
|
||
scrollToSel();
|
||
};
|
||
|
||
const scrollToSel = () => {
|
||
const el = list.children[sel];
|
||
if (el && el.scrollIntoView) el.scrollIntoView({ block: "nearest" });
|
||
};
|
||
|
||
const move = (delta) => {
|
||
if (!rows.length) return;
|
||
const prev = list.children[sel];
|
||
if (prev) prev.classList.remove("sel");
|
||
sel = (sel + delta + rows.length) % rows.length;
|
||
const next = list.children[sel];
|
||
if (next) next.classList.add("sel");
|
||
scrollToSel();
|
||
};
|
||
|
||
input.addEventListener("input", render);
|
||
input.addEventListener("keydown", (e) => {
|
||
if (e.key === "ArrowDown") { e.preventDefault(); move(1); }
|
||
else if (e.key === "ArrowUp") { e.preventDefault(); move(-1); }
|
||
else if (e.key === "PageDown") { e.preventDefault(); move(10); }
|
||
else if (e.key === "PageUp") { e.preventDefault(); move(-10); }
|
||
else if (e.key === "Enter") {
|
||
e.preventDefault();
|
||
if (rows[sel]) commit(rows[sel].name);
|
||
} else if (e.key === "Escape") {
|
||
e.preventDefault();
|
||
close();
|
||
}
|
||
e.stopPropagation(); // keep ComfyUI's global hotkeys out of the box
|
||
});
|
||
|
||
render();
|
||
const grabFocus = () => {
|
||
if (root.isConnected && document.activeElement !== input) input.focus();
|
||
};
|
||
// Defer by a full task, not a microtask: the pointerdown that opened this is
|
||
// still being dispatched, and a capture listener added now would see it and
|
||
// close the picker immediately.
|
||
setTimeout(() => {
|
||
document.addEventListener("pointerdown", onOutside, true);
|
||
grabFocus();
|
||
}, 0);
|
||
// LiteGraph pulls focus back to the canvas on a delay in some paths; re-assert
|
||
// once so the caret really is in the filter box and typing goes there.
|
||
setTimeout(grabFocus, 60);
|
||
}
|
||
|
||
function loraWidgets(node) {
|
||
return (node.widgets || []).filter((w) => w.type === "H3_LORA");
|
||
}
|
||
|
||
/** Shape test for a serialized stack row, as it appears in widgets_values. */
|
||
function isRowValue(v) {
|
||
return !!v && typeof v === "object" && !Array.isArray(v) && "lora" in v;
|
||
}
|
||
|
||
function copyRow(v) {
|
||
return { ...v };
|
||
}
|
||
|
||
function rowsFromList(values) {
|
||
if (Array.isArray(values)) return values.filter(isRowValue).map(copyRow);
|
||
if (values && typeof values === "object") {
|
||
return Object.values(values).filter(isRowValue).map(copyRow);
|
||
}
|
||
return [];
|
||
}
|
||
|
||
/** Named slots are `lora_N`. Anything else that still looks like a row sorts last. */
|
||
function rowsFromNamed(named) {
|
||
if (!named || typeof named !== "object" || Array.isArray(named)) return [];
|
||
const order = (key) => {
|
||
const match = /^lora_(\d+)$/.exec(key);
|
||
return match ? Number(match[1]) : Number.MAX_SAFE_INTEGER;
|
||
};
|
||
return Object.keys(named)
|
||
.filter((key) => isRowValue(named[key]))
|
||
.sort((a, b) => order(a) - order(b) || (a < b ? -1 : a > b ? 1 : 0))
|
||
.map((key) => copyRow(named[key]));
|
||
}
|
||
|
||
/**
|
||
* Row list to rebuild after configure. properties.h3Rows wins once a graph has
|
||
* been saved by this version; older graphs still carry the rows in
|
||
* widgets_values, which must be read before migrateWidgetsValues compacts them.
|
||
*/
|
||
function extractRows(info) {
|
||
if (!info || typeof info !== "object") return [];
|
||
const stored = info.properties?.[ROWS_KEY];
|
||
if (Array.isArray(stored)) return stored.filter(isRowValue).map(copyRow);
|
||
const positional = rowsFromList(info.widgets_values);
|
||
if (positional.length) return positional;
|
||
return rowsFromNamed(info.widgets_values_named);
|
||
}
|
||
|
||
function hasRowPayload(info) {
|
||
if (!info || typeof info !== "object") return false;
|
||
return (
|
||
Array.isArray(info.properties?.[ROWS_KEY]) ||
|
||
info.widgets_values != null ||
|
||
info.widgets_values_named != null
|
||
);
|
||
}
|
||
|
||
function syncProperties(node) {
|
||
node.properties = node.properties || {};
|
||
node.properties[ROWS_KEY] = loraWidgets(node).map((w) => copyRow(w.value));
|
||
}
|
||
|
||
/** Persist the row list and tell LiteGraph the graph changed, so autosave runs. */
|
||
function markChanged(node) {
|
||
syncProperties(node);
|
||
node.graph?.change?.();
|
||
}
|
||
|
||
/* --------------------------------------------------------------- balance -- */
|
||
|
||
function activeRows(node) {
|
||
return loraWidgets(node).filter(
|
||
(w) => w.value.on && w.value.lora && w.value.lora !== "None"
|
||
);
|
||
}
|
||
|
||
/**
|
||
* Whether the stack is in auto-balance mode. Derived from the rows rather than
|
||
* held as node state: ``manual`` is stashed on every balanced row and rows are
|
||
* already serialized, so the mode survives a save/reload for free.
|
||
*/
|
||
function isBalanced(node) {
|
||
return loraWidgets(node).some((w) => w.value.manual !== undefined);
|
||
}
|
||
|
||
function updateBalanceLabel(node) {
|
||
const button = (node.widgets || []).find((w) => w.h3Role === "balance");
|
||
if (!button) return;
|
||
if (!isBalanced(node)) {
|
||
button.name = BALANCE_LABEL;
|
||
} else {
|
||
const applied = loraWidgets(node).filter((w) => w.value.autoApplied);
|
||
const trimmed = applied.filter((w) => (w.value.factor ?? 1) < 0.995).length;
|
||
button.name = `⚖ Balanced — ${trimmed}/${applied.length} trimmed`;
|
||
}
|
||
node.setDirtyCanvas(true, true);
|
||
}
|
||
|
||
/**
|
||
* Put every active row onto one strength unit.
|
||
*
|
||
* The measured factor multiplies the strength the user already chose, so their
|
||
* relative intent between rows survives; what changes is that a LoRA which
|
||
* perturbs the model 18x harder than usual stops arriving at full force. The
|
||
* factor only ever trims (server-side it is clamped to <= 1): a LoRA measuring
|
||
* below the reference may be quiet deliberately -- distillation adapters sit an
|
||
* order of magnitude down and are correct at 1.0 -- so boosting is not safe,
|
||
* while a LoRA measuring far above it essentially never is.
|
||
*
|
||
* ``force`` re-measures every active row, which is what the button does. The
|
||
* implicit calls -- made when a row is switched on or repointed inside an
|
||
* already-balanced stack -- only touch rows that carry no stash yet, so a row
|
||
* the user has since edited by hand is not quietly pulled back to its
|
||
* calibrated value.
|
||
*/
|
||
async function applyBalance(node, { force = false } = {}) {
|
||
const rows = activeRows(node).filter((w) => force || w.value.manual === undefined);
|
||
if (!rows.length) {
|
||
updateBalanceLabel(node);
|
||
return;
|
||
}
|
||
let data;
|
||
try {
|
||
data = await fetchBalance([...new Set(rows.map((w) => w.value.lora))]);
|
||
} catch (err) {
|
||
console.error("[H3PowerLoraStack] auto-balance failed", err);
|
||
const button = (node.widgets || []).find((w) => w.h3Role === "balance");
|
||
if (button) {
|
||
button.name = "⚖ Balance failed";
|
||
button.tooltip = err?.message || "Could not measure the selected LoRAs";
|
||
node.setDirtyCanvas(true, true);
|
||
}
|
||
return;
|
||
}
|
||
const results = data?.results ?? {};
|
||
for (const widget of rows) {
|
||
const result = results[widget.value.lora];
|
||
if (!result) continue;
|
||
const manual =
|
||
widget.value.autoApplied && widget.value.manual !== undefined
|
||
? widget.value.manual
|
||
: widget.value.strength;
|
||
const rawFactor = Number(result.factor ?? 1);
|
||
if (!Number.isFinite(rawFactor) || rawFactor < 0) continue;
|
||
const factor = Math.min(1, rawFactor);
|
||
widget.value = {
|
||
...widget.value,
|
||
manual,
|
||
strength: manual * factor,
|
||
factor,
|
||
rel: result.rel ?? null,
|
||
note: result.note ?? "",
|
||
autoApplied: true,
|
||
};
|
||
if (result.note) {
|
||
console.warn(`[H3PowerLoraStack] ${widget.value.lora}: ${result.note}`);
|
||
}
|
||
}
|
||
updateBalanceLabel(node);
|
||
node.setDirtyCanvas(true, true);
|
||
markChanged(node);
|
||
}
|
||
|
||
/** Hand every row back the strength it had before auto-balance touched it. */
|
||
function restoreManual(node) {
|
||
let restored = 0;
|
||
for (const widget of loraWidgets(node)) {
|
||
if (widget.value.manual === undefined) continue;
|
||
const next = { ...widget.value, strength: widget.value.manual };
|
||
delete next.manual;
|
||
delete next.factor;
|
||
delete next.rel;
|
||
delete next.note;
|
||
delete next.autoApplied;
|
||
widget.value = next;
|
||
restored += 1;
|
||
}
|
||
updateBalanceLabel(node);
|
||
node.setDirtyCanvas(true, true);
|
||
if (restored) markChanged(node);
|
||
return restored;
|
||
}
|
||
|
||
function renumber(node) {
|
||
loraWidgets(node).forEach((w, i) => {
|
||
w.name = `lora_${i + 1}`;
|
||
});
|
||
}
|
||
|
||
function addLoraWidget(node, value, { notify = true } = {}) {
|
||
const widget = makeLoraWidget(node, `lora_${loraWidgets(node).length + 1}`, value);
|
||
node.widgets = node.widgets || [];
|
||
// keep the "Add LoRA" button last
|
||
const buttonIndex = node.widgets.findIndex((w) => w.h3Role === "add");
|
||
if (buttonIndex >= 0) node.widgets.splice(buttonIndex, 0, widget);
|
||
else node.widgets.push(widget);
|
||
renumber(node);
|
||
resize(node);
|
||
if (notify) markChanged(node);
|
||
return widget;
|
||
}
|
||
|
||
function removeLoraWidget(node, widget, { notify = true } = {}) {
|
||
const i = node.widgets.indexOf(widget);
|
||
if (i >= 0) node.widgets.splice(i, 1);
|
||
renumber(node);
|
||
resize(node);
|
||
node.setDirtyCanvas(true, true);
|
||
if (notify) markChanged(node);
|
||
}
|
||
|
||
/**
|
||
* Rebuild the stack from a snapshot taken before configure. notify stays off:
|
||
* loading a graph is not an edit, and graph.change() here would dirty a fresh load.
|
||
*/
|
||
function restoreRows(node, rows) {
|
||
for (const w of [...loraWidgets(node)]) removeLoraWidget(node, w, { notify: false });
|
||
for (const row of rows) addLoraWidget(node, row, { notify: false });
|
||
updateBalanceLabel(node);
|
||
syncProperties(node);
|
||
resize(node);
|
||
}
|
||
|
||
function resize(node) {
|
||
const [minWidth, minHeight] = node.computeSize();
|
||
// Width stays the user's choice; height tracks the row count so removing a
|
||
// row hands the space back instead of leaving a hole.
|
||
node.setSize([Math.max(node.size[0], minWidth), minHeight]);
|
||
node.setDirtyCanvas(true, true);
|
||
}
|
||
|
||
app.registerExtension({
|
||
name: "h3.power.lora.stack",
|
||
|
||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||
if (nodeData.name !== NODE_TYPE) return;
|
||
|
||
const onCreated = nodeType.prototype.onNodeCreated;
|
||
nodeType.prototype.onNodeCreated = function () {
|
||
onCreated?.apply(this, arguments);
|
||
this.serialize_widgets = true;
|
||
|
||
// (value, canvas, node, pos, event) -- the event is the pointerdown that
|
||
// started the click, so the picker opens right under the button.
|
||
const add = this.addWidget("button", "➕ Add LoRA", null, (_v, _canvas, _node, _pos, event) => {
|
||
const widget = addLoraWidget(this);
|
||
// open the picker immediately: the row is added to be filled in, so the
|
||
// filter box takes the keyboard straight away
|
||
showLoraMenu(this, widget, event, { removeOnCancel: true });
|
||
});
|
||
add.h3Role = "add";
|
||
|
||
const balance = this.addWidget("button", BALANCE_LABEL, null, () => {
|
||
applyBalance(this, { force: true });
|
||
});
|
||
balance.h3Role = "balance";
|
||
|
||
const restore = this.addWidget("button", RESTORE_LABEL, null, () => {
|
||
if (!restoreManual(this)) {
|
||
console.info("[H3PowerLoraStack] nothing to restore, no balance applied");
|
||
}
|
||
});
|
||
restore.h3Role = "restore";
|
||
|
||
// The buttons hold no state. Leaving them in widgets_values made the stack
|
||
// grow: positional restore wrote saved rows into the buttons, and the next
|
||
// save echoed them back as extra rows. serialize:false stops that, but
|
||
// ComfyUI's migrateWidgetsValues then treats those button slots as holes
|
||
// and drops the first rows. Restore happens in configure, from a snapshot
|
||
// taken before that migrator runs — not from the stripped array.
|
||
for (const button of [add, balance, restore]) button.serialize = false;
|
||
|
||
resize(this);
|
||
};
|
||
|
||
// Wrap configure, not onConfigure. ComfyNode.configure compacts
|
||
// widgets_values before it calls super, and super calls onConfigure.
|
||
const origConfigure = nodeType.prototype.configure;
|
||
nodeType.prototype.configure = function (info) {
|
||
const captured = hasRowPayload(info) ? extractRows(info) : null;
|
||
const result = origConfigure?.apply(this, arguments);
|
||
if (captured !== null) restoreRows(this, captured);
|
||
return result;
|
||
};
|
||
|
||
const origSerialize = nodeType.prototype.onSerialize;
|
||
nodeType.prototype.onSerialize = function (info) {
|
||
origSerialize?.apply(this, arguments);
|
||
syncProperties(this);
|
||
if (info && typeof info === "object") {
|
||
info.properties = info.properties || {};
|
||
info.properties[ROWS_KEY] = loraWidgets(this).map((w) => copyRow(w.value));
|
||
}
|
||
};
|
||
},
|
||
});
|