diff --git a/CHANGELOG.md b/CHANGELOG.md index ac88333a..3218144f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,7 +24,7 @@ A rewrite from scratch, not an update of 4.x. Same idea (a studio for video insi - **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. -- **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. +- **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, Pause and Resume, and previews: every N steps training waits while your own pipeline makes a picture or clip with the LoRA so far, then carries on. Results: past runs. It trains everything in the set. H3 (picture and sound together) uses the module's own training loop with ComfyUI carrying the batches; other models use ComfyUI's trainer, with the loaders in your pipeline. **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/generate/index.js b/app/ui/features/generate/index.js index 9f482ec1..374bfb5c 100644 --- a/app/ui/features/generate/index.js +++ b/app/ui/features/generate/index.js @@ -138,6 +138,36 @@ export default { } finally { busy = false; draw(); } } app.runner.units = (units, times) => runUnits(units, times); + /** 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 (busy || held || g.state().phase === QUEUED || g.state().phase === RUNNING) return { error: "A run is already going: the preview waits for it." }; + 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 === "FunPackLoraLoader" && !s.bypassed) || frozen.find((s) => s.node === "FunPackLoraLoader"); + 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 } = await buildInputs({ project: snap, scene: root, slots: frozen, expand, frames, hooks: app.inputHooks, prefix: (app.promptPrefix || []).flatMap((f) => { try { return f(root, snap); } catch { return []; } }) }); + 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))]; diff --git a/app/ui/features/training/index.js b/app/ui/features/training/index.js index 255abb1c..973859c8 100644 --- a/app/ui/features/training/index.js +++ b/app/ui/features/training/index.js @@ -7,6 +7,8 @@ import { monitor, 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", label: "1 Data" }, { value: "tags", label: "2 Tags" }, { value: "train", label: "3 Train" }, { value: "results", label: "Results" }]; const POLL = 1000; const QUEUE_WAIT = 60000; @@ -21,7 +23,7 @@ const smooth = (values) => { // a gentle moving average, s export default { id: "training", mount: "timeline.status", - needs: ["api", "dom", "pipeline"], + needs: ["api", "dom", "pipeline", "project", "generate"], setup({ host, app }) { const api = app.api; const session = makeSession(api); @@ -29,7 +31,9 @@ export default { 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: [], smooth: [], step: 0, total: 0, elapsed: 0, message: "", error: "", lora: "", before: "" }, + run: { shown: false, active: false, state: "idle", loss: [], 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), 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. */ @@ -103,25 +107,31 @@ export default { // --- a run const setRun = (patch) => Object.assign(ui.run, patch); - async function start() { + const start = () => launch(""); + async function launch(resumeId) { const doc = session.doc; if (!doc) return; await session.flush(); const settings = settingsOf(doc); settings.trigger = doc.trigger || ""; settings.postfix = doc.postfix || ""; - setRun({ shown: true, active: true, state: "queued", loss: [], smooth: [], step: 0, total: settings.steps, elapsed: 0, message: "Checking the pipeline", error: "", lora: "" }); + 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: [], smooth: [], step: 0, total: settings.steps, elapsed: 0, message: "Checking the pipeline", error: "", lora: "", runId: "", paused: null }); draw(); let before = ""; - try { before = (await api.trainRun(0)).run || ""; } catch { /* first run */ } + if (!resumeId) { 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." }); + 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: "error", error: err.message, message: "" }); draw(); return; } + } 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 watch() { @@ -135,9 +145,16 @@ export default { } let snap = null; try { snap = await api.trainRun(ui.run.loss.length); } catch { /* try again */ } - if (snap && snap.run && snap.run !== ui.before) { + if (snap && snap.run && snap.run !== ui.before && !(ui.awaitStart && snap.state === "paused")) { + ui.awaitStart = false; 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 || "" }); + 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 || "", 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}` }); @@ -152,16 +169,46 @@ export default { runPoll = setTimeout(tick, POLL); } async function stop() { + const waiting = ui.run.state === "paused" || ui.run.state === "previewing"; setRun({ message: "Stopping" }); + if (ui.run.state === "previewing") { ui.stopAfterPreview = true; try { app.generate.cancel(); } catch { /* it was over */ } } 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(); } + 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 */ } } + if (ui.tab === "train" && ui.monitorBox) ui.monitorBox.set([monitor(ui)]); + } + /** A picture or clip made by the pipeline with the LoRA so far; when `auto`, training carries on afterwards. */ + async function previewNow(auto) { + const paused = ui.run.paused; + if (!paused || !paused.preview) return; + 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 }) : { 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 === "training") { + if (snap.state === "paused") { + ui.before = ""; + setRun({ shown: true, active: false, state: "paused", loss: snap.loss, smooth: smooth(snap.loss), step: snap.step, total: snap.total, message: "", runId: snap.run, paused: snap.paused }); + ui.holdPause = true; // found after a reload: the person decides + ui.tab = "train"; + } else 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"; diff --git a/app/ui/features/training/train.js b/app/ui/features/training/train.js index 8dd96a6b..3f12b75b 100644 --- a/app/ui/features/training/train.js +++ b/app/ui/features/training/train.js @@ -1,9 +1,10 @@ // Step 3, Train: the settings, Start and Stop, and the loss graph while it runs. import { composer as c } from "../../composer/composer.js"; +import { stillOf } from "./state.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.0001, rank: 16, algorithm: "LoRA", optimizer: "AdamW", loss: "MSE", + name: "funpack_lora", use: "all", preview_every: 0, steps: 500, learning_rate: 0.0001, 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, swap_blocks: 0, }; @@ -12,15 +13,10 @@ export const PRESETS = [ { 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) || {}) }); +export const settingsOf = (doc) => ({ ...DEFAULTS, ...((doc && doc.settings) || {}), use: "all" }); // the Data step decides what is in; a set made before that remembers an old choice 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`; }; @@ -36,7 +32,7 @@ export function trainStep({ app, session, go, ui }) { 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 used = items; 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; @@ -66,20 +62,28 @@ export function trainStep({ app, session, go, ui }) { 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 waiting = run.state === "paused" || run.state === "previewing"; // between steps: nothing is training, the run is still open + const busy = run.active || run.state === "previewing"; + const start = c.button.lg({ label: busy ? "Training…" : "Start training", tone: "primary", disabled: busy || waiting || !used.length, onClick: () => ui.start() }); + start.node.hidden = waiting; const stop = c.button.lg({ label: "Stop", tone: "danger", onClick: () => ui.stop() }); - stop.node.hidden = !busy; + stop.node.hidden = !busy && !waiting; + const pause = c.button.lg({ label: run.state === "previewing" ? "Stay paused" : "Pause", tone: "neutral", disabled: ui.holdPause && run.state !== "paused", onClick: () => ui.pause() }); + pause.node.hidden = !(run.state === "training" || run.state === "previewing"); + const resume = c.button.lg({ label: "Resume", tone: "primary", onClick: () => ui.resume() }); + resume.node.hidden = run.state !== "paused"; + const previewNow = c.button.lg({ label: "Preview now", tone: "neutral", title: "Make a picture or clip with the LoRA so far", onClick: () => ui.previewNow() }); + previewNow.node.hidden = !(run.state === "paused" && run.paused && run.paused.preview); 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 a caption are learned from with only the prefix and postfix." : ""}`.trim() - : `Nothing in "${session.name}" fits "${(USES.find((u) => u.value === v.use) || {}).label}". Add media or change the choice below.` }), + : `"${session.name}" is empty. Add media in the Data step.` }), 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. musubi-tuner's H3 recipe: learning rate 0.0001, rank 16, batch 1, bf16, gradient checkpointing, plus block swap. Its plain flow-matching loss undoes the few-step quality of a distilled (turbo) model; musubi adds a training adapter, guidance loss or teacher matching for that, which this trainer does not have." }) : 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 })] }), + num("rank", "Rank", "How much it can learn. Higher holds more detail and makes a bigger file.", { min: 1, max: 128, step: 1 }), + num("preview_every", "Preview every (steps)", "0 never. Otherwise training waits every this many steps while your pipeline makes a picture or clip with the LoRA so far (the picked scene, the project's own settings), then carries on. Needs a LoRA loader in the pipeline; each wait loads the models again.", { min: 0, max: 100000, step: 10 })] }), 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."), @@ -87,7 +91,7 @@ export function trainStep({ app, session, go, ui }) { 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] }), + c.toolbar.default({ items: [c.button.md({ label: "Back", tone: "ghost", onClick: () => go("tags") })], trailing: [stop, pause, previewNow, resume, start] }), live, ].filter(Boolean) }); } @@ -103,11 +107,17 @@ export function monitor(ui) { 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: [ + const previews = ui.previews && ui.previews.length ? c.gallery.adaptive({ id: "training-previews", cols: 3, selection: [], + items: ui.previews.map((p) => ({ id: p.id, label: `Step ${p.step}`, thumb: stillOf(p), badge: p.media_kind })), + onActivate: (item) => ui.peek(ui.previews.find((p) => p.id === item.id)) }) : null; + const heading = { done: "Finished", error: "Stopped with an error", stopped: "Stopped", paused: `Paused at step ${done}`, previewing: `Making a preview at step ${done}` }[run.state] || "Training"; + return c.group.default({ label: heading, rows: [ + run.state === "previewing" ? c.progress.indeterminate({ label: "Making a preview" }) : null, 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, + previews, 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 index d5cf6c00..368cee7a 100644 --- a/app/ui/features/training/training.test.js +++ b/app/ui/features/training/training.test.js @@ -285,3 +285,63 @@ test("Tags step: a clip's frame count and rate show beside its label", async () assert.ok(text().includes("124 f · 24 fps")); off(); }); + +test("a preview step pauses the run, the pipeline makes a preview with the LoRA so far, and training carries on", async () => { + const f = fake(); + const previews = []; + f.app.runner = { preview: async (o) => { previews.push(o); return { image: { filename: "pv.mp4", subfolder: "", type: "output" } }; } }; + let queued = 0, polls = 0; + f.api.queueGraph = async (g) => { queued += 1; f.calls.push(["queue", g]); return "pid"; }; + f.api.trainRun = async () => { + polls += 1; + if (queued === 0) return { run: "", state: "idle", loss: [], loss_from: 0 }; + if (queued === 1) return polls % 2 ? { run: "r1", state: "training", step: 2, total: 6, elapsed: 3, message: "", loss: [0.9, 0.8], loss_from: 0 } + : { run: "r1", state: "paused", step: 2, total: 6, elapsed: 3, message: "", loss: [], loss_from: 2, paused: { reason: "preview", step: 2, preview: "cats_preview.safetensors", checkpoint: true } }; + if (polls % 3 === 0) return { run: "r1", state: "paused", step: 2, total: 6, message: "", loss: [], loss_from: 2, paused: { reason: "preview", step: 2, preview: "cats_preview.safetensors", checkpoint: true } }; // the record still shows the pause until the new run begins + return { run: "r1", state: "done", step: 6, total: 6, elapsed: 9, message: "", loss: [0.7, 0.6, 0.5, 0.4], loss_from: 2, 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(3600); + assert.equal(previews.length, 1); + assert.equal(previews[0].lora, "cats_preview.safetensors"); + assert.equal(queued, 2, "training was queued again to carry on"); + const plans = f.calls.filter((c) => c[0] === "plan"); + assert.equal(plans[1][2].resume, "r1"); + assert.ok(text().includes("Step 2"), "the preview is shown with its step"); + await wait(3000); + assert.ok(text().includes("cats_lora.safetensors")); + off(); +}); + +test("Pause asks the server, a paused run offers Resume, Preview now and Stop, and Stop ends it", async () => { + const f = fake(); + f.app.runner = { preview: async () => ({ error: "x" }) }; + let phase = 0; + f.api.trainPause = async () => { f.calls.push(["pause"]); phase = 2; return { paused: true }; }; + f.api.queueGraph = async () => { phase = 1; return "pid"; }; + f.api.trainRun = async () => phase === 0 ? { run: "", state: "idle", loss: [], loss_from: 0 } + : phase === 1 ? { run: "r2", state: "training", step: 1, total: 6, elapsed: 1, message: "", loss: [0.9], loss_from: 0 } + : { run: "r2", state: "paused", step: 2, total: 6, message: "", loss: [0.8], loss_from: 1, paused: { reason: "user", step: 2, preview: "cats_preview.safetensors", checkpoint: true } }; + 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(1400); + click("Pause").click(); + await wait(1400); + assert.ok(f.calls.some((c) => c[0] === "pause")); + assert.ok(text().includes("Paused at step 2")); + assert.ok(click("Resume") && !click("Resume").hidden && !click("Preview now").hidden); + click("Stop").click(); + await wait(40); + assert.ok(f.calls.some((c) => c[0] === "stop")); + assert.ok(text().includes("Stopped")); + off(); +}); diff --git a/app/ui/shell/api.js b/app/ui/shell/api.js index ea4f648a..db585a44 100644 --- a/app/ui/shell/api.js +++ b/app/ui/shell/api.js @@ -97,6 +97,7 @@ export const api = { 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}`), + trainPause: () => call("POST", `${TRAIN}/run/pause`, {}), trainStop: () => call("POST", `${TRAIN}/run/stop`, {}), trainRuns: () => call("GET", `${TRAIN}/runs`).then((d) => d.runs || []), trainRunRecord: (id) => call("GET", `${TRAIN}/runs/${encodeURIComponent(id)}`),