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..1d399f353 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,10 +1020,42 @@ 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 - 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) - return samples + 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 = 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={}): 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"