415 lines
12 KiB
JavaScript
415 lines
12 KiB
JavaScript
import { app } from "../../scripts/app.js";
|
|
import { api } from "../../scripts/api.js";
|
|
|
|
/**
|
|
* Workflow Iterator - ComfyUI JavaScript Extension
|
|
*
|
|
* Provides:
|
|
* 1. Node target browser for WI Parameter nodes
|
|
* 2. Progress bar overlay for WorkflowIterator node
|
|
* 3. Value preview for WI Parameter nodes
|
|
*/
|
|
|
|
const WI_PARAM_TYPES = [
|
|
"WIParameterInt",
|
|
"WIParameterFloat",
|
|
"WIParameterString",
|
|
"WIParameterCombo",
|
|
"WIParameterSeed",
|
|
];
|
|
|
|
app.registerExtension({
|
|
name: "WorkflowIterator",
|
|
|
|
async beforeRegisterNodeDef(nodeType, nodeData, app) {
|
|
if (!nodeData.name) return;
|
|
|
|
// WorkflowIterator main node - add progress overlay
|
|
if (nodeData.name === "WorkflowIterator") {
|
|
const origOnDrawForeground = nodeType.prototype.onDrawForeground;
|
|
|
|
nodeType.prototype.onDrawForeground = function (ctx) {
|
|
if (origOnDrawForeground) {
|
|
origOnDrawForeground.apply(this, arguments);
|
|
}
|
|
drawProgressOverlay(this, ctx);
|
|
};
|
|
|
|
const origOnExecuted = nodeType.prototype.onExecuted;
|
|
nodeType.prototype.onExecuted = function (output) {
|
|
if (origOnExecuted) {
|
|
origOnExecuted.apply(this, arguments);
|
|
}
|
|
// Start polling for batch status
|
|
startBatchPolling(this);
|
|
};
|
|
}
|
|
|
|
// WI Parameter nodes - add browse button and value preview
|
|
if (WI_PARAM_TYPES.includes(nodeData.name)) {
|
|
const origOnNodeCreated = nodeType.prototype.onNodeCreated;
|
|
|
|
nodeType.prototype.onNodeCreated = function () {
|
|
if (origOnNodeCreated) {
|
|
origOnNodeCreated.apply(this, arguments);
|
|
}
|
|
addBrowseButton(this);
|
|
addValuePreview(this);
|
|
};
|
|
}
|
|
},
|
|
});
|
|
|
|
// ============================================================
|
|
// Progress Overlay for WorkflowIterator
|
|
// ============================================================
|
|
|
|
const batchStatusCache = new Map();
|
|
|
|
function drawProgressOverlay(node, ctx) {
|
|
const status = batchStatusCache.get(node.id);
|
|
if (!status || status.total <= 1) return;
|
|
|
|
const progress = status.completed / status.total;
|
|
const barHeight = 6;
|
|
const barY = node.size[1] - barHeight - 4;
|
|
const barWidth = node.size[0] - 20;
|
|
const barX = 10;
|
|
|
|
// Background
|
|
ctx.fillStyle = "rgba(0, 0, 0, 0.5)";
|
|
ctx.fillRect(barX, barY, barWidth, barHeight);
|
|
|
|
// Progress fill
|
|
ctx.fillStyle = status.cancelled
|
|
? "#ff4444"
|
|
: progress >= 1
|
|
? "#44ff44"
|
|
: "#4488ff";
|
|
ctx.fillRect(barX, barY, barWidth * progress, barHeight);
|
|
|
|
// Border
|
|
ctx.strokeStyle = "rgba(255, 255, 255, 0.3)";
|
|
ctx.strokeRect(barX, barY, barWidth, barHeight);
|
|
|
|
// Text
|
|
const text = status.cancelled
|
|
? `Cancelled (${status.completed}/${status.total})`
|
|
: `${status.completed}/${status.total}`;
|
|
ctx.fillStyle = "#ffffff";
|
|
ctx.font = "11px sans-serif";
|
|
ctx.textAlign = "center";
|
|
ctx.fillText(text, node.size[0] / 2, barY - 4);
|
|
ctx.textAlign = "left";
|
|
}
|
|
|
|
let pollInterval = null;
|
|
|
|
function startBatchPolling(node) {
|
|
if (pollInterval) return;
|
|
|
|
pollInterval = setInterval(async () => {
|
|
try {
|
|
const resp = await api.fetchApi("/wi/active_batches");
|
|
if (resp.ok) {
|
|
const batches = await resp.json();
|
|
const entries = Object.entries(batches);
|
|
|
|
if (entries.length === 0) {
|
|
clearInterval(pollInterval);
|
|
pollInterval = null;
|
|
batchStatusCache.delete(node.id);
|
|
app.graph.setDirtyCanvas(true);
|
|
return;
|
|
}
|
|
|
|
// Use the most recent active batch
|
|
const [batchId, status] = entries[entries.length - 1];
|
|
batchStatusCache.set(node.id, status);
|
|
app.graph.setDirtyCanvas(true);
|
|
|
|
// Stop polling if complete or cancelled
|
|
if (status.completed >= status.total || status.cancelled) {
|
|
setTimeout(() => {
|
|
clearInterval(pollInterval);
|
|
pollInterval = null;
|
|
}, 3000);
|
|
}
|
|
}
|
|
} catch (e) {
|
|
// Silently ignore polling errors
|
|
}
|
|
}, 1000);
|
|
}
|
|
|
|
// ============================================================
|
|
// Browse Button for WI Parameter Nodes
|
|
// ============================================================
|
|
|
|
function addBrowseButton(node) {
|
|
const browseWidget = node.addWidget(
|
|
"button",
|
|
"Browse Target Node",
|
|
null,
|
|
() => {
|
|
showNodeBrowser(node);
|
|
}
|
|
);
|
|
browseWidget.serialize = false;
|
|
}
|
|
|
|
function showNodeBrowser(paramNode) {
|
|
// Get all nodes in the graph
|
|
const allNodes = app.graph._nodes;
|
|
if (!allNodes || allNodes.length === 0) return;
|
|
|
|
// Build list of unique node titles/types
|
|
const nodeOptions = [];
|
|
const seen = new Set();
|
|
|
|
for (const node of allNodes) {
|
|
if (node === paramNode) continue;
|
|
// Skip other WI nodes
|
|
if (
|
|
node.type &&
|
|
(node.type.startsWith("WIParameter") ||
|
|
node.type === "WorkflowIterator")
|
|
)
|
|
continue;
|
|
|
|
const title = node.title || node.type || `Node #${node.id}`;
|
|
const key = `${title}__${node.id}`;
|
|
if (seen.has(key)) continue;
|
|
seen.add(key);
|
|
|
|
nodeOptions.push({
|
|
title: title,
|
|
type: node.type || "",
|
|
id: node.id,
|
|
node: node,
|
|
});
|
|
}
|
|
|
|
// Sort by title
|
|
nodeOptions.sort((a, b) => a.title.localeCompare(b.title));
|
|
|
|
// Create modal dialog
|
|
const dialog = document.createElement("dialog");
|
|
dialog.style.cssText = `
|
|
background: #1a1a2e; color: #eee; border: 1px solid #444;
|
|
border-radius: 8px; padding: 16px; min-width: 350px;
|
|
max-height: 70vh; font-family: sans-serif;
|
|
`;
|
|
|
|
const title = document.createElement("h3");
|
|
title.textContent = "Select Target Node";
|
|
title.style.cssText = "margin: 0 0 12px 0; font-size: 14px;";
|
|
dialog.appendChild(title);
|
|
|
|
// Search input
|
|
const searchInput = document.createElement("input");
|
|
searchInput.type = "text";
|
|
searchInput.placeholder = "Filter nodes...";
|
|
searchInput.style.cssText = `
|
|
width: 100%; padding: 6px; margin-bottom: 8px;
|
|
background: #2a2a4e; color: #eee; border: 1px solid #555;
|
|
border-radius: 4px; box-sizing: border-box;
|
|
`;
|
|
dialog.appendChild(searchInput);
|
|
|
|
// Node list
|
|
const listContainer = document.createElement("div");
|
|
listContainer.style.cssText = `
|
|
max-height: 300px; overflow-y: auto; margin-bottom: 12px;
|
|
`;
|
|
|
|
function renderList(filter) {
|
|
listContainer.innerHTML = "";
|
|
const filtered = filter
|
|
? nodeOptions.filter(
|
|
(n) =>
|
|
n.title.toLowerCase().includes(filter.toLowerCase()) ||
|
|
n.type.toLowerCase().includes(filter.toLowerCase())
|
|
)
|
|
: nodeOptions;
|
|
|
|
for (const opt of filtered) {
|
|
const item = document.createElement("div");
|
|
item.style.cssText = `
|
|
padding: 6px 8px; cursor: pointer; border-radius: 4px;
|
|
margin: 2px 0; font-size: 13px;
|
|
`;
|
|
item.innerHTML = `<strong>${opt.title}</strong> <span style="color:#888">(${opt.type} #${opt.id})</span>`;
|
|
|
|
item.onmouseenter = () => (item.style.background = "#333");
|
|
item.onmouseleave = () => (item.style.background = "");
|
|
item.onclick = () => {
|
|
dialog.close();
|
|
selectNodeTarget(paramNode, opt);
|
|
};
|
|
|
|
listContainer.appendChild(item);
|
|
}
|
|
}
|
|
|
|
renderList("");
|
|
searchInput.addEventListener("input", (e) => renderList(e.target.value));
|
|
|
|
dialog.appendChild(listContainer);
|
|
|
|
// Cancel button
|
|
const cancelBtn = document.createElement("button");
|
|
cancelBtn.textContent = "Cancel";
|
|
cancelBtn.style.cssText = `
|
|
padding: 6px 16px; background: #444; color: #eee;
|
|
border: none; border-radius: 4px; cursor: pointer;
|
|
`;
|
|
cancelBtn.onclick = () => dialog.close();
|
|
dialog.appendChild(cancelBtn);
|
|
|
|
document.body.appendChild(dialog);
|
|
dialog.showModal();
|
|
searchInput.focus();
|
|
|
|
dialog.addEventListener("close", () => dialog.remove());
|
|
}
|
|
|
|
function selectNodeTarget(paramNode, selectedNode) {
|
|
// Set the target_node_title widget
|
|
const titleWidget = paramNode.widgets?.find(
|
|
(w) => w.name === "target_node_title"
|
|
);
|
|
if (titleWidget) {
|
|
titleWidget.value = selectedNode.title;
|
|
}
|
|
|
|
// Now show widget picker for the selected node
|
|
showWidgetPicker(paramNode, selectedNode.node);
|
|
}
|
|
|
|
function showWidgetPicker(paramNode, targetNode) {
|
|
if (!targetNode.widgets || targetNode.widgets.length === 0) return;
|
|
|
|
const dialog = document.createElement("dialog");
|
|
dialog.style.cssText = `
|
|
background: #1a1a2e; color: #eee; border: 1px solid #444;
|
|
border-radius: 8px; padding: 16px; min-width: 300px;
|
|
font-family: sans-serif;
|
|
`;
|
|
|
|
const title = document.createElement("h3");
|
|
title.textContent = `Select Widget on "${targetNode.title || targetNode.type}"`;
|
|
title.style.cssText = "margin: 0 0 12px 0; font-size: 14px;";
|
|
dialog.appendChild(title);
|
|
|
|
const listContainer = document.createElement("div");
|
|
listContainer.style.cssText = "max-height: 300px; overflow-y: auto;";
|
|
|
|
for (const widget of targetNode.widgets) {
|
|
if (!widget.name || widget.name.startsWith("_")) continue;
|
|
|
|
const item = document.createElement("div");
|
|
item.style.cssText = `
|
|
padding: 6px 8px; cursor: pointer; border-radius: 4px;
|
|
margin: 2px 0; font-size: 13px;
|
|
`;
|
|
|
|
const widgetType = widget.type || "unknown";
|
|
const currentValue =
|
|
widget.value !== undefined ? ` = ${widget.value}` : "";
|
|
item.innerHTML = `<strong>${widget.name}</strong> <span style="color:#888">(${widgetType}${currentValue})</span>`;
|
|
|
|
item.onmouseenter = () => (item.style.background = "#333");
|
|
item.onmouseleave = () => (item.style.background = "");
|
|
item.onclick = () => {
|
|
const widgetNameWidget = paramNode.widgets?.find(
|
|
(w) => w.name === "target_widget_name"
|
|
);
|
|
if (widgetNameWidget) {
|
|
widgetNameWidget.value = widget.name;
|
|
}
|
|
dialog.close();
|
|
};
|
|
|
|
listContainer.appendChild(item);
|
|
}
|
|
|
|
dialog.appendChild(listContainer);
|
|
|
|
const cancelBtn = document.createElement("button");
|
|
cancelBtn.textContent = "Cancel";
|
|
cancelBtn.style.cssText = `
|
|
padding: 6px 16px; background: #444; color: #eee;
|
|
border: none; border-radius: 4px; cursor: pointer; margin-top: 8px;
|
|
`;
|
|
cancelBtn.onclick = () => dialog.close();
|
|
dialog.appendChild(cancelBtn);
|
|
|
|
document.body.appendChild(dialog);
|
|
dialog.showModal();
|
|
dialog.addEventListener("close", () => dialog.remove());
|
|
}
|
|
|
|
// ============================================================
|
|
// Value Preview for WI Parameter Nodes
|
|
// ============================================================
|
|
|
|
function addValuePreview(node) {
|
|
const previewWidget = node.addWidget("text", "preview", "", () => {}, {
|
|
serialize: false,
|
|
multiline: false,
|
|
});
|
|
previewWidget.computeSize = () => [0, 20];
|
|
|
|
// Override draw to show parsed value count
|
|
const origDraw = node.onDrawForeground;
|
|
node.onDrawForeground = function (ctx) {
|
|
if (origDraw) origDraw.apply(this, arguments);
|
|
|
|
const valuesWidget = this.widgets?.find((w) => w.name === "values");
|
|
if (!valuesWidget) return;
|
|
|
|
const raw = valuesWidget.value || "";
|
|
const typeWidget = this.type;
|
|
let count = 0;
|
|
|
|
try {
|
|
if (raw.includes("-") && raw.includes(",") && !raw.includes("\n")) {
|
|
// Range syntax - estimate count
|
|
const parts = raw.split(",").map((s) => s.trim());
|
|
if (parts.length === 2) {
|
|
const rangeParts = parts[0].split("-");
|
|
if (rangeParts.length === 2) {
|
|
const start = parseFloat(rangeParts[0]);
|
|
const end = parseFloat(rangeParts[1]);
|
|
const step = parseFloat(parts[1]);
|
|
if (!isNaN(start) && !isNaN(end) && !isNaN(step) && step > 0) {
|
|
count = Math.floor(Math.abs(end - start) / step) + 1;
|
|
}
|
|
}
|
|
}
|
|
} else if (raw.includes("\n")) {
|
|
count = raw
|
|
.split("\n")
|
|
.filter((l) => l.trim()).length;
|
|
} else {
|
|
count = raw
|
|
.split(",")
|
|
.filter((p) => p.trim()).length;
|
|
}
|
|
} catch (e) {
|
|
count = 0;
|
|
}
|
|
|
|
if (count > 0) {
|
|
const text = `${count} value${count !== 1 ? "s" : ""}`;
|
|
const y = this.size[1] - 18;
|
|
ctx.fillStyle = "rgba(100, 180, 255, 0.7)";
|
|
ctx.font = "10px sans-serif";
|
|
ctx.textAlign = "right";
|
|
ctx.fillText(text, this.size[0] - 10, y);
|
|
ctx.textAlign = "left";
|
|
}
|
|
};
|
|
}
|