diff --git a/app/ui/features/training/index.js b/app/ui/features/training/index.js index 60eab916..6269e2cd 100644 --- a/app/ui/features/training/index.js +++ b/app/ui/features/training/index.js @@ -157,7 +157,7 @@ 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 && !(ui.awaitStart && (snap.state === "paused" || (ui.staleStarted && snap.started === ui.staleStarted)))) { + if (snap && snap.run && snap.run !== ui.before && !(ui.awaitStart && (ui.staleStarted ? snap.started === ui.staleStarted : snap.state === "paused"))) { ui.awaitStart = false; const times = snap.loss_from === ui.run.loss.length ? ui.run.times.concat(snap.times || []) : ui.run.times; const loss = snap.loss_from === ui.run.loss.length ? ui.run.loss.concat(snap.loss) : ui.run.loss; @@ -182,6 +182,29 @@ export default { }; runPoll = setTimeout(tick, POLL); } + // The server is the truth. A window that sits idle on a pause (a tab left in the background, a missed poll) asks now and then and takes over what it finds. + let syncPoll = 0; + async function reconcile() { + clearTimeout(syncPoll); + if (!win) return; + syncPoll = setTimeout(reconcile, POLL * 5); + if (ui.run.active || ui.run.state !== "paused") return; + let snap = null; + try { snap = await api.trainRun(0); } catch { return; } + if (!win || ui.run.active || ui.run.state !== "paused" || !snap || !snap.run) return; + const moved = snap.state === "training" || (snap.state === "paused" && (snap.run !== ui.run.runId || snap.step !== ui.run.step)); + if (!moved) return; + ui.before = ""; + if (snap.state === "training") { + setRun({ shown: true, active: true, state: "training", runId: snap.run, loss: snap.loss, times: snap.times || [], stepSeconds: snap.step_seconds, lastStep: snap.last_step_seconds, smooth: smooth(snap.loss), step: snap.step, total: snap.total, elapsed: snap.elapsed, message: snap.message, error: "", paused: null }); + draw(); + watch(); + } else { + setRun({ shown: true, active: false, state: "paused", runId: snap.run, loss: snap.loss, times: snap.times || [], stepSeconds: snap.step_seconds, lastStep: snap.last_step_seconds, smooth: smooth(snap.loss), step: snap.step, total: snap.total, elapsed: snap.elapsed, message: "", paused: snap.paused }); + draw(); + if (snap.paused && snap.paused.reason === "preview" && !ui.holdPause) previewNow(true); + } + } async function stop() { const waiting = ui.run.state === "paused" || ui.run.state === "previewing"; setRun({ message: "Stopping" }); @@ -237,7 +260,8 @@ export default { 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); } }); + onClose: () => { win = null; session.flush(); clearTimeout(tagPoll); clearTimeout(syncPoll); } }); + syncPoll = setTimeout(reconcile, POLL * 5); offSession = session.on(draw); session.refresh().then(async () => { if (!session.name) { @@ -254,6 +278,6 @@ export default { } 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(); }; + return () => { clearTimeout(tagPoll); clearTimeout(runPoll); clearTimeout(syncPoll); timers.forEach(clearTimeout); if (offSession) offSession(); offOpen(); if (win) win.close(); }; }, }; diff --git a/app/ui/features/training/training.test.js b/app/ui/features/training/training.test.js index daa285db..95011c3c 100644 --- a/app/ui/features/training/training.test.js +++ b/app/ui/features/training/training.test.js @@ -331,6 +331,27 @@ test("a preview step pauses the run, the pipeline makes a preview with the LoRA off(); }); +test("a pause that the window never saw training towards is still shown at the step the server is at", async () => { + const f = fake(); + f.app.runner = { preview: async () => ({ image: { filename: "pv.png", subfolder: "", type: "output" } }) }; + let queued = 0; + f.api.queueGraph = async () => { queued += 1; return "pid"; }; + const pause = (step, started) => ({ run: "r1", state: "paused", step, total: 6, started, elapsed: 3, message: "", loss: [0.9, 0.8], loss_from: 0, paused: { reason: "user", step, preview: "p.safetensors", checkpoint: true } }); + f.api.trainRun = async () => (queued === 0 ? { run: "", state: "idle", loss: [], loss_from: 0 } : queued === 1 ? pause(2, 100) : pause(4, 200)); + 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(1100); + assert.ok(text().includes("paused at step 2")); + click("Resume").click(); + await wait(1400); + assert.ok(text().includes("paused at step 4"), "the second pause is not mistaken for the old one"); + 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" }) }; diff --git a/modules/training/lora/nodes.py b/modules/training/lora/nodes.py index a2e75bee..5d5d1c00 100644 --- a/modules/training/lora/nodes.py +++ b/modules/training/lora/nodes.py @@ -56,6 +56,11 @@ class FunPackTrainLora(io.ComfyNode): outputs=[io.String.Output(display_name="lora")], ) + @classmethod + def fingerprint_inputs(cls, **_): + """Always different: a run is work to do, never a result to reuse. Without this ComfyUI answers a second 'carry on' from a pause (same inputs as the first) with its cached output in 0 seconds.""" + return float("nan") + @classmethod def execute(cls, clip, vae, set_name, name, use, steps, learning_rate, rank, algorithm, optimizer, loss, batch_size, grad_accumulation, resolution, aspect_buckets, frames, fps, start, trigger, seed, training_dtype, lora_dtype, diff --git a/modules/training/lora/tests/test_plan_and_routes.py b/modules/training/lora/tests/test_plan_and_routes.py index b3391528..ee45c5b8 100644 --- a/modules/training/lora/tests/test_plan_and_routes.py +++ b/modules/training/lora/tests/test_plan_and_routes.py @@ -313,3 +313,9 @@ def test_a_rated_result_is_captioned_from_its_own_file_and_an_older_rating_on_a_ got["results"] = [r for r in got["results"] if r["media"]["filename"] != "bare.png"] assert [(r["media"]["filename"], r["caption"], r["doubtful"]) for r in got["results"]] == [("now.png", "a red fox runs through the snow", False), ("reused.png", "a city street at night", True)] assert got["doubtful"] == 1 + + +def test_the_training_node_is_never_answered_from_comfys_cache(): + from modules.training.lora.nodes import FunPackTrainLora + a, b = FunPackTrainLora.fingerprint_inputs(resume="r1"), FunPackTrainLora.fingerprint_inputs(resume="r1") + assert a != b, "NaN never equals itself: ComfyUI runs the node again even with identical inputs"