From 3c421b2c1ddbc3921ed5b76bc72ccdf74898a97d Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 17 Oct 2026 20:54:49 +0000 Subject: [PATCH] LoRA training: fix resume in inference mode, Restore and continue, live graph in Results The saved optimizer state was read in ComfyUI's inference mode, so the first optimizer step after a resume failed. It is copied to normal tensors first. A failed or stopped run keeps its saved state and Results offers Restore and continue. The graph, progress and previews of a running training moved from Train to Results; Train says where to look. Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P --- app/ui/features/training/index.js | 24 ++++++++++++----- app/ui/features/training/results.js | 10 +++++--- app/ui/features/training/train.js | 14 ++++++++-- app/ui/features/training/training.test.js | 13 ++++++++-- modules/training/lora/__init__.py | 12 ++++++--- modules/training/lora/loop.py | 17 ++++++++++++ modules/training/lora/tests/test_train_h3.py | 27 ++++++++++++++++++++ 7 files changed, 99 insertions(+), 18 deletions(-) diff --git a/app/ui/features/training/index.js b/app/ui/features/training/index.js index 973859c8..98af78b9 100644 --- a/app/ui/features/training/index.js +++ b/app/ui/features/training/index.js @@ -3,7 +3,7 @@ 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 { monitor, notice, settingsOf, trainStep } from "./train.js"; import { resultsStep } from "./results.js"; import { makeSession, srcKey } from "./state.js"; @@ -33,7 +33,7 @@ export default { 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: "", runId: "", paused: null }, previews: [], holdPause: false, awaitStart: false, peek: (item) => peek(item), - pause: () => pause(), resume: () => launch(ui.run.runId), previewNow: () => previewNow(), + pause: () => pause(), resume: () => launch(ui.run.runId), restore: (id) => { ui.run.loss = []; 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. */ @@ -122,7 +122,12 @@ export default { : { 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 = ""; - if (!resumeId) { try { before = (await api.trainRun(0)).run || ""; } catch { /* first run */ } } + 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("; ")); @@ -134,6 +139,11 @@ export default { 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)); + } function watch() { clearTimeout(runPoll); const queuedAt = Date.now(); @@ -141,11 +151,11 @@ export default { if (!ui.run.active) 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." }); - if (ui.tab === "train" && ui.monitorBox) ui.monitorBox.set([monitor(ui)]); + refreshLive(); } let snap = null; try { snap = await api.trainRun(ui.run.loss.length); } catch { /* try again */ } - if (snap && snap.run && snap.run !== ui.before && !(ui.awaitStart && snap.state === "paused")) { + if (snap && snap.run && snap.run !== ui.before && !(ui.awaitStart && (snap.state === "paused" || (ui.staleStarted && snap.started === ui.staleStarted)))) { 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 || "", runId: snap.run, paused: snap.paused || null }); @@ -162,7 +172,7 @@ export default { loadRuns(); return; } - if (ui.tab === "train" && ui.monitorBox) ui.monitorBox.set([monitor(ui)]); + refreshLive(); } runPoll = setTimeout(tick, POLL); }; @@ -180,7 +190,7 @@ export default { 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)]); + refreshLive(); } /** A picture or clip made by the pipeline with the LoRA so far; when `auto`, training carries on afterwards. */ async function previewNow(auto) { diff --git a/app/ui/features/training/results.js b/app/ui/features/training/results.js index c22cd13f..0eb832c2 100644 --- a/app/ui/features/training/results.js +++ b/app/ui/features/training/results.js @@ -1,12 +1,14 @@ // 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"; +import { clock, fmt, monitor } 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 live = c.region.stack({ gap: "md", children: [monitor(ui)].filter(Boolean) }); // the run going on now, redrawn on its own + ui.monitorBox = live; + if (!runs.length) return ui.run.shown ? live : 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(); } }); @@ -19,13 +21,15 @@ export function resultsStep({ app, ui }) { 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: [ + live, 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, + rec && rec.resumable && ["error", "stopped"].includes(rec.state) ? c.banner.info({ text: `The state at step ${rec.steps || 0} or earlier is kept. Restore and continue picks the training up from it with the Train step's current settings.` }) : 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(); } })] }), + c.toolbar.default({ items: rec && rec.resumable && !ui.run.active ? [c.button.sm({ label: "Restore and continue", tone: "primary", title: "Carry on from the saved state", onClick: async () => { await ui.restore(chosen.run); } })] : [], 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/train.js b/app/ui/features/training/train.js index 670ef154..12ac692d 100644 --- a/app/ui/features/training/train.js +++ b/app/ui/features/training/train.js @@ -1,4 +1,4 @@ -// Step 3, Train: the settings, Start and Stop, and the loss graph while it runs. +// Step 3, Train: the settings, Start and Stop. The graph, previews and progress of a run are in the Results step. import { composer as c } from "../../composer/composer.js"; import { stillOf } from "./state.js"; @@ -60,7 +60,7 @@ 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 + const live = c.region.stack({ gap: "md", children: [notice(ui, go)].filter(Boolean) }); // redrawn on its own while a run goes on: the settings above stay as they are ui.monitorBox = live; 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"; @@ -97,6 +97,16 @@ export function trainStep({ app, session, go, ui }) { ].filter(Boolean) }); } +/** What the Train step says about a run: where to look, and the error if it failed. */ +export function notice(ui, go) { + const run = ui.run; + if (!run.shown) return null; + const where = c.button.sm({ label: "Open Results", tone: "ghost", onClick: () => go("results") }); + if (run.error) return c.region.stack({ gap: "sm", children: [c.inlineError.default({ text: run.error }), where] }); + const text = { done: "Training finished.", stopped: "Training stopped.", paused: `Training paused at step ${run.step || 0}.`, previewing: "Making a preview." }[run.state] || "Training started."; + return c.region.stack({ gap: "sm", children: [c.banner.info({ text: `${text} Check graphs and previews in the Results tab.` }), where] }); +} + /** Progress, the message, and the loss graph. */ export function monitor(ui) { const run = ui.run; diff --git a/app/ui/features/training/training.test.js b/app/ui/features/training/training.test.js index 368cee7a..983d7c85 100644 --- a/app/ui/features/training/training.test.js +++ b/app/ui/features/training/training.test.js @@ -47,6 +47,7 @@ const open = async (f) => { return { host, button, off }; }; const text = () => document.body.textContent; +const tab = (label) => [...document.querySelectorAll("button, [role=tab]")].find((b) => b.textContent.trim().endsWith(label)); 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", () => { @@ -199,11 +200,18 @@ test("a run in progress draws its loss, then reports the saved file; Stop asks t await wait(20); click("Start training").click(); await wait(1200); + assert.ok(text().includes("Check graphs and previews in the Results tab"), "Train only says where to look"); + assert.ok(!document.querySelector(".cx-chart"), "the graph is not on the Train tab"); + tab("Results").click(); + await wait(20); 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"); + tab("Train").click(); + await wait(20); click("Stop").click(); await wait(20); assert.ok(f.calls.some((c) => c[0] === "stop")); + tab("Results").click(); await wait(1200); assert.ok(text().includes("cats_lora.safetensors")); off(); @@ -306,6 +314,7 @@ test("a preview step pauses the run, the pipeline makes a preview with the LoRA click("Skip to training").click(); await wait(20); click("Start training").click(); + tab("Results").click(); await wait(3600); assert.equal(previews.length, 1); assert.equal(previews[0].lora, "cats_preview.safetensors"); @@ -337,11 +346,11 @@ test("Pause asks the server, a paused run offers Resume, Preview now and Stop, a 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(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")); + assert.ok(text().includes("Training stopped.")); off(); }); diff --git a/modules/training/lora/__init__.py b/modules/training/lora/__init__.py index f4b4ee76..23e01c92 100644 --- a/modules/training/lora/__init__.py +++ b/modules/training/lora/__init__.py @@ -9,7 +9,7 @@ H3 trains picture and sound together; the repairs its latent and its rotary op n used when this module asks. UNTESTED on real weights or hardware: proven only on tiny stand-in models on CPU. """ -from . import captions, jobs, media_io, plan, rated, store, taggers, train +from . import captions, jobs, loop, media_io, plan, rated, store, taggers, train from .nodes import FunPackTrainLora ID = "training_lora" @@ -226,16 +226,20 @@ def routes(table, base, web): @table.get(base + "/runs") async def _runs(_req): - return web.json_response({"runs": jobs.history()}) + return web.json_response({"runs": [{**r, "resumable": loop.has_checkpoint(r["run"])} for r in jobs.history()]}) @table.get(base + "/runs/{run}") async def _record(req): rec = jobs.record(req.match_info["run"]) - return web.json_response(rec) if rec else _bad(web, "There is no such run.", 404) + return web.json_response({**rec, "resumable": loop.has_checkpoint(rec.get("run"))}) if rec else _bad(web, "There is no such run.", 404) @table.delete(base + "/runs/{run}") async def _forget(req): - return web.json_response({"deleted": jobs.forget(req.match_info["run"])}) + run = req.match_info["run"] + gone = jobs.forget(run) + if gone: + loop.forget_checkpoint(run) + return web.json_response({"deleted": gone}) PROVIDES = {"routes": routes} diff --git a/modules/training/lora/loop.py b/modules/training/lora/loop.py index d37d7f23..d2256653 100644 --- a/modules/training/lora/loop.py +++ b/modules/training/lora/loop.py @@ -47,6 +47,22 @@ def load_checkpoint(run, torch): return torch.load(path, map_location="cpu", weights_only=False) +def has_checkpoint(run): + import re + return bool(re.fullmatch(r"[0-9a-f]{12}", str(run))) and _checkpoint_file(run).is_file() + + +def _normal(value, torch): + """A copy made outside inference mode: ComfyUI runs nodes in it, and what it creates there cannot be updated in place by the optimizer.""" + if isinstance(value, torch.Tensor): + return value.clone() + if isinstance(value, dict): + return {k: _normal(v, torch) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return type(value)(_normal(v, torch) for v in value) + return value + + def forget_checkpoint(run): try: _checkpoint_file(run).unlink() @@ -216,6 +232,7 @@ def train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim, resume= raise Stopped() first = 0 if resume: + resume = _normal(resume, torch) lora.restore(resume["lora"]) optimizer.load_state_dict(resume["optimizer"]) first = int(resume["step"]) diff --git a/modules/training/lora/tests/test_train_h3.py b/modules/training/lora/tests/test_train_h3.py index db6c2df1..e221e180 100644 --- a/modules/training/lora/tests/test_train_h3.py +++ b/modules/training/lora/tests/test_train_h3.py @@ -376,6 +376,33 @@ def test_a_run_pauses_for_a_preview_and_carries_on_to_the_same_end(run_env): assert not (tmp_path / "training" / "checkpoints" / f"{run}.pt").exists() +def test_a_failed_run_is_restored_from_its_saved_state_in_comfyuis_inference_mode(run_env, monkeypatch): + """The optimizer's memory is read back in inference mode there; updating it in place used to fail on the first step.""" + import torch + from modules.training.lora import jobs as j, loop, train + tmp_path, jobs = run_env + with torch.inference_mode(): + train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2)) + run = jobs.snapshot()["run"] + assert loop.has_checkpoint(run) + real = j.step + + def fail(loss): + real(loss) + raise RuntimeError("the card fell over") + monkeypatch.setattr(j, "step", fail) + with torch.inference_mode(), pytest.raises(RuntimeError, match="fell over"): + train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run)) + assert jobs.snapshot()["state"] == "error" and loop.has_checkpoint(run), "a failed run keeps its saved state" + monkeypatch.setattr(j, "step", real) + with torch.inference_mode(): + train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run)) + done = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run)) + snap = jobs.snapshot() + assert done == "tiny_lora.safetensors" and snap["state"] == "done" and len(snap["loss"]) == 6 + assert not loop.has_checkpoint(run) + + def test_pause_button_waits_at_the_next_step_and_stop_ends_a_paused_run(run_env, monkeypatch): from modules.training.lora import jobs as j, train tmp_path, jobs = run_env