Files

286 lines
24 KiB
JavaScript

// Generate / Selected / Stop: queue generated units one after another and put each result on its clips.
import { composer as c } from "../../composer/composer.js";
import { isGenerative, isVideoClip, genUnitId, unitRoot, effFrames, effFps } from "../../shell/scenes.js";
import { QUEUED, RUNNING } from "../../shell/run.js";
import { jobsRunning, onJobs, stopJobs } from "../../shell/renders.js";
import { inPlace } from "../../shell/place.js";
import { buildInputs, rolesAt } from "./inputs.js";
import { offer } from "../../shell/actions.js";
import { names } from "../../shell/bin.js";
import { showCompare } from "./compare.js";
const MAX_TAKES = 24;
const expand = (body) => fetch("/fun/api/prompt/expand", { method: "POST", headers: { "Content-Type": "application/json" },
body: JSON.stringify({ ...body, seed: Math.floor(Math.random() * 2 ** 31) || 1 }) }).then((r) => (r.ok ? r.json() : null));
/** The render a scene gets from its unit's one clip: a cut half plays from where its half begins. */
export const renderFor = (sc, p, media, unitSec, promptId, text) => ({ media, ...(promptId ? { promptId } : {}), ...(typeof text === "string" ? { text } : {}), inSec: (sc.cut_offset_frames || 0) / effFps(sc, p),
...((sc.frames_mode || "project") === "project" ? { durationSec: unitSec } : {}) });
export default {
id: "generate",
mount: "timeline.actions",
needs: ["project", "pipeline", "generate", "api", "selection", "dom"],
setup({ host, app }) {
const p = app.project, g = app.generate;
let said = false;
const tell = (text) => { if (!text) return; said = true; c.toast.warn({ text }); }; // no words, no toast: an empty one read as a message gone missing
let current = 0; // the scene being queued: a refusal said meanwhile names it
g.on("say", (t) => tell(t && (current ? `Scene ${current}: ${t}` : t))); g.on("warn", tell);
g.on("hold", () => { held = true; draw(); }); g.on("release", () => { held = false; draw(); }); // while the page asks ComfyUI whether a run is already going // the pipeline check's refusals: said where the person is looking
let busy = false, stopped = false, held = false;
const record = (pid, unit, image, frames, promptId, text) => { app.lastRun.n += 1; return p.editFor(pid, (pr) => {
const group = pr.scenes.filter((s) => genUnitId(s) === unit);
if (!group.length) return false;
const secs = (frames || group.reduce((t, s) => t + effFrames(s, pr), 0) - (group.length - 1)) / (pr.frame_rate || 25); // as queued, not as the project reads now
const media = { filename: image.filename, subfolder: image.subfolder || "", type: image.type || "output" };
const head = group.find((s) => !(s.cut_offset_frames > 0)) || group[0];
const takes = ((pr.scene_variants ||= {})[head.id] ||= []); // every render stays as a take: the one that was on the clip keeps its rating
const was = (pr.scene_renders || {})[head.id];
const sameTake = (t) => was && (was.promptId ? t.promptId === was.promptId : t.media.filename === (was.media || {}).filename);
const old = was && was.media && takes.find(sameTake);
if (old) Object.assign(old, { rating: head.rating || "", rated: head.rated || {} });
else if (was && was.media) takes.push({ media: was.media, ...(was.promptId ? { promptId: was.promptId } : {}), rating: head.rating || "", rated: head.rated || {}, ...(was.durationSec ? { secs: was.durationSec } : {}), ...(was.text !== undefined ? { text: was.text } : {}) }); // a render from before takes is a take too
takes.push({ media, ...(promptId ? { promptId } : {}), rating: "", secs, ...(typeof text === "string" ? { text } : {}) }); // the prompt that made it: what a rating of it is about
while (takes.length > MAX_TAKES) { const at = takes.findIndex((t) => !t.rating); takes.splice(at < 0 || at === takes.length - 1 ? 0 : at, 1); } // the oldest unrated goes first; a rated take stays as long as anything else can go
group.forEach((s) => {
(pr.scene_renders ||= {})[s.id] = renderFor(s, pr, media, secs, promptId, text);
if (!isVideoClip(s)) s.source_in = 0;
if (!(s.cut_offset_frames > 0)) Object.assign(s, { rating: "", rated: {} }); // a new render has not been rated yet // a fresh render was made at this length: an earlier trim's window no longer applies
});
return true;
}).catch(() => tell("The result could not be saved to its project.")); };
const unwiredText = (unwired, left) => `${unwired} reference(s) did not fit this pipeline and are not used: ${(left || []).map((l) => `${names.get(l.id) || l.id} (looked for ${l.what ? `${l.what}, ${l.at.replace("assets.", "")}` : l.at.replace("assets.", "")}: ${l.why})`).join("; ")}.`;
// `times` > 1: the same shots again, only the seed differing (the prompt's random picks are drawn once), so the takes can be compared.
// `once`: for a single run only - {seed} uses that seed on every seed input; onQueued runs once ComfyUI has taken (or refused) the run.
// This page's own upscale, or one the server still runs for a reloaded page or another tab.
const upscaling = async () => {
if (app.has && app.has.has("upscale.model")) return true;
let timer;
try { return Boolean(await Promise.race([Promise.resolve(app.api.upscaleBusy ? app.api.upscaleBusy() : null).then((r) => r && r.busy), new Promise((r) => { timer = setTimeout(() => r(false), 5000); })])); }
catch { return false; } // a server that does not answer, or errs, is not "busy"
finally { clearTimeout(timer); }
};
async function runUnits(units, times = 1, once = {}) {
if (busy) { tell("A run is already going: wait for it, or Stop it."); return; }
if (await upscaling()) { tell("An upscale model is running. It shares ComfyUI's model memory, so Generate when it has finished."); return; }
if (busy) { tell("A run is already going: wait for it, or Stop it."); return; } // a second click that came in while the server was asked
busy = true; stopped = false; draw();
const pid = p.project.id;
let made = 0;
said = false;
try {
app.dom.blur(); // a field still being typed in commits first, however Generate was started
await Promise.all((app.beforeRun || []).map((f) => f().catch(() => {}))); // an edit still on its way (the Story box) lands first
// A project's pipeline still going in is the one this run uses, not the last one's.
// a typed value is already live and is never waited on; only a pipeline still going in, or a node change, is
const ready = app.pipeline.readyToRun || app.pipeline.settled;
if (ready && !(await ready())) { tell("Not started: this project's pipeline, or a node change in it, is still on its way. ComfyUI is slow to answer; try again in a moment."); return; }
if (app.pipelineOwned && !app.pipelineOwned()) { tell(`Not started: ${(app.pipelineWhy && app.pipelineWhy()) || "this project's pipeline is not loaded yet."}`); return; }
// The pipeline as it is NOW, at the click: every shot of this run uses it, whatever is opened or edited meanwhile.
const frozen = structuredClone(app.pipeline.slots() || []);
if (times > 1 && !rolesAt(frozen, "generation.seed").length) { tell("This pipeline has no seed input the app controls, so the takes would all come out the same."); return; }
// The project too: text, lengths and references edited after the click reach the next Generate, not this one's later shots.
const snap = structuredClone(p.project);
// A picture deleted from the bin (in another project, or on another machine) is still named by id here: it would fail the run in ComfyUI.
const bin = await Promise.resolve().then(() => app.api.media()).then((r) => new Set((r.media || []).map((m) => m.id)), () => null);
const dead = (id) => bin && id && !bin.has(id);
const prune = (pr) => {
if (dead(pr.fixed_source_image)) pr.fixed_source_image = null;
const fx = pr.fixed_references || {}, deadFx = Object.keys(fx).filter(dead); // a fixed slot whose picture is gone goes too, or every scene sends it
for (const id of deadFx) delete fx[id];
if (deadFx.length) pr.fixed_references = fx;
return pr.scenes.reduce((hit, s) => {
const had = dead(s.source_image) || (s.references || []).some(dead);
if (dead(s.source_image)) s.source_image = "";
if (had) s.references = (s.references || []).filter((r) => !dead(r));
return hit || had;
}, deadFx.length > 0);
};
if (prune(snap)) { tell("Some pictures this project used are no longer in the media bin; they were removed from their scenes."); p.edit(prune); }
const fixedSeed = times === 1 && p.pref("random_seed", true) === false; // takes are compared across seeds: always fresh ones
const drawn = new Map();
const same = times > 1 ? (body) => { const key = JSON.stringify(body); if (!drawn.has(key)) drawn.set(key, expand(body)); return drawn.get(key); } : expand;
let told = false; // the wait is said once per Generate
const waitForUpscale = async () => { // an Always upscale of the shot before may be running: the next shot waits for it
while (!stopped && await upscaling()) {
if (!told) { told = true; tell("Waiting for the upscale model to finish before the next shot."); }
await new Promise((r) => setTimeout(r, 2000));
}
};
for (const unit of Array.from({ length: times }, () => units).flat()) {
if (stopped) break;
await waitForUpscale();
if (stopped) break;
if (!p.project || p.project.id !== pid) { tell("Stopped: another project was opened."); break; }
const root = unitRoot(snap, unit);
const group = snap.scenes.filter((s) => genUnitId(s) === unit);
if (!root || !p.project.scenes.some((s) => s.id === root.id)) continue; // deleted since the click
if (group.every((s) => s.excluded) || !isGenerative(root)) continue; // all left out, or not made by the model
// The unit is made once at the length of its clips together (each cut adds one shared frame).
const frames = group.reduce((t, s) => t + effFrames(s, snap), 0) - (group.length - 1);
const { inputs, unwired, unusedPicture, left, noPrompt, notes, placed } = await buildInputs({ project: snap, scene: root, slots: frozen, expand: same, nodeInfo: app.api.wireTargets, frames, hooks: app.inputHooks, ...(Number.isInteger(once.seed) ? { seed: () => once.seed } : fixedSeed ? { seed: null } : {}), prefix: (app.promptPrefix || []).flatMap((f) => { try { return f(root, snap); } catch { return []; } }) });
app.lastRun.typed = (root.text || "").trim(); // what a Chat comment made now would be about
const seeded = rolesAt(frozen, "generation.seed")[0];
if (seeded) app.lastRun.seed = (inputs[seeded.slot.id] || {})[seeded.role.input] ?? (seeded.slot.inputs || {})[seeded.role.input]; // shown in ⏱ Sampler, to keep a liked one
if (stopped) break;
if (noPrompt && !made) tell("This pipeline has no prompt input, so the scene text is not sent.");
notes.forEach(tell);
for (const x of placed || []) app.api.carryLog(pid, { scene_id: root.id, ...x }).catch(() => {}); // the Log records where a carried picture went; no toast
if (unusedPicture) tell(`The scene's picture (${names.get(unusedPicture) || unusedPicture}) is not used: this pipeline starts from a video, not a picture.`);
if (unwired) tell(unwiredText(unwired, left));
await waitForUpscale(); // asked again right before queueing: the shot's inputs took time to build
if (stopped) break;
if (!p.project || p.project.id !== pid) { tell("Stopped: another project was opened."); break; } // the wait can take minutes
// The page's own guard does not see another tab: ComfyUI's queue does. Asked per shot: between this run's shots it is empty.
if (app.api.projectQueued && (await app.api.projectQueued(pid))) { tell("This project already has a run in ComfyUI's queue (another tab or window?). Wait for it, or Stop it there."); break; }
if (!p.project || p.project.id !== pid) { tell("Stopped: another project was opened."); break; } // the queue answer can take seconds
const done = g.waitForTerminal(); // listening before the run starts, so a fast one is not missed
said = false;
current = snap.scenes.findIndex((s) => s.id === root.id) + 1;
const queued = await g.generate({ sceneId: root.id, projectId: pid, inputs, slots: frozen });
current = 0;
if (once.onQueued) await once.onQueued();
if (!queued) {
done.cancel();
const n = snap.scenes.findIndex((s) => s.id === root.id) + 1; // which clip: in a run of several, the toast alone does not say
if (!said) tell(`Scene ${n} was not started: ${(g.run.state.error && g.run.state.error.message) || "could not queue the run. Is ComfyUI running, and is a run already going?"}`);
break;
}
if (stopped) g.cancel(); // Stop landed while this one was being queued
if (!made && (!app.hasModule || app.hasModule("taste"))) app.api.newTasteGeneration(g.run.state.promptId).catch(() => {}); // a run really started: the last run's unrated clips can no longer be paired with what it learned
made += 1;
const end = await done;
if (end === g.CANCELLED) break;
const images = g.run.state.images;
if (end !== g.DONE) { // ComfyUI's own words: which node, and why
const e = g.run.state.error;
tell(e && e.message ? `Generation failed${e.node ? ` in ${e.node}` : ""}: ${e.message}` : "Generation failed. The log has ComfyUI's message.");
break;
}
if (!images.length) { tell("ComfyUI finished without a result. Try again."); break; }
const kept = await record(pid, unit, images[images.length - 1], frames, g.run.state.promptId, root.text || ""); // the next shot may continue from this one: it must be on the clip first
if (kept === false) tell(`Scene ${snap.scenes.findIndex((s) => s.id === root.id) + 1} was removed while it was being made, so its result is not on the timeline. It is in ComfyUI's output folder as ${images[images.length - 1].filename}.`);
snap.scene_renders = structuredClone(p.project.scene_renders || {}); // ...and in what the next shot reads
}
if (!made && !stopped && !said) tell("Nothing to generate: every scene is left out or a video clip.");
} finally { busy = false; draw(); }
}
app.runner.units = (units, times) => runUnits(units, times);
/** Quick A/B: the picked scene(s) once more with every enhancement off and the last run's seed, so the two clips differ only by the enhancements. Nothing stays changed afterwards. */
app.runner.quickAB = async () => {
const seed = app.lastRun && app.lastRun.seed;
if (busy) { tell("A run is already going: wait for it, or Stop it."); return; }
if (!p.project || !Number.isInteger(seed)) { tell("Quick A/B needs a run to compare with: generate once first."); return; }
const units = unitsOf(pickedScenes());
if (!units.length) return;
const headOf = (u) => { const g = p.project.scenes.filter((s) => genUnitId(s) === u); return g.find((s) => !(s.cut_offset_frames > 0)) || g[0]; };
const renderOf = (u) => { const h = headOf(u), r = h && (p.project.scene_renders || {})[h.id]; return r && r.media ? { media: r.media, key: r.promptId || r.media.filename } : null; };
const before = new Map(units.map((u) => [u, renderOf(u)]));
const was = app.pipeline.allOff();
let restored = was;
const restore = async () => { if (!restored) { restored = true; await app.pipeline.setAllOff(false); } }; // as soon as ComfyUI has the run: later edits are the person's again
try {
if (!was) { restored = false; await app.pipeline.setAllOff(true); }
await runUnits(units, 1, { seed, onQueued: restore });
const pairs = units.map((u) => [before.get(u), renderOf(u)]).filter(([b, a]) => b && a && b.key !== a.key).map(([b, a]) => ({ before: b.media, after: a.media }));
if (pairs.length) showCompare({ pairs, seed });
} finally { await restore().catch(() => tell("Enhancements could not be switched back on: use Enable enhancements on the wheel.")); }
};
/** One run for the Training window: the pipeline as it is, its LoRA loader pointed at `lora`, the picked scene (else the first one) and the project's own settings.
* Nothing is put on the timeline; resolves to {image} or {error}. */
app.runner.preview = async ({ lora, strength = 1 }) => {
if (await upscaling()) return { error: "An upscale model is running. It shares ComfyUI's model memory, so the preview waits for it." };
if (busy || held || g.state().phase === QUEUED || g.state().phase === RUNNING) return { error: "A run is already going: the preview waits for it." }; // after the wait: a Generate may have started meanwhile
if (!p.project) return { error: "Open a project first: the preview uses its settings." };
const frozen = structuredClone(app.pipeline.slots() || []);
const slot = frozen.find((s) => s.node === "FunUILoraLoader" && !s.bypassed) || frozen.find((s) => s.node === "FunUILoraLoader");
if (!slot) return { error: "This pipeline has no LoRA loader to try the LoRA in. Add one in Models & Pipeline." };
slot.bypassed = false;
slot.inputs = { ...slot.inputs, lora_name: lora, strength_model: strength };
const snap = structuredClone(p.project);
const first = [p.selected, ...p.scenes].find((s) => s && isGenerative(s) && !s.excluded);
const root = first && unitRoot(snap, genUnitId(first));
if (!root) return { error: "The project has no scene the model makes, so there is nothing to preview with." };
const group = snap.scenes.filter((s) => genUnitId(s) === genUnitId(root));
const frames = group.reduce((t, s) => t + effFrames(s, snap), 0) - (group.length - 1);
busy = true; draw();
try {
await Promise.all((app.beforeRun || []).map((f) => f().catch(() => {})));
const { inputs, unwired, left } = await buildInputs({ project: snap, scene: root, slots: frozen, expand, nodeInfo: app.api.wireTargets, frames, hooks: app.inputHooks, prefix: (app.promptPrefix || []).flatMap((f) => { try { return f(root, snap); } catch { return []; } }) });
if (unwired) tell(unwiredText(unwired, left));
if (await upscaling()) return { error: "An upscale model is running. It shares ComfyUI's model memory, so the preview waits for it." };
const done = g.waitForTerminal();
const queued = await g.generate({ sceneId: root.id, projectId: snap.id, inputs, slots: frozen });
if (!queued) { done.cancel(); return { error: (g.run.state.error && g.run.state.error.message) || "The preview could not be queued. Is ComfyUI running?" }; }
const end = await done;
const images = g.run.state.images;
if (end === g.DONE && images.length) return { image: images[images.length - 1] };
const e = g.run.state.error;
return { error: end === g.CANCELLED ? "The preview was stopped." : e && e.message ? `The preview failed${e.node ? ` in ${e.node}` : ""}: ${e.message}` : "The preview finished without a result." };
} finally { busy = false; draw(); }
};
const pickedScenes = () => { const ids = new Set(app.selection.ids); const hit = p.scenes.filter((s) => ids.has(s.id)); return hit.length ? hit : p.selected ? [p.selected] : []; }; // every picked clip, else the one in focus
const unitsOf = (scenes) => [...new Set(scenes.map(genUnitId))];
const put = inPlace(host);
let drawn = "";
function draw() {
const open = p.project;
const picked = open ? pickedScenes().length : 0, key = `${busy}|${held}|${g.state().phase}|${Boolean(open)}|${picked}|${jobsRunning()}`;
if (key === drawn) return; // progress ticks must not rebuild the buttons
drawn = key;
const working = busy || held || g.state().phase === QUEUED || g.state().phase === RUNNING; // also a run this page did not start
const go = c.button.sm({ label: working ? "Generating…" : "▶ Generate", tone: "primary", disabled: working || !open, title: "Generate every scene on the timeline, except left-out scenes and video clips. An earlier render of a scene stays as a take.",
onClick: () => runUnits(unitsOf(p.scenes)) });
const one = c.button.sm({ label: picked > 1 ? `Selected (${picked})` : "Selected", title: "Generate only the picked scenes", disabled: working || !picked, onClick: () => runUnits(unitsOf(pickedScenes())) });
const stop = c.button.sm({ label: "■ Stop", tone: "danger", title: "Stop the run that is going, and any Upscale or Interpolate in progress", onClick: () => {
stopJobs();
if (working) { stopped = true; g.cancel(); }
} });
stop.node.hidden = !working && !jobsRunning();
put(go.node, one.node, stop.node);
}
// A run found after a reload belongs to this page's last session: its result still goes on its clips.
const claim = ({ sceneId, projectId }) => {
const sc = p.project && p.project.id === projectId && p.project.scenes.find((s) => s.id === sceneId);
if (!sc) return false;
const unit = genUnitId(sc);
const finish = (end) => {
const im = g.run.state.images;
if (end === g.DONE && im.length) record(projectId, unit, im[im.length - 1], undefined, g.run.state.promptId);
else if (end === g.FAILED) tell((g.run.state.error && g.run.state.error.message) || "A run from before the reload failed.");
else if (end !== g.CANCELLED) tell("A run from before the reload ended without a result that could be attached. Generate again.");
busy = false; draw();
};
const now = g.run.state.phase; // a run that ended while the page reloaded is over already: nothing to wait for
if (now === g.DONE || now === g.FAILED || now === g.CANCELLED) { finish(now); return true; }
busy = true; draw(); // Stop works on it, Generate waits
g.waitForTerminal().then(finish);
return true;
};
// The page may know the run before it has opened a project: wait for the project, then claim or give up.
g.on("adopt", (a) => {
if (claim(a)) return;
const off = app.on(() => { // until its project opens, or the run is over
if (!p.project) return;
if (claim(a)) return off();
const phase = g.run.state.phase;
if (phase === g.DONE || phase === g.FAILED || phase === g.CANCELLED) {
off();
tell("A run from before the reload finished, but its project or scene is not open here, so its result was not attached.");
}
});
});
draw();
const offRun = g.subscribe(draw), offJobs = onJobs(draw); // the Stop button also shows while an Upscale or Interpolate runs
const offApp = app.on((what) => { if (what === "generate.scene" && p.selected) runUnits(unitsOf([p.selected]));
else if (what === "generate.selected" && pickedScenes().length) runUnits(unitsOf(pickedScenes())); else draw(); });
const offers = [offer(app, { id: "generate-all", label: "Generate", icon: "▶", run: () => { if (p.project) runUnits(unitsOf(p.scenes)); } }),
offer(app, { id: "generate-selected", label: "Generate selected", icon: "▶", run: () => { if (pickedScenes().length) runUnits(unitsOf(pickedScenes())); } })];
return () => { offRun(); offJobs(); offApp(); offers.forEach((f) => f()); };
},
};