From 7dfbcf3c1bbdabbe60071a950606944480ece669 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 17 Oct 2026 15:18:38 +0000 Subject: [PATCH] Training window: data, tags, train and results steps; docs and changelog Built only from kit elements. UNTESTED on real models or GPUs; backend proven on tiny stand-ins on CPU. Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P --- CHANGELOG.md | 5 +- app/ui/features/training/captions.js | 44 ++++++ app/ui/features/training/data.js | 99 ++++++++++++ app/ui/features/training/index.js | 181 ++++++++++++++++++++++ app/ui/features/training/results.js | 31 ++++ app/ui/features/training/state.js | 76 +++++++++ app/ui/features/training/tags.js | 149 ++++++++++++++++++ app/ui/features/training/train.js | 113 ++++++++++++++ app/ui/features/training/training.test.js | 155 ++++++++++++++++++ app/ui/modules.js | 1 + app/ui/shell/api.js | 33 ++++ docs/module-reference.md | 20 +++ docs/testing.md | 1 + 13 files changed, 906 insertions(+), 2 deletions(-) create mode 100644 app/ui/features/training/captions.js create mode 100644 app/ui/features/training/data.js create mode 100644 app/ui/features/training/index.js create mode 100644 app/ui/features/training/results.js create mode 100644 app/ui/features/training/state.js create mode 100644 app/ui/features/training/tags.js create mode 100644 app/ui/features/training/train.js create mode 100644 app/ui/features/training/training.test.js diff --git a/CHANGELOG.md b/CHANGELOG.md index 8cd7b877..ac88333a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,11 +20,12 @@ A rewrite from scratch, not an update of 4.x. Same idea (a studio for video insi ### Features - **Proven**: prompt markup, empty latent, second pass (the H3 preset for it is held back), loaders (diffusion model, CLIP, VAE, LoRA, checkpoint, media), cut opening, decode, save video, anchor de-staticking (ALG), randomizer, DynaShift, first-step seed search, modifiers, sampler, detail, shot camera, prompt enhancer, Render, taste modules (slider, guidance, prompt nudge, attention, REINS, early guidance), shot memory, decisiveness, late-branch guidance, momentum, velocity bias. -- **Experimental**: continuity, identity transfer, negative erase, camera move, context windows, block influence, block repeat, phrase probe, shadow negative, video detail, temporal style, massive-activation steering, video-to-audio coupling. +- **Experimental**: continuity, identity transfer, negative erase, camera move, context windows, block influence, block repeat, phrase probe, shadow negative, video detail, temporal style, massive-activation steering, video-to-audio coupling, **LoRA training**. - **Prompt enhancer**: a language model rewrites each finished prompt. Has its own encoder pick (Composer ▸ Enhance ▸ Encoder), because H3's own encoder cannot write prose. Verified with Gemma 4. - **Auto Montage**: builds a new cut at the end of the timeline from your rendered clips; originals stay untouched, and the pieces are never regenerated or rated. - **Upscale** a finished render with an upscale model (`.safetensors`) or NVIDIA RTX Super Resolution. -- **Add-ons** (spaCy for automatic camera and focus, nvidia-vfx) install on the next start after you switch them on. +- **LoRA training** (add-on, Settings > Modules): a Training window with four steps. Data (optional): drop files, pick from your Bin, or take your rated results (Liked, Disliked, Bad composition, Bad image, All). Tags: a caption editor beside every picture, a tag cloud, bulk edits and automatic tagging with downloadable tagger models. Train: all the settings, a live loss graph, Stop. Results: past runs. It trains on ComfyUI's own trainer, using the loaders in your pipeline, and supports H3 (picture and sound together). **Experimental, and UNTESTED on real models or GPUs**: proven only on tiny stand-in models on CPU. +- **Add-ons** (spaCy for automatic camera and focus, nvidia-vfx, onnxruntime for LoRA training) install on the next start after you switch them on. ### Not carried over diff --git a/app/ui/features/training/captions.js b/app/ui/features/training/captions.js new file mode 100644 index 00000000..3e40a536 --- /dev/null +++ b/app/ui/features/training/captions.js @@ -0,0 +1,44 @@ +// Captions as tags, the same rules as the server's modules/training/lora/captions.py. + +export const splitTags = (text) => { + const seen = new Set(), out = []; + for (const part of String(text ?? "").replace(/\n/g, ",").split(",")) { + const tag = part.split(/\s+/).filter(Boolean).join(" "); + if (tag && !seen.has(tag.toLowerCase())) { seen.add(tag.toLowerCase()); out.push(tag); } + } + return out; +}; +export const joinTags = (tags) => splitTags((tags || []).join(", ")).join(", ").slice(0, 4000); + +/** Every tag with how many captions hold it, most common first. */ +export function frequency(items) { + const counts = new Map(); + for (const item of items) for (const tag of splitTags(item.caption)) { + const key = tag.toLowerCase(), have = counts.get(key); + if (have) have.count += 1; else counts.set(key, { tag, count: 1 }); + } + return [...counts.values()].sort((a, b) => b.count - a.count || a.tag.localeCompare(b.tag)); +} + +/** Change captions in place; returns how many changed. op: add, prepend, remove, replace. `ids` (a Set) limits it. */ +export function bulk(items, op, tag, to, ids) { + tag = String(tag ?? "").split(/\s+/).filter(Boolean).join(" "); + if (!tag || !["add", "prepend", "remove", "replace"].includes(op)) return 0; + const low = tag.toLowerCase(); + to = String(to ?? "").split(/\s+/).filter(Boolean).join(" "); + if (op === "replace" && !to) op = "remove"; + let changed = 0; + for (const item of items) { + if (ids && !ids.has(item.id)) continue; + let tags = splitTags(item.caption); + const has = tags.some((t) => t.toLowerCase() === low); + if (op === "add" && !has) tags.push(tag); + else if (op === "prepend") tags = [tag, ...tags.filter((t) => t.toLowerCase() !== low)]; + else if (op === "remove" && has) tags = tags.filter((t) => t.toLowerCase() !== low); + else if (op === "replace" && has) tags = tags.map((t) => (t.toLowerCase() === low ? to : t)); + else continue; + const text = joinTags(tags); + if (text !== item.caption) { item.caption = text; changed += 1; } + } + return changed; +} diff --git a/app/ui/features/training/data.js b/app/ui/features/training/data.js new file mode 100644 index 00000000..26c094a7 --- /dev/null +++ b/app/ui/features/training/data.js @@ -0,0 +1,99 @@ +// Step 1, Training data: pick or make a set, then put media in it from any mix of sources. +// Skippable: a set that already has what you want goes straight on to tags or to training. +import { composer as c } from "../../composer/composer.js"; +import { LABEL_WORDS } from "./state.js"; + +export const RATED = [ + { value: "liked", label: "Liked results" }, + { value: "disliked", label: "Disliked results" }, + { value: "bad_composition", label: "Bad composition results" }, + { value: "bad_image", label: "Bad image results" }, + { value: "all", label: "All" }, +]; +const ALL_RATED = ["liked", "disliked", "bad_composition", "bad_image"]; + +export const binEntry = (m) => ({ src: { kind: "bin", id: m.id }, media_kind: m.kind === "video" ? "video" : "image", label: "own", name: m.name }); +export const ratedEntry = (r) => ({ src: { kind: "output", type: r.media.type, subfolder: r.media.subfolder, filename: r.media.filename }, + media_kind: /\.(mp4|webm|mov|mkv|avi)$/i.test(r.media.filename) ? "video" : "image", label: r.label, caption: r.caption || "", name: r.media.filename }); + +const countsLine = (counts) => Object.keys(LABEL_WORDS).filter((k) => counts[k]).map((k) => `${counts[k]} ${LABEL_WORDS[k].toLowerCase()}`).join(", "); + +/** Builds the step. `ui` keeps what is picked between redraws. */ +export function dataStep({ app, session, go, ui }) { + const added = async (entries, what) => { + if (!entries.length) { c.toast.info({ text: `Nothing to add from ${what}.` }); return; } + try { + const n = await session.add(entries); + c.toast.info({ text: n ? `Added ${n} to "${session.name}".` : "They were all in the set already." }); + } catch (err) { c.toast.warn({ text: err.message }); } + }; + + const names = session.sets.map((x) => ({ value: x.name, label: `${x.name} (${x.count})` })); + const setSelect = c.select.md({ label: "Training set", options: [{ value: "", label: names.length ? "Pick a set" : "No sets yet" }, ...names], value: session.name, onChange: (v) => session.open(v) }); + const make = c.button.sm({ label: "New set", tone: "neutral", onClick: async () => { + const answer = await c.modal.prompt({ title: "New training set", label: "Name", placeholder: "my_style", confirmLabel: "Create" }).result; + if (!answer) return; + try { await session.create(String(answer).trim()); } catch (err) { c.toast.warn({ text: err.message }); } + } }); + const drop = c.button.sm({ label: "Delete set", tone: "danger", onClick: async () => { + if (!(await c.modal.dialogue({ title: "Delete training set", message: `Delete "${session.name}"? Your media and results stay where they are.`, tone: "danger", confirmLabel: "Delete" }).result)) return; + await session.remove(); + } }); + drop.node.hidden = !session.name; + + const open = !!session.name; + const counts = session.counts(); + const total = ((session.doc || {}).items || []).length; + const status = c.hint.default({ text: open ? (total ? `${total} in "${session.name}": ${countsLine(counts)}.` : `"${session.name}" is empty. Add something below, or go on and add later.`) : "Make a set first. It is only a saved list, nothing is copied or changed." }); + + const dropzone = c.dropzone.default({ label: "Drop pictures or clips here", hint: "They are kept in your Bin and added to this set", accept: "image/*,video/*", onFiles: async (files) => { + if (!open) { c.toast.warn({ text: "Make or pick a set first." }); return; } + const out = await app.api.uploadMedia(files); + if (out.problems && out.problems.length) c.toast.warn({ text: out.problems.join("; ") }); + if (app.say) app.say("media"); + await added((out.media || []).filter((m) => m.kind === "image" || m.kind === "video").map(binEntry), "the drop"); + } }); + + const picked = ui.picked; + const pickedNote = c.text.sm({ text: `${picked.size} picked` }); + const usable = ui.bin.filter((m) => m.kind === "image" || m.kind === "video"); + const gallery = c.gallery.adaptive({ id: "training-bin", items: usable.map((m) => ({ id: m.id, label: m.name, thumb: `/funpack/api/media/${encodeURIComponent(m.id)}/thumb`, badge: m.kind })), selection: [...picked], + empty: "Your Bin has no pictures or clips yet", onActivate: (item) => { + if (picked.has(item.id)) picked.delete(item.id); else picked.add(item.id); + gallery.setValue([...picked]); + pickedNote.setText(`${picked.size} picked`); + addBin.node.disabled = !picked.size; + } }); + const addBin = c.button.sm({ label: "Add picked", tone: "primary", disabled: !picked.size || !open, onClick: async () => { + const chosen = usable.filter((m) => picked.has(m.id)); + await added(chosen.map(binEntry), "the Bin"); + picked.clear(); + } }); + + const want = () => (ui.ratedKind === "all" ? ALL_RATED : [ui.ratedKind]); + const ratedCount = () => want().reduce((a, k) => a + ((ui.rated.counts || {})[k] || 0), 0); + const ratedNote = c.text.sm({ text: "" }); + const addRated = c.button.sm({ label: "Add these results", tone: "primary", disabled: true, onClick: async () => { + await added((ui.rated.results || []).filter((r) => want().includes(r.label)).map(ratedEntry), "your results"); + } }); + const showRated = () => { + const n = ratedCount(); + ratedNote.setText(n ? `${n} rated result${n > 1 ? "s" : ""} found in your projects.` : "None found yet. Rate results in your projects and they appear here."); + addRated.node.disabled = !n || !open; + }; + const ratedSel = c.segmented.sm({ label: "Which results", options: RATED, value: ui.ratedKind, onChange: (v) => { ui.ratedKind = v; showRated(); } }); + showRated(); + + const next = c.toolbar.default({ items: [c.hint.default({ text: "This step is optional: a set you made earlier is ready as it is." })], trailing: [ + c.button.md({ label: "Skip to training", tone: "ghost", onClick: () => go("train") }), + c.button.md({ label: "Next: tags", tone: "primary", disabled: !open, onClick: () => go("tags") })] }); + + return c.region.stack({ gap: "md", children: [ + c.field.default({ label: "Training set", hint: "A saved list of what to learn from.", control: c.region.stack({ gap: "sm", children: [setSelect, c.toolbar.default({ items: [make, drop] })] }) }), + status, + c.group.default({ label: "Drop files", rows: [dropzone] }), + c.group.default({ label: "Choose in your Bin", rows: [gallery, c.toolbar.default({ items: [pickedNote], trailing: [addBin] })] }), + c.group.default({ label: "Choose your own results", hint: "Results you rated in your projects. Liked ones teach what you want; disliked ones can teach what to avoid.", rows: [ratedSel, c.toolbar.default({ items: [ratedNote], trailing: [addRated] })] }), + next, + ] }); +} diff --git a/app/ui/features/training/index.js b/app/ui/features/training/index.js new file mode 100644 index 00000000..0825bd73 --- /dev/null +++ b/app/ui/features/training/index.js @@ -0,0 +1,181 @@ +// "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 { monitor, settingsOf, trainStep } from "./train.js"; +import { resultsStep } from "./results.js"; +import { makeSession } from "./state.js"; + +const TABS = [{ value: "data", label: "1 Data" }, { value: "tags", label: "2 Tags" }, { value: "train", label: "3 Train" }, { value: "results", label: "Results" }]; +const POLL = 1000; + +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"], + setup({ host, app }) { + const api = app.api; + const session = makeSession(api); + const ui = { + tab: "data", picked: new Set(), bin: [], rated: { results: [], counts: {} }, ratedKind: "liked", 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: [], smooth: [], step: 0, total: 0, elapsed: 0, message: "", error: "", lora: "", before: "" }, + monitorBox: null, redraw: () => draw(), pollTaggers: () => pollTaggers(true), loadRecord: (id) => loadRecord(id), loadRuns: () => loadRuns(), + start: () => start(), stop: () => stop(), + }; + 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" ? `/funpack/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) { ui.tab = tab; if (tab === "results") loadRuns(); draw(); } + function draw() { + if (!win) return; + const tabs = c.tabs.underline({ label: "Training steps", tabs: TABS, value: ui.tab, onChange: go }); + const note = session.error ? c.inlineError.default({ text: session.error }) : null; + body.set([tabs, note, steps[ui.tab]()].filter(Boolean)); + } + + // --- 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(); } 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); + async function start() { + const doc = session.doc; + if (!doc) return; + await session.flush(); + const settings = settingsOf(doc); + settings.trigger = doc.trigger || ""; + setRun({ shown: true, active: true, state: "queued", loss: [], smooth: [], step: 0, total: settings.steps, elapsed: 0, message: "Checking the pipeline", error: "", lora: "" }); + draw(); + let before = ""; + try { before = (await api.trainRun(0)).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: "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: "error", error: err.message, message: "" }); draw(); return; } + ui.before = before; + watch(); + } + function watch() { + clearTimeout(runPoll); + const tick = async () => { + if (!ui.run.active) return; + let snap = null; + try { snap = await api.trainRun(ui.run.loss.length); } catch { /* try again */ } + if (snap && snap.run && snap.run !== ui.before) { + const loss = snap.loss_from === ui.run.loss.length ? ui.run.loss.concat(snap.loss) : ui.run.loss; + setRun({ 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 || "" }); + if (["done", "error", "stopped"].includes(snap.state)) { + setRun({ active: false }); + if (snap.state === "done") c.toast.info({ text: `Done: ${snap.lora}` }); + draw(); + loadRuns(); + return; + } + if (ui.tab === "train" && ui.monitorBox) ui.monitorBox.set([monitor(ui)]); + } + runPoll = setTimeout(tick, POLL); + }; + runPoll = setTimeout(tick, POLL); + } + async function stop() { + setRun({ message: "Stopping" }); + try { await api.trainStop(); } catch { /* it ends anyway */ } + try { await fetch("/interrupt", { method: "POST" }); } catch { /* the training step check ends it */ } + if (ui.run.state === "queued") { setRun({ active: false, state: "stopped", message: "" }); 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 === "training") { + ui.before = ""; + setRun({ shown: true, active: true, state: "training", loss: snap.loss, 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: "md", fill: true }); + win = c.floating.window({ id: "training", title: "Training", subtitle: "make a LoRA", body, width: 780, height: 720, x: 240, y: 40, + onClose: () => { win = null; session.flush(); clearTimeout(tagPoll); } }); + 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); timers.forEach(clearTimeout); if (offSession) offSession(); offOpen(); if (win) win.close(); }; + }, +}; diff --git a/app/ui/features/training/results.js b/app/ui/features/training/results.js new file mode 100644 index 00000000..c22cd13f --- /dev/null +++ b/app/ui/features/training/results.js @@ -0,0 +1,31 @@ +// Step 4, Results: the runs that finished, with their loss graphs. +import { composer as c } from "../../composer/composer.js"; +import { clock, fmt } from "./train.js"; + +const STATES = { done: "Finished", error: "Failed", stopped: "Stopped", training: "Running" }; + +export function resultsStep({ app, ui }) { + const runs = ui.runs; + if (!runs.length) return c.emptyState.default({ title: "No runs yet", hint: "A run you start in the Train step is kept here, with its loss graph." }); + const chosen = runs.find((r) => r.run === ui.chosenRun) || runs[0]; + const rows = runs.map((r) => ({ id: r.run, label: `${r.name || "LoRA"} · ${STATES[r.state] || r.state}`, hint: `${r.set} · ${r.steps || 0} steps${r.final_loss != null ? ` · loss ${fmt(r.final_loss)}` : ""}` })); + const list = c.gallery.rows({ items: rows, selection: [chosen.run], onActivate: (item) => { ui.chosenRun = item.id; ui.loadRecord(item.id); ui.redraw(); } }); + const rec = ui.record && ui.record.run === chosen.run ? ui.record : null; + const chart = c.chart.line({ series: rec && rec.loss ? [{ name: "loss", values: rec.loss, tone: "accent" }] : [], xLabel: "step", yLabel: "loss", height: "md", empty: rec ? "No loss was recorded" : "Loading" }); + const facts = rec ? [ + `${STATES[rec.state] || rec.state}`, `${rec.steps || 0} of ${rec.total || 0} steps`, rec.seconds ? clock(rec.seconds) : null, + rec.final_loss != null ? `final loss ${fmt(rec.final_loss)}` : null, rec.set ? `from "${rec.set}"` : null, + ].filter(Boolean).join(" · ") : ""; + const s = (rec && rec.settings) || {}; + const settings = rec ? ["algorithm", "rank", "learning_rate", "optimizer", "loss", "resolution", "batch_size", "use"].filter((k) => s[k] !== undefined).map((k) => `${k.replace("_", " ")} ${s[k]}`).join(" · ") : ""; + return c.region.stack({ gap: "md", children: [ + list, + c.group.default({ label: chosen.name || "LoRA", rows: [ + c.hint.default({ text: facts }), + rec && rec.lora ? c.banner.info({ text: `File: ${rec.lora}. Pick it in a LoRA loader to use it.` }) : null, + rec && rec.error ? c.inlineError.default({ text: rec.error }) : null, + chart, c.hint.default({ text: settings }), + c.toolbar.default({ items: [], trailing: [c.button.sm({ label: "Delete this record", tone: "ghost", title: "Removes the record only, not the LoRA file", onClick: async () => { await app.api.trainForgetRun(chosen.run); ui.chosenRun = ""; await ui.loadRuns(); ui.redraw(); } })] }), + ].filter(Boolean) }), + ] }); +} diff --git a/app/ui/features/training/state.js b/app/ui/features/training/state.js new file mode 100644 index 00000000..2574c464 --- /dev/null +++ b/app/ui/features/training/state.js @@ -0,0 +1,76 @@ +// What the Training window shares between its steps: which set is open, the set itself, and saving it. +// A caption edited on screen is the truth: the server's copy follows a moment later, and tagger results are merged in by the window, not written behind its back. + +const TRAIN = "/funpack/api/m/training_lora"; +const SAVE_AFTER = 600; + +export const LABEL_WORDS = { own: "Your media", liked: "Liked", disliked: "Disliked", bad_composition: "Bad composition", bad_image: "Bad image" }; +export const LABEL_TONES = { own: "neutral", liked: "good", disliked: "bad", bad_composition: "bad", bad_image: "bad" }; + +/** A picture to show for an item: the bin's thumbnail, or the output file's still. */ +export function stillOf(item) { + const s = item.src; + if (s.kind === "bin") return `/funpack/api/media/${encodeURIComponent(s.id)}/thumb`; + const q = `filename=${encodeURIComponent(s.filename)}&subfolder=${encodeURIComponent(s.subfolder || "")}&type=${encodeURIComponent(s.type || "output")}`; + return item.media_kind === "video" ? `/funpack/api/m/render/poster?${q}` : `/view?${q}`; +} + +export function makeSession(api) { + const listeners = new Set(); + const s = { name: "", doc: null, sets: [], error: "", saving: false, state: null }; + let timer = null, queued = null; + + const emit = () => listeners.forEach((f) => f()); + s.on = (f) => { listeners.add(f); return () => listeners.delete(f); }; + s.emit = emit; + + async function push() { + if (!queued) return; + const { name, doc } = queued; + queued = null; + s.saving = true; + try { await api.trainSaveSet(name, doc); s.error = ""; } catch (err) { s.error = `Could not save the set: ${err.message}`; } + s.saving = false; + emit(); + } + s.flush = async () => { clearTimeout(timer); timer = null; await push(); }; + /** Change the open set in place and save it soon. Does not redraw: the controls on screen already show the edit. */ + s.edit = (change) => { + if (!s.doc) return; + change(s.doc); + queued = { name: s.name, doc: s.doc }; + clearTimeout(timer); + timer = setTimeout(push, SAVE_AFTER); + }; + s.refresh = async () => { + try { s.sets = (await api.trainSets()).sets || []; s.error = ""; } catch (err) { s.error = err.message; } + emit(); + }; + s.open = async (name) => { + await s.flush(); + s.name = name || ""; + s.doc = null; + if (name) { + try { s.doc = await api.trainSet(name); s.error = ""; } catch (err) { s.error = err.message; s.name = ""; } + } + emit(); + }; + /** Take the server's version of the open set (after it changed there). */ + s.reload = async () => { await s.flush(); if (s.name) { try { s.doc = await api.trainSet(s.name); } catch (err) { s.error = err.message; } } await s.refresh(); }; + s.add = async (entries) => { + if (!s.name) throw new Error("Make or pick a set first."); + await s.flush(); + const out = await api.trainAdd(s.name, entries); + s.doc = out.set; + await s.refresh(); + return out.added; + }; + s.create = async (name) => { await api.trainCreateSet(name); await s.refresh(); await s.open(name); }; + s.remove = async () => { if (s.name) { await api.trainDeleteSet(s.name); } await s.open(""); await s.refresh(); }; + s.counts = () => { + const out = {}; + for (const item of (s.doc && s.doc.items) || []) out[item.label] = (out[item.label] || 0) + 1; + return out; + }; + return s; +} diff --git a/app/ui/features/training/tags.js b/app/ui/features/training/tags.js new file mode 100644 index 00000000..0e3ff252 --- /dev/null +++ b/app/ui/features/training/tags.js @@ -0,0 +1,149 @@ +// Step 2, Tags: every item with its caption beside its picture, bulk edits, a tag cloud, and automatic tagging. +import { composer as c } from "../../composer/composer.js"; +import { LABEL_TONES, LABEL_WORDS, stillOf } from "./state.js"; +import { bulk, frequency, joinTags, splitTags } from "./captions.js"; + +const PAGE = 48; +const FILTERS = [ + { value: "all", label: "All" }, + { value: "untagged", label: "No tags yet" }, + { value: "selected", label: "Selected" }, +]; + +export function tagsStep({ app, session, go, ui, peek }) { + const doc = session.doc; + if (!doc) return c.emptyState.default({ title: "No set open", hint: "Pick or make a training set in the Data step.", action: c.button.sm({ label: "Back to data", tone: "primary", onClick: () => go("data") }) }); + const all = doc.items; + const sel = ui.selected; + for (const id of [...sel]) if (!all.some((i) => i.id === id)) sel.delete(id); + const cloud = frequency(all); + + // --- what is shown + const shown = () => all.filter((i) => (ui.filter === "all" ? true : ui.filter === "untagged" ? !i.caption.trim() : sel.has(i.id)) + && (!ui.tag || splitTags(i.caption).some((t) => t.toLowerCase() === ui.tag.toLowerCase()))); + const list = shown(); + const count = c.text.sm({ text: "" }); + const sync = () => count.setText(`${sel.size} selected of ${all.length}${ui.tag ? ` · showing "${ui.tag}"` : ""}`); + sync(); + + const cards = list.slice(0, ui.limit).map((item) => { + const card = c.captioned.cell({ title: item.name || item.src.filename || item.src.id, src: stillOf(item), kind: item.media_kind, caption: item.caption, + badge: LABEL_WORDS[item.label], tone: LABEL_TONES[item.label], selected: sel.has(item.id), suggestions: cloud.map((x) => x.tag).slice(0, 400), + onChange: (text) => session.edit(() => { item.caption = text; }), + onSelect: (on) => { if (on) sel.add(item.id); else sel.delete(item.id); sync(); }, + onOpen: () => peek(item) }); + return card; + }); + const grid = cards.length ? c.grid.cards({ children: cards }) : c.hint.default({ text: ui.filter === "untagged" ? "Everything has tags." : "Nothing to show." }); + const more = list.length > ui.limit ? c.button.sm({ label: `Show ${Math.min(PAGE, list.length - ui.limit)} more (${list.length - ui.limit} left)`, tone: "ghost", onClick: () => { ui.limit += PAGE; ui.redraw(); } }) : null; + + // --- the bar over the grid + const filter = c.segmented.sm({ label: "Show", options: FILTERS, value: ui.filter, onChange: (v) => { ui.filter = v; ui.limit = PAGE; ui.redraw(); } }); + const everyone = c.button.sm({ label: "Select all shown", tone: "ghost", onClick: () => { list.forEach((i) => sel.add(i.id)); ui.redraw(); } }); + const none = c.button.sm({ label: "Clear", tone: "ghost", onClick: () => { sel.clear(); ui.redraw(); } }); + + // --- bulk edits on the selected (or all, when nothing is selected) + const tagBox = c.input.md({ label: "Tag", placeholder: "tag" }); + const toBox = c.input.md({ label: "Replace with", placeholder: "replace with" }); + const target = () => (sel.size ? sel : null); + const apply = (op) => { + const word = tagBox.value; + let n = 0; + session.edit((d) => { n = bulk(d.items, op, word, toBox.value, target()); }); + c.toast.info({ text: n ? `Changed ${n} caption${n > 1 ? "s" : ""}${sel.size ? " (selected)" : " (all)"}.` : "Nothing changed." }); + ui.redraw(); + }; + const bar = c.toolbar.default({ label: "Bulk tag edits", items: [tagBox, toBox, c.button.sm({ label: "Add", tone: "neutral", title: "Add the tag at the end", onClick: () => apply("add") }), + c.button.sm({ label: "Add first", tone: "neutral", title: "Put the tag at the start", onClick: () => apply("prepend") }), + c.button.sm({ label: "Remove", tone: "neutral", onClick: () => apply("remove") }), + c.button.sm({ label: "Replace", tone: "neutral", onClick: () => apply("replace") })] }); + + // --- tag cloud: click to show only those, x to take the tag off everything + const cloudEl = c.tags.cloud({ label: "Tags in this set", items: cloud, selected: ui.tag ? [ui.tag] : [], empty: "No tags yet. Type some, or use automatic tagging below.", + onPick: (t) => { ui.tag = ui.tag && ui.tag.toLowerCase() === t.toLowerCase() ? "" : t; ui.limit = PAGE; ui.redraw(); }, + onRemove: async (t) => { + if (!(await c.modal.dialogue({ title: "Remove tag", message: `Remove "${t}" from every caption in this set?`, tone: "danger", confirmLabel: "Remove" }).result)) return; + session.edit((d) => { bulk(d.items, "remove", t, "", null); }); + ui.redraw(); + } }); + + // --- trigger word + const trigger = c.input.md({ label: "Trigger word", placeholder: "none", value: doc.trigger || "", + onCommit: (v) => session.edit((d) => { d.trigger = v.split(/\s+/).filter(Boolean).join(" ").slice(0, 80); }) }); + + return c.region.stack({ gap: "md", children: [ + c.field.default({ label: "Trigger word", hint: "Put first in every caption when training, so you can call the LoRA up later by writing it. Optional.", control: trigger }), + c.group.default({ label: "Automatic tagging", rows: [autoTag({ app, session, ui })] }), + c.group.default({ label: "Tags in this set", hint: "Click one to show only the items that have it. The x removes it from all.", rows: [cloudEl] }), + c.toolbar.default({ items: [filter, count], trailing: [everyone, none] }), + bar, + grid, more, + c.toolbar.default({ items: [], trailing: [c.button.md({ label: "Back", tone: "ghost", onClick: () => go("data") }), c.button.md({ label: "Next: train", tone: "primary", onClick: () => go("train") })] }), + ] }); +} + +// --- automatic tagging ------------------------------------------------------- + +function autoTag({ app, session, ui }) { + const t = ui.tagger; + const row = (t.list || []).find((x) => x.id === t.id) || (t.list || [])[0]; + const options = (t.list || []).map((x) => ({ value: x.id, label: `${x.title}${x.state === "ready" ? "" : " (not downloaded)"}` })); + const pick = c.select.md({ label: "Tagger", options: options.length ? options : [{ value: "", label: "Loading" }], value: row ? row.id : "", onChange: (v) => { t.id = v; ui.redraw(); } }); + const state = row ? row.state : "absent"; + const fetching = state === "downloading"; + const pct = row && row.total ? Math.round((100 * row.done) / row.total) : 0; + const download = c.button.sm({ label: state === "ready" ? "Downloaded" : fetching ? "Downloading" : state === "error" ? "Try the download again" : "Download", tone: "primary", + disabled: !row || state === "ready" || fetching, onClick: async () => { try { await app.api.trainDownloadTagger(row.id); } catch (err) { c.toast.warn({ text: err.message }); } ui.pollTaggers(); } }); + const remove = c.button.sm({ label: "Delete download", tone: "ghost", onClick: async () => { await app.api.trainRemoveTagger(row.id); ui.pollTaggers(); } }); + remove.node.hidden = state !== "ready"; + const custom = c.button.sm({ label: "Use another model", tone: "ghost", title: "Any Hugging Face repo with model.onnx and selected_tags.csv, like owner/name", onClick: async () => { + const answer = await c.modal.prompt({ title: "Another tagger", label: "Hugging Face repo (owner/name)", placeholder: "SmilingWolf/wd-eva02-large-tagger-v3", confirmLabel: "Add" }).result; + if (!answer) return; + t.id = String(answer).trim(); + try { await app.api.trainDownloadTagger(t.id); } catch (err) { c.toast.warn({ text: err.message }); } + ui.pollTaggers(); + } }); + const info = c.hint.default({ text: row ? (fetching ? `Downloading ${pct}%` : state === "error" ? row.error : row.hint || "") : "" }); + + const general = c.number.md({ label: "Confidence", value: t.general, min: 0.05, max: 0.95, step: 0.05, precision: 2, onChange: (v) => { t.general = v; } }); + const limit = c.number.md({ label: "Most tags per item", value: t.limit, min: 0, max: 200, step: 1, onChange: (v) => { t.limit = v; } }); + const characters = c.toggle.default({ label: "Include character names", checked: t.characters, onChange: (v) => { t.characters = v; } }); + const mode = c.segmented.sm({ label: "With existing tags", options: [{ value: "add", label: "Keep and add" }, { value: "replace", label: "Replace" }], value: t.mode, onChange: (v) => { t.mode = v; } }); + const scope = c.segmented.sm({ label: "Which items", options: [{ value: "untagged", label: "No tags yet" }, { value: "selected", label: "Selected" }, { value: "all", label: "All" }], value: t.scope, onChange: (v) => { t.scope = v; } }); + + const job = t.job || { state: "idle" }; + const running = job.state === "running"; + const bar = running ? c.progress.bar({ value: job.done || 0, max: job.total || 1, label: "Tagging" }) : null; + const run = c.button.md({ label: running ? "Stop tagging" : "Tag now", tone: running ? "danger" : "primary", disabled: !running && state !== "ready", onClick: async () => { + if (running) { await app.api.trainAutotagCancel(); return; } + await session.flush(); + const items = session.doc.items; + const ids = items.filter((i) => (t.scope === "all" ? true : t.scope === "selected" ? ui.selected.has(i.id) : !i.caption.trim())).map((i) => i.id); + if (!ids.length) { c.toast.info({ text: "No items match that choice." }); return; } + try { + await app.api.trainAutotag({ set: session.name, tagger: t.id || (row && row.id), ids, general: t.general, characters: t.characters, limit: t.limit }); + ui.pollTaggers(true); + } catch (err) { c.toast.warn({ text: err.message }); } + } }); + const why = job.state === "error" ? c.inlineError.default({ text: job.error }) : null; + const problems = (job.problems || []).length ? c.hint.default({ text: `Skipped: ${job.problems.slice(0, 3).join("; ")}${job.problems.length > 3 ? "…" : ""}` }) : null; + return c.region.stack({ gap: "sm", children: [ + c.toolbar.default({ items: [pick], trailing: [download, remove, custom] }), info, + c.field.row({ fields: [c.field.default({ label: "Confidence", hint: "Higher keeps only sure tags.", control: general }), c.field.default({ label: "Most tags per item", hint: "0 means no limit.", control: limit })] }), + characters, mode, scope, bar, why, problems, c.toolbar.default({ items: [], trailing: [run] }), + ].filter(Boolean) }); +} + +/** Merge a finished tagging job into the open set. */ +export function applyTagging(session, job, mode) { + let n = 0; + session.edit((d) => { + for (const item of d.items) { + const found = (job.results || {})[item.id]; + if (!found) continue; + const text = joinTags(mode === "replace" ? found : [...splitTags(item.caption), ...found]); + if (text !== item.caption) { item.caption = text; n += 1; } + } + }); + return n; +} diff --git a/app/ui/features/training/train.js b/app/ui/features/training/train.js new file mode 100644 index 00000000..c4f9d12e --- /dev/null +++ b/app/ui/features/training/train.js @@ -0,0 +1,113 @@ +// Step 3, Train: the settings, Start and Stop, and the loss graph while it runs. +import { composer as c } from "../../composer/composer.js"; + +/** The node's own defaults. A set remembers what was changed from them. */ +export const DEFAULTS = { + name: "funpack_lora", use: "normal", steps: 500, learning_rate: 0.0005, rank: 16, algorithm: "LoRA", optimizer: "AdamW", loss: "MSE", + batch_size: 1, grad_accumulation: 1, resolution: 512, aspect_buckets: false, frames: 33, fps: 0, start: "start", seed: 0, + training_dtype: "bf16", lora_dtype: "bf16", gradient_checkpointing: true, offloading: false, bypass: false, +}; +export const PRESETS = [ + { value: "check", label: "Quick check", steps: 100, rank: 8 }, + { value: "standard", label: "Standard", steps: 500, rank: 16 }, + { value: "long", label: "Thorough", steps: 2000, rank: 32 }, +]; +const USES = [ + { value: "normal", label: "Learn what I like" }, + { value: "avoid", label: "Learn what to avoid" }, + { value: "all", label: "Learn from everything" }, +]; +const opts = (list) => list.map((v) => ({ value: v, label: v })); + +/** The settings for a run: the defaults, then what the set remembers. */ +export const settingsOf = (doc) => ({ ...DEFAULTS, ...((doc && doc.settings) || {}) }); + +const fmt = (n) => (Number.isFinite(n) ? (Math.abs(n) >= 0.01 ? n.toFixed(4) : n.toPrecision(3)) : ""); +const clock = (s) => { s = Math.round(s || 0); return s >= 3600 ? `${Math.floor(s / 3600)} h ${Math.floor((s % 3600) / 60)} min` : s >= 60 ? `${Math.floor(s / 60)} min ${s % 60} s` : `${s} s`; }; +export { fmt, clock }; + +export function trainStep({ app, session, go, ui }) { + const doc = session.doc; + if (!doc) return c.emptyState.default({ title: "No set open", hint: "Pick or make a training set in the Data step.", action: c.button.sm({ label: "Back to data", tone: "primary", onClick: () => go("data") }) }); + const v = settingsOf(doc); + const set = (key) => (value) => { session.edit((d) => { d.settings = { ...(d.settings || {}), [key]: value }; }); }; + const num = (key, label, hint, o = {}) => c.field.default({ label, hint, control: c.number.md({ label, value: v[key], onChange: set(key), ...o }) }); + const pickOf = (key, label, hint, options) => c.field.default({ label, hint, control: c.select.md({ label, options, value: v[key], onChange: set(key) }) }); + const flag = (key, label, hint) => c.toggle.default({ label, hint, checked: v[key], onChange: set(key) }); + + const items = doc.items; + const used = items.filter((i) => (v.use === "normal" ? ["own", "liked"] : v.use === "avoid" ? ["disliked", "bad_composition", "bad_image"] : Object.keys({ own: 1, liked: 1, disliked: 1, bad_composition: 1, bad_image: 1 })).includes(i.label)); + const slots = (app.pipeline && app.pipeline.slots && app.pipeline.slots()) || []; + const hasAudioVae = slots.some((s) => s.id === "audio_vae"); + const caption = used.filter((i) => i.caption.trim()).length; + + const name = c.field.default({ label: "Name of the LoRA", hint: "Saved in your loras folder. An earlier file is never overwritten.", control: c.input.md({ label: "Name", value: v.name, onCommit: set("name") }) }); + const preset = c.segmented.sm({ label: "Preset", options: PRESETS, value: (PRESETS.find((p) => p.steps === v.steps && p.rank === v.rank) || {}).value, onChange: (key) => { + const p = PRESETS.find((x) => x.value === key); + session.edit((d) => { d.settings = { ...(d.settings || {}), steps: p.steps, rank: p.rank }; }); + ui.redraw(); + } }); + + const advanced = c.collapsible.default({ label: "More settings", open: ui.advanced, body: c.region.stack({ gap: "sm", children: [ + pickOf("algorithm", "Adapter", "LoRA is the common one.", opts(["LoRA", "LoHa", "LoKr", "OFT"])), + pickOf("optimizer", "Optimizer", "How the weights are updated.", opts(["AdamW", "Adam", "SGD", "RMSprop"])), + pickOf("loss", "Loss", "How the error is measured.", opts(["MSE", "L1", "Huber", "SmoothL1"])), + num("batch_size", "Batch size", "Samples per step. H3 always uses 1.", { min: 1, max: 64, step: 1 }), + num("grad_accumulation", "Gradient accumulation", "Steps added up before the weights change: a bigger batch for less memory.", { min: 1, max: 256, step: 1 }), + pickOf("training_dtype", "Training precision", "bf16 is right for nearly everything.", opts(["bf16", "fp32", "none"])), + pickOf("lora_dtype", "LoRA precision", "How the file is stored.", opts(["bf16", "fp32"])), + num("seed", "Seed", "The same seed repeats the same run.", { min: 0, step: 1 }), + flag("gradient_checkpointing", "Save memory (gradient checkpointing)", "Slower, but needs far less memory."), + flag("offloading", "Keep weights in system memory", "For cards that cannot hold the model. Slower."), + flag("bypass", "Bypass mode", "Applies the LoRA by hooks, for quantized models."), + ] }) }); + + const run = ui.run; + const live = c.region.stack({ gap: "md", children: [monitor(ui)] }); // redrawn on its own while a run goes on: the settings above stay as they are + ui.monitorBox = live; + const busy = run.active; + const start = c.button.lg({ label: busy ? "Training…" : "Start training", tone: "primary", disabled: busy || !used.length, onClick: () => ui.start() }); + const stop = c.button.lg({ label: "Stop", tone: "danger", onClick: () => ui.stop() }); + stop.node.hidden = !busy; + + return c.region.stack({ gap: "md", children: [ + c.hint.default({ text: used.length ? `${used.length} item${used.length > 1 ? "s" : ""} will be learned from (${caption} with tags). ${used.length - caption ? "Items without tags are learned from with only the trigger word." : ""}`.trim() + : `Nothing in "${session.name}" fits "${(USES.find((u) => u.value === v.use) || {}).label}". Add media or change the choice below.` }), + hasAudioVae ? c.banner.info({ text: "This pipeline has an audio VAE, so it trains picture and sound together (H3). One clip at a time, so training is slow and needs a lot of memory." }) : null, + c.group.default({ label: "Basics", rows: [name, preset, + pickOf("use", "What to learn", "Liked results and your own media teach what you want. Disliked results teach what to avoid: apply that LoRA at a negative strength.", USES), + num("steps", "Steps", "More steps learn more and can overdo it.", { min: 1, max: 100000, step: 10 }), + num("learning_rate", "Learning rate", "How big each step is. Lower is safer and slower.", { min: 0.0000001, max: 1, step: 0.0001, precision: 7 }), + num("rank", "Rank", "How much it can learn. Higher holds more detail and makes a bigger file.", { min: 1, max: 128, step: 1 })] }), + c.group.default({ label: "Pictures and clips", rows: [ + num("resolution", "Size", "Pixels per side. Everything is fitted to about this area.", { min: 64, max: 2048, step: 32 }), + flag("aspect_buckets", "Keep each shape", "A few size groups instead of one shape for all. Not used for H3."), + num("frames", "Frames per clip", "For video models. H3 rounds to its own grid (17 x n + 5).", { min: 1, max: 257, step: 1 }), + num("fps", "Frames per second", "0 uses each clip's own rate (H3: 24).", { min: 0, max: 120, step: 0.5 }), + pickOf("start", "Take frames from", "", [{ value: "start", label: "The start of the clip" }, { value: "middle", label: "The middle" }])] }), + advanced, + c.toolbar.default({ items: [c.button.md({ label: "Back", tone: "ghost", onClick: () => go("tags") })], trailing: [stop, start] }), + live, + ].filter(Boolean) }); +} + +/** Progress, the message, and the loss graph. */ +export function monitor(ui) { + const run = ui.run; + if (!run.shown) return null; + const chart = c.chart.line({ series: [{ name: "loss", values: run.loss, tone: "accent" }, ...(run.smooth.length > 1 ? [{ name: "smoothed", values: run.smooth, tone: "good" }] : [])], + xLabel: "step", yLabel: "loss", height: "md", empty: "The first steps are on their way" }); + const done = run.step || 0, total = run.total || 1; + const lines = [ + run.state === "training" || run.state === "queued" ? `Step ${done} of ${total}${run.elapsed ? ` · ${clock(run.elapsed)}` : ""}${done && run.elapsed ? ` · about ${clock((run.elapsed / done) * (total - done))} left` : ""}` : null, + run.message || null, + ].filter(Boolean); + return c.group.default({ label: run.state === "done" ? "Finished" : run.state === "error" ? "Stopped with an error" : run.state === "stopped" ? "Stopped" : "Training", rows: [ + run.state === "training" ? c.progress.bar({ value: done, max: total, label: "Training progress" }) : run.state === "queued" ? c.progress.indeterminate({ label: "Waiting for ComfyUI" }) : null, + lines.length ? c.hint.default({ text: lines.join(" · ") }) : null, + run.error ? c.inlineError.default({ text: run.error }) : null, + run.lora ? c.banner.info({ text: `Saved as ${run.lora}. Pick it in a LoRA loader to try it.` }) : null, + chart, + run.loss.length ? c.hint.default({ text: `Latest loss ${fmt(run.loss[run.loss.length - 1])}. It should fall at first, then wobble; a LoRA is rarely better after the wobble stops shrinking.` }) : null, + ].filter(Boolean) }); +} diff --git a/app/ui/features/training/training.test.js b/app/ui/features/training/training.test.js new file mode 100644 index 00000000..c80da2a8 --- /dev/null +++ b/app/ui/features/training/training.test.js @@ -0,0 +1,155 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { setupDom, teardownDom } from "../../composer/tests/_dom.js"; +import { bulk, frequency, joinTags, splitTags } from "./captions.js"; + +let feature; +test.before(async () => { setupDom(); feature = (await import("./index.js")).default; }); +test.after(() => teardownDom()); + +const wait = (ms = 20) => new Promise((r) => setTimeout(r, ms)); + +function fake({ enabled = true } = {}) { + const calls = []; + const doc = { name: "cats", trigger: "", settings: {}, items: [ + { id: "aaaaaaaaaaaa", src: { kind: "bin", id: "bbbbbbbbbbbb" }, media_kind: "image", label: "own", caption: "cat, grey", name: "a.png" }, + { id: "cccccccccccc", src: { kind: "output", type: "output", subfolder: "", filename: "v.mp4" }, media_kind: "video", label: "liked", caption: "", name: "v.mp4" }, + ] }; + const run = { run: "", state: "idle", loss: [], loss_from: 0, step: 0, total: 0 }; + const api = { + trainState: async () => ({ enabled }), + trainSets: async () => ({ sets: [{ name: "cats", count: 2 }] }), + trainSet: async () => structuredClone(doc), + trainSaveSet: async (name, d) => { calls.push(["save", name, d]); return d; }, + trainAdd: async () => ({ set: doc, added: 0 }), + trainRated: async () => ({ results: [], counts: { liked: 3 } }), + trainTaggers: async () => ({ taggers: [{ id: "wd_vit_v3", title: "WD ViT v3", state: "absent", hint: "", done: 0, total: 0 }] }), + trainAutotagState: async () => ({ state: "idle" }), + trainRun: async () => ({ ...run }), + trainRuns: async () => [], + trainPlan: async (set, slots, settings) => { calls.push(["plan", set, settings]); return { prompt: { 1: { class_type: "FunPackTrainLora", inputs: {} } }, problems: [] }; }, + queueGraph: async (g) => { calls.push(["queue", g]); return "pid"; }, + trainStop: async () => { calls.push(["stop"]); return { stopped: true }; }, + media: async () => ({ media: [{ id: "bbbbbbbbbbbb", name: "a.png", kind: "image" }, { id: "dddddddddddd", name: "s.wav", kind: "audio" }] }), + uploadMedia: async () => ({ media: [], problems: [] }), + }; + const app = { api, dom: {}, pipeline: { slots: () => [{ id: "model" }] }, say: () => {}, on: () => () => {} }; + return { api, app, calls, doc, run }; +} + +const open = async (f) => { + const host = document.createElement("div"); + document.body.append(host); + const off = feature.setup({ host, app: f.app }); + await wait(); + const button = [...host.querySelectorAll("button")].find((b) => b.textContent.includes("Training")); + return { host, button, off }; +}; +const text = () => document.body.textContent; +const click = (label, root = document) => [...root.querySelectorAll("button")].find((b) => b.textContent.trim() === label); + +test("captions: tags are trimmed and unique, counts are per caption, bulk edits only the chosen", () => { + assert.deepEqual(splitTags(" a,B ,, a\nc "), ["a", "B", "c"]); + assert.equal(joinTags(["x", "X", " y "]), "x, y"); + assert.deepEqual(frequency([{ caption: "a, b" }, { caption: "a, a" }, { caption: "c" }]).map((x) => [x.tag, x.count]), [["a", 2], ["b", 1], ["c", 1]]); + const items = [{ id: "x", caption: "cat, dog" }, { id: "y", caption: "cat" }]; + assert.equal(bulk(items, "remove", "cat", "", new Set(["x"])), 1); + assert.deepEqual(items.map((i) => i.caption), ["dog", "cat"]); + assert.equal(bulk(items, "replace", "cat", "lion"), 1); + assert.equal(bulk(items, "prepend", "zq"), 2); + assert.equal(items[1].caption, "zq, lion"); + assert.equal(bulk(items, "add", ""), 0); +}); + +test("the window is hidden while the add-on is off", async () => { + const f = fake({ enabled: false }); + const { button, off } = await open(f); + assert.ok(button.hidden); + off(); +}); + +test("the Data step lists the sets, offers the three sources and can be skipped", async () => { + const f = fake(); + const { button, off } = await open(f); + assert.ok(!button.hidden); + button.click(); + await wait(60); + const t = text(); + for (const word of ["Drop pictures or clips here", "Choose in your Bin", "Liked results", "Bad composition results", "Bad image results", "Skip to training", "Next: tags"]) assert.ok(t.includes(word), word); + assert.ok(t.includes("2 in \"cats\""), "the open set is summarised"); + assert.ok(document.querySelectorAll(".cx-cell").length === 1, "only pictures and clips are offered from the Bin, not sound"); + off(); +}); + +test("Tags step: a card per item; editing a caption saves the set", async () => { + const f = fake(); + const { button, off } = await open(f); + button.click(); + await wait(60); + click("Next: tags").click(); + await wait(20); + assert.equal(document.querySelectorAll(".cx-captioned").length, 2); + document.querySelector(".cx-captioned .cx-tag-x").click(); // takes "cat" off the first item + await wait(700); + const saved = f.calls.filter((c) => c[0] === "save").pop(); + assert.equal(saved[2].items[0].caption, "grey"); + off(); +}); + +test("Train step: Start sends the settings and the pipeline to the plan and queues the graph", async () => { + const f = fake(); + const { button, off } = await open(f); + button.click(); + await wait(60); + click("Skip to training").click(); + await wait(20); + assert.ok(text().includes("Start training")); + click("Start training").click(); + await wait(40); + const plan = f.calls.find((c) => c[0] === "plan"); + assert.equal(plan[1], "cats"); + assert.equal(plan[2].steps, 500); + assert.ok(f.calls.some((c) => c[0] === "queue")); + off(); +}); + +test("a refused plan shows the reason and starts nothing", async () => { + const f = fake(); + f.api.trainPlan = async () => ({ prompt: {}, problems: ["the pipeline has no clip loader"] }); + const { button, off } = await open(f); + button.click(); + await wait(60); + click("Skip to training").click(); + await wait(20); + click("Start training").click(); + await wait(40); + assert.ok(text().includes("no clip loader")); + assert.ok(!f.calls.some((c) => c[0] === "queue")); + off(); +}); + +test("a run in progress draws its loss, then reports the saved file; Stop asks the server", async () => { + const f = fake(); + let n = 0; + f.api.trainRun = async (since) => { + n += 1; + if (n < 3) return { run: "", state: "idle", loss: [], loss_from: 0 }; + if (n === 3) return { run: "r1", state: "training", step: 2, total: 5, elapsed: 4, message: "", loss: [0.9, 0.7], loss_from: since }; + return { run: "r1", state: "done", step: 5, total: 5, elapsed: 9, message: "", loss: [0.5, 0.4, 0.3].slice(0, 5 - since), loss_from: since, lora: "cats_lora.safetensors" }; + }; + const { button, off } = await open(f); + button.click(); + await wait(60); + click("Skip to training").click(); + await wait(20); + click("Start training").click(); + await wait(1200); + assert.ok(text().includes("Step 2 of 5"), "progress is shown"); + assert.ok(document.querySelector(".cx-chart path.cx-chart-line"), "the loss is drawn"); + click("Stop").click(); + await wait(20); + assert.ok(f.calls.some((c) => c[0] === "stop")); + await wait(1200); + assert.ok(text().includes("cats_lora.safetensors")); + off(); +}); diff --git a/app/ui/modules.js b/app/ui/modules.js index add96b57..56b8389e 100644 --- a/app/ui/modules.js +++ b/app/ui/modules.js @@ -44,4 +44,5 @@ export default [ "./features/welcome/index.js", "./features/addons/index.js", "./features/temp/index.js", + "./features/training/index.js", ].map((path) => new URL(path, import.meta.url).href); diff --git a/app/ui/shell/api.js b/app/ui/shell/api.js index c87e03bb..9ecc4b71 100644 --- a/app/ui/shell/api.js +++ b/app/ui/shell/api.js @@ -1,5 +1,6 @@ // The server's routes, one line each. An error carries the sentence the server gave, not a bare status. const BASE = "/funpack"; +const TRAIN = "/api/m/training_lora"; function why(res, payload) { const p = payload && typeof payload === "object" ? payload : {}; @@ -77,6 +78,38 @@ export const api = { if (!res.ok) throw new Error((payload && payload.why) || `HTTP ${res.status}`); return payload; }, + // LoRA training (the add-on module): training sets, rated results, taggers, and a run. + trainState: () => call("GET", `${TRAIN}/state`), + trainSets: () => call("GET", `${TRAIN}/sets`), + trainCreateSet: (name) => call("POST", `${TRAIN}/sets`, { name }), + trainSet: (name) => call("GET", `${TRAIN}/sets/${encodeURIComponent(name)}`), + trainSaveSet: (name, doc) => call("PUT", `${TRAIN}/sets/${encodeURIComponent(name)}`, doc), + trainDeleteSet: (name) => call("DELETE", `${TRAIN}/sets/${encodeURIComponent(name)}`), + trainAdd: (name, entries) => call("POST", `${TRAIN}/sets/${encodeURIComponent(name)}/add`, { entries }), + trainBulk: (name, body) => call("POST", `${TRAIN}/sets/${encodeURIComponent(name)}/bulk`, body), + trainRated: (label) => call("GET", `${TRAIN}/rated${label ? `?label=${encodeURIComponent(label)}` : ""}`), + trainTaggers: () => call("GET", `${TRAIN}/taggers`), + trainDownloadTagger: (tagger) => call("POST", `${TRAIN}/taggers/download`, { tagger }), + trainRemoveTagger: (tagger) => call("POST", `${TRAIN}/taggers/remove`, { tagger }), + trainAutotag: (body) => call("POST", `${TRAIN}/autotag`, body), + trainAutotagState: () => call("GET", `${TRAIN}/autotag`), + trainAutotagCancel: () => call("POST", `${TRAIN}/autotag/cancel`, {}), + trainPlan: (set, slots, settings) => call("POST", `${TRAIN}/plan`, { set, slots, settings }), + trainRun: (since) => call("GET", `${TRAIN}/run?since=${since || 0}`), + trainStop: () => call("POST", `${TRAIN}/run/stop`, {}), + trainRuns: () => call("GET", `${TRAIN}/runs`).then((d) => d.runs || []), + trainRunRecord: (id) => call("GET", `${TRAIN}/runs/${encodeURIComponent(id)}`), + trainForgetRun: (id) => call("DELETE", `${TRAIN}/runs/${encodeURIComponent(id)}`), + /** Any graph as a ComfyUI job. -> the prompt id. */ + async queueGraph(graph) { + const res = await fetch("/prompt", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ prompt: graph }) }); + const body = await res.json().catch(() => ({})); + if (!res.ok || !body.prompt_id) { + const nodes = Object.values(body.node_errors || {}).flatMap((n) => (n.errors || []).map((x) => x.message)); + throw new Error([(body.error || {}).message, ...nodes].filter(Boolean).join("; ") || `ComfyUI refused the job (HTTP ${res.status})`); + } + return body.prompt_id; + }, upscaleModels: () => call("GET", "/api/m/render/upscale_models"), /** The upscale is an ordinary ComfyUI job: queued on /prompt, read back from /history. -> the prompt id. */ async queueUpscale(media, model, rtx) { diff --git a/docs/module-reference.md b/docs/module-reference.md index feaeb9b9..031f437b 100644 --- a/docs/module-reference.md +++ b/docs/module-reference.md @@ -148,6 +148,7 @@ A capability is a name in `PROVIDES`. Core defines none: the contract is between | `anchor_pin` | `(kwargs, transform) -> kwargs` or `None` | | | `rescale_conditioning` | `(conditioning, height, width) -> (conditioning, count)` or `None` | | | `latent_upscaler` | `(state_dict) -> object with funpack_latent_upscale(x, scale)` or `None` | | +| `training_adapt` | `(model) -> {name, step, video, batch_limit, fps, frames, pack, shims, silence}` or `None` | what LoRA training must know about a model whose latent is not a plain tensor (H3 packs picture and sound). `shims()` is a context manager that applies the repairs only for the length of a run | ### Learning @@ -158,6 +159,25 @@ A capability is a name in `PROVIDES`. Core defines none: the contract is between | `taste_latest_key` | `taste` | `() -> the key the most recent capture went to, or None` | | `on_rating` | `taste` calls every module that provides it | `(prompt_id, rating, axis) -> None` | +## An add-on module: LoRA training + +`modules/training/lora/` is announced as `training_lora` and switched on in Settings > Modules (the `lora_training` add-on, which asks for `onnxruntime`, needed only by the automatic tagger). It is a node-and-routes module with no settings of its own: the window passes every value to the node. + +| File | What it holds | +|---|---| +| `store.py` | Training sets as JSON under `training/sets/`. An item points at a Bin file or an output file; the path is resolved at run time, inside its root. | +| `rated.py` | Finds rated results in the saved projects and maps a rating to `liked`, `disliked`, `bad_composition` or `bad_image`. | +| `captions.py` | Tag text helpers, shared in spirit with `app/ui/features/training/captions.js`. | +| `taggers.py` | WD tagger models (ONNX) downloaded from Hugging Face on request, or any repo with `model.onnx` and `selected_tags.csv`; one tagging job at a time. | +| `media_io.py`, `frames.py` | Pictures and clips as tensors, fitted to a size bucket. | +| `plan.py` | Builds the graph a run is queued as: the pipeline's own loaders (ids `model`, `clip`, `vae`, `audio_vae`, or the first of each kind) feed `FunPackTrainLora`. | +| `train.py`, `nodes.py` | The run, on ComfyUI's own Train LoRA node. Latents and conditioning are built from the set; each step's loss goes to `jobs.py`; Stop and ComfyUI's interrupt end it at the next step. | +| `jobs.py` | The run in progress and finished run records under `training/runs/`. | + +"Learn what I like" trains on liked results and your own media; "learn what to avoid" trains on the disliked ones, for a LoRA applied at negative strength. This is not VADER (reward gradients), which stays parked. + +The diffusion model loader sets `cached_patcher_init` for a plain file, because the trainer clones the model from it. An H3 run uses the repairs in `modules/models/minimax_h3/training.py`. **UNTESTED on real weights or hardware**: proven only on a tiny H3 and a stand-in text encoder and VAEs, on CPU. + ## Helpers a module may import from `_core` `from ..._core import ...` offers exactly these: `addons`, `chain`, `config`, `contract`, `control`, `dit_hooks`, `graph`, `input_steer`, `log`, `media`, `patching`, `projects`, `rated_dial`, `registry`, `relations`, `run`, `schema`, `shortcuts`, `streams`, `temp_files`, `traits`, `vae_tiles`. Adding one means editing `modules/_core.py` in both of its branches. diff --git a/docs/testing.md b/docs/testing.md index 3e6c627e..d4a5b596 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -38,6 +38,7 @@ These run with the rest and fail the build if you break them. | `core/tests/test_loads_as_comfyui_does.py` | The pack loads the way ComfyUI loads it. | | `core/tests/test_log_call_sites.py` | Log calls use the right level and shape. | | `core/tests/test_control.py` | Quarantine and `Unavailable` behave as described in [Architecture](architecture.md). | +| `modules/training/lora/tests/test_train_h3.py` | A whole LoRA training run on a tiny real H3 on CPU, with a stand-in text encoder and VAEs. It proves our side (sets, latents, repairs, loss stream, saved file), nothing about real weights or hardware. | ### Environment-dependent failures