nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

LoRA training: fix resume in inference mode, Restore and continue, live graph in Results

Browse Source 
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P
This commit is contained in:
Claude
2026-10-17 20:54:49 +00:00
parent b19f7dec2f
commit 3c421b2c1d
7 changed files with 99 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/ui/features/training/index.js
+17 -7
View File
@@ -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) {
app/ui/features/training/results.js
+7 -3
View File
@@ -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) }),
] });
}
app/ui/features/training/train.js
+12 -2
View File
@@ -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;
app/ui/features/training/training.test.js
+11 -2
View File
@@ -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();
});
modules/training/lora/__init__.py
+8 -4
View File
@@ -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}
modules/training/lora/loop.py
+17
View File
@@ -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"])
modules/training/lora/tests/test_train_h3.py
+27
View File
@@ -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