Claude
3adebb731f
Old project/workflow names and browser settings migrate on load; /funpack redirects to /fun. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016acRss69eH1235eRdFTZHY
306 lines
20 KiB
JavaScript
306 lines
20 KiB
JavaScript
// "Training": make a LoRA from your own media and rated results. An add-on (Settings > Modules); the window is absent until it is on.
|
|
// Four steps in one window: Data (optional), Tags, Train, Results. A training set is saved, so each step can be left and come back to.
|
|
import { composer as c } from "../../composer/composer.js";
|
|
import { dataStep } from "./data.js";
|
|
import { applyTagging, tagsStep } from "./tags.js";
|
|
import { footNote, monitor, notice, settingsOf, trainStep } from "./train.js";
|
|
import { resultsStep } from "./results.js";
|
|
import { makeSession, srcKey } from "./state.js";
|
|
|
|
const VIDEO = /\.(mp4|webm|mov|mkv)$/i;
|
|
|
|
const TABS = [{ value: "data" }, { value: "tags" }, { value: "train" }, { value: "results" }]; // the steps, in order: Back and Next walk them
|
|
const STEPS = {
|
|
data: { title: "Pick what to learn from", sub: "Drop pictures and clips, choose from your Bin, or take results you rated. A set you filled earlier can go straight on; an empty one cannot be trained." },
|
|
tags: { title: "Describe each one", sub: "The words under each picture are what the LoRA learns to connect with it. Fix them by hand, or let the tagger suggest." },
|
|
train: { title: "Train the LoRA", sub: "Choose how thorough it should be and press Start. Progress and previews appear in Results." },
|
|
results: { title: "Results", sub: "Watch the training, look at the previews and download a LoRA, the final one or any step." },
|
|
};
|
|
const POLL = 1000;
|
|
const QUEUE_WAIT = 60000;
|
|
|
|
const smooth = (values) => { // a gentle moving average, so the trend shows through a noisy loss
|
|
const w = Math.max(3, Math.round(values.length / 40)), out = [];
|
|
let sum = 0;
|
|
for (let i = 0; i < values.length; i += 1) { sum += values[i]; if (i >= w) sum -= values[i - w]; out.push(sum / Math.min(i + 1, w)); }
|
|
return out;
|
|
};
|
|
|
|
export default {
|
|
id: "training",
|
|
mount: "timeline.status",
|
|
needs: ["api", "dom", "pipeline", "project", "generate"],
|
|
setup({ host, app }) {
|
|
const api = app.api;
|
|
const session = makeSession(api);
|
|
const ui = {
|
|
tab: "data", openBin: false, openRated: false, clips: {}, picked: new Set(), ratedPicked: new Set(), bin: [], rated: { results: [], counts: {} }, ratedKind: "liked", include: ["main"], selected: new Set(), filter: "all", tag: "", limit: 48,
|
|
tagger: { id: "wd_vit_v3", list: [], job: null, general: 0.35, limit: 0, characters: false, mode: "add", scope: "untagged" },
|
|
advanced: false, runs: [], chosenRun: "", record: null,
|
|
run: { shown: false, active: false, state: "idle", loss: [], times: [], stepSeconds: null, lastStep: null, smooth: [], step: 0, total: 0, elapsed: 0, message: "", error: "", lora: "", before: "", runId: "", paused: null },
|
|
previews: [], holdPause: false, awaitStart: false, peek: (item) => peek(item),
|
|
pause: () => pause(), resume: () => launch(ui.run.runId), restore: (id) => { ui.run.loss = []; ui.run.times = []; ui.run.runId = id; ui.record = null; return launch(id); }, previewNow: () => previewNow(),
|
|
monitorBox: null, redraw: () => draw(), pollTaggers: () => pollTaggers(true), loadRecord: (id) => loadRecord(id), loadRuns: () => loadRuns(),
|
|
start: () => start(), stop: () => stop(),
|
|
/** Ask for the frame counts of clips not known yet; `then` runs when new numbers arrived. */
|
|
async needClips(srcs, then) {
|
|
const fresh = srcs.filter((x) => !(srcKey(x) in ui.clips));
|
|
if (!fresh.length) return;
|
|
fresh.forEach((x) => { ui.clips[srcKey(x)] = null; });
|
|
try { Object.assign(ui.clips, (await api.trainClips(fresh)).clips || {}); } catch { fresh.forEach((x) => { delete ui.clips[srcKey(x)]; }); return; }
|
|
if (then) then();
|
|
},
|
|
};
|
|
let foot = () => {};
|
|
let win = null, body = null, timers = [], offSession = null, tagPoll = null, runPoll = null;
|
|
|
|
const button = c.button.sm({ label: "🎓 Training", tone: "neutral", title: "Make a LoRA from your media and rated results", onClick: () => (win ? win.close() : open()) });
|
|
button.node.hidden = true;
|
|
host.append(button.node);
|
|
api.trainState().then((s) => { button.node.hidden = !(s && s.enabled); }).catch(() => {});
|
|
|
|
const peek = (item) => {
|
|
const s = item.src;
|
|
const src = s.kind === "bin" ? `/fun/api/media/${encodeURIComponent(s.id)}/file` : `/view?filename=${encodeURIComponent(s.filename)}&subfolder=${encodeURIComponent(s.subfolder || "")}&type=${encodeURIComponent(s.type || "output")}`;
|
|
const viewer = c.viewer.media({ src, kind: item.media_kind === "video" ? "video" : "image" });
|
|
const m = c.modal.generic({ title: item.name || "Preview", body: viewer, onClose: () => viewer.destroy && viewer.destroy() });
|
|
return m;
|
|
};
|
|
|
|
// --- the steps
|
|
const steps = {
|
|
data: () => dataStep({ app, session, go, ui }),
|
|
tags: () => tagsStep({ app, session, go, ui, peek }),
|
|
train: () => trainStep({ app, session, go, ui }),
|
|
results: () => resultsStep({ app, ui }),
|
|
};
|
|
function go(tab) { if (tab !== "data" && tab !== "results" && !((session.doc || {}).items || []).length) tab = "data"; ui.tab = tab; if (tab === "results") loadRuns(); draw(); }
|
|
function draw() {
|
|
if (!win) return;
|
|
const note = session.error ? c.inlineError.default({ text: session.error }) : null;
|
|
const info = STEPS[ui.tab], order = TABS.map((t) => t.value), at = order.indexOf(ui.tab);
|
|
const ready = ((session.doc || {}).items || []).length > 0; // nothing to learn from: no way on to training
|
|
body.set([c.label.section({ text: `${at + 1} / ${order.length} · ${info.title}` }), c.hint.default({ text: info.sub }), note, steps[ui.tab]()].filter(Boolean));
|
|
const results = (ui.tab === "data" || ui.tab === "train") ? c.button.sm({ label: "See results", tone: "ghost", title: "Earlier runs, the live graph and downloads", onClick: () => go("results") }) : null;
|
|
const back = ui.tab === "results" ? c.button.sm({ label: "Back to training", tone: "ghost", onClick: () => go("train") }) // a run carries on while another set is made or tags are fixed
|
|
: at ? c.button.sm({ label: "Back", tone: "ghost", onClick: () => go(order[at - 1]) }) : null;
|
|
const skip = ui.tab === "data" ? c.button.sm({ label: "Skip to training", tone: "ghost", disabled: !ready, title: ready ? "" : "Add pictures or clips to the set first", onClick: () => go("train") }) : null;
|
|
const next = ui.tab === "data" ? c.button.sm({ label: "Next: tags", tone: "primary", disabled: !ready, title: ready ? "" : "Add pictures or clips to the set first", onClick: () => go("tags") })
|
|
: ui.tab === "tags" ? c.button.sm({ label: "Next: train", tone: "primary", disabled: !ready, onClick: () => go("train") }) : null;
|
|
const elsewhere = ui.tab === "results" ? [c.button.sm({ label: "Prepare a set", tone: "ghost", title: "Make or fill a set; a training going on carries on", onClick: () => go("data") }),
|
|
c.button.sm({ label: "Fix tags", tone: "ghost", title: "Correct the captions; a training going on carries on", onClick: () => go("tags") })] : [];
|
|
const actions = [results, ...elsewhere, back, skip, next];
|
|
foot = () => win && win.setFooter({ actions: [...actions, ...(ui.tab === "train" ? ui.controls || [] : [])].filter(Boolean), note: ui.tab === "train" ? footNote(ui) : "" });
|
|
foot();
|
|
}
|
|
|
|
// --- the Bin and rated results, read when the window opens
|
|
async function loadSources() {
|
|
try { ui.bin = (await api.media()).media || []; } catch { ui.bin = []; }
|
|
try { ui.rated = await api.trainRated(ui.include); } catch { ui.rated = { results: [], counts: {} }; }
|
|
draw();
|
|
}
|
|
async function loadRuns() { try { ui.runs = await api.trainRuns(); } catch { ui.runs = []; } if (ui.runs.length) loadRecord(ui.chosenRun || ui.runs[0].run); }
|
|
async function loadRecord(id) { try { ui.record = await api.trainRunRecord(id); } catch { ui.record = null; } if (ui.tab === "results") draw(); }
|
|
|
|
// --- tagger downloads and the tagging job
|
|
async function pollTaggers(fast) {
|
|
clearTimeout(tagPoll);
|
|
let again = false;
|
|
try {
|
|
const d = await api.trainTaggers();
|
|
ui.tagger.list = d.taggers || [];
|
|
const was = ui.tagger.job && ui.tagger.job.state;
|
|
const job = await api.trainAutotagState();
|
|
ui.tagger.job = job;
|
|
if (was === "running" && job.state === "done") {
|
|
const n = applyTagging(session, job, ui.tagger.mode);
|
|
c.toast.info({ text: `Tagged ${n} item${n === 1 ? "" : "s"}.` });
|
|
await session.flush();
|
|
} else if (was === "running" && job.state === "error") c.toast.warn({ text: job.error });
|
|
again = job.state === "running" || ui.tagger.list.some((t) => t.state === "downloading");
|
|
} catch { /* the page shows what it had */ }
|
|
if (win && ui.tab === "tags") draw();
|
|
if (win && (again || fast === true)) tagPoll = setTimeout(() => pollTaggers(again ? undefined : false), again ? POLL : POLL * 2);
|
|
}
|
|
|
|
// --- a run
|
|
const setRun = (patch) => Object.assign(ui.run, patch);
|
|
const start = () => launch("");
|
|
async function launch(resumeId) {
|
|
const doc = session.doc;
|
|
if (!doc) return;
|
|
await session.flush();
|
|
const settings = settingsOf(doc);
|
|
delete settings.preset; // a starting point the person picked, not a setting of the trainer
|
|
delete settings.preview_strength; // used by the previews only
|
|
settings.trigger = doc.trigger || "";
|
|
settings.postfix = doc.postfix || "";
|
|
if (resumeId) settings.resume = resumeId;
|
|
else { ui.previews = []; settings.resume = ""; }
|
|
ui.holdPause = false;
|
|
setRun(resumeId ? { shown: true, active: true, state: "queued", message: "Carrying on. The models load again first.", error: "", paused: null }
|
|
: { shown: true, active: true, state: "queued", loss: [], times: [], stepSeconds: null, lastStep: null, smooth: [], step: 0, total: settings.steps, elapsed: 0, message: "Checking the pipeline", error: "", lora: "", runId: "", paused: null });
|
|
draw();
|
|
let before = "";
|
|
ui.staleStarted = 0;
|
|
try {
|
|
const last = await api.trainRun(0);
|
|
if (resumeId) ui.staleStarted = last.started || 0; // a failed or stopped run being restored still shows as the current record until the new one begins
|
|
else before = last.run || "";
|
|
} catch { /* first run */ }
|
|
try {
|
|
const plan = await api.trainPlan(session.name, app.pipeline.slots(), settings);
|
|
if (plan.problems && plan.problems.length) throw new Error(plan.problems.join("; "));
|
|
setRun({ message: resumeId ? "Sent to ComfyUI. It loads the models again, which can take a while." : "Sent to ComfyUI. It loads the models first, which can take a while." });
|
|
draw();
|
|
await api.queueGraph(plan.prompt);
|
|
} catch (err) { setRun({ active: false, state: resumeId ? "paused" : "error", error: err.message, message: "" }); draw(); return; }
|
|
ui.before = before;
|
|
ui.awaitStart = Boolean(resumeId); // until the new run begins, the record still shows the pause it is carrying on from
|
|
watch();
|
|
}
|
|
function refreshLive() { // the graph lives in Results, a short notice in Train
|
|
if (!ui.monitorBox) return;
|
|
if (ui.tab === "results") ui.monitorBox.set([monitor(ui)].filter(Boolean));
|
|
else if (ui.tab === "train") { ui.monitorBox.set([notice(ui, go)].filter(Boolean)); foot(); }
|
|
}
|
|
function watch() {
|
|
clearTimeout(runPoll);
|
|
const queuedAt = Date.now(), mine = (ui.watching = (ui.watching || 0) + 1); // one loop at a time: an older one still waiting on the server stops when it hears back
|
|
const tick = async () => {
|
|
if (!ui.run.active || mine !== ui.watching) return;
|
|
if (ui.run.state === "queued" && Date.now() - queuedAt > QUEUE_WAIT && !ui.run.message.includes("log")) {
|
|
setRun({ message: "ComfyUI has not started the run yet. It may be busy with another job or loading models; if this stays, look at ComfyUI's log." });
|
|
refreshLive();
|
|
}
|
|
let snap = null;
|
|
try { snap = await api.trainRun(ui.run.loss.length); } catch { /* try again */ }
|
|
if (mine !== ui.watching) return;
|
|
if (snap && ui.run.state === "previewing") { runPoll = setTimeout(tick, POLL); return; } // the pause is being dealt with: its snapshot must not set the state back to paused
|
|
if (snap && snap.run && snap.run !== ui.before && !(ui.awaitStart && (ui.staleStarted ? snap.started === ui.staleStarted : snap.state === "paused"))) {
|
|
ui.awaitStart = false;
|
|
const times = snap.loss_from === ui.run.loss.length ? ui.run.times.concat(snap.times || []) : ui.run.times;
|
|
const loss = snap.loss_from === ui.run.loss.length ? ui.run.loss.concat(snap.loss) : ui.run.loss;
|
|
setRun({ times, stepSeconds: snap.step_seconds, lastStep: snap.last_step_seconds, state: snap.state, step: snap.step, total: snap.total || ui.run.total, elapsed: snap.elapsed, message: snap.message, loss, smooth: smooth(loss), error: snap.error || "", lora: snap.lora || "", runId: snap.run, paused: snap.paused || null });
|
|
if (snap.state === "paused") {
|
|
setRun({ active: false });
|
|
draw();
|
|
if (snap.paused && snap.paused.reason === "preview" && !ui.holdPause) previewNow(true);
|
|
return;
|
|
}
|
|
if (["done", "error", "stopped"].includes(snap.state)) {
|
|
setRun({ active: false });
|
|
if (snap.state === "done") c.toast.info({ text: `Done: ${snap.lora}` });
|
|
draw();
|
|
loadRuns();
|
|
loadRecord(snap.run); // a failure offers Restore and continue again
|
|
return;
|
|
}
|
|
refreshLive();
|
|
}
|
|
runPoll = setTimeout(tick, POLL);
|
|
};
|
|
runPoll = setTimeout(tick, POLL);
|
|
}
|
|
// The server is the truth. A window that sits idle on a pause (a tab left in the background, a missed poll) asks now and then and takes over what it finds.
|
|
let syncPoll = 0;
|
|
async function reconcile() {
|
|
clearTimeout(syncPoll);
|
|
if (!win) return;
|
|
syncPoll = setTimeout(reconcile, POLL * 5);
|
|
if (ui.run.active || ui.run.state !== "paused") return;
|
|
let snap = null;
|
|
try { snap = await api.trainRun(0); } catch { return; }
|
|
if (!win || ui.run.active || ui.run.state !== "paused" || !snap || !snap.run) return;
|
|
const moved = snap.state === "training" || (snap.state === "paused" && (snap.run !== ui.run.runId || snap.step !== ui.run.step));
|
|
if (!moved) return;
|
|
ui.before = "";
|
|
if (snap.state === "training") {
|
|
setRun({ shown: true, active: true, state: "training", runId: snap.run, loss: snap.loss, times: snap.times || [], stepSeconds: snap.step_seconds, lastStep: snap.last_step_seconds, smooth: smooth(snap.loss), step: snap.step, total: snap.total, elapsed: snap.elapsed, message: snap.message, error: "", paused: null });
|
|
draw();
|
|
watch();
|
|
} else {
|
|
setRun({ shown: true, active: false, state: "paused", runId: snap.run, loss: snap.loss, times: snap.times || [], stepSeconds: snap.step_seconds, lastStep: snap.last_step_seconds, smooth: smooth(snap.loss), step: snap.step, total: snap.total, elapsed: snap.elapsed, message: "", paused: snap.paused });
|
|
draw();
|
|
if (snap.paused && snap.paused.reason === "preview" && !ui.holdPause) previewNow(true);
|
|
}
|
|
}
|
|
async function stop() {
|
|
const waiting = ui.run.state === "paused" || ui.run.state === "previewing";
|
|
setRun({ message: "Stopping" });
|
|
refreshLive();
|
|
if (ui.run.state === "previewing") { ui.stopAfterPreview = true; try { app.generate.cancel(); } catch { /* it was over */ } }
|
|
try { await api.trainStop(); } catch { /* it ends anyway */ }
|
|
if (!waiting) { try { await fetch("/interrupt", { method: "POST" }); } catch { /* the training step check ends it */ } }
|
|
if (ui.run.state === "queued" || waiting) { setRun({ active: false, state: "stopped", message: "" }); draw(); }
|
|
}
|
|
async function pause() {
|
|
ui.holdPause = true; // a pause the person asked for stays: no preview step then carries on by itself
|
|
setRun({ message: ui.run.state === "previewing" ? "It will wait after this preview" : "Pausing at the end of this step" });
|
|
if (ui.run.state === "training") { try { await api.trainPause(); } catch { /* the run goes on */ } }
|
|
refreshLive();
|
|
}
|
|
/** A picture or clip made by the pipeline with the LoRA so far; when `auto`, training carries on afterwards. */
|
|
const previewStrength = () => { const v = Number(settingsOf(session.doc).preview_strength); return Number.isFinite(v) ? v : 1; };
|
|
async function previewNow(auto) {
|
|
const paused = ui.run.paused;
|
|
if (!paused || !paused.preview || ui.run.state === "previewing") return; // one preview at a time
|
|
ui.stopAfterPreview = false;
|
|
setRun({ state: "previewing", message: `Making a preview with the LoRA at step ${paused.step}` });
|
|
draw();
|
|
let out;
|
|
try { out = app.runner && app.runner.preview ? await app.runner.preview({ lora: paused.preview, strength: previewStrength() }) : { error: "Previews need the app's Generate." }; } catch (err) { out = { error: err.message }; }
|
|
if (out.image) {
|
|
const im = out.image;
|
|
ui.previews.push({ id: `p${ui.previews.length + 1}`, step: paused.step, src: { kind: "output", type: im.type || "output", subfolder: im.subfolder || "", filename: im.filename }, media_kind: VIDEO.test(im.filename) ? "video" : "image" });
|
|
} else if (!ui.stopAfterPreview) c.toast.warn({ text: out.error || "The preview could not be made." });
|
|
if (ui.stopAfterPreview) { draw(); return; }
|
|
setRun({ state: "paused", message: "" });
|
|
if (auto && !ui.holdPause) launch(ui.run.runId); else draw();
|
|
}
|
|
// A run already going when the window opens (page reloaded): pick it up.
|
|
async function resume() {
|
|
try {
|
|
const snap = await api.trainRun(0);
|
|
if (snap.state === "paused") {
|
|
ui.before = "";
|
|
setRun({ shown: true, active: false, state: "paused", loss: snap.loss, times: snap.times || [], stepSeconds: snap.step_seconds, lastStep: snap.last_step_seconds, elapsed: snap.elapsed, smooth: smooth(snap.loss), step: snap.step, total: snap.total, message: "", runId: snap.run, paused: snap.paused });
|
|
ui.holdPause = !(snap.paused && snap.paused.reason === "preview"); // a pause the person asked for waits for them; one for a preview step was never theirs, so the preview and the carrying on go ahead
|
|
ui.tab = "train";
|
|
if (!ui.holdPause) previewNow(true);
|
|
} else if (snap.state === "training") {
|
|
ui.before = "";
|
|
setRun({ shown: true, active: true, state: "training", runId: snap.run, loss: snap.loss, times: snap.times || [], stepSeconds: snap.step_seconds, lastStep: snap.last_step_seconds, smooth: smooth(snap.loss), step: snap.step, total: snap.total, elapsed: snap.elapsed, message: snap.message });
|
|
ui.tab = "train";
|
|
watch();
|
|
}
|
|
} catch { /* nothing to resume */ }
|
|
}
|
|
|
|
function open() {
|
|
if (win) return;
|
|
body = c.region.stack({ gap: "lg" });
|
|
win = c.modal.generic({ title: "Make a LoRA", subtitle: "teach the model your style, in four steps", size: "xl", closeOnOutside: false, body, origin: () => button.node,
|
|
onClose: () => { win = null; session.flush(); clearTimeout(tagPoll); clearTimeout(syncPoll); } });
|
|
syncPoll = setTimeout(reconcile, POLL * 5);
|
|
offSession = session.on(draw);
|
|
session.refresh().then(async () => {
|
|
if (!session.name) {
|
|
const remembered = (() => { try { return localStorage.getItem("fp.training.set") || ""; } catch { return ""; } })();
|
|
const pick = session.sets.find((x) => x.name === remembered) || session.sets[0];
|
|
if (pick) await session.open(pick.name);
|
|
}
|
|
draw();
|
|
});
|
|
loadSources();
|
|
pollTaggers();
|
|
resume();
|
|
draw();
|
|
}
|
|
const offOpen = session.on(() => { try { if (session.name) localStorage.setItem("fp.training.set", session.name); } catch { /* private mode */ } });
|
|
|
|
return () => { clearTimeout(tagPoll); clearTimeout(runPoll); clearTimeout(syncPoll); timers.forEach(clearTimeout); if (offSession) offSession(); offOpen(); if (win) win.close(); };
|
|
},
|
|
};
|