Report restart replay steps in hybrid sampler

docs/FunPackHybridEuler2SSampler.md

@@ -31,7 +31,7 @@ Example: `0.35` means only the last 35% of steps use the ODE refinement path.
 - `0.0` = pure late-step Euler ODE
 - `1.0` = full late-step DPM++(2S)-style correction
 
-**restart_steps**: How many previous late-phase sigma steps are revisited during Restart replay.
+**restart_steps**: Size of the restarted sigma interval. This behaves like the paper's restart segment length, so each replay adds `restart_steps - 1` extra denoise transitions.
 
 **restart_repeats**: How many Restart loops to run. `0` disables Restart.
 
@@ -56,6 +56,8 @@ For a first Restart test, try:
 - `restart_trigger_pct = 0.85`
 - `restart_noise = 1.0`
 
+Example: with `8` base sampling steps, `restart_steps = 3`, and `restart_repeats = 1`, the Restart section adds `2` extra replay denoise steps, so you should observe the equivalent of `10` denoise callbacks rather than `11`.
+
 ## Expected behavior
 
 Compared to plain `euler_ancestral`, this sampler should usually:

samplers.py

@@ -57,21 +57,39 @@ def _renoise_to_sigma(x, current_sigma, target_sigma, restart_noise, noise_sampl
     return x + noise_sampler(current_sigma, target_sigma) * (restart_noise * sigma_delta)
 
 
-def _run_restart_replay(model, x, restart_sigmas, s_in, extra_args, correction_blend):
+def _run_restart_replay(model, x, restart_sigmas, s_in, extra_args, correction_blend,
+                        callback=None, callback_step_start=0):
     if restart_sigmas.numel() <= 1:
-        return x
+        return x, callback_step_start
+
+    callback_step = int(callback_step_start)
 
     for idx in range(int(restart_sigmas.shape[0]) - 1):
+        sigma = restart_sigmas[idx]
+        sigma_next = restart_sigmas[idx + 1]
+        denoised = model(x, sigma * s_in, **extra_args)
+
+        if callback is not None:
+            callback({
+                "x": x,
+                "i": callback_step,
+                "sigma": sigma,
+                "sigma_hat": sigma,
+                "denoised": denoised,
+            })
+            callback_step += 1
+
         x, _ = _hybrid_ode_step(
             model,
             x,
-            restart_sigmas[idx],
-            restart_sigmas[idx + 1],
+            sigma,
+            sigma_next,
             s_in,
             extra_args,
             correction_blend,
+            denoised=denoised,
         )
-    return x
+    return x, callback_step
 
 
 def sample_funpack_hybrid_euler_2s(model, x, sigmas, extra_args=None, callback=None,
@@ -110,6 +128,7 @@ def sample_funpack_hybrid_euler_2s(model, x, sigmas, extra_args=None, callback=N
     s_in = x.new_ones([x.shape[0]])
     restart_anchor = min(total_steps - 1, max(late_start, int(math.floor(total_steps * restart_trigger_pct))))
     restart_has_run = False
+    callback_step = 0
 
     for i in comfy.utils.model_trange(total_steps, disable=disable):
         sigma = sigmas[i]
@@ -119,7 +138,14 @@ def sample_funpack_hybrid_euler_2s(model, x, sigmas, extra_args=None, callback=N
             denoised = model(x, sigma * s_in, **extra_args)
 
             if callback is not None:
-                callback({"x": x, "i": i, "sigma": sigma, "sigma_hat": sigma, "denoised": denoised})
+                callback({
+                    "x": x,
+                    "i": callback_step,
+                    "sigma": sigma,
+                    "sigma_hat": sigma,
+                    "denoised": denoised,
+                })
+                callback_step += 1
 
             sigma_down, sigma_up = k_diffusion_sampling.get_ancestral_step(sigma, sigma_next, eta=eta)
 
@@ -145,20 +171,29 @@ def sample_funpack_hybrid_euler_2s(model, x, sigmas, extra_args=None, callback=N
                             restart_noise,
                             noise_sampler,
                         )
-                        x = _run_restart_replay(
+                        x, callback_step = _run_restart_replay(
                             model,
                             x,
                             restart_sigmas,
                             s_in,
                             extra_args,
                             correction_blend,
+                            callback=callback,
+                            callback_step_start=callback_step,
                         )
                     restart_has_run = True
 
             denoised = model(x, sigma * s_in, **extra_args)
 
             if callback is not None:
-                callback({"x": x, "i": i, "sigma": sigma, "sigma_hat": sigma, "denoised": denoised})
+                callback({
+                    "x": x,
+                    "i": callback_step,
+                    "sigma": sigma,
+                    "sigma_hat": sigma,
+                    "denoised": denoised,
+                })
+                callback_step += 1
 
             x, _ = _hybrid_ode_step(
                 model,