Add restart logic to hybrid sampler

docs/FunPackHybridEuler2SSampler.md

@@ -7,9 +7,15 @@ This node outputs a ComfyUI `SAMPLER` object that can be connected directly into
 It is designed as a quality/speed compromise between:
 
 - `euler_ancestral`: fast, reliable, one model evaluation per step
-- full 2S-style samplers: higher quality, but more expensive because they need an extra evaluation every step
+- deterministic 2S-style ODE refinement: higher quality, but more expensive because it needs extra evaluations
 
-This sampler keeps classic Euler ancestral for most of the schedule, then applies a DPM-Solver++(2S)-style correction only on the late denoise steps where fine detail and cleanup usually matter most.
+This sampler keeps classic Euler ancestral for the early "structure building" stage where motion, anatomy and main composition are forming. On the late denoise steps it switches to a deterministic Euler / DPM-Solver++(2S) ODE refinement path where fine detail and cleanup usually matter most.
+
+It can also run an optional paper-style Restart pass during that late quality stage:
+
+- re-noise the current latent up to a higher sigma
+- replay a short late-step sigma interval with the deterministic refinement path
+- repeat that replay multiple times if requested
 
 ## Parameters
 
@@ -17,13 +23,21 @@ This sampler keeps classic Euler ancestral for most of the schedule, then applie
 
 **s_noise**: Noise multiplier for ancestral noise injection.
 
-**high_quality_pct**: Fraction of the *late* denoise steps that receive the extra 2S correction.  
-Example: `0.35` means only the last 35% of steps pay the extra compute cost.
+**high_quality_pct**: Fraction of the *late* denoise steps that switch from Euler ancestral into deterministic quality refinement.  
+Example: `0.35` means only the last 35% of steps use the ODE refinement path.
 
-**correction_blend**: Blend between pure Euler ancestral and the late-step 2S correction.
+**correction_blend**: Blend between late-step Euler ODE and late-step 2S correction.
 
-- `0.0` = behaves like Euler ancestral
-- `1.0` = full late-step correction
+- `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_repeats**: How many Restart loops to run. `0` disables Restart.
+
+**restart_trigger_pct**: Sampling progress point where Restart is triggered. This value is clamped so Restart only happens inside the late quality phase.
+
+**restart_noise**: Noise strength used when re-noising the latent up to the restart interval's higher sigma.
 
 ## Recommended starting values
 
@@ -31,11 +45,16 @@ Example: `0.35` means only the last 35% of steps pay the extra compute cost.
 - `s_noise = 1.0`
 - `high_quality_pct = 0.30` to `0.40`
 - `correction_blend = 1.0`
+- `restart_repeats = 0` for baseline testing
 
-For a safer first test, try:
+For a first Restart test, try:
 
-- `high_quality_pct = 0.25`
-- `correction_blend = 0.7`
+- `high_quality_pct = 0.35`
+- `correction_blend = 1.0`
+- `restart_steps = 3`
+- `restart_repeats = 1`
+- `restart_trigger_pct = 0.85`
+- `restart_noise = 1.0`
 
 ## Expected behavior
 
@@ -43,8 +62,11 @@ Compared to plain `euler_ancestral`, this sampler should usually:
 
 - preserve more late-stage detail
 - clean up texture and edge quality better
-- cost noticeably less than a full 2S sampler
+- keep the early motion/anatomy formation more lively
+- cost less than running a heavier deterministic solver for the whole schedule
+
+With Restart enabled, it may further improve late cleanup and detail consistency, but it will increase runtime and may be too aggressive for some video workflows if overused.
 
 ## Limitation
 
-This node improves the **sampler-side quality/speed tradeoff**. It does **not** reduce the underlying cost of CFG++ guidance itself. If the goal becomes specifically “CFG++ quality at lower guidance cost”, that likely requires a custom `GUIDER`, not just a custom `SAMPLER`.
+This node improves the **sampler-side quality/speed tradeoff**. It does **not** reduce the underlying cost of CFG++ guidance itself. If the goal becomes specifically "CFG++ quality at lower guidance cost", that likely requires a custom `GUIDER`, not just a custom `SAMPLER`.

samplers.py

@@ -7,17 +7,83 @@ import comfy.model_sampling
 import comfy.samplers
 import comfy.utils
 
+
+def _sigma_fn(t):
+    return t.neg().exp()
+
+
+def _t_fn(sigma):
+    return sigma.log().neg()
+
+
+def _hybrid_ode_step(model, x, sigma, sigma_next, s_in, extra_args, correction_blend, denoised=None):
+    if denoised is None:
+        denoised = model(x, sigma * s_in, **extra_args)
+
+    if sigma_next == 0:
+        return denoised, denoised
+
+    d = k_diffusion_sampling.to_d(x, sigma, denoised)
+    dt = sigma_next - sigma
+    x_euler = x + d * dt
+
+    if correction_blend <= 0.0:
+        return x_euler, denoised
+
+    t = _t_fn(sigma)
+    t_next = _t_fn(sigma_next)
+    h = t_next - t
+    r = 0.5
+    s = t + r * h
+
+    x_mid = (_sigma_fn(s) / _sigma_fn(t)) * x - torch.expm1(-h * r) * denoised
+    denoised_mid = model(x_mid, _sigma_fn(s) * s_in, **extra_args)
+    x_2s = (_sigma_fn(t_next) / _sigma_fn(t)) * x - torch.expm1(-h) * denoised_mid
+
+    if correction_blend >= 1.0:
+        return x_2s, denoised
+    return x_euler.lerp(x_2s, correction_blend), denoised
+
+
+def _renoise_to_sigma(x, current_sigma, target_sigma, restart_noise, noise_sampler):
+    if target_sigma <= current_sigma or restart_noise <= 0.0:
+        return x
+
+    sigma_delta_sq = max(0.0, float(target_sigma * target_sigma - current_sigma * current_sigma))
+    if sigma_delta_sq <= 0.0:
+        return x
+
+    sigma_delta = math.sqrt(sigma_delta_sq)
+    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):
+    if restart_sigmas.numel() <= 1:
+        return x
+
+    for idx in range(int(restart_sigmas.shape[0]) - 1):
+        x, _ = _hybrid_ode_step(
+            model,
+            x,
+            restart_sigmas[idx],
+            restart_sigmas[idx + 1],
+            s_in,
+            extra_args,
+            correction_blend,
+        )
+    return x
+
+
 def sample_funpack_hybrid_euler_2s(model, x, sigmas, extra_args=None, callback=None,
                                    disable=None, eta=1.0, s_noise=1.0,
-                                   high_quality_pct=0.35, correction_blend=1.0):
+                                   high_quality_pct=0.35, correction_blend=1.0,
+                                   restart_steps=3, restart_repeats=0,
+                                   restart_trigger_pct=0.85, restart_noise=1.0):
     """
     Hybrid sampler:
-    - Early schedule: classic Euler ancestral (fast).
-    - Late schedule: DPM-Solver++(2S)-style correction (higher quality) blended onto
-      the Euler ancestral proposal only where late-step detail matters most.
-
-    This keeps the cost close to Euler ancestral while paying the extra model
-    evaluation only on the tail of the schedule.
+    - Early schedule: Euler ancestral for motion/anatomy buildup.
+    - Late schedule: deterministic Euler / DPM-Solver++(2S) ODE refinement for detail.
+    - Restart: paper-style re-noise and short ODE replay inside the late quality window.
     """
     if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
         return k_diffusion_sampling.sample_euler_ancestral(
@@ -35,47 +101,75 @@ def sample_funpack_hybrid_euler_2s(model, x, sigmas, extra_args=None, callback=N
 
     high_quality_pct = max(0.0, min(1.0, float(high_quality_pct)))
     correction_blend = max(0.0, min(1.0, float(correction_blend)))
+    restart_steps = max(2, int(restart_steps))
+    restart_repeats = max(0, int(restart_repeats))
+    restart_trigger_pct = max(0.0, min(1.0, float(restart_trigger_pct)))
+    restart_noise = max(0.0, float(restart_noise))
     late_steps = max(1, int(math.ceil(total_steps * high_quality_pct))) if high_quality_pct > 0.0 else 0
     late_start = max(0, total_steps - late_steps)
-
-    sigma_fn = lambda t: t.neg().exp()
-    t_fn = lambda sigma: sigma.log().neg()
     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
 
     for i in comfy.utils.model_trange(total_steps, disable=disable):
         sigma = sigmas[i]
         sigma_next = sigmas[i + 1]
-        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})
+        if i < late_start:
+            denoised = model(x, sigma * s_in, **extra_args)
 
-        sigma_down, sigma_up = k_diffusion_sampling.get_ancestral_step(sigma, sigma_next, eta=eta)
+            if callback is not None:
+                callback({"x": x, "i": i, "sigma": sigma, "sigma_hat": sigma, "denoised": denoised})
 
-        if sigma_down == 0:
-            x = denoised
-        else:
-            d = k_diffusion_sampling.to_d(x, sigma, denoised)
-            dt = sigma_down - sigma
-            x_euler = x + d * dt
+            sigma_down, sigma_up = k_diffusion_sampling.get_ancestral_step(sigma, sigma_next, eta=eta)
 
-            if i >= late_start and correction_blend > 0.0:
-                t = t_fn(sigma)
-                t_next = t_fn(sigma_down)
-                h = t_next - t
-                r = 0.5
-                s = t + r * h
-
-                x_mid = (sigma_fn(s) / sigma_fn(t)) * x - torch.expm1(-h * r) * denoised
-                denoised_mid = model(x_mid, sigma_fn(s) * s_in, **extra_args)
-                x_2s = (sigma_fn(t_next) / sigma_fn(t)) * x - torch.expm1(-h) * denoised_mid
-
-                x = x_euler.lerp(x_2s, correction_blend)
+            if sigma_down == 0:
+                x = denoised
             else:
-                x = x_euler
+                d = k_diffusion_sampling.to_d(x, sigma, denoised)
+                dt = sigma_down - sigma
+                x = x + d * dt
 
-            if sigma_next > 0 and eta > 0 and s_noise > 0:
-                x = x + noise_sampler(sigma, sigma_next) * s_noise * sigma_up
+                if sigma_next > 0 and eta > 0 and s_noise > 0:
+                    x = x + noise_sampler(sigma, sigma_next) * s_noise * sigma_up
+        else:
+            if restart_repeats > 0 and not restart_has_run and i >= restart_anchor:
+                restart_from = max(late_start, i - restart_steps + 1)
+                restart_sigmas = sigmas[restart_from:i + 1]
+                if restart_sigmas.numel() > 1 and restart_sigmas[0] > restart_sigmas[-1]:
+                    for _ in range(restart_repeats):
+                        x = _renoise_to_sigma(
+                            x,
+                            restart_sigmas[-1],
+                            restart_sigmas[0],
+                            restart_noise,
+                            noise_sampler,
+                        )
+                        x = _run_restart_replay(
+                            model,
+                            x,
+                            restart_sigmas,
+                            s_in,
+                            extra_args,
+                            correction_blend,
+                        )
+                    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})
+
+            x, _ = _hybrid_ode_step(
+                model,
+                x,
+                sigma,
+                sigma_next,
+                s_in,
+                extra_args,
+                correction_blend,
+                denoised=denoised,
+            )
 
     return x
 
@@ -113,6 +207,34 @@ class FunPackHybridEuler2SSampler:
                     "step": 0.01,
                     "tooltip": "Blend between Euler ancestral (0.0) and late-step 2S correction (1.0)."
                 }),
+                "restart_steps": ("INT", {
+                    "default": 3,
+                    "min": 2,
+                    "max": 32,
+                    "step": 1,
+                    "tooltip": "How many previous late-phase sigma steps are revisited during Restart replay."
+                }),
+                "restart_repeats": ("INT", {
+                    "default": 0,
+                    "min": 0,
+                    "max": 8,
+                    "step": 1,
+                    "tooltip": "How many paper-style re-noise/replay loops to run in the late quality phase. Set to 0 to disable Restart."
+                }),
+                "restart_trigger_pct": ("FLOAT", {
+                    "default": 0.85,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01,
+                    "tooltip": "Sampling progress point where Restart is triggered. Clamped into the late quality phase."
+                }),
+                "restart_noise": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 2.0,
+                    "step": 0.01,
+                    "tooltip": "Noise strength used when re-noising up to the Restart interval's higher sigma."
+                }),
             }
         }
 
@@ -120,9 +242,10 @@ class FunPackHybridEuler2SSampler:
     RETURN_NAMES = ("sampler",)
     FUNCTION = "get_sampler"
     CATEGORY = "FunPack/Sampling"
-    DESCRIPTION = "Hybrid sampler: Euler ancestral for most steps, DPM-Solver++(2S)-style correction only on late steps for a better quality/speed tradeoff."
+    DESCRIPTION = "Hybrid sampler: early Euler ancestral for motion, late DPM-Solver++(2S) ODE refinement for quality, with optional paper-style Restart replay."
 
-    def get_sampler(self, eta, s_noise, high_quality_pct, correction_blend):
+    def get_sampler(self, eta, s_noise, high_quality_pct, correction_blend,
+                    restart_steps, restart_repeats, restart_trigger_pct, restart_noise):
         sampler = comfy.samplers.KSAMPLER(
             sample_funpack_hybrid_euler_2s,
             extra_options={
@@ -130,6 +253,10 @@ class FunPackHybridEuler2SSampler:
                 "s_noise": s_noise,
                 "high_quality_pct": high_quality_pct,
                 "correction_blend": correction_blend,
+                "restart_steps": restart_steps,
+                "restart_repeats": restart_repeats,
+                "restart_trigger_pct": restart_trigger_pct,
+                "restart_noise": restart_noise,
             }
         )
         return (sampler,)