From f128eb2d74ef73a0bb6d1cf63a9e96f45cdca721 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 29 Aug 2026 18:27:31 +0300 Subject: [PATCH] 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 --- modules/sampling/sampler/nodes.py | 23 ++-- modules/sampling/tests/test_cancel.py | 179 ++++++++++++++++++++++++++ 2 files changed, 194 insertions(+), 8 deletions(-) create mode 100644 modules/sampling/tests/test_cancel.py diff --git a/modules/sampling/sampler/nodes.py b/modules/sampling/sampler/nodes.py index b1f8a196..eb91744f 100644 --- a/modules/sampling/sampler/nodes.py +++ b/modules/sampling/sampler/nodes.py @@ -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. diff --git a/modules/sampling/tests/test_cancel.py b/modules/sampling/tests/test_cancel.py new file mode 100644 index 00000000..4d83529f --- /dev/null +++ b/modules/sampling/tests/test_cancel.py @@ -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")