1018 lines
34 KiB
JavaScript
1018 lines
34 KiB
JavaScript
import { app } from "/scripts/app.js";
|
|
import { api } from "/scripts/api.js";
|
|
|
|
/* =============================================================================
|
|
Prompt Database — ComfyUI Web Extension
|
|
========================================================================== */
|
|
|
|
console.log("[PromptDB] Extension loading...");
|
|
|
|
// ─── API helpers ────────────────────────────────────────────────────────────
|
|
const API = {
|
|
async getCategories() {
|
|
const r = await api.fetchApi("/prompt_db/categories");
|
|
const data = await r.json();
|
|
return data.categories || [];
|
|
},
|
|
async getPrompts(category) {
|
|
const url = category
|
|
? `/prompt_db/prompts?category=${encodeURIComponent(category)}`
|
|
: "/prompt_db/prompts";
|
|
const r = await api.fetchApi(url);
|
|
const data = await r.json();
|
|
return data.prompts || [];
|
|
},
|
|
};
|
|
|
|
function getNodeClass(node) {
|
|
return node.comfyClass || node.type || node.constructor?.type || "";
|
|
}
|
|
|
|
// ─── Setup PromptDBManager ──────────────────────────────────────────────────
|
|
function setupManagerNode(node) {
|
|
if (node._pdb_manager_setup) return;
|
|
node._pdb_manager_setup = true;
|
|
node.addWidget("button", "📝 Open Prompt Manager", null, () => {
|
|
window.open("/prompt_db/manager", "_blank");
|
|
});
|
|
node.setSize?.(node.computeSize?.() || [260, 58]);
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
}
|
|
|
|
// ─── Setup PromptSelector ───────────────────────────────────────────────────
|
|
function setupSelectorNode(node) {
|
|
if (node._pdb_selector_setup) return;
|
|
node._pdb_selector_setup = true;
|
|
|
|
const catWidget = node.widgets?.find((w) => w.name === "category");
|
|
const promptWidget = node.widgets?.find((w) => w.name === "prompt_name");
|
|
if (!catWidget || !promptWidget) return;
|
|
|
|
// Strip 📖 marker before sending to Python backend
|
|
promptWidget.serializeValue = () =>
|
|
(promptWidget.value || "").replace(/\s*📖$/, "");
|
|
|
|
// ── Refresh categories + prompts from DB ──
|
|
async function refreshFromDB() {
|
|
try {
|
|
const categories = await API.getCategories();
|
|
if (categories.length > 0) {
|
|
catWidget.options.values = categories;
|
|
if (!categories.includes(catWidget.value)) catWidget.value = categories[0];
|
|
}
|
|
const cat = catWidget.value;
|
|
if (cat && cat !== "(no categories)") {
|
|
const prompts = await API.getPrompts(cat);
|
|
const names = prompts.map((p) => {
|
|
const hasLoras = p.loras && p.loras.length > 0;
|
|
return hasLoras ? `${p.name} 📖` : p.name;
|
|
});
|
|
if (names.length > 0) {
|
|
promptWidget.options.values = names;
|
|
if (!names.includes(promptWidget.value)) promptWidget.value = names[0];
|
|
} else {
|
|
promptWidget.options.values = ["(no prompts)"];
|
|
promptWidget.value = "(no prompts)";
|
|
}
|
|
}
|
|
await updateLoraDisplay();
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
} catch (e) {
|
|
console.error("[PromptDB] Refresh failed:", e);
|
|
}
|
|
}
|
|
|
|
// ── LoRA display widgets (read-only info) ──
|
|
async function updateLoraDisplay() {
|
|
if (node.widgets) {
|
|
node.widgets = node.widgets.filter((w) => !w._pdb_lora_display);
|
|
}
|
|
const rawName = promptWidget.value || "";
|
|
const cleanName = rawName.replace(/\s*📖$/, "");
|
|
const cat = catWidget.value;
|
|
if (!cleanName || cleanName === "(no prompts)") return;
|
|
|
|
try {
|
|
const prompts = await API.getPrompts(cat);
|
|
const prompt = prompts.find((p) => p.name === cleanName);
|
|
if (!prompt || !prompt.loras || prompt.loras.length === 0) return;
|
|
|
|
for (const lora of prompt.loras) {
|
|
const displayName = (lora.name || "").replace(/\.safetensors$/, "");
|
|
const strength = lora.strength != null ? lora.strength : 1.0;
|
|
const w = node.addWidget(
|
|
"text",
|
|
`🔗 ${displayName}`,
|
|
`${strength.toFixed(2)}`,
|
|
() => {},
|
|
{ serialize: false }
|
|
);
|
|
w._pdb_lora_display = true;
|
|
w.disabled = true;
|
|
}
|
|
} catch (e) {
|
|
console.error("[PromptDB] Failed to load LoRA info:", e);
|
|
}
|
|
const computed = node.computeSize?.();
|
|
if (computed) node.setSize?.(computed);
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
}
|
|
|
|
// ── Add Refresh button ──
|
|
node.addWidget("button", "🔄 Refresh", null, () => refreshFromDB());
|
|
|
|
// ── Category change callback ──
|
|
const origCatCallback = catWidget.callback;
|
|
catWidget.callback = async function (value) {
|
|
if (origCatCallback) origCatCallback.call(this, value);
|
|
try {
|
|
const prompts = await API.getPrompts(value);
|
|
const names = prompts.map((p) => {
|
|
const hasLoras = p.loras && p.loras.length > 0;
|
|
return hasLoras ? `${p.name} 📖` : p.name;
|
|
});
|
|
if (names.length === 0) {
|
|
promptWidget.options.values = ["(no prompts)"];
|
|
promptWidget.value = "(no prompts)";
|
|
} else {
|
|
promptWidget.options.values = names;
|
|
if (!names.includes(promptWidget.value)) promptWidget.value = names[0];
|
|
}
|
|
await updateLoraDisplay();
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
} catch (e) {
|
|
console.error("[PromptDB] Failed to fetch prompts:", e);
|
|
}
|
|
};
|
|
|
|
// ── Prompt change callback ──
|
|
const origPromptCallback = promptWidget.callback;
|
|
promptWidget.callback = async function (value) {
|
|
if (origPromptCallback) origPromptCallback.call(this, value);
|
|
await updateLoraDisplay();
|
|
};
|
|
|
|
refreshFromDB();
|
|
}
|
|
|
|
// ─── Setup PromptCombiner ───────────────────────────────────────────────────
|
|
const DEFAULT_WIDTH = 500;
|
|
|
|
function setupCombinerNode(node) {
|
|
if (node._pdb_combiner_setup) return;
|
|
node._pdb_combiner_setup = true;
|
|
console.log("[PromptDB] Setting up PromptCombiner node", node.id);
|
|
|
|
const selectionsWidget = node.widgets?.find(
|
|
(w) => w.name === "prompt_selections"
|
|
);
|
|
if (selectionsWidget) {
|
|
selectionsWidget.type = "converted-widget";
|
|
selectionsWidget.computeSize = () => [0, -4];
|
|
selectionsWidget.serializeValue = () => selectionsWidget.value;
|
|
}
|
|
|
|
let userWidth = Math.max(node.size?.[0] || DEFAULT_WIDTH, DEFAULT_WIDTH);
|
|
|
|
const origOnResize = node.onResize;
|
|
node.onResize = function (size) {
|
|
if (origOnResize) origOnResize.apply(this, arguments);
|
|
if (size && size[0]) userWidth = size[0];
|
|
};
|
|
|
|
let rebuildPending = false;
|
|
let rebuildTimer = null;
|
|
|
|
function getMaxSelectionIndex() {
|
|
let max = 0;
|
|
for (const inp of node.inputs || []) {
|
|
const m = inp.name.match(/^selection_(\d+)$/);
|
|
if (m) max = Math.max(max, parseInt(m[1]));
|
|
}
|
|
return max;
|
|
}
|
|
|
|
function syncSelectionsJSON() {
|
|
if (!selectionsWidget) return;
|
|
const obj = {};
|
|
if (node.widgets) {
|
|
for (const w of node.widgets) {
|
|
if (w._pdb_dynamic && w._pdb_inputName) {
|
|
// Skip muted entries
|
|
if (w._pdb_muted) {
|
|
obj[w._pdb_inputName] = "";
|
|
} else {
|
|
const val = (w.value || "").replace(/\s*📖$/, "");
|
|
obj[w._pdb_inputName] = val;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
selectionsWidget.value = JSON.stringify(obj);
|
|
}
|
|
|
|
function getSourceCategory(slotIndex) {
|
|
const input = node.inputs?.[slotIndex];
|
|
if (!input || input.link == null) return null;
|
|
const linkInfo = app.graph.links[input.link];
|
|
if (!linkInfo) return null;
|
|
const sourceNode = app.graph.getNodeById(linkInfo.origin_id);
|
|
if (!sourceNode) return null;
|
|
const catWidget = sourceNode.widgets?.find((w) => w.name === "category");
|
|
return catWidget?.value || null;
|
|
}
|
|
|
|
async function fullRebuild() {
|
|
// Remove dynamic widgets
|
|
if (node.widgets) {
|
|
node.widgets = node.widgets.filter(
|
|
(w) => !w._pdb_dynamic && !w._pdb_lora_display
|
|
);
|
|
}
|
|
|
|
// Cleanup disconnected slots
|
|
if (node.inputs) {
|
|
for (let i = node.inputs.length - 1; i >= 0; i--) {
|
|
const inp = node.inputs[i];
|
|
if (inp.name.startsWith("selection_") && inp.link == null) {
|
|
node.removeInput(i);
|
|
}
|
|
}
|
|
}
|
|
|
|
const nextIdx = getMaxSelectionIndex() + 1;
|
|
node.addInput(`selection_${nextIdx}`, "DB_SELECTION");
|
|
|
|
let savedSelections = {};
|
|
try {
|
|
savedSelections = JSON.parse(selectionsWidget?.value || "{}");
|
|
} catch (e) {}
|
|
|
|
for (let i = 0; i < (node.inputs?.length || 0); i++) {
|
|
const inp = node.inputs[i];
|
|
if (!inp.name.startsWith("selection_")) continue;
|
|
if (inp.link == null) continue;
|
|
|
|
const category = getSourceCategory(i);
|
|
if (!category) continue;
|
|
|
|
let allPrompts = [];
|
|
let promptNames = [];
|
|
try {
|
|
allPrompts = await API.getPrompts(category);
|
|
promptNames = allPrompts.map((p) => {
|
|
const hasLoras = p.loras && p.loras.length > 0;
|
|
return hasLoras ? `${p.name} 📖` : p.name;
|
|
});
|
|
} catch (e) {
|
|
console.error("[PromptDB] fetch failed for", category, e);
|
|
}
|
|
if (promptNames.length === 0) promptNames = ["(no prompts)"];
|
|
|
|
const saved = savedSelections[inp.name];
|
|
let defaultVal = promptNames[0];
|
|
if (saved) {
|
|
const match = promptNames.find(
|
|
(n) =>
|
|
n === saved ||
|
|
n === `${saved} 📖` ||
|
|
n.replace(/\s*📖$/, "") === saved
|
|
);
|
|
if (match) defaultVal = match;
|
|
}
|
|
|
|
// Check saved mute state
|
|
const muteKey = `pdb_mute_${node.id}_${inp.name}`;
|
|
const isMuted = localStorage.getItem(muteKey) === "1";
|
|
|
|
// ── Custom single-line widget: ● 📂 Category ◀ PromptName ▶ ──
|
|
const closureCat = category;
|
|
const closureNames = promptNames;
|
|
const closureInpName = inp.name;
|
|
const widget = {
|
|
type: "pdb_combo",
|
|
name: `📂 ${category}`,
|
|
value: defaultVal,
|
|
_pdb_dynamic: true,
|
|
_pdb_inputName: closureInpName,
|
|
_pdb_category: closureCat,
|
|
_pdb_allPrompts: allPrompts,
|
|
_pdb_muted: isMuted,
|
|
_pdb_values: closureNames,
|
|
options: { serialize: false },
|
|
last_y: 0,
|
|
computeSize: function () {
|
|
return [node.size[0], 20];
|
|
},
|
|
// ── Custom draw ──
|
|
draw: function (ctx, nodeRef, widgetWidth, posY, H) {
|
|
this.last_y = posY;
|
|
const margin = 15;
|
|
const innerW = widgetWidth - margin * 2;
|
|
|
|
// Background bar (standard LiteGraph color)
|
|
const bgColor = LiteGraph.WIDGET_BGCOLOR || "#29293e";
|
|
ctx.fillStyle = bgColor;
|
|
ctx.beginPath();
|
|
ctx.roundRect
|
|
? ctx.roundRect(margin, posY, innerW, H, H * 0.3)
|
|
: ctx.rect(margin, posY, innerW, H);
|
|
ctx.fill();
|
|
|
|
// Toggle circle (left side)
|
|
const toggleX = margin + 14;
|
|
const toggleY = posY + H * 0.5;
|
|
const r = 5;
|
|
ctx.beginPath();
|
|
ctx.arc(toggleX, toggleY, r, 0, Math.PI * 2);
|
|
ctx.fillStyle = this._pdb_muted ? "#555" : "#7cacf8";
|
|
ctx.fill();
|
|
if (!this._pdb_muted) {
|
|
ctx.beginPath();
|
|
ctx.arc(toggleX - 1, toggleY - 1, r * 0.35, 0, Math.PI * 2);
|
|
ctx.fillStyle = "rgba(255,255,255,0.25)";
|
|
ctx.fill();
|
|
}
|
|
|
|
// Category name (after circle)
|
|
ctx.fillStyle = this._pdb_muted ? "#777" : "#ddd";
|
|
ctx.font = "12px sans-serif";
|
|
ctx.textAlign = "left";
|
|
ctx.textBaseline = "middle";
|
|
const catText = `📂 ${closureCat}`;
|
|
ctx.fillText(catText, margin + 28, posY + H * 0.5);
|
|
const catTextW = ctx.measureText(catText).width;
|
|
const catEnd = margin + 28 + catTextW + 8; // where category text ends
|
|
|
|
// Right side layout: ◀ PromptName ▶ (right-aligned)
|
|
const rightEdge = margin + innerW - 10;
|
|
const displayVal = (this.value || "").replace(/\s*📖$/, "");
|
|
|
|
// Measure prompt name and truncate if overlaps category
|
|
ctx.font = "12px sans-serif";
|
|
const maxTextW = rightEdge - catEnd - 36; // space for both arrows + gaps
|
|
let valText = displayVal;
|
|
if (ctx.measureText(valText).width > maxTextW) {
|
|
while (ctx.measureText(valText + "…").width > maxTextW && valText.length > 1) {
|
|
valText = valText.slice(0, -1);
|
|
}
|
|
valText = valText + "…";
|
|
}
|
|
|
|
const textW = ctx.measureText(valText).width;
|
|
// Positions: ...◀ [text] ▶
|
|
const arrowRight = rightEdge;
|
|
const textRight = arrowRight - 14; // text right edge (before ▶)
|
|
const textLeft = textRight - textW; // text left edge
|
|
const arrowLeftX = textLeft - 14; // ◀ position
|
|
|
|
// Cache for mouse handler
|
|
this._pdb_arrowLeft = arrowLeftX;
|
|
this._pdb_catEnd = catEnd;
|
|
|
|
// ◀ button
|
|
ctx.fillStyle = "#aaa";
|
|
ctx.textAlign = "center";
|
|
ctx.font = "10px sans-serif";
|
|
ctx.fillText("◀", arrowLeftX, posY + H * 0.5);
|
|
|
|
// Prompt name (right-aligned)
|
|
ctx.fillStyle = this._pdb_muted ? "#777" : "#ddd";
|
|
ctx.font = "12px sans-serif";
|
|
ctx.textAlign = "right";
|
|
ctx.fillText(valText, textRight, posY + H * 0.5);
|
|
|
|
// ▶ button
|
|
ctx.fillStyle = "#aaa";
|
|
ctx.textAlign = "center";
|
|
ctx.font = "10px sans-serif";
|
|
ctx.fillText("▶", arrowRight, posY + H * 0.5);
|
|
|
|
ctx.textBaseline = "alphabetic";
|
|
},
|
|
// ── Custom mouse ──
|
|
mouse: function (event, pos, nodeRef) {
|
|
if (event.type !== "pointerdown" && event.type !== "mousedown")
|
|
return false;
|
|
|
|
const x = pos[0];
|
|
const widgetWidth = nodeRef.size[0];
|
|
const margin = 15;
|
|
const innerW = widgetWidth - margin * 2;
|
|
const rightEdge = margin + innerW - 10;
|
|
const arrowLeftX = this._pdb_arrowLeft || (margin + 120);
|
|
|
|
// Toggle area (left 45px — circle click)
|
|
if (x < margin + 45) {
|
|
this._pdb_muted = !this._pdb_muted;
|
|
localStorage.setItem(muteKey, this._pdb_muted ? "1" : "0");
|
|
syncSelectionsJSON();
|
|
updateCombinerLoras();
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
return true;
|
|
}
|
|
|
|
// ◀ button area
|
|
if (x > arrowLeftX - 10 && x < arrowLeftX + 10) {
|
|
const idx = this._pdb_values.indexOf(this.value);
|
|
const newIdx = (idx - 1 + this._pdb_values.length) % this._pdb_values.length;
|
|
this.value = this._pdb_values[newIdx];
|
|
syncSelectionsJSON();
|
|
updateCombinerLoras();
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
return true;
|
|
}
|
|
|
|
// ▶ button area
|
|
if (x > rightEdge - 10) {
|
|
const idx = this._pdb_values.indexOf(this.value);
|
|
const newIdx = (idx + 1) % this._pdb_values.length;
|
|
this.value = this._pdb_values[newIdx];
|
|
syncSelectionsJSON();
|
|
updateCombinerLoras();
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
return true;
|
|
}
|
|
|
|
// Text area click — open dropdown
|
|
if (x > arrowLeftX + 10 && x < rightEdge - 10) {
|
|
const that = this;
|
|
const canvas = app.canvas;
|
|
const ref_window = canvas.getCanvasWindow();
|
|
// Clean display names for the menu
|
|
const cleanValues = this._pdb_values.map(v => v.replace(/\s*📖$/, ""));
|
|
new LiteGraph.ContextMenu(
|
|
this._pdb_values,
|
|
{
|
|
scale: Math.max(1, canvas.ds?.scale || 1),
|
|
event: event,
|
|
className: "dark",
|
|
callback: function (v) {
|
|
if (v) {
|
|
that.value = v;
|
|
syncSelectionsJSON();
|
|
updateCombinerLoras();
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
}
|
|
},
|
|
},
|
|
ref_window
|
|
);
|
|
return true;
|
|
}
|
|
|
|
return false;
|
|
},
|
|
};
|
|
node.widgets.push(widget);
|
|
|
|
// Add LoRA info strips for the currently selected prompt
|
|
const cleanVal = defaultVal.replace(/\s*📖$/, "");
|
|
const selectedPrompt = allPrompts.find((p) => p.name === cleanVal);
|
|
if (selectedPrompt && selectedPrompt.loras && selectedPrompt.loras.length > 0) {
|
|
for (const lora of selectedPrompt.loras) {
|
|
const dName = (lora.name || "").replace(/\.safetensors$/, "");
|
|
const str = lora.strength != null ? lora.strength : 1.0;
|
|
const loraW = node.addWidget(
|
|
"text",
|
|
` 🔗 ${dName}`,
|
|
`${str.toFixed(2)}`,
|
|
() => {},
|
|
{ serialize: false }
|
|
);
|
|
loraW._pdb_lora_display = true;
|
|
loraW.disabled = true;
|
|
}
|
|
}
|
|
}
|
|
|
|
syncSelectionsJSON();
|
|
const computed = node.computeSize?.() || [userWidth, 100];
|
|
node.setSize?.([userWidth, computed[1]]);
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
}
|
|
|
|
// Update LoRA display strips when combo selection changes
|
|
async function updateCombinerLoras() {
|
|
// Remove old lora display widgets
|
|
if (node.widgets) {
|
|
node.widgets = node.widgets.filter((w) => !w._pdb_lora_display);
|
|
}
|
|
|
|
// Re-add LoRA strips after each combo widget
|
|
const combos = (node.widgets || []).filter((w) => w._pdb_dynamic);
|
|
// We need to insert in order, so collect all widgets to add
|
|
const toInsert = []; // { afterWidget, loraWidgets[] }
|
|
|
|
for (const combo of combos) {
|
|
const cleanVal = (combo.value || "").replace(/\s*📖$/, "");
|
|
const cat = combo._pdb_category;
|
|
if (!cleanVal || cleanVal === "(no prompts)" || !cat) continue;
|
|
|
|
let prompts = combo._pdb_allPrompts;
|
|
if (!prompts) {
|
|
try {
|
|
prompts = await API.getPrompts(cat);
|
|
} catch (e) {
|
|
continue;
|
|
}
|
|
}
|
|
|
|
const prompt = prompts.find((p) => p.name === cleanVal);
|
|
if (!prompt || !prompt.loras || prompt.loras.length === 0) continue;
|
|
|
|
toInsert.push({ afterWidget: combo, loras: prompt.loras });
|
|
}
|
|
|
|
// Insert LoRA widgets after their combo (in reverse to keep indices stable)
|
|
for (const item of toInsert.reverse()) {
|
|
const idx = node.widgets.indexOf(item.afterWidget);
|
|
if (idx === -1) continue;
|
|
|
|
const newWidgets = [];
|
|
for (const lora of item.loras) {
|
|
const dName = (lora.name || "").replace(/\.safetensors$/, "");
|
|
const str = lora.strength != null ? lora.strength : 1.0;
|
|
const w = {
|
|
type: "text",
|
|
name: ` 🔗 ${dName}`,
|
|
value: `${str.toFixed(2)}`,
|
|
options: { serialize: false },
|
|
_pdb_lora_display: true,
|
|
disabled: true,
|
|
callback: () => {},
|
|
computeSize: function (width) {
|
|
return [width, 20];
|
|
},
|
|
};
|
|
newWidgets.push(w);
|
|
}
|
|
// Insert after the combo widget
|
|
node.widgets.splice(idx + 1, 0, ...newWidgets);
|
|
}
|
|
|
|
const computed = node.computeSize?.() || [userWidth, 100];
|
|
node.setSize?.([userWidth, computed[1]]);
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
}
|
|
|
|
function scheduleRebuild() {
|
|
if (rebuildTimer) clearTimeout(rebuildTimer);
|
|
rebuildTimer = setTimeout(async () => {
|
|
if (rebuildPending) return;
|
|
rebuildPending = true;
|
|
try {
|
|
await fullRebuild();
|
|
} finally {
|
|
rebuildPending = false;
|
|
}
|
|
}, 150);
|
|
}
|
|
|
|
// Refresh button
|
|
node.addWidget("button", "🔄 Refresh", null, () => scheduleRebuild());
|
|
|
|
const origOnConnectionsChange = node.onConnectionsChange;
|
|
node.onConnectionsChange = function (
|
|
type,
|
|
slotIndex,
|
|
isConnected,
|
|
link,
|
|
ioSlot
|
|
) {
|
|
if (origOnConnectionsChange) origOnConnectionsChange.apply(this, arguments);
|
|
if (type === 1) scheduleRebuild();
|
|
};
|
|
|
|
// Periodically check for category changes
|
|
node._pdb_refreshInterval = setInterval(() => {
|
|
if (!node.graph) {
|
|
clearInterval(node._pdb_refreshInterval);
|
|
return;
|
|
}
|
|
let needsRebuild = false;
|
|
for (let i = 0; i < (node.inputs?.length || 0); i++) {
|
|
const inp = node.inputs[i];
|
|
if (!inp.name.startsWith("selection_") || inp.link == null) continue;
|
|
const category = getSourceCategory(i);
|
|
const existing = node.widgets?.find(
|
|
(w) => w._pdb_dynamic && w._pdb_inputName === inp.name
|
|
);
|
|
if (category && (!existing || existing._pdb_category !== category)) {
|
|
needsRebuild = true;
|
|
break;
|
|
}
|
|
}
|
|
if (needsRebuild) scheduleRebuild();
|
|
}, 1000);
|
|
|
|
node.setSize?.([userWidth, node.size?.[1] || 100]);
|
|
scheduleRebuild();
|
|
}
|
|
|
|
// ─── Setup ApplyLoraStack ───────────────────────────────────────────────────
|
|
const APPLY_LORA_WIDTH = 450;
|
|
|
|
function setupApplyLoraNode(node) {
|
|
if (node._pdb_apply_setup) return;
|
|
node._pdb_apply_setup = true;
|
|
console.log("[PromptDB] Setting up ApplyLoraStack node", node.id);
|
|
|
|
// Find and hide the lora_config text widget
|
|
const configWidget = node.widgets?.find((w) => w.name === "lora_config");
|
|
if (configWidget) {
|
|
configWidget.type = "converted-widget";
|
|
configWidget.computeSize = () => [0, -4];
|
|
configWidget.serializeValue = () => configWidget.value;
|
|
}
|
|
|
|
let rebuildPending = false;
|
|
let rebuildTimer = null;
|
|
|
|
// Get current config overrides
|
|
function getConfig() {
|
|
try {
|
|
return JSON.parse(configWidget?.value || "{}");
|
|
} catch (e) {
|
|
return {};
|
|
}
|
|
}
|
|
|
|
// Save config back to hidden widget
|
|
function saveConfig(config) {
|
|
if (configWidget) configWidget.value = JSON.stringify(config);
|
|
}
|
|
|
|
// Trace back through lora_stack input to discover LoRAs
|
|
async function discoverLoras() {
|
|
const loraStackInput = node.inputs?.find((inp) => inp.name === "lora_stack");
|
|
if (!loraStackInput || loraStackInput.link == null) return [];
|
|
|
|
const linkInfo = app.graph.links[loraStackInput.link];
|
|
if (!linkInfo) return [];
|
|
|
|
const sourceNode = app.graph.getNodeById(linkInfo.origin_id);
|
|
if (!sourceNode) return [];
|
|
|
|
const sourceClass = getNodeClass(sourceNode);
|
|
|
|
if (sourceClass === "PromptSelector") {
|
|
return await getLorasFromSelector(sourceNode);
|
|
} else if (sourceClass === "PromptCombiner") {
|
|
return await getLorasFromCombiner(sourceNode);
|
|
}
|
|
|
|
return [];
|
|
}
|
|
|
|
// Get LoRAs from a PromptSelector node
|
|
async function getLorasFromSelector(selectorNode) {
|
|
const catWidget = selectorNode.widgets?.find((w) => w.name === "category");
|
|
const promptWidget = selectorNode.widgets?.find(
|
|
(w) => w.name === "prompt_name"
|
|
);
|
|
if (!catWidget || !promptWidget) return [];
|
|
|
|
const cat = catWidget.value;
|
|
const rawName = promptWidget.value || "";
|
|
const cleanName = rawName.replace(/\s*📖$/, "");
|
|
|
|
try {
|
|
const prompts = await API.getPrompts(cat);
|
|
const prompt = prompts.find((p) => p.name === cleanName);
|
|
if (prompt && prompt.loras && prompt.loras.length > 0) {
|
|
return prompt.loras.map((l) => ({
|
|
name: l.name || "",
|
|
strength: l.strength != null ? l.strength : 1.0,
|
|
}));
|
|
}
|
|
} catch (e) {
|
|
console.error("[PromptDB] Failed to get LoRAs from selector:", e);
|
|
}
|
|
return [];
|
|
}
|
|
|
|
// Get LoRAs from a PromptCombiner node
|
|
async function getLorasFromCombiner(combinerNode) {
|
|
const allLoras = [];
|
|
const seen = new Set();
|
|
|
|
// Read all dynamic combo widgets to find selected prompts
|
|
const combos = (combinerNode.widgets || []).filter((w) => w._pdb_dynamic);
|
|
|
|
for (const combo of combos) {
|
|
const category = combo._pdb_category;
|
|
const rawVal = combo.value || "";
|
|
const cleanName = rawVal.replace(/\s*📖$/, "");
|
|
|
|
if (!category || !cleanName || cleanName === "(no prompts)") continue;
|
|
// Skip muted categories
|
|
if (combo._pdb_muted) continue;
|
|
|
|
try {
|
|
const prompts = await API.getPrompts(category);
|
|
const prompt = prompts.find((p) => p.name === cleanName);
|
|
if (prompt && prompt.loras) {
|
|
for (const l of prompt.loras) {
|
|
const name = l.name || "";
|
|
if (name && !seen.has(name)) {
|
|
seen.add(name);
|
|
allLoras.push({
|
|
name,
|
|
strength: l.strength != null ? l.strength : 1.0,
|
|
});
|
|
}
|
|
}
|
|
}
|
|
} catch (e) {
|
|
console.error("[PromptDB] Failed to get LoRAs from combiner:", e);
|
|
}
|
|
}
|
|
|
|
return allLoras;
|
|
}
|
|
|
|
// Full rebuild of LoRA display widgets
|
|
async function fullRebuild() {
|
|
console.log("[PromptDB] ApplyLoraStack rebuild", node.id);
|
|
|
|
// Remove all dynamic LoRA widgets
|
|
if (node.widgets) {
|
|
node.widgets = node.widgets.filter((w) => !w._pdb_lora_widget);
|
|
}
|
|
|
|
const loras = await discoverLoras();
|
|
if (loras.length === 0) {
|
|
const emptyW = node.addWidget(
|
|
"text",
|
|
"Status",
|
|
"No LoRAs connected",
|
|
() => {},
|
|
{ serialize: false }
|
|
);
|
|
emptyW._pdb_lora_widget = true;
|
|
emptyW.disabled = true;
|
|
|
|
const computed = node.computeSize?.() || [APPLY_LORA_WIDTH, 100];
|
|
node.setSize?.([Math.max(computed[0], APPLY_LORA_WIDTH), computed[1]]);
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
return;
|
|
}
|
|
|
|
const config = getConfig();
|
|
|
|
for (const lora of loras) {
|
|
const displayName = (lora.name || "").replace(/\.safetensors$/, "");
|
|
const override = config[lora.name] || {};
|
|
const isEnabled = override.enabled !== false;
|
|
const str =
|
|
override.strength != null ? override.strength : lora.strength;
|
|
|
|
// ── Custom single-line widget: ● name ◀ 1.00 ▶ ──
|
|
const loraFullName = lora.name; // closure capture
|
|
const widget = {
|
|
type: "pdb_lora",
|
|
name: displayName,
|
|
value: str,
|
|
_pdb_lora_widget: true,
|
|
_pdb_lora_name: loraFullName,
|
|
_pdb_enabled: isEnabled,
|
|
options: { serialize: false },
|
|
last_y: 0,
|
|
computeSize: function () {
|
|
return [node.size[0], 20];
|
|
},
|
|
// ── Custom draw ──
|
|
draw: function (ctx, nodeRef, widgetWidth, posY, H) {
|
|
this.last_y = posY;
|
|
const margin = 15;
|
|
const innerW = widgetWidth - margin * 2;
|
|
|
|
// Background bar
|
|
ctx.fillStyle = "#232333";
|
|
ctx.beginPath();
|
|
ctx.roundRect
|
|
? ctx.roundRect(margin, posY, innerW, H, H * 0.3)
|
|
: ctx.rect(margin, posY, innerW, H);
|
|
ctx.fill();
|
|
|
|
// Toggle circle
|
|
const toggleX = margin + 14;
|
|
const toggleY = posY + H * 0.5;
|
|
const r = 5;
|
|
ctx.beginPath();
|
|
ctx.arc(toggleX, toggleY, r, 0, Math.PI * 2);
|
|
ctx.fillStyle = this._pdb_enabled ? "#7cacf8" : "#555";
|
|
ctx.fill();
|
|
if (this._pdb_enabled) {
|
|
ctx.beginPath();
|
|
ctx.arc(toggleX - 1, toggleY - 1, r * 0.35, 0, Math.PI * 2);
|
|
ctx.fillStyle = "rgba(255,255,255,0.25)";
|
|
ctx.fill();
|
|
}
|
|
|
|
// LoRA name
|
|
ctx.fillStyle = this._pdb_enabled ? "#ddd" : "#777";
|
|
ctx.font = "12px sans-serif";
|
|
ctx.textAlign = "left";
|
|
ctx.textBaseline = "middle";
|
|
ctx.fillText(this.name, margin + 28, posY + H * 0.5);
|
|
|
|
// ◀ value ▶ on right side
|
|
const rightEdge = margin + innerW - 8;
|
|
const valStr = this.value.toFixed(2);
|
|
|
|
// ▶ button
|
|
ctx.fillStyle = "#aaa";
|
|
ctx.textAlign = "center";
|
|
ctx.font = "10px sans-serif";
|
|
ctx.fillText("▶", rightEdge, posY + H * 0.5);
|
|
|
|
// Value
|
|
ctx.fillStyle = this._pdb_enabled ? "#ddd" : "#777";
|
|
ctx.font = "12px sans-serif";
|
|
ctx.fillText(valStr, rightEdge - 28, posY + H * 0.5);
|
|
|
|
// ◀ button
|
|
ctx.fillStyle = "#aaa";
|
|
ctx.font = "10px sans-serif";
|
|
ctx.fillText("◀", rightEdge - 56, posY + H * 0.5);
|
|
|
|
ctx.textBaseline = "alphabetic";
|
|
},
|
|
// ── Custom mouse ──
|
|
mouse: function (event, pos, nodeRef) {
|
|
if (
|
|
event.type !== "pointerdown" &&
|
|
event.type !== "mousedown"
|
|
)
|
|
return false;
|
|
|
|
const x = pos[0];
|
|
const widgetWidth = nodeRef.size[0];
|
|
const margin = 15;
|
|
const innerW = widgetWidth - margin * 2;
|
|
const rightEdge = margin + innerW - 8;
|
|
|
|
// Toggle area (left 50px)
|
|
if (x < margin + 45) {
|
|
this._pdb_enabled = !this._pdb_enabled;
|
|
const cfg = getConfig();
|
|
if (!cfg[this._pdb_lora_name]) cfg[this._pdb_lora_name] = {};
|
|
cfg[this._pdb_lora_name].enabled = this._pdb_enabled;
|
|
saveConfig(cfg);
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
return true;
|
|
}
|
|
|
|
// ◀ button (decrease strength)
|
|
if (x > rightEdge - 70 && x < rightEdge - 42) {
|
|
this.value = Math.max(-2, +(this.value - 0.05).toFixed(2));
|
|
const cfg = getConfig();
|
|
if (!cfg[this._pdb_lora_name]) cfg[this._pdb_lora_name] = {};
|
|
cfg[this._pdb_lora_name].strength = this.value;
|
|
saveConfig(cfg);
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
return true;
|
|
}
|
|
|
|
// ▶ button (increase strength)
|
|
if (x > rightEdge - 14 && x < rightEdge + 14) {
|
|
this.value = Math.min(2, +(this.value + 0.05).toFixed(2));
|
|
const cfg = getConfig();
|
|
if (!cfg[this._pdb_lora_name]) cfg[this._pdb_lora_name] = {};
|
|
cfg[this._pdb_lora_name].strength = this.value;
|
|
saveConfig(cfg);
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
return true;
|
|
}
|
|
|
|
return false;
|
|
},
|
|
};
|
|
|
|
node.widgets = node.widgets || [];
|
|
node.widgets.push(widget);
|
|
}
|
|
|
|
// ── Add Reset Values button ──
|
|
const resetBtn = node.addWidget(
|
|
"button",
|
|
"🔄 Reset Values",
|
|
null,
|
|
async () => {
|
|
// Clear all overrides
|
|
if (configWidget) configWidget.value = "{}";
|
|
await fullRebuild();
|
|
},
|
|
{ serialize: false }
|
|
);
|
|
resetBtn._pdb_lora_widget = true;
|
|
|
|
const computed = node.computeSize?.() || [APPLY_LORA_WIDTH, 100];
|
|
node.setSize?.([Math.max(computed[0], APPLY_LORA_WIDTH), computed[1]]);
|
|
app.graph?.setDirtyCanvas?.(true);
|
|
}
|
|
|
|
function scheduleRebuild() {
|
|
if (rebuildTimer) clearTimeout(rebuildTimer);
|
|
rebuildTimer = setTimeout(async () => {
|
|
if (rebuildPending) return;
|
|
rebuildPending = true;
|
|
try {
|
|
await fullRebuild();
|
|
} finally {
|
|
rebuildPending = false;
|
|
}
|
|
}, 200);
|
|
}
|
|
|
|
// ── Refresh button ──
|
|
node.addWidget("button", "🔄 Refresh", null, () => scheduleRebuild());
|
|
|
|
// ── React to connection changes ──
|
|
const origOnConnectionsChange = node.onConnectionsChange;
|
|
node.onConnectionsChange = function (
|
|
type,
|
|
slotIndex,
|
|
isConnected,
|
|
link,
|
|
ioSlot
|
|
) {
|
|
if (origOnConnectionsChange) origOnConnectionsChange.apply(this, arguments);
|
|
scheduleRebuild();
|
|
};
|
|
|
|
// Periodically check for changes in connected source nodes
|
|
let lastLoraFingerprint = "";
|
|
node._pdb_refreshInterval = setInterval(async () => {
|
|
if (!node.graph) {
|
|
clearInterval(node._pdb_refreshInterval);
|
|
return;
|
|
}
|
|
|
|
try {
|
|
const loras = await discoverLoras();
|
|
const fingerprint = JSON.stringify(loras.map((l) => l.name).sort());
|
|
if (fingerprint !== lastLoraFingerprint) {
|
|
lastLoraFingerprint = fingerprint;
|
|
scheduleRebuild();
|
|
}
|
|
} catch (e) {}
|
|
}, 1500);
|
|
|
|
// Initial setup
|
|
node.setSize?.([APPLY_LORA_WIDTH, node.size?.[1] || 100]);
|
|
scheduleRebuild();
|
|
}
|
|
|
|
// ─── Register extension ─────────────────────────────────────────────────────
|
|
app.registerExtension({
|
|
name: "comfy.promptdb",
|
|
|
|
async beforeRegisterNodeDef(nodeType, nodeData, _app) {
|
|
if (nodeData.name === "PromptDBManager") {
|
|
const orig = nodeType.prototype.onNodeCreated;
|
|
nodeType.prototype.onNodeCreated = function () {
|
|
orig?.apply(this, arguments);
|
|
setTimeout(() => setupManagerNode(this), 100);
|
|
};
|
|
}
|
|
if (nodeData.name === "PromptSelector") {
|
|
const orig = nodeType.prototype.onNodeCreated;
|
|
nodeType.prototype.onNodeCreated = function () {
|
|
orig?.apply(this, arguments);
|
|
setTimeout(() => setupSelectorNode(this), 100);
|
|
};
|
|
}
|
|
if (nodeData.name === "PromptCombiner") {
|
|
const orig = nodeType.prototype.onNodeCreated;
|
|
nodeType.prototype.onNodeCreated = function () {
|
|
orig?.apply(this, arguments);
|
|
setTimeout(() => setupCombinerNode(this), 100);
|
|
};
|
|
}
|
|
if (nodeData.name === "ApplyLoraStack") {
|
|
const orig = nodeType.prototype.onNodeCreated;
|
|
nodeType.prototype.onNodeCreated = function () {
|
|
orig?.apply(this, arguments);
|
|
setTimeout(() => setupApplyLoraNode(this), 100);
|
|
};
|
|
}
|
|
},
|
|
|
|
nodeCreated(node) {
|
|
const cls = getNodeClass(node);
|
|
if (cls === "PromptDBManager")
|
|
setTimeout(() => setupManagerNode(node), 150);
|
|
if (cls === "PromptSelector")
|
|
setTimeout(() => setupSelectorNode(node), 150);
|
|
if (cls === "PromptCombiner")
|
|
setTimeout(() => setupCombinerNode(node), 150);
|
|
if (cls === "ApplyLoraStack")
|
|
setTimeout(() => setupApplyLoraNode(node), 150);
|
|
},
|
|
});
|
|
|
|
console.log("[PromptDB] Extension registered.");
|