443 lines
18 KiB
JavaScript
443 lines
18 KiB
JavaScript
import { app } from "../../scripts/app.js";
|
|
import { collectChain } from "./widgets/generation-data.js";
|
|
import { serializeModelOverrides } from "./utils.js";
|
|
|
|
export function getComboFlakes(node) {
|
|
return node.properties?._combo_flakes || [];
|
|
}
|
|
|
|
export function getComboPresets(node) {
|
|
return node.properties?._combo_presets || [];
|
|
}
|
|
|
|
// Active (non-bypassed) flakes for a FlakeCombo node.
|
|
function activeComboFlakes(node) {
|
|
return getComboFlakes(node).filter(f => !f.bypassed);
|
|
}
|
|
|
|
// Active (non-bypassed) presets for a FlakeModelCombo node. Bypass is keyed by
|
|
// preset name in node.properties._combo_bypassed.
|
|
function activeComboPresets(node) {
|
|
const bypassed = new Set(node.properties?._combo_bypassed || []);
|
|
return getComboPresets(node).filter(p => !bypassed.has(p));
|
|
}
|
|
|
|
// Number of prompts a generation will queue: the product of each
|
|
// FlakeCombo's active-flake count and each FlakeModelCombo's preset count.
|
|
// Returns 1 when there are combo nodes but a single combination, or 1 when
|
|
// there are no combo nodes at all (a normal single queue).
|
|
export function computeJobCount(graph) {
|
|
const g = graph || app.graph;
|
|
if (!g) return 1;
|
|
let count = 1;
|
|
for (const node of g.nodes) {
|
|
if (node.type === "FlakeCombo") {
|
|
const n = activeComboFlakes(node).length;
|
|
if (n > 0) count *= n;
|
|
} else if (node.type === "FlakeModelCombo") {
|
|
const n = activeComboPresets(node).length;
|
|
if (n > 0) count *= n;
|
|
}
|
|
}
|
|
return count;
|
|
}
|
|
|
|
export function cartesianProduct(arrays) {
|
|
if (arrays.length === 0) return [[]];
|
|
const result = [];
|
|
const head = arrays[0];
|
|
const tail = cartesianProduct(arrays.slice(1));
|
|
for (const h of head) {
|
|
for (const t of tail) {
|
|
result.push([h, ...t]);
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
const _originalQueuePrompt = app.queuePrompt;
|
|
|
|
let _batchCleanup = null;
|
|
|
|
function getPromptIdFromEvent(e) {
|
|
const detail = e.detail;
|
|
return typeof detail === "string" ? detail : detail?.prompt_id;
|
|
}
|
|
|
|
function clearAllHighlights() {
|
|
for (const n of app.graph.nodes) {
|
|
if (n.type === "FlakeCombo") {
|
|
delete n._combo_generating_index;
|
|
n._combo_render?.();
|
|
}
|
|
if (n.type === "FlakeModelCombo") {
|
|
delete n._model_combo_generating_index;
|
|
n._model_combo_render?.();
|
|
}
|
|
}
|
|
}
|
|
|
|
// ComfyUI exposes the websocket-event-emitting `api` singleton in slightly
|
|
// different places depending on frontend version. Try several spots so we
|
|
// don't silently lose batch tracking on a newer/older frontend.
|
|
function _resolveComfyApi() {
|
|
return (
|
|
window.comfyAPI?.api?.api
|
|
|| window.comfyAPI?.api
|
|
|| app.api
|
|
|| window.app?.api
|
|
|| null
|
|
);
|
|
}
|
|
|
|
// Set by the outer batch loop before each combo's queue call; read by the
|
|
// patched api.queuePrompt so the captured prompt_id can be tagged with the
|
|
// right comboKey. Cleared between combos.
|
|
let _currentBatchComboKey = null;
|
|
|
|
// Active batch promptId → comboKey map, plus a live "running" pointer.
|
|
// Exported so the Generation Data overlay can show per-combo progress (#230).
|
|
const _activeBatch = { promptIds: [], comboKeys: [], runningPromptId: null };
|
|
export function getActiveBatch() {
|
|
return _activeBatch;
|
|
}
|
|
|
|
function setupBatchTracking(combinations, comboKeys) {
|
|
const comfyApi = _resolveComfyApi();
|
|
if (!comfyApi || typeof comfyApi.addEventListener !== "function" || combinations.length === 0) {
|
|
if (!comfyApi) {
|
|
// eslint-disable-next-line no-console
|
|
console.warn("[flakes] could not resolve ComfyUI api singleton; combo highlight/per-combo image storage disabled.");
|
|
}
|
|
return null;
|
|
}
|
|
|
|
const promptIds = []; // insertion order = combo index
|
|
let completedCount = 0;
|
|
|
|
// Pending event buffer: websocket events for a prompt_id can arrive BEFORE
|
|
// the REST response that gives us that prompt_id (different transport
|
|
// channels). Buffer events we don't yet know about and replay them when
|
|
// the patched api.queuePrompt finally registers the id. Without this, the
|
|
// first generation in a batch consistently lost its execution_start /
|
|
// executed events (#241).
|
|
const pendingByPid = new Map(); // promptId -> [{type, e}]
|
|
function applyEventForKnownPid(type, e, pid) {
|
|
if (type === "execution_start") {
|
|
const idx = promptIds.indexOf(pid);
|
|
if (idx >= 0) { setComboHighlight(idx); _activeBatch.runningPromptId = pid; }
|
|
} else if (type === "execution_success" || type === "execution_error" || type === "execution_interrupted") {
|
|
const idx = promptIds.indexOf(pid);
|
|
if (idx < 0) return;
|
|
completedCount++;
|
|
for (const n of app.graph.nodes) {
|
|
if (n.type === "FlakeGenerate") {
|
|
n._batch_completed_count = completedCount;
|
|
n._batch_total_count = combinations.length;
|
|
n._batch_progress_render?.();
|
|
}
|
|
}
|
|
maybeFinish();
|
|
} else if (type === "executed") {
|
|
handleExecuted(e);
|
|
} else if (type === "progress") {
|
|
// Progress events are handled by the overlay subscription; nothing
|
|
// to replay batch-side. We only buffer the lifecycle events.
|
|
}
|
|
}
|
|
function bufferOrApply(type, e) {
|
|
const pid = type === "executed" ? (typeof e?.detail === "string" ? null : e?.detail?.prompt_id) : getPromptIdFromEvent(e);
|
|
if (!pid) return;
|
|
if (promptIds.includes(pid)) {
|
|
applyEventForKnownPid(type, e, pid);
|
|
} else {
|
|
let arr = pendingByPid.get(pid);
|
|
if (!arr) { arr = []; pendingByPid.set(pid, arr); }
|
|
arr.push({ type, e });
|
|
}
|
|
}
|
|
function flushPendingFor(pid) {
|
|
const arr = pendingByPid.get(pid);
|
|
if (!arr) return;
|
|
pendingByPid.delete(pid);
|
|
for (const { type, e } of arr) applyEventForKnownPid(type, e, pid);
|
|
}
|
|
|
|
// Patch api.queuePrompt for the duration of this batch so we capture the
|
|
// returned prompt_id. (app.queuePrompt in current frontend versions
|
|
// returns void; the low-level api.queuePrompt still returns the
|
|
// {prompt_id, number, node_errors} REST response.)
|
|
const _origApiQueuePrompt = comfyApi.queuePrompt?.bind(comfyApi);
|
|
if (_origApiQueuePrompt) {
|
|
comfyApi.queuePrompt = async function(...args) {
|
|
// Snapshot the combo key now — _currentBatchComboKey may already be
|
|
// cleared by the time the await below resolves.
|
|
const comboKeyForThisCall = _currentBatchComboKey;
|
|
const result = await _origApiQueuePrompt(...args);
|
|
if (result?.prompt_id && comboKeyForThisCall !== null) {
|
|
promptIds.push(result.prompt_id);
|
|
comboKeys.push(comboKeyForThisCall);
|
|
_activeBatch.promptIds.push(result.prompt_id);
|
|
_activeBatch.comboKeys.push(comboKeyForThisCall);
|
|
// Replay any websocket events that arrived before this push.
|
|
flushPendingFor(result.prompt_id);
|
|
}
|
|
return result;
|
|
};
|
|
}
|
|
|
|
function setComboHighlight(comboIndex) {
|
|
const combination = combinations[comboIndex];
|
|
if (!combination) return;
|
|
for (const item of combination) {
|
|
if (item.type === "combo") {
|
|
item.node._combo_generating_index = item.index;
|
|
item.node._combo_render?.();
|
|
} else {
|
|
item.node._model_combo_generating_index = item.index;
|
|
item.node._model_combo_render?.();
|
|
}
|
|
}
|
|
}
|
|
|
|
function maybeFinish() {
|
|
if (completedCount >= combinations.length) {
|
|
if (_batchCleanup) {
|
|
_batchCleanup();
|
|
_batchCleanup = null;
|
|
}
|
|
}
|
|
}
|
|
|
|
// The `executed` event carries the per-node output payload. Used to stamp
|
|
// generated images onto FlakeGenerate nodes keyed by their comboKey.
|
|
function handleExecuted(e) {
|
|
const detail = e?.detail || {};
|
|
const pid = typeof detail === "string" ? null : detail.prompt_id;
|
|
const output = detail?.output;
|
|
const nodeId = detail?.node;
|
|
if (!pid || !output) return;
|
|
const idx = promptIds.indexOf(pid);
|
|
if (idx < 0) return;
|
|
const comboKey = (comboKeys && comboKeys[idx]) ?? "";
|
|
const images = output.flake_images;
|
|
if (!Array.isArray(images) || images.length === 0) return;
|
|
const candidates = nodeId
|
|
? [app.graph.getNodeById(nodeId)].filter(n => n && n.type === "FlakeGenerate")
|
|
: app.graph.nodes.filter(n => n.type === "FlakeGenerate");
|
|
for (const fg of candidates) {
|
|
fg.properties = fg.properties || {};
|
|
fg.properties._images_by_combo = fg.properties._images_by_combo || {};
|
|
fg.properties._images_by_combo[comboKey] = images[0];
|
|
}
|
|
}
|
|
|
|
// All websocket lifecycle/output events route through bufferOrApply so
|
|
// events arriving before their prompt_id is registered get queued.
|
|
const onExecStart = (e) => bufferOrApply("execution_start", e);
|
|
const onExecSuccess = (e) => bufferOrApply("execution_success", e);
|
|
const onExecError = (e) => bufferOrApply("execution_error", e);
|
|
const onExecInterrupted = (e) => bufferOrApply("execution_interrupted", e);
|
|
const onExecuted = (e) => bufferOrApply("executed", e);
|
|
|
|
comfyApi.addEventListener("execution_start", onExecStart);
|
|
comfyApi.addEventListener("execution_success", onExecSuccess);
|
|
comfyApi.addEventListener("execution_error", onExecError);
|
|
comfyApi.addEventListener("execution_interrupted", onExecInterrupted);
|
|
comfyApi.addEventListener("executed", onExecuted);
|
|
|
|
const cleanup = () => {
|
|
comfyApi.removeEventListener("execution_start", onExecStart);
|
|
comfyApi.removeEventListener("execution_success", onExecSuccess);
|
|
comfyApi.removeEventListener("execution_error", onExecError);
|
|
comfyApi.removeEventListener("execution_interrupted", onExecInterrupted);
|
|
comfyApi.removeEventListener("executed", onExecuted);
|
|
if (_origApiQueuePrompt) {
|
|
comfyApi.queuePrompt = _origApiQueuePrompt;
|
|
}
|
|
_currentBatchComboKey = null;
|
|
_activeBatch.promptIds = [];
|
|
_activeBatch.comboKeys = [];
|
|
_activeBatch.runningPromptId = null;
|
|
pendingByPid.clear();
|
|
clearAllHighlights();
|
|
};
|
|
|
|
return { promptIds, cleanup };
|
|
}
|
|
|
|
// Order the combinatorial nodes (FlakeCombo + FlakeModelCombo) by their
|
|
// position in the pipeline connection chain (source -> sink), so the loop
|
|
// nesting follows the graph flow rather than node creation order. The
|
|
// FlakeModelCombo is the chain source (it loads the heavy checkpoint), so it
|
|
// lands first = outermost loop = reloaded least; each subsequent FlakeCombo
|
|
// toward FlakeGenerate becomes a deeper inner loop. Combo nodes not reachable
|
|
// from any FlakeGenerate (orphans) are appended in creation order so
|
|
// disconnected graphs keep working.
|
|
function orderedComboNodes() {
|
|
const isCombo = (n) => n.type === "FlakeCombo" || n.type === "FlakeModelCombo";
|
|
const ordered = [];
|
|
const seen = new Set();
|
|
for (const gen of app.graph.nodes) {
|
|
if (gen.type !== "FlakeGenerate") continue;
|
|
const fdInput = gen.inputs?.find(i => i.name === "flake_data");
|
|
if (fdInput?.link == null) continue;
|
|
const link = app.graph.links?.[fdInput.link];
|
|
if (!link) continue;
|
|
const startNode = app.graph.getNodeById(link.origin_id);
|
|
if (!startNode) continue;
|
|
let chain = [];
|
|
try { chain = collectChain(startNode); } catch { chain = []; }
|
|
for (const { node } of chain) {
|
|
if (isCombo(node) && !seen.has(node.id)) { seen.add(node.id); ordered.push(node); }
|
|
}
|
|
}
|
|
for (const node of app.graph.nodes) {
|
|
if (isCombo(node) && !seen.has(node.id)) { seen.add(node.id); ordered.push(node); }
|
|
}
|
|
return ordered;
|
|
}
|
|
|
|
app.queuePrompt = async function(number, batchCount = 1) {
|
|
const orderedNodes = orderedComboNodes();
|
|
|
|
if (orderedNodes.length === 0) {
|
|
return _originalQueuePrompt.call(this, number, batchCount);
|
|
}
|
|
|
|
const optionsArrays = [];
|
|
|
|
for (const node of orderedNodes) {
|
|
if (node.type === "FlakeCombo") {
|
|
const flakes = getComboFlakes(node);
|
|
if (flakes.length === 0) {
|
|
window.alert("FlakeCombo node has no flakes selected.");
|
|
return;
|
|
}
|
|
const activeFlakes = flakes
|
|
.map((flake, i) => ({ flake, i }))
|
|
.filter(({ flake }) => !flake.bypassed);
|
|
if (activeFlakes.length === 0) {
|
|
window.alert("FlakeCombo node has no active (non-bypassed) flakes.");
|
|
return;
|
|
}
|
|
optionsArrays.push(activeFlakes.map(({ flake, i }) => ({
|
|
node,
|
|
type: "combo",
|
|
value: flake,
|
|
index: i,
|
|
})));
|
|
} else {
|
|
const presets = getComboPresets(node);
|
|
if (presets.length === 0) {
|
|
window.alert("FlakeModelCombo node has no presets selected.");
|
|
return;
|
|
}
|
|
const bypassed = new Set(node.properties?._combo_bypassed || []);
|
|
const activePresets = presets
|
|
.map((preset, i) => ({ preset, i }))
|
|
.filter(({ preset }) => !bypassed.has(preset));
|
|
if (activePresets.length === 0) {
|
|
window.alert("FlakeModelCombo node has no active (non-bypassed) presets.");
|
|
return;
|
|
}
|
|
// Keep the original preset index so combo highlighting maps correctly.
|
|
optionsArrays.push(activePresets.map(({ preset, i }) => ({
|
|
node,
|
|
type: "model_combo",
|
|
value: preset,
|
|
index: i,
|
|
})));
|
|
}
|
|
}
|
|
|
|
const combinations = cartesianProduct(optionsArrays);
|
|
if (combinations.length === 0) {
|
|
return _originalQueuePrompt.call(this, number, batchCount);
|
|
}
|
|
|
|
// Save original widget values (once per unique node)
|
|
const nodeOriginals = new Map();
|
|
for (const item of combinations[0]) {
|
|
if (!nodeOriginals.has(item.node.id)) {
|
|
if (item.type === "combo") {
|
|
const w = item.node.widgets?.find(w => w.name === "flakes_json");
|
|
nodeOriginals.set(item.node.id, { node: item.node, widget: w, value: w?.value });
|
|
} else {
|
|
const w = item.node.widgets?.find(w => w.name === "preset");
|
|
const ow = item.node.widgets?.find(w => w.name === "overrides_json");
|
|
nodeOriginals.set(item.node.id, {
|
|
node: item.node, widget: w, value: w?.value,
|
|
overrideWidget: ow, overrideValue: ow?.value,
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
const flakeGenerateNodes = app.graph.nodes.filter(n => n.type === "FlakeGenerate");
|
|
|
|
// Clean up any previous batch tracking
|
|
if (_batchCleanup) {
|
|
_batchCleanup();
|
|
_batchCleanup = null;
|
|
}
|
|
|
|
const comboKeys = []; // parallel to promptIds; comboKeys[idx] = key for promptIds[idx]
|
|
const tracker = setupBatchTracking(combinations, comboKeys);
|
|
if (tracker) {
|
|
_batchCleanup = tracker.cleanup;
|
|
}
|
|
// Initialize progress counters on every FlakeGenerate node so the "[i/N] done"
|
|
// label resets to "[0/N]" at batch start (#227).
|
|
for (const n of flakeGenerateNodes) {
|
|
n._batch_completed_count = 0;
|
|
n._batch_total_count = combinations.length;
|
|
n._batch_progress_render?.();
|
|
}
|
|
|
|
try {
|
|
for (let comboIdx = 0; comboIdx < combinations.length; comboIdx++) {
|
|
const combination = combinations[comboIdx];
|
|
for (const item of combination) {
|
|
if (item.type === "combo") {
|
|
const w = item.node.widgets?.find(w => w.name === "flakes_json");
|
|
if (w) w.value = JSON.stringify([item.value]);
|
|
} else {
|
|
const w = item.node.widgets?.find(w => w.name === "preset");
|
|
if (w) w.value = item.value;
|
|
// Per-instance preset overrides for this combination (#279).
|
|
const ow = item.node.widgets?.find(w => w.name === "overrides_json");
|
|
if (ow) {
|
|
const ovr = (item.node.properties?._combo_overrides || [])[item.index] || {};
|
|
ow.value = serializeModelOverrides(ovr);
|
|
}
|
|
}
|
|
}
|
|
const comboKey = combination
|
|
.map(it => [it.node.id, it.index])
|
|
.sort((a, b) => a[0] - b[0])
|
|
.map(([id, idx]) => `${id}:${idx}`)
|
|
.join("|");
|
|
|
|
// Tell our patched api.queuePrompt which comboKey to attach to the
|
|
// upcoming prompt_id. app.queuePrompt in modern frontends returns
|
|
// void; the prompt_id is captured inside the patch instead.
|
|
if (tracker) _currentBatchComboKey = comboKey;
|
|
try {
|
|
await _originalQueuePrompt.call(this, number, 1);
|
|
} finally {
|
|
if (tracker) _currentBatchComboKey = null;
|
|
}
|
|
}
|
|
} finally {
|
|
for (const orig of nodeOriginals.values()) {
|
|
if (orig.widget) orig.widget.value = orig.value;
|
|
if (orig.overrideWidget) orig.overrideWidget.value = orig.overrideValue;
|
|
}
|
|
// If no tracker is available, clear highlights immediately as fallback
|
|
if (!tracker) {
|
|
clearAllHighlights();
|
|
}
|
|
}
|
|
};
|