609 lines
20 KiB
JavaScript
609 lines
20 KiB
JavaScript
import { app } from "../../scripts/app.js";
|
|
import { api } from "../../scripts/api.js";
|
|
|
|
const WIDGETS_PER_ENTRY = 6; // header, ckpt, cfg_min, cfg_max, weight, enabled
|
|
const NODE_CLASS = "ModelRandomizer";
|
|
const ANCHOR_WIDGET = "action_separator";
|
|
|
|
let cachedCheckpoints = null;
|
|
|
|
async function fetchCheckpoints(forceRefresh = false) {
|
|
if (cachedCheckpoints && !forceRefresh) return cachedCheckpoints;
|
|
try {
|
|
const resp = await api.fetchApi("/models/checkpoints");
|
|
if (!resp.ok) throw new Error(`HTTP ${resp.status}`);
|
|
cachedCheckpoints = await resp.json();
|
|
} catch (e) {
|
|
console.warn("[ModelRandomizer] Failed to fetch checkpoints:", e);
|
|
cachedCheckpoints = [];
|
|
}
|
|
return cachedCheckpoints;
|
|
}
|
|
|
|
// ── Drawing Helpers ────────────────────────────────────────────────
|
|
|
|
function isLowQuality() {
|
|
return (app.canvas?.ds?.scale || 1) <= 0.5;
|
|
}
|
|
|
|
function drawRoundedRect(ctx, x, y, w, h, radius, fill, stroke) {
|
|
ctx.beginPath();
|
|
ctx.roundRect(x, y, w, h, [radius]);
|
|
if (fill) {
|
|
ctx.fillStyle = fill;
|
|
ctx.fill();
|
|
}
|
|
if (stroke) {
|
|
ctx.strokeStyle = stroke;
|
|
ctx.lineWidth = 1;
|
|
ctx.stroke();
|
|
}
|
|
}
|
|
|
|
function drawStyledButton(ctx, x, y, w, h, label, isPressed, colors) {
|
|
const radius = 4;
|
|
const lowQ = isLowQuality();
|
|
|
|
// Shadow
|
|
if (!lowQ && !isPressed) {
|
|
drawRoundedRect(ctx, x + 1, y + 1, w, h, radius, "#00000055");
|
|
}
|
|
|
|
// Main fill + border
|
|
const offsetY = isPressed ? 1 : 0;
|
|
drawRoundedRect(
|
|
ctx, x, y + offsetY, w, h, radius,
|
|
isPressed ? colors.pressed : colors.bg,
|
|
lowQ ? null : colors.border
|
|
);
|
|
|
|
if (lowQ) return;
|
|
|
|
// Text
|
|
ctx.save();
|
|
ctx.textAlign = "center";
|
|
ctx.textBaseline = "middle";
|
|
ctx.fillStyle = colors.text;
|
|
ctx.fillText(label, x + w / 2, y + h / 2 + offsetY);
|
|
ctx.restore();
|
|
}
|
|
|
|
function drawSmallButton(ctx, x, y, w, h, label, isPressed, disabled, colors) {
|
|
const radius = 3;
|
|
const lowQ = isLowQuality();
|
|
|
|
const c = colors || {};
|
|
const bgColor = c.bg || "#4a4a4a";
|
|
const pressedColor = c.pressed || "#3a3a3a";
|
|
const borderColor = c.border || "#666";
|
|
const textColor = c.text || "#ddd";
|
|
const textPressedColor = c.textPressed || "#ccc";
|
|
const disabledBg = c.disabledBg || "#2a2a2a";
|
|
const disabledBorder = c.disabledBorder || "#3a3a3a";
|
|
const disabledText = c.disabledText || "#555";
|
|
|
|
if (disabled) {
|
|
drawRoundedRect(ctx, x, y, w, h, radius, disabledBg, lowQ ? null : disabledBorder);
|
|
if (!lowQ) {
|
|
ctx.save();
|
|
ctx.textAlign = "center";
|
|
ctx.textBaseline = "middle";
|
|
ctx.fillStyle = disabledText;
|
|
ctx.font = "11px Arial";
|
|
ctx.fillText(label, x + w / 2, y + h / 2);
|
|
ctx.restore();
|
|
}
|
|
return;
|
|
}
|
|
|
|
if (!lowQ && !isPressed) {
|
|
drawRoundedRect(ctx, x + 0.5, y + 0.5, w, h, radius, "#00000033");
|
|
}
|
|
|
|
const offsetY = isPressed ? 1 : 0;
|
|
drawRoundedRect(
|
|
ctx, x, y + offsetY, w, h, radius,
|
|
isPressed ? pressedColor : bgColor,
|
|
lowQ ? null : borderColor
|
|
);
|
|
|
|
if (lowQ) return;
|
|
|
|
ctx.save();
|
|
ctx.textAlign = "center";
|
|
ctx.textBaseline = "middle";
|
|
ctx.fillStyle = isPressed ? textPressedColor : textColor;
|
|
ctx.font = "11px Arial";
|
|
ctx.fillText(label, x + w / 2, y + h / 2 + offsetY);
|
|
ctx.restore();
|
|
}
|
|
|
|
// ── Custom Widgets ─────────────────────────────────────────────────
|
|
|
|
function createDividerWidget(name, opts = {}) {
|
|
const marginTop = opts.marginTop ?? 4;
|
|
const marginBottom = opts.marginBottom ?? 4;
|
|
const thickness = opts.thickness ?? 1;
|
|
const color = opts.color ?? LiteGraph.WIDGET_OUTLINE_COLOR;
|
|
|
|
return {
|
|
name,
|
|
type: "custom",
|
|
value: "",
|
|
y: 0,
|
|
last_y: 0,
|
|
options: { serialize: false },
|
|
|
|
draw(ctx, node, width, y, height) {
|
|
this.last_y = y;
|
|
if (isLowQuality() || !thickness) return;
|
|
const margin = 15;
|
|
ctx.save();
|
|
ctx.strokeStyle = color;
|
|
ctx.lineWidth = thickness;
|
|
ctx.beginPath();
|
|
ctx.moveTo(margin, y + marginTop);
|
|
ctx.lineTo(width - margin, y + marginTop);
|
|
ctx.stroke();
|
|
ctx.restore();
|
|
},
|
|
|
|
mouse() { return false; },
|
|
|
|
computeSize(width) {
|
|
return [width, marginTop + marginBottom + thickness];
|
|
},
|
|
|
|
serializeValue() { return null; },
|
|
};
|
|
}
|
|
|
|
function createEntryHeaderWidget(index, onMoveUp, onMoveDown, onDelete) {
|
|
return {
|
|
name: `header_${index}`,
|
|
type: "custom",
|
|
value: "",
|
|
y: 0,
|
|
last_y: 0,
|
|
options: { serialize: false },
|
|
_entryIndex: index,
|
|
_upPressed: false,
|
|
_downPressed: false,
|
|
_deletePressed: false,
|
|
_upBounds: null,
|
|
_downBounds: null,
|
|
_deleteBounds: null,
|
|
_mouseDownTarget: null,
|
|
|
|
draw(ctx, node, width, y, height) {
|
|
this.last_y = y;
|
|
const margin = 15;
|
|
const lowQ = isLowQuality();
|
|
|
|
// Top divider line
|
|
if (!lowQ) {
|
|
ctx.strokeStyle = LiteGraph.WIDGET_OUTLINE_COLOR;
|
|
ctx.lineWidth = 1;
|
|
ctx.beginPath();
|
|
ctx.moveTo(margin, y + 4);
|
|
ctx.lineTo(width - margin, y + 4);
|
|
ctx.stroke();
|
|
}
|
|
|
|
if (lowQ) return;
|
|
|
|
// Vertical center for label and buttons (shifted down past divider)
|
|
const centerY = y + height * 0.5 + 4;
|
|
|
|
// Label
|
|
ctx.save();
|
|
ctx.fillStyle = "#999";
|
|
ctx.font = "11px Arial";
|
|
ctx.textAlign = "left";
|
|
ctx.textBaseline = "middle";
|
|
ctx.fillText(`Model ${this._entryIndex}`, margin + 2, centerY);
|
|
ctx.restore();
|
|
|
|
// Button layout (right-aligned): [▲] [▼] [✕]
|
|
const btnH = 16;
|
|
const btnW = 20;
|
|
const btnGap = 4;
|
|
const deleteSeparation = 10;
|
|
const btnY = centerY - btnH / 2;
|
|
|
|
const isFirst = this._entryIndex === 1;
|
|
const isLast = this._entryIndex === node._modelCount;
|
|
const isOnly = node._modelCount <= 1;
|
|
|
|
const deleteColors = {
|
|
bg: "#5a2a2a",
|
|
pressed: "#4a1a1a",
|
|
border: "#8a3a3a",
|
|
text: "#e66",
|
|
textPressed: "#c55",
|
|
};
|
|
|
|
// ✕ Delete (rightmost)
|
|
const delX = width - margin - btnW;
|
|
this._deleteBounds = [delX, btnY, btnW, btnH];
|
|
drawSmallButton(ctx, delX, btnY, btnW, btnH, "\u2715", this._deletePressed, isOnly, deleteColors);
|
|
|
|
// ▼ Move Down
|
|
const downX = delX - btnW - deleteSeparation;
|
|
this._downBounds = [downX, btnY, btnW, btnH];
|
|
drawSmallButton(ctx, downX, btnY, btnW, btnH, "\u25BC", this._downPressed, isLast);
|
|
|
|
// ▲ Move Up
|
|
const upX = downX - btnW - btnGap;
|
|
this._upBounds = [upX, btnY, btnW, btnH];
|
|
drawSmallButton(ctx, upX, btnY, btnW, btnH, "\u25B2", this._upPressed, isFirst);
|
|
},
|
|
|
|
_hitTest(pos) {
|
|
const tests = [
|
|
["up", this._upBounds],
|
|
["down", this._downBounds],
|
|
["delete", this._deleteBounds],
|
|
];
|
|
for (const [name, bounds] of tests) {
|
|
if (bounds &&
|
|
pos[0] >= bounds[0] && pos[0] <= bounds[0] + bounds[2] &&
|
|
pos[1] >= bounds[1] && pos[1] <= bounds[1] + bounds[3]) {
|
|
return name;
|
|
}
|
|
}
|
|
return null;
|
|
},
|
|
|
|
_isDisabled(target, node) {
|
|
if (target === "up") return this._entryIndex === 1;
|
|
if (target === "down") return this._entryIndex === node._modelCount;
|
|
if (target === "delete") return node._modelCount <= 1;
|
|
return false;
|
|
},
|
|
|
|
mouse(event, pos, node) {
|
|
if (event.type === "pointerdown") {
|
|
const target = this._hitTest(pos);
|
|
if (!target || this._isDisabled(target, node)) return false;
|
|
this._mouseDownTarget = target;
|
|
if (target === "up") this._upPressed = true;
|
|
if (target === "down") this._downPressed = true;
|
|
if (target === "delete") this._deletePressed = true;
|
|
app.graph.setDirtyCanvas(true, false);
|
|
return true;
|
|
}
|
|
|
|
if (event.type === "pointerup") {
|
|
const wasTarget = this._mouseDownTarget;
|
|
this._upPressed = false;
|
|
this._downPressed = false;
|
|
this._deletePressed = false;
|
|
this._mouseDownTarget = null;
|
|
app.graph.setDirtyCanvas(true, false);
|
|
if (!wasTarget) return false;
|
|
|
|
const target = this._hitTest(pos);
|
|
if (target === wasTarget && !this._isDisabled(target, node)) {
|
|
if (target === "up") onMoveUp.call(node, this._entryIndex);
|
|
if (target === "down") onMoveDown.call(node, this._entryIndex);
|
|
if (target === "delete") onDelete.call(node, this._entryIndex);
|
|
}
|
|
return true;
|
|
}
|
|
|
|
if (event.type === "pointermove") {
|
|
if (!this._mouseDownTarget) return false;
|
|
const target = this._hitTest(pos);
|
|
this._upPressed = (this._mouseDownTarget === "up" && target === "up");
|
|
this._downPressed = (this._mouseDownTarget === "down" && target === "down");
|
|
this._deletePressed = (this._mouseDownTarget === "delete" && target === "delete");
|
|
app.graph.setDirtyCanvas(true, false);
|
|
return true;
|
|
}
|
|
|
|
return false;
|
|
},
|
|
|
|
computeSize(width) {
|
|
return [width, LiteGraph.NODE_WIDGET_HEIGHT + 12];
|
|
},
|
|
|
|
serializeValue() { return null; },
|
|
};
|
|
}
|
|
|
|
function createAddButtonWidget(addCallback) {
|
|
return {
|
|
name: "model_actions",
|
|
type: "custom",
|
|
value: "",
|
|
y: 0,
|
|
last_y: 0,
|
|
options: { serialize: false },
|
|
_pressed: false,
|
|
_bounds: null,
|
|
_mouseDown: false,
|
|
|
|
draw(ctx, node, width, y, height) {
|
|
this.last_y = y;
|
|
const margin = 15;
|
|
const w = width - margin * 2;
|
|
this._bounds = [margin, y, w, height];
|
|
|
|
drawStyledButton(ctx, margin, y, w, height,
|
|
"\u2795 Add Model", this._pressed,
|
|
{ bg: "#2a4a2a", pressed: "#1a3a1a", border: "#3a6a3a", text: "#9fcf9f" });
|
|
},
|
|
|
|
_isInside(pos) {
|
|
const b = this._bounds;
|
|
return b && pos[0] >= b[0] && pos[0] <= b[0] + b[2] &&
|
|
pos[1] >= b[1] && pos[1] <= b[1] + b[3];
|
|
},
|
|
|
|
mouse(event, pos, node) {
|
|
if (event.type === "pointerdown") {
|
|
if (!this._isInside(pos)) return false;
|
|
this._mouseDown = true;
|
|
this._pressed = true;
|
|
app.graph.setDirtyCanvas(true, false);
|
|
return true;
|
|
}
|
|
|
|
if (event.type === "pointerup") {
|
|
const wasDown = this._mouseDown;
|
|
this._pressed = false;
|
|
this._mouseDown = false;
|
|
app.graph.setDirtyCanvas(true, false);
|
|
if (!wasDown) return false;
|
|
if (this._isInside(pos)) addCallback.call(node);
|
|
return true;
|
|
}
|
|
|
|
if (event.type === "pointermove") {
|
|
if (!this._mouseDown) return false;
|
|
this._pressed = this._isInside(pos);
|
|
app.graph.setDirtyCanvas(true, false);
|
|
return true;
|
|
}
|
|
|
|
return false;
|
|
},
|
|
|
|
computeSize(width) {
|
|
return [width, LiteGraph.NODE_WIDGET_HEIGHT];
|
|
},
|
|
|
|
serializeValue() { return null; },
|
|
};
|
|
}
|
|
|
|
// ── Resize Helper ──────────────────────────────────────────────────
|
|
|
|
function resizeNodeToFit(node) {
|
|
const sz = node.computeSize();
|
|
node.setSize([Math.max(node.size[0], sz[0]), sz[1]]);
|
|
}
|
|
|
|
// ── Entry Manipulation ─────────────────────────────────────────────
|
|
|
|
function swapModelEntries(node, idxA, idxB) {
|
|
const fields = ["ckpt", "cfg_min", "cfg_max", "weight", "enabled"];
|
|
for (const field of fields) {
|
|
const wA = node.widgets.find((w) => w.name === `${field}_${idxA}`);
|
|
const wB = node.widgets.find((w) => w.name === `${field}_${idxB}`);
|
|
if (wA && wB) {
|
|
const tmp = wA.value;
|
|
wA.value = wB.value;
|
|
wB.value = tmp;
|
|
}
|
|
}
|
|
}
|
|
|
|
function removeModelEntryAt(node, targetIndex) {
|
|
if (node._modelCount <= 1) return false;
|
|
|
|
const prefixes = ["header", "ckpt", "cfg_min", "cfg_max", "weight", "enabled"];
|
|
|
|
for (const prefix of prefixes) {
|
|
const idx = node.widgets.findIndex((w) => w.name === `${prefix}_${targetIndex}`);
|
|
if (idx >= 0) node.widgets.splice(idx, 1);
|
|
}
|
|
|
|
for (let i = targetIndex + 1; i <= node._modelCount; i++) {
|
|
const newIdx = i - 1;
|
|
for (const prefix of prefixes) {
|
|
const w = node.widgets.find((w) => w.name === `${prefix}_${i}`);
|
|
if (w) {
|
|
w.name = `${prefix}_${newIdx}`;
|
|
if (prefix === "header") {
|
|
w._entryIndex = newIdx;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
node._modelCount--;
|
|
node.properties = node.properties || {};
|
|
node.properties.modelCount = node._modelCount;
|
|
|
|
return true;
|
|
}
|
|
|
|
// ── Entry Action Callbacks (called with `this` = node) ─────────────
|
|
|
|
function handleMoveUp(entryIndex) {
|
|
if (entryIndex <= 1) return;
|
|
swapModelEntries(this, entryIndex, entryIndex - 1);
|
|
app.graph.setDirtyCanvas(true, true);
|
|
}
|
|
|
|
function handleMoveDown(entryIndex) {
|
|
if (entryIndex >= this._modelCount) return;
|
|
swapModelEntries(this, entryIndex, entryIndex + 1);
|
|
app.graph.setDirtyCanvas(true, true);
|
|
}
|
|
|
|
function handleDeleteEntry(entryIndex) {
|
|
if (removeModelEntryAt(this, entryIndex)) {
|
|
resizeNodeToFit(this);
|
|
app.graph.setDirtyCanvas(true, true);
|
|
}
|
|
}
|
|
|
|
// ── Core Widget Management ─────────────────────────────────────────
|
|
|
|
function addModelEntry(node, index, checkpoints, defaults) {
|
|
const ckptDefault = defaults?.ckpt ?? checkpoints[0] ?? "";
|
|
const cfgMinDefault = defaults?.cfg_min ?? 5.0;
|
|
const cfgMaxDefault = defaults?.cfg_max ?? 8.0;
|
|
const weightDefault = defaults?.weight ?? 1.0;
|
|
const enabledDefault = defaults?.enabled ?? true;
|
|
|
|
const anchorIdx = node.widgets.findIndex((w) => w.name === ANCHOR_WIDGET);
|
|
|
|
const newWidgets = [];
|
|
|
|
newWidgets.push(
|
|
node.addCustomWidget(
|
|
createEntryHeaderWidget(index, handleMoveUp, handleMoveDown, handleDeleteEntry)
|
|
)
|
|
);
|
|
|
|
newWidgets.push(
|
|
node.addWidget("combo", `ckpt_${index}`, ckptDefault, () => {}, {
|
|
values: checkpoints,
|
|
})
|
|
);
|
|
newWidgets.push(
|
|
node.addWidget("number", `cfg_min_${index}`, cfgMinDefault, () => {}, {
|
|
min: 0.0,
|
|
max: 100.0,
|
|
step: 5,
|
|
precision: 1,
|
|
})
|
|
);
|
|
newWidgets.push(
|
|
node.addWidget("number", `cfg_max_${index}`, cfgMaxDefault, () => {}, {
|
|
min: 0.0,
|
|
max: 100.0,
|
|
step: 5,
|
|
precision: 1,
|
|
})
|
|
);
|
|
newWidgets.push(
|
|
node.addWidget("number", `weight_${index}`, weightDefault, () => {}, {
|
|
min: 0.0,
|
|
max: 10.0,
|
|
step: 1,
|
|
precision: 1,
|
|
})
|
|
);
|
|
newWidgets.push(
|
|
node.addWidget("toggle", `enabled_${index}`, enabledDefault, () => {}, {})
|
|
);
|
|
|
|
if (anchorIdx >= 0) {
|
|
for (let i = 0; i < newWidgets.length; i++) {
|
|
const w = newWidgets[i];
|
|
const currentIdx = node.widgets.indexOf(w);
|
|
if (currentIdx >= 0) {
|
|
node.widgets.splice(currentIdx, 1);
|
|
node.widgets.splice(anchorIdx + i, 0, w);
|
|
}
|
|
}
|
|
}
|
|
|
|
return newWidgets;
|
|
}
|
|
|
|
// ── Extension Registration ─────────────────────────────────────────
|
|
|
|
app.registerExtension({
|
|
name: "comfyui.ModelRandomizer",
|
|
|
|
async beforeRegisterNodeDef(nodeType, nodeData, app) {
|
|
if (nodeData.name !== NODE_CLASS) return;
|
|
|
|
const checkpoints = await fetchCheckpoints();
|
|
|
|
// --- onNodeCreated: set up initial widgets ---
|
|
const origOnNodeCreated = nodeType.prototype.onNodeCreated;
|
|
nodeType.prototype.onNodeCreated = function () {
|
|
origOnNodeCreated?.apply(this, arguments);
|
|
|
|
this._modelCount = 0;
|
|
|
|
this.addCustomWidget(createDividerWidget(ANCHOR_WIDGET));
|
|
this.addCustomWidget(createAddButtonWidget(
|
|
// Add Model callback
|
|
async function () {
|
|
const freshCheckpoints = await fetchCheckpoints(true);
|
|
this._modelCount++;
|
|
addModelEntry(this, this._modelCount, freshCheckpoints);
|
|
this.properties = this.properties || {};
|
|
this.properties.modelCount = this._modelCount;
|
|
resizeNodeToFit(this);
|
|
app.graph.setDirtyCanvas(true, true);
|
|
}
|
|
));
|
|
|
|
// Default first entry
|
|
this._modelCount++;
|
|
addModelEntry(this, this._modelCount, checkpoints);
|
|
this.properties = this.properties || {};
|
|
this.properties.modelCount = this._modelCount;
|
|
this.setSize(this.computeSize());
|
|
};
|
|
|
|
const origOnSerialize = nodeType.prototype.onSerialize;
|
|
nodeType.prototype.onSerialize = function (o) {
|
|
origOnSerialize?.apply(this, arguments);
|
|
const entries = [];
|
|
for (let i = 1; i <= this._modelCount; i++) {
|
|
entries.push({
|
|
ckpt: this.widgets.find(w => w.name === `ckpt_${i}`)?.value,
|
|
cfg_min: this.widgets.find(w => w.name === `cfg_min_${i}`)?.value,
|
|
cfg_max: this.widgets.find(w => w.name === `cfg_max_${i}`)?.value,
|
|
weight: this.widgets.find(w => w.name === `weight_${i}`)?.value,
|
|
enabled: this.widgets.find(w => w.name === `enabled_${i}`)?.value,
|
|
});
|
|
}
|
|
o.properties = o.properties || {};
|
|
o.properties.modelEntries = entries;
|
|
o.properties.modelCount = this._modelCount;
|
|
};
|
|
|
|
const origConfigure = nodeType.prototype.configure;
|
|
nodeType.prototype.configure = function (info) {
|
|
// Nuke model widgets BEFORE super.configure so LiteGraph's
|
|
// widgets_values restoration has nothing to mismatch against.
|
|
// Includes both header_ (current) and sep_ (legacy) prefixes.
|
|
for (let i = this.widgets.length - 1; i >= 0; i--) {
|
|
const n = this.widgets[i]?.name;
|
|
if (n && (n.startsWith("header_") || n.startsWith("sep_") ||
|
|
n.startsWith("ckpt_") || n.startsWith("cfg_min_") ||
|
|
n.startsWith("cfg_max_") || n.startsWith("weight_") ||
|
|
n.startsWith("enabled_"))) {
|
|
this.widgets.splice(i, 1);
|
|
}
|
|
}
|
|
this._modelCount = 0;
|
|
|
|
origConfigure?.apply(this, arguments);
|
|
|
|
const entries = info.properties?.modelEntries || [];
|
|
const savedCount = entries.length || info.properties?.modelCount || 0;
|
|
|
|
for (let i = 0; i < savedCount; i++) {
|
|
this._modelCount++;
|
|
addModelEntry(this, this._modelCount, checkpoints, entries[i]);
|
|
}
|
|
|
|
resizeNodeToFit(this);
|
|
app.graph.setDirtyCanvas(true, true);
|
|
};
|
|
},
|
|
});
|