LoRA training: pause, resume and a preview LoRA at chosen steps (backend)

A run can wait at a step boundary (Pause, or every N steps for a preview): its state is saved,
the LoRA so far is written as <name>_preview.safetensors, and the card is let go. Resume carries on
from the saved state under the same run. UNTESTED on real weights.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P

core/tests/test_shipped_nodes.py

@@ -74,7 +74,7 @@ FROZEN = {
         "model", "clip", "vae", "audio_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", "gradient_checkpointing", "offloading", "bypass",
-        "model_name", "weight_dtype", "compute_dtype", "postfix", "swap_blocks",
+        "model_name", "weight_dtype", "compute_dtype", "postfix", "swap_blocks", "preview_every", "resume",
     ],
     "FunPackEmptyLatent": ["model", "width", "height", "length", "batch_size", "audio_vae", "frame_rate"],
     "FunPackDiffusionModelLoader": [

modules/training/lora/__init__.py

@@ -204,9 +204,18 @@ def routes(table, base, web):
             since = 0
         return web.json_response(jobs.snapshot(since))
 
+    @table.post(base + "/run/pause")
+    async def _pause(_req):
+        return web.json_response({"paused": jobs.request_pause()})
+
     @table.post(base + "/run/stop")
     async def _stop(_req):
+        now = jobs.snapshot()
+        was_paused = now["state"] == "paused"
         stopped = jobs.request_stop()
+        if stopped and was_paused:
+            train.discard(now["run"], now["name"])
+            return web.json_response({"stopped": True})
         if stopped:
             try:
                 import comfy.model_management

modules/training/lora/jobs.py

@@ -19,10 +19,11 @@ MAX_KEPT_LOSSES = 4000          # a saved run keeps a thinned curve: a 100000 st
 _LOCK = threading.Lock()
 _ID = re.compile(r"\A[0-9a-f]{12}\Z")
 IDLE = {"state": "idle", "run": "", "set": "", "name": "", "step": 0, "total": 0, "started": 0.0, "message": "", "error": "",
-        "lora": "", "samples": [], "settings": {}}
+        "lora": "", "samples": [], "settings": {}, "paused": None}
 _now: dict = dict(IDLE)
 _loss: list[float] = []
 _stop = False
+_pause = False
 
 
 def _dir():
@@ -42,14 +43,19 @@ def trail(message: str) -> None:
     log.info("FunPack LoRA training", message + rss)
 
 
-def begin(set_name: str, name: str, total: int, settings: dict, prompt_id: str = "") -> str:
-    """Start a record. A run that was still marked running (ComfyUI restarted under it) is simply replaced."""
-    global _stop
+def begin(set_name: str, name: str, total: int, settings: dict, prompt_id: str = "", resume: dict | None = None) -> str:
+    """Start a record. A run that was still marked running (ComfyUI restarted under it) is simply replaced.
+
+    `resume` ({"run", "loss"}) carries on a paused run: same record, same loss curve, the steps counting on.
+    """
+    global _stop, _pause
     with _LOCK:
-        _stop = False
+        _stop = _pause = False
         _loss.clear()
         _now.clear()
-        _now.update(IDLE, state="training", run=uuid.uuid4().hex[:12], set=set_name, name=name, total=int(total), started=time.time(),
+        if resume:
+            _loss.extend(float(v) for v in resume.get("loss") or [])
+        _now.update(IDLE, state="training", run=(resume or {}).get("run") or uuid.uuid4().hex[:12], step=len(_loss), set=set_name, name=name, total=int(total), started=time.time(),
                     settings=dict(settings), prompt_id=prompt_id, message="Loading the model and the training data")
         return _now["run"]
 
@@ -84,12 +90,40 @@ def thin(values, keep: int = MAX_KEPT_LOSSES) -> list[float]:
     return out + [values[-1]]
 
 
+def losses() -> list[float]:
+    with _LOCK:
+        return list(_loss)
+
+
+def pause(reason: str, step: int, preview: str, checkpoint: bool = True) -> None:
+    """The run waits: its state is on disk, and the models are let go so the pipeline can use the card. `reason`: user or preview."""
+    with _LOCK:
+        if _now["state"] != "training":
+            return
+        _now.update(state="paused", message="", paused={"reason": reason, "step": int(step), "preview": preview, "checkpoint": bool(checkpoint)})
+
+
+def request_pause() -> bool:
+    """Ask the running training to wait at its next step. False when nothing is training."""
+    global _pause
+    with _LOCK:
+        if _now["state"] != "training":
+            return False
+        _pause = True
+        _now["message"] = "Pausing"
+        return True
+
+
+def pause_requested() -> bool:
+    return _pause
+
+
 def finish(state: str, lora: str = "", error: str = "") -> None:
     """state: done, stopped (the person stopped it) or error. Writes the run's record."""
     with _LOCK:
         if _now["run"] == "" or _now["state"] in ("done", "error", "stopped"):
             return
-        _now.update(state=state, lora=lora, error=error, message="")
+        _now.update(state=state, lora=lora, error=error, message="", paused=None)
         record = {k: _now.get(k) for k in ("run", "set", "name", "total", "started", "settings", "lora", "samples", "error")}
         record.update(state=state, steps=len(_loss), seconds=round(time.time() - _now["started"], 1),
                       final_loss=_loss[-1] if _loss else None, loss=thin(_loss))
@@ -144,11 +178,17 @@ def request_stop() -> bool:
     """Ask the running training to end at its next step. False when nothing is training."""
     global _stop
     with _LOCK:
-        if _now["state"] != "training":
+        if _now["state"] == "paused":
+            ended = True                                  # nothing is running: the run simply ends
+        elif _now["state"] != "training":
             return False
-        _stop = True
-        _now["message"] = "Stopping"
-        return True
+        else:
+            ended = False
+            _stop = True
+            _now["message"] = "Stopping"
+    if ended:
+        finish("stopped")
+    return True
 
 
 def stop_requested() -> bool:

modules/training/lora/loop.py

@@ -15,6 +15,45 @@ from ..._core import log
 from . import jobs
 
 
+class Paused(Exception):
+    """The loop stops at a step boundary and keeps its state: `reason` is user (Pause pressed) or preview (the step asked for)."""
+
+    def __init__(self, reason, step):
+        super().__init__(f"paused at step {step}")
+        self.reason, self.step = reason, step
+
+
+def _checkpoint_file(run):
+    from ..._core import config
+    folder = config.TRAINING_DIR / "checkpoints"
+    folder.mkdir(parents=True, exist_ok=True)
+    return folder / f"{run}.pt"
+
+
+def save_checkpoint(run, lora, optimizer, step, torch):
+    """Everything a resumed run needs besides the data: the LoRA, the optimizer's memory, the loss so far and where it stopped."""
+    torch.save({"lora": {n: tuple(p.detach().cpu() for p in pair) for n, pair in lora.pairs.items()}, "optimizer": optimizer.state_dict(),
+                "step": int(step), "loss": jobs.losses()}, _checkpoint_file(run))
+
+
+def load_checkpoint(run, torch):
+    """The saved state of a paused run, or a ValueError saying why it cannot be carried on."""
+    import re
+    if not re.fullmatch(r"[0-9a-f]{12}", str(run)):
+        raise ValueError("that is not the name of a paused run")
+    path = _checkpoint_file(run)
+    if not path.is_file():
+        raise ValueError("the paused run's saved state is gone: start the training again")
+    return torch.load(path, map_location="cpu", weights_only=False)
+
+
+def forget_checkpoint(run):
+    try:
+        _checkpoint_file(run).unlink()
+    except OSError:
+        pass
+
+
 class Lora:
     """A low-rank pair per target layer, applied by a forward hook so the layer, its ops and its weight casting stay ComfyUI's own."""
 
@@ -41,6 +80,17 @@ class Lora:
             return out + delta.to(out.dtype)
         return add
 
+    def restore(self, saved):
+        """Put a checkpoint's weights back; the layers must be the ones it was saved from."""
+        if set(saved) != set(self.pairs):
+            raise ValueError("the paused run was made for a different model or LoRA layout: start the training again")
+        with self.torch.no_grad():
+            for name, (down, up) in self.pairs.items():
+                if saved[name][0].shape != down.shape:
+                    raise ValueError("the paused run used another rank: start the training again")
+                down.copy_(saved[name][0])
+                up.copy_(saved[name][1])
+
     def parameters(self):
         return [p for pair in self.pairs.values() for p in pair]
 
@@ -88,7 +138,7 @@ class Sampler:
     """Handed to ComfyUI's guider in place of a sampler: instead of denoising, it trains. Built lazily so this file imports without comfy."""
 
     @staticmethod
-    def make(nt, torch, comfy_samplers, custom_sampler, lora, optimizer, criterion, cfg, dtype, stopped):
+    def make(nt, torch, comfy_samplers, custom_sampler, lora, optimizer, criterion, cfg, dtype, stopped, first=0):
         class OwnSampler(comfy_samplers.Sampler):
             def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
                 model_wrap.conds = nt.process_cond_list(model_wrap.conds)
@@ -96,7 +146,9 @@ class Sampler:
                 size = sigmas.size(0)
                 accumulate, seed = cfg["grad_accumulation"], cfg["seed"]
                 sampling = model_wrap.inner_model.model_sampling
-                for i in range(cfg["steps"] * accumulate):
+                total = cfg["steps"] * accumulate
+                every = -(-int(cfg.get("preview_every") or 0) // accumulate) * accumulate        # a preview waits for a whole optimizer step
+                for i in range(first, total):
                     gen = torch.Generator().manual_seed(seed + i)
                     pick = torch.randperm(size, generator=gen)[:1].tolist()
                     latent = torch.stack([latent_image[j] for j in pick])
@@ -121,12 +173,17 @@ class Sampler:
                     if (i + 1) % accumulate == 0:
                         optimizer.step()
                         optimizer.zero_grad()
+                        if (i + 1) < total:
+                            if jobs.pause_requested():
+                                raise Paused("user", i + 1)
+                            if every and (i + 1) % every == 0:
+                                raise Paused("preview", i + 1)
                 torch.cuda.empty_cache()
                 return torch.zeros_like(latent_image)
         return OwnSampler()
 
 
-def train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim):
+def train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim, resume=None, run=None):
     """Train a LoRA on `model` with ComfyUI's guider carrying each batch. Returns the file's tensors."""
     import comfy.samplers
     import comfy_extras.nodes_custom_sampler as custom_sampler
@@ -154,19 +211,30 @@ def train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim):
             if jobs.stop_requested():
                 from .train import Stopped
                 raise Stopped()
-        sampler = Sampler.make(nt, torch, comfy.samplers, custom_sampler, lora, optimizer, criterion, cfg, dtype, stopped)
+        first = 0
+        if resume:
+            lora.restore(resume["lora"])
+            optimizer.load_state_dict(resume["optimizer"])
+            first = int(resume["step"])
+            log.info("FunPack LoRA training", f"carrying on from step {first}")
+        sampler = Sampler.make(nt, torch, comfy.samplers, custom_sampler, lora, optimizer, criterion, cfg, dtype, stopped, first)
         if cfg["offloading"]:
             log.warning("FunPack LoRA training", "ComfyUI's own 'keep weights in system memory' is ignored by this training loop: it has brought ComfyUI down on H3. Use 'Blocks kept in system memory' instead")
         guider = nt.TrainGuider(model, offloading=False)
         guider.set_conds(positive)
+        paused = None
         try:
             mm.in_training = True
             with watch:
-                nt._run_training_loop(guider, sampler, latents, count, cfg["seed"], False, multi)
+                try:
+                    nt._run_training_loop(guider, sampler, latents, count, cfg["seed"], False, multi)
+                except Paused as stop:
+                    paused = stop
+                    save_checkpoint(run, lora, optimizer, stop.step, torch)
         finally:
             mm.in_training = False
             restore()
             state = lora.state() if lora else {}
             lora.remove()
             del optimizer, sampler, guider                # a failed run leaves nothing on the card for the next one
-        return state
+        return state, paused

modules/training/lora/nodes.py

@@ -49,6 +49,8 @@ class FunPackTrainLora(io.ComfyNode):
                 io.Combo.Input("compute_dtype", options=["default", "fp16", "bf16", "fp32"], default="default", tooltip="What the maths runs in when the model is loaded here."),
                 io.String.Input("postfix", default="", tooltip="Text put at the end of every caption."),
                 io.Int.Input("swap_blocks", default=0, min=0, max=128, tooltip="How many transformer blocks keep their weights in system memory, fetched one layer at a time: slower, and each one frees about a block's weight on the card. 0 keeps the whole model on the card."),
+                io.Int.Input("preview_every", default=0, min=0, max=100000, tooltip="Pause every this many steps so the pipeline can make a preview with the LoRA so far. 0 never pauses."),
+                io.String.Input("resume", default="", tooltip="The name of a paused run to carry on. Empty starts a new run."),
             ],
             outputs=[io.String.Output(display_name="lora")],
         )
@@ -56,7 +58,7 @@ class FunPackTrainLora(io.ComfyNode):
     @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,
-                gradient_checkpointing, offloading, bypass, model=None, audio_vae=None, model_name="", weight_dtype="default", compute_dtype="default", postfix="", swap_blocks=0) -> io.NodeOutput:
+                gradient_checkpointing, offloading, bypass, model=None, audio_vae=None, model_name="", weight_dtype="default", compute_dtype="default", postfix="", swap_blocks=0, preview_every=0, resume="") -> io.NodeOutput:
         if not addons.enabled("lora_training"):
             raise RuntimeError("LoRA training is switched off. Turn it on in Settings > Modules.")
         cfg = {"name": name, "use": use, "steps": steps, "learning_rate": learning_rate, "rank": rank, "algorithm": algorithm,
@@ -64,5 +66,5 @@ class FunPackTrainLora(io.ComfyNode):
                "resolution": resolution, "aspect_buckets": aspect_buckets, "frames": frames, "fps": fps, "start": start,
                "trigger": trigger, "seed": seed, "training_dtype": training_dtype, "lora_dtype": lora_dtype,
                "gradient_checkpointing": gradient_checkpointing, "offloading": offloading, "bypass": bypass,
-               "model_name": model_name, "weight_dtype": weight_dtype, "compute_dtype": compute_dtype, "postfix": postfix, "swap_blocks": swap_blocks}
+               "model_name": model_name, "weight_dtype": weight_dtype, "compute_dtype": compute_dtype, "postfix": postfix, "swap_blocks": swap_blocks, "preview_every": preview_every, "resume": resume}
         return io.NodeOutput(train.run(model, clip, vae, audio_vae, set_name, cfg))

modules/training/lora/tests/test_train_h3.py

@@ -351,3 +351,49 @@ def test_block_modulation_is_out_of_place_while_training_and_put_back_after():
         assert torch.allclose(m._mod_gate(x0, gate, other, segs), want_x)
         assert torch.equal(x0, x)
     assert (m._mod_scale_shift, m._mod_gate) == before
+
+
+def test_a_run_pauses_for_a_preview_and_carries_on_to_the_same_end(run_env):
+    import safetensors.torch
+    from modules.training.lora import train
+    tmp_path, jobs = run_env
+    first = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2))
+    snap = jobs.snapshot()
+    assert first == "" and snap["state"] == "paused"
+    assert snap["paused"]["reason"] == "preview" and snap["paused"]["step"] == 2 and snap["step"] == 2
+    preview = tmp_path / "loras" / snap["paused"]["preview"]
+    assert preview.name == "tiny_lora_preview.safetensors"
+    assert any(k.endswith(".lora_up.weight") for k in safetensors.torch.load_file(str(preview)))
+    run = snap["run"]
+    # the person looks at the preview, and training carries on from where it stopped
+    second = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run))
+    snap = jobs.snapshot()
+    assert second == "" and snap["run"] == run and snap["paused"]["step"] == 4 and len(snap["loss"]) == 4
+    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 preview.exists(), "the preview file stays after the run is over"
+    assert not (tmp_path / "training" / "checkpoints" / f"{run}.pt").exists()
+
+
+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
+    real = j.step
+
+    def press(loss):
+        real(loss)
+        j.request_pause()
+    monkeypatch.setattr(j, "step", press)
+    train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6))
+    snap = jobs.snapshot()
+    assert snap["state"] == "paused" and snap["paused"]["reason"] == "user" and snap["paused"]["step"] == 1
+    assert j.request_stop() and jobs.snapshot()["state"] == "stopped"
+
+
+def test_a_run_that_cannot_be_resumed_says_why(run_env):
+    from modules.training.lora import train
+    _, jobs = run_env
+    with pytest.raises(ValueError, match="saved state is gone"):
+        train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(resume="0123456789ab"))
+    assert jobs.snapshot()["state"] == "error"

modules/training/lora/train.py

@@ -128,7 +128,10 @@ def run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str =
         torch.cuda.empty_cache()
     _card("at the start of the run", mm, torch)
     try:
-        return _run(model, clip, vae, audio_vae, set_name, cfg, prompt_id)
+        result = _run(model, clip, vae, audio_vae, set_name, cfg, prompt_id)
+        model = clip = vae = audio_vae = None             # a paused or finished run lets the card go: the pipeline's preview or the next run needs it
+        _free_the_card(mm, torch)
+        return result
     except BaseException as exc:
         traceback.clear_frames(exc.__traceback__)
         model = clip = vae = audio_vae = None
@@ -144,8 +147,16 @@ def _run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str =
     import safetensors.torch
 
     from ..._core import registry
-    run_id = jobs.begin(set_name, cfg["name"], cfg["steps"], cfg, prompt_id)        # first, so a refusal below is shown in the window and not only in ComfyUI's log
+    resume, failed = None, None
+    if cfg.get("resume"):
+        try:
+            resume = loop.load_checkpoint(cfg["resume"], torch)
+        except ValueError as exc:
+            failed = exc
+    run_id = jobs.begin(set_name, cfg["name"], cfg["steps"], cfg, prompt_id, {"run": cfg["resume"], "loss": resume["loss"]} if resume else None)        # first, so a refusal below is shown in the window and not only in ComfyUI's log
     try:
+        if failed:
+            raise failed
         if model is None:
             model = load_model(cfg.get("model_name"), cfg.get("weight_dtype", "default"), cfg.get("compute_dtype", "default"))
         adapt = registry.current().ask("training_adapt", model)
@@ -205,7 +216,12 @@ def _run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str =
         shim = adapt["shims"]() if adapt and adapt.get("shims") else _nothing()
         watch = _watch(nt, mm, torch, _swapper(adapt, cfg, mm, torch))
         if adapt and adapt.get("lora_targets"):
-            lora = loop.train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim)       # our own loop, ComfyUI carries the batches
+            lora, paused = loop.train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim, resume, run_id)       # our own loop, ComfyUI carries the batches
+            if paused:
+                preview = _save(lora, f"{cfg['name']}_preview", folder_paths, safetensors.torch, set_name, cfg, keep=True)
+                jobs.pause(paused.reason, paused.step, preview)
+                log.info("FunPack LoRA training", f"paused at step {paused.step} ({paused.reason}); the LoRA so far is {preview}")
+                return ""
         else:
             losses = _loss_hook(nt, cfg["steps"] * cfg["grad_accumulation"])
             with shim, losses, watch:
@@ -218,6 +234,7 @@ def _run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str =
                     existing_lora=["[None]"], bucket_mode=[False], bypass_mode=[cfg["bypass"]])
             lora = (out.result if hasattr(out, "result") else out)[0]
         file = _save(lora, cfg["name"], folder_paths, safetensors.torch, set_name, cfg)
+        discard(run_id, cfg["name"])
         jobs.finish("done", lora=file)
         return file
     except Stopped:
@@ -467,13 +484,28 @@ class _loss_hook:
         return False
 
 
-def _save(lora, name, folder_paths, safetensors_torch, set_name, cfg) -> str:
-    """Into the first loras folder, as <name>.safetensors, or <name>_2 and so on rather than over an earlier one."""
+def _preview_name(name) -> str:
+    stem = "".join(c for c in f"{name}_preview" if c.isalnum() or c in "-_ .").strip() or "funpack_lora_preview"
+    return f"{stem}.safetensors"
+
+
+def discard(run, name) -> None:
+    """A run that is over (finished, stopped) keeps neither its saved state nor its preview file."""
+    import folder_paths
+    loop.forget_checkpoint(run)
+    try:
+        (Path(folder_paths.get_folder_paths("loras")[0]) / _preview_name(name.removesuffix("_preview"))).unlink()
+    except OSError:
+        pass
+
+
+def _save(lora, name, folder_paths, safetensors_torch, set_name, cfg, keep=False) -> str:
+    """Into the first loras folder, as <name>.safetensors, or <name>_2 and so on rather than over an earlier one. `keep`: replace the file of that name (the preview)."""
     folder = Path(folder_paths.get_folder_paths("loras")[0])
     folder.mkdir(parents=True, exist_ok=True)
     stem = "".join(c for c in name if c.isalnum() or c in "-_ .") .strip() or "funpack_lora"
     n, path = 1, folder / f"{stem}.safetensors"
-    while path.exists():
+    while path.exists() and not keep:
         n += 1
         path = folder / f"{stem}_{n}.safetensors"
     meta = {"software": "ComfyUI-FunPack", "training_set": set_name, "algorithm": cfg["algorithm"], "rank": str(cfg["rank"]),