From 72d2be807919e1de0614a7e7b8108d24e05561e2 Mon Sep 17 00:00:00 2001 From: Masahiro Date: Wed, 27 May 2026 02:59:51 +1200 Subject: [PATCH 1/3] Add sampler lifecycle callbacks --- comfy/patcher_extension.py | 3 +++ comfy/samplers.py | 46 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/comfy/patcher_extension.py b/comfy/patcher_extension.py index 189ee84ca..b869cdef9 100644 --- a/comfy/patcher_extension.py +++ b/comfy/patcher_extension.py @@ -13,6 +13,9 @@ class CallbacksMP: ON_REGISTER_ALL_HOOK_PATCHES = "on_register_all_hook_patches" ON_INJECT_MODEL = "on_inject_model" ON_EJECT_MODEL = "on_eject_model" + ON_SAMPLER_START = "on_sampler_start" + ON_SAMPLER_STEP = "on_sampler_step" + ON_SAMPLER_END = "on_sampler_end" # callbacks dict is in the format: # {"call_type": {"key": [Callable1, Callable2, ...]} } diff --git a/comfy/samplers.py b/comfy/samplers.py index 1d6a4e104..de2aa1a1a 100755 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -993,6 +993,25 @@ class KSAMPLER(Sampler): noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas)) total_steps = len(sigmas) - 1 + model_options = extra_args.get("model_options", {}) + callback_types = comfy.patcher_extension.CallbacksMP + get_callbacks = comfy.patcher_extension.get_all_callbacks + sampler_function_name = getattr(self.sampler_function, "__name__", self.sampler_function.__class__.__name__) + sampler_start_callbacks = get_callbacks(callback_types.ON_SAMPLER_START, model_options, is_model_options=True) + sampler_step_callbacks = get_callbacks(callback_types.ON_SAMPLER_STEP, model_options, is_model_options=True) + sampler_end_callbacks = get_callbacks(callback_types.ON_SAMPLER_END, model_options, is_model_options=True) + + if len(sampler_start_callbacks) > 0: + sampler_info = { + "total_steps": total_steps, + "sample_sigmas": sigmas, + "noise_shape": tuple(noise.shape), + "latent_shape": tuple(latent_image.shape) if latent_image is not None else None, + "sampler_function": sampler_function_name, + } + for sampler_callback in sampler_start_callbacks: + sampler_callback(sampler_info) + first_step = True def k_callback(x): nonlocal first_step @@ -1001,9 +1020,36 @@ class KSAMPLER(Sampler): first_step = False if callback is not None: callback(x["i"], x["denoised"], x["x"], total_steps) + if len(sampler_step_callbacks) == 0: + return + + step = x["i"] + sigma_next = sigmas[step + 1] if step + 1 < len(sigmas) else None + sampler_info = { + "step": step, + "total_steps": total_steps, + "sigma": x.get("sigma", sigmas[step] if step < len(sigmas) else None), + "sigma_next": sigma_next, + "sigma_hat": x.get("sigma_hat", None), + "sample_sigmas": sigmas, + "x_shape": tuple(x["x"].shape) if "x" in x else None, + "denoised_shape": tuple(x["denoised"].shape) if "denoised" in x else None, + "sampler_function": sampler_function_name, + } + for sampler_callback in sampler_step_callbacks: + sampler_callback(sampler_info) samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options) samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples) + if len(sampler_end_callbacks) > 0: + sampler_info = { + "total_steps": total_steps, + "sample_sigmas": sigmas, + "samples_shape": tuple(samples.shape), + "sampler_function": sampler_function_name, + } + for sampler_callback in sampler_end_callbacks: + sampler_callback(sampler_info) return samples From 4d12c753a9b9073b3601873d8a196b76b21a0ad0 Mon Sep 17 00:00:00 2001 From: Masahiro Date: Wed, 27 May 2026 12:39:10 +1200 Subject: [PATCH 2/3] Ensure sampler end callbacks run on errors --- comfy/samplers.py | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/comfy/samplers.py b/comfy/samplers.py index de2aa1a1a..1d399f353 100755 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -1039,18 +1039,23 @@ class KSAMPLER(Sampler): for sampler_callback in sampler_step_callbacks: sampler_callback(sampler_info) - samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options) - samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples) - if len(sampler_end_callbacks) > 0: - sampler_info = { - "total_steps": total_steps, - "sample_sigmas": sigmas, - "samples_shape": tuple(samples.shape), - "sampler_function": sampler_function_name, - } - for sampler_callback in sampler_end_callbacks: - sampler_callback(sampler_info) - return samples + samples = None + sampling_succeeded = False + try: + samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options) + samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples) + sampling_succeeded = True + return samples + finally: + if len(sampler_end_callbacks) > 0: + sampler_info = { + "total_steps": total_steps, + "sample_sigmas": sigmas, + "samples_shape": tuple(samples.shape) if sampling_succeeded else None, + "sampler_function": sampler_function_name, + } + for sampler_callback in sampler_end_callbacks: + sampler_callback(sampler_info) def ksampler(sampler_name, extra_options={}, inpaint_options={}): From f273baa94d58780d780554ba4eca575691b79285 Mon Sep 17 00:00:00 2001 From: Masahiro Date: Sun, 16 Aug 2026 09:27:58 +1200 Subject: [PATCH 3/3] Add tests for sampler lifecycle callbacks Covers the three paths raised in review: no registered callbacks (legacy per-step callback unchanged), start/step/end delivery with expected payloads, and end delivery with samples_shape=None when sampling raises. Co-Authored-By: Claude Opus 5 (1M context) --- .../comfy_test/sampler_callbacks_test.py | 118 ++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 tests-unit/comfy_test/sampler_callbacks_test.py diff --git a/tests-unit/comfy_test/sampler_callbacks_test.py b/tests-unit/comfy_test/sampler_callbacks_test.py new file mode 100644 index 000000000..530dcc69d --- /dev/null +++ b/tests-unit/comfy_test/sampler_callbacks_test.py @@ -0,0 +1,118 @@ +import pytest +import torch + +import comfy.patcher_extension as patcher_extension +from comfy.samplers import KSAMPLER + +SIGMAS = torch.tensor([14.6, 7.0, 0.0]) +NOISE = torch.zeros(1, 4, 8, 8) +LATENT = torch.zeros(1, 4, 8, 8) + + +class _ModelSampling: + sigma_max = 14.6 + + def noise_scaling(self, sigma, noise, latent_image, max_denoise=False): + return noise + + def inverse_noise_scaling(self, sigma, latent): + return latent + + +class _InnerModel: + def __init__(self): + self.model_sampling = _ModelSampling() + + +class _ModelPatcher: + def __init__(self): + self.model = _InnerModel() + + +class _ModelWrap: + """Minimal stand-in for CFGGuider: only what KSAMPLER.sample and the first-step log touch.""" + + def __init__(self): + self.inner_model = _InnerModel() + self.model_patcher = _ModelPatcher() + self.cfg = 8.0 + + +def _sampler_function(model, noise, sigmas, extra_args=None, callback=None, disable=False, **kwargs): + for i in range(len(sigmas) - 1): + if callback is not None: + callback({"i": i, "denoised": noise, "x": noise, "sigma": sigmas[i]}) + return noise + + +def _failing_sampler_function(model, noise, sigmas, extra_args=None, callback=None, disable=False, **kwargs): + raise RuntimeError("sampling failed") + + +@pytest.fixture +def extra_args(): + return {"model_options": {}} + + +def _register(extra_args, call_type, callback): + patcher_extension.add_callback(call_type, callback, extra_args["model_options"], is_model_options=True) + + +def test_sampling_unchanged_without_lifecycle_callbacks(extra_args): + """No registered lifecycle callbacks: the legacy per-step callback still fires and the result is returned.""" + legacy_steps = [] + sampler = KSAMPLER(_sampler_function) + + samples = sampler.sample( + _ModelWrap(), SIGMAS, extra_args, + lambda i, denoised, x, total_steps: legacy_steps.append((i, total_steps)), + NOISE, latent_image=LATENT, + ) + + assert samples.shape == NOISE.shape + assert legacy_steps == [(0, 2), (1, 2)] + + +def test_start_step_end_callbacks_are_delivered(extra_args): + started, stepped, ended = [], [], [] + _register(extra_args, patcher_extension.CallbacksMP.ON_SAMPLER_START, started.append) + _register(extra_args, patcher_extension.CallbacksMP.ON_SAMPLER_STEP, stepped.append) + _register(extra_args, patcher_extension.CallbacksMP.ON_SAMPLER_END, ended.append) + + KSAMPLER(_sampler_function).sample( + _ModelWrap(), SIGMAS, extra_args, None, NOISE, latent_image=LATENT, + ) + + assert len(started) == 1 + assert started[0]["total_steps"] == 2 + assert started[0]["noise_shape"] == tuple(NOISE.shape) + assert started[0]["latent_shape"] == tuple(LATENT.shape) + assert started[0]["sampler_function"] == "_sampler_function" + + assert [s["step"] for s in stepped] == [0, 1] + assert [s["total_steps"] for s in stepped] == [2, 2] + assert float(stepped[0]["sigma"]) == pytest.approx(float(SIGMAS[0])) + assert float(stepped[0]["sigma_next"]) == pytest.approx(float(SIGMAS[1])) + assert float(stepped[1]["sigma_next"]) == pytest.approx(float(SIGMAS[2])) + assert stepped[0]["x_shape"] == tuple(NOISE.shape) + assert stepped[0]["denoised_shape"] == tuple(NOISE.shape) + + assert len(ended) == 1 + assert ended[0]["total_steps"] == 2 + assert ended[0]["samples_shape"] == tuple(NOISE.shape) + assert ended[0]["sampler_function"] == "_sampler_function" + + +def test_end_callback_runs_when_sampling_raises(extra_args): + ended = [] + _register(extra_args, patcher_extension.CallbacksMP.ON_SAMPLER_END, ended.append) + + with pytest.raises(RuntimeError, match="sampling failed"): + KSAMPLER(_failing_sampler_function).sample( + _ModelWrap(), SIGMAS, extra_args, None, NOISE, latent_image=LATENT, + ) + + assert len(ended) == 1 + assert ended[0]["samples_shape"] is None + assert ended[0]["total_steps"] == 2 + assert ended[0]["sampler_function"] == "_failing_sampler_function"