Merge f273baa94d into 37ac9ff44f

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, ...]} }

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={}):

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"