sampling: cancelling is ordinary, so it is checked like everything else

Stopping a generation is a normal thing to do -- a bad seed, a wrong prompt, a rental minute -- so it is not an error path that can be left approximate. A leak here is invisible: the run AFTER a cancel just behaves slightly wrong. The sampler says a run was cancelled and re-raises. It is not ours to absorb; ComfyUI ends the queue item. But a run that simply stops with nothing said is indistinguishable from one that broke. Verified against the real engine: a 40-step generation interrupted mid-way reports execution_interrupted, and the next one completes as run 2. Tests cover what must be left behind, which is nothing -- no wrappers on the shared model after five cancelled runs, the sharp anchor put back on a denoiser that outlives the run, and no modifier marked as dropped, because a cancelled modifier did not fail and disabling it next time would punish the user for stopping. Temp previews needed no change: core's PreviewImage already writes to the temp directory under a unique prefix, so where the app SHOWS them is a UI decision and is recorded as one rather than built here. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-29 18:27:31 +03:00
parent 165bc6e9b6
commit f128eb2d74
2 changed files with 194 additions and 8 deletions
@@ -126,14 +126,21 @@ class FunPackSampler(io.ComfyNode):
        # is missing, and already reports steps to whatever is watching.
        callback = latent_preview.prepare_callback(model, max(0, len(sigmas) - 1))

-        samples = guider.sample(
-            Noise_RandomNoise(seed).generate_noise(latent),
-            samples_in, sampler, sigmas,
-            denoise_mask=latent.get("noise_mask"),
-            callback=callback,
-            disable_pbar=not comfy.utils.PROGRESS_BAR_ENABLED,
-            seed=seed,
-        )
+        try:
+            samples = guider.sample(
+                Noise_RandomNoise(seed).generate_noise(latent),
+                samples_in, sampler, sigmas,
+                denoise_mask=latent.get("noise_mask"),
+                callback=callback,
+                disable_pbar=not comfy.utils.PROGRESS_BAR_ENABLED,
+                seed=seed,
+            )
+        except comfy.model_management.InterruptProcessingException:
+            # Said, then re-raised. Cancelling is a normal thing to do and the
+            # log should show it happening rather than a run simply stopping --
+            # but it is not ours to absorb: ComfyUI ends the queue item.
+            log.alert("FunPack Sampler", f"{run} was cancelled part way through")
+            raise
        # Back to the device every other sampler returns on. Leaving it on the
        # GPU works until something downstream is on the CPU, and then it is a
        # device mismatch a long way from here.
@@ -0,0 +1,179 @@
+"""Cancelling a run, and what it must leave behind: nothing.
+
+Stopping a generation is ordinary -- a bad seed, a wrong prompt, a rental
+minute -- so it is not an error path that can be left approximate. What matters
+is that the next run starts clean, because a leak here is invisible: the run
+after a cancel just behaves slightly wrong.
+"""
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def _needs_comfy(comfyui):
+    """Imports comfy."""
+
+
+@pytest.fixture
+def patcher():
+    import torch
+    from comfy.model_patcher import ModelPatcher
+
+    class Stub(torch.nn.Module):
+        pass
+
+    return ModelPatcher(Stub(), load_device=torch.device("cpu"),
+                        offload_device=torch.device("cpu"))
+
+
+@pytest.fixture
+def registry(monkeypatch):
+    from core import registry as registry_mod
+    fake = registry_mod.Registry()
+    monkeypatch.setattr(registry_mod, "current", lambda rescan=False: fake)
+    return fake
+
+
+def _spec(module_id, install):
+    from core.contract import ModuleSpec
+    return ModuleSpec(id=module_id, title=module_id, mount="", stage="sampling",
+                      provides={"modifier": install})
+
+
+def test_a_cancel_reaches_the_top_and_is_not_absorbed(patcher, registry):
+    """The guard catches Exception, and ComfyUI's cancel is a BaseException, so
+    it passes through. If that ever changed, stop would stop nothing."""
+    from comfy.model_management import InterruptProcessingException
+    from comfy.patcher_extension import WrappersMP
+    from modules.sampling.modifiers.nodes import FunPackLoadModifiers
+
+    def install(target, values, key):
+        def cancels(args):
+            raise InterruptProcessingException()
+        target.set_model_sampler_pre_cfg_function(cancels)
+        return "on"
+
+    registry.specs["m"] = _spec("m", install)
+    patched, _ = FunPackLoadModifiers.execute(patcher).result
+    hook = patched.model_options["sampler_pre_cfg_function"][0]
+
+    with pytest.raises(InterruptProcessingException):
+        hook({"conds_out": ["a", "b"]})
+
+
+def test_a_cancelled_run_leaves_nothing_on_the_shared_model(patcher, registry):
+    """The model everyone else holds must be exactly as it was. Installing on a
+    clone is what makes this true, and a cancel must not change that."""
+    from comfy.model_management import InterruptProcessingException
+    from comfy.patcher_extension import WrappersMP
+    from modules.sampling.modifiers.nodes import FunPackLoadModifiers
+
+    def install(target, values, key):
+        target.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, key,
+                                    lambda e, *a, **k: e(*a, **k))
+        return "on"
+
+    registry.specs["m"] = _spec("m", install)
+
+    for _ in range(5):
+        patched, _ = FunPackLoadModifiers.execute(patcher).result
+        try:
+            raise InterruptProcessingException()      # cancelled mid-run
+        except InterruptProcessingException:
+            pass
+
+    assert patcher.wrappers.get(WrappersMP.SAMPLER_SAMPLE, {}) == {}
+    assert not [k for k in patcher.model_options.get("sampler_pre_cfg_function", [])]
+
+
+def test_a_cancel_does_not_mark_a_modifier_as_dropped(patcher, registry):
+    """A cancelled modifier did not fail. Recording it would disable it on the
+    next run for something the user did on purpose."""
+    from comfy.model_management import InterruptProcessingException
+    from core import patching
+
+    dropped = patching.Dropped()
+
+    def cancels(args):
+        raise InterruptProcessingException()
+
+    guarded = patching.guard(cancels, "funpack.m",
+                             lambda args: args["conds_out"], dropped)
+    with pytest.raises(InterruptProcessingException):
+        guarded({"conds_out": "x"})
+
+    assert "funpack.m" not in dropped and not dropped
+
+
+def test_the_next_run_starts_with_a_clean_record(patcher, registry):
+    """The record travels on the model, which ComfyUI may hand back from its
+    cache -- so a run cancelled after a modifier failed must not leave that
+    modifier disabled for the next one."""
+    from core import patching
+    from modules.sampling.sampler.nodes import FunPackSampler
+
+    dropped = patching.Dropped()
+    dropped.record("funpack.m", RuntimeError("failed last time"))
+    patcher.funpack_dropped = dropped
+
+    # What the sampler does at the start of every generation.
+    found = getattr(patcher, "funpack_dropped", None)
+    found.clear()
+    assert "funpack.m" not in found
+
+
+def test_alg_puts_the_sharp_anchor_back_when_a_run_is_cancelled():
+    """ALG swaps the pinned anchor per step. A cancel between steps must not
+    leave the blurred copy pinned on a denoiser that outlives the run."""
+    import torch
+    import comfy.samplers
+    from comfy.model_management import InterruptProcessingException
+    from modules.sampling import alg
+
+    sharp = torch.randn(1, 3, 4, 16, 16)
+    holder = {}
+
+    def sampler_function(model_k, x, sigmas, **kwargs):
+        model_k(x, torch.tensor([0.9]), denoise_mask=None)
+        holder["denoiser"] = model_k._inner
+        raise InterruptProcessingException()
+
+    captured = {}
+
+    class Patcher:
+        def add_wrapper_with_key(self, kind, key, wrapper):
+            captured["wrapper"] = wrapper
+
+    alg.install(Patcher(), {"enabled": True, "strength": 4.0, "until_sigma": 0.6},
+                key="funpack.alg")
+
+    class Guider:
+        def __init__(self):
+            class ModelSampling:
+                sigma_max = 1.0
+                def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
+                    return noise
+                def inverse_noise_scaling(self, sigma, x):
+                    return x
+
+            class InnerModel:
+                model_sampling = ModelSampling()
+                def scale_latent_inpaint(self, **kw):
+                    return kw["latent_image"]
+
+            self.inner_model = InnerModel()
+
+        def __call__(self, x, sigma, model_options={}, seed=None):
+            return x
+
+    class Executor:
+        class_obj = comfy.samplers.KSAMPLER(sampler_function)
+        def __call__(self, *a, **k):
+            raise AssertionError("should not fall back")
+
+    with pytest.raises(InterruptProcessingException):
+        captured["wrapper"](Executor(), Guider(), torch.tensor([0.9]), {}, None,
+                            torch.zeros_like(sharp), sharp, torch.ones_like(sharp), True)
+
+    assert torch.equal(holder["denoiser"].latent_image, sharp), (
+        "a cancelled run left the loosened anchor pinned")