Turbo sampler: stage-aware flow step, step spacing per stage, optional polish step; funui_turbo sampler and scheduler registered in Comfy

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Tw9fjUCoaPNZ8uLoL2xACR

README.md

@@ -169,6 +169,7 @@ I express my deepest gratitude to:
 - NVIDIA and the NVIDIA Video Effects SDK (`nvidia-vfx`, NVIDIA's own licence) — RTX Super Resolution in Upscale and RTX Frame Generation in Interpolate run on it; FunUI installs it only when you switch the add-on on;
 - NVIDIA and the authors of Sol-H3 (arXiv 2609.35110, NVlabs/Sana, Apache-2.0) — the MXFP8 loader option (which blocks stay in bf16) and the Two stages setting (the cross-resolution schedule) follow its recipe;
 - LBH-123-AI ([Comfyui_Minimax_h3_latent_Upscaler](https://github.com/LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler)) and TripleHeadedMonkey ([ComfyUI-MiniMaxH3_LatentUpscaler](https://github.com/TripleHeadedMonkey/ComfyUI-MiniMaxH3_LatentUpscaler)) — the H3 latent upscaler network (rebuilt from the published weights, not copied) and its use in Second pass and as the Two stages enlarge option;
+- The Turbo sampler (stochastic flow step, polish step) rests on consistency models (Song, Dhariwal, Chen, Sutskever; arXiv 2303.01469), latent consistency models (Luo, Tan, Huang, Li, Zhao; arXiv 2310.04378) and the ancestral noise injection of k-diffusion by Katherine Crowson; its second-order late steps reuse Distilled Flow;
 
 Without all of you, this project would've been impossible.
 

core/dit_hooks.py

@@ -195,6 +195,10 @@ def current_step(transformer_options):
 STAGE = "funui_stage"
 
 
+# Set on the guider by Two stages while it runs: the whole schedule, of which a stage sees only its part.
+FULL = "funui_full_sigmas"
+
+
 def stage_of(guider):
     """1 or 2 while a Two stages run is in that stage, None for a plain run."""
     return getattr(guider, STAGE, None)

modules/sampling/turbo/__init__.py

@@ -0,0 +1,253 @@
+"""Turbo sampler: a step rule made for few-step turbo runs, aware of Two stages, for any step count.
+
+Replaces the chosen sampler for the run (the model calls are the same: one per step, no extra cost) with a
+flow step that has three dials:
+
+* Fresh noise per step. After each model answer the next point is `(1 - s) * picture + s * noise`; the
+  noise is the model's own implied noise, blended with fresh noise by `eta` (0 = plain Euler, 1 = a full
+  re-draw as in consistency sampling). `eta` runs from "Noise at the start" to "Noise at the end" across the
+  WHOLE schedule, so it means the same for 3+2 steps, 4+4 or a run without Two stages.
+* Second order late in the run: the picture guesses of the last steps are extrapolated from the one before
+  (as in Distilled Flow's `dflow_ord2_ab2`), faded in with the run's progress. History starts fresh in each
+  stage, because a Two stages hand-over changes the size.
+* Step spacing per stage: the steps inside a stage are moved earlier or later while the stage's first and
+  last sigma stay put, so the hand-over point never moves. 1 = the schedule as given.
+* Polish step (optional): after the last step, noise the finished picture back to a small sigma and let the
+  model answer once more at full size. One extra model call, in the last stage only.
+
+The sound is stepped without fresh noise and without second order. Samplers that call the model twice per
+step are replaced, not wrapped; the other sampler wrappers still run.
+
+Ideas and credit: consistency models (Song et al., arXiv 2303.01469) and latent consistency models (Luo et
+al., arXiv 2310.04378) for the re-noise-and-answer multistep rule; k-diffusion (Katherine Crowson) for
+ancestral-style noise injection; NVIDIA Sol-H3 (arXiv 2609.35110) for the cross-resolution schedule this is
+tuned around. See the README Thank you section.
+"""
+
+import math
+
+import torch
+from comfy.patcher_extension import WrappersMP
+
+from ..._core import dit_hooks, log
+
+ID = "turbo_sampler"
+TITLE = "Turbo sampler"
+MOUNT = "generation.sampling"
+STAGE = "load"
+CATEGORY = "sampling"
+STATUS = "experimental"
+REQUIRES = ["audio_stream"]
+
+SETTINGS = {
+    "enabled": {
+        "type": "bool", "default": False,
+        "label": "Turbo sampler",
+        "hint": "Replaces the chosen sampler with a step rule made for few-step turbo runs. No extra model calls "
+                "(except the optional polish step). Works with any step count and with Two stages.",
+    },
+    "noise_start": {
+        "type": "float", "default": 0.4, "min": 0.0, "max": 1.0, "step": 0.05,
+        "label": "Noise at the start", "ui": "slider",
+        "hint": "How much of each early step's noise is drawn fresh. 0 = plain Euler, 1 = a full re-draw.",
+        "when": {"enabled": True},
+    },
+    "noise_end": {
+        "type": "float", "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05,
+        "label": "Noise at the end", "ui": "slider",
+        "hint": "The same for the last steps; it fades from the start value to this one across the whole run.",
+        "when": {"enabled": True},
+    },
+    "second_order": {
+        "type": "float", "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
+        "label": "Second order late", "ui": "slider",
+        "hint": "How much the late steps lean on the step before. 0 = off. It fades in over the run.",
+        "when": {"enabled": True},
+    },
+    "spacing_first": {
+        "type": "float", "default": 1.0, "min": 0.5, "max": 2.0, "step": 0.05,
+        "label": "Step spacing, first stage", "ui": "slider",
+        "hint": "Moves the steps inside the first stage (with Two stages) or the whole run (without). Above 1 "
+                "pulls them toward the end of the stage, below 1 toward its start. 1 = as scheduled.",
+        "when": {"enabled": True},
+    },
+    "spacing_last": {
+        "type": "float", "default": 1.0, "min": 0.5, "max": 2.0, "step": 0.05,
+        "label": "Step spacing, last stage", "ui": "slider",
+        "hint": "The same for the stage at the full size. Ignored without Two stages.",
+        "when": {"enabled": True},
+    },
+    "polish": {
+        "type": "float", "default": 0.0, "min": 0.0, "max": 0.3, "step": 0.01,
+        "label": "Polish step sigma", "ui": "slider",
+        "hint": "0 = off. Otherwise one extra full-size model call after the last step, from this small noise "
+                "level. Costs one step (about 15 to 25% of a run). 0.05 is a start.",
+        "when": {"enabled": True},
+    },
+}
+
+
+def warp(sigmas, k):
+    """The inner sigmas of a stage moved by `k`, the first and last kept. u runs 0..1 along the stage; u' = k u / (1 + (k - 1) u)."""
+    if k == 1.0 or len(sigmas) < 3:
+        return sigmas
+    first, last = sigmas[0], sigmas[-1]
+    span = first - last
+    if float(span.abs()) < 1e-9:
+        return sigmas
+    u = ((first - sigmas) / span).clamp(0.0, 1.0)
+    moved = k * u / (1.0 + (k - 1.0) * u)
+    out = first - span * moved
+    out[0], out[-1] = first, last
+    return out
+
+
+def _parts(t):
+    return list(t.unbind()) if getattr(t, "is_nested", False) else [t]
+
+
+def _join(like, parts):
+    if getattr(like, "is_nested", False):
+        from comfy.nested_tensor import NestedTensor
+        return NestedTensor(tuple(parts))
+    return parts[0]
+
+
+def _noise(like, generator):
+    return torch.randn(like.shape, generator=generator, device=like.device, dtype=like.dtype)
+
+
+def step_to(x, x0, sigma, to, eta, noise):
+    """One flow step from `sigma` to `to` given the picture guess `x0`; `eta` of the next point's noise is fresh `noise`."""
+    sigma = sigma.to(device=x.device, dtype=x.dtype)
+    to = to.to(device=x.device, dtype=x.dtype)
+    implied = (x - (1 - sigma) * x0) / sigma
+    if eta > 0.0:
+        implied = math.sqrt(max(0.0, 1.0 - eta * eta)) * implied + eta * noise
+    return (1 - to) * x0 + to * implied
+
+
+def make(settings):
+    """The sampler function. It reads the stage from the guider behind `model` (set by Two stages), so the same
+    function serves the registered `funui_turbo` sampler and the settings module."""
+    n_start, n_end = settings["noise_start"], settings["noise_end"]
+    order, tail = settings["second_order"], settings["polish"]
+
+    @torch.no_grad()
+    def sample(model, x, sigmas, extra_args=None, callback=None, disable=None, **_ignored):
+        extra_args = {} if extra_args is None else extra_args
+        guider = getattr(model, "inner_model", None)
+        stage = dit_hooks.stage_of(guider)
+        full = getattr(guider, dit_hooks.FULL, None)
+        total = (len(full) - 1) if full is not None else len(sigmas) - 1
+        n0 = total - (len(sigmas) - 1) if stage == 2 else 0                  # global index of this run's first step
+        last_stage = stage != 1
+        sigmas = warp(sigmas, settings["spacing_last"] if stage == 2 else settings["spacing_first"])
+        steps = len(sigmas) - 1
+        log.once(f"turbo:{stage}:{steps}:{n0}", log.INFO, TITLE,
+                 f"{'stage ' + str(stage) if stage else 'run'}: {steps} step(s) from step {n0 + 1} of {total}"
+                 + (", polish step after the last" if tail > 0 and last_stage else ""))
+        base = int(extra_args.get("seed") or 0)
+        parts_like = _parts(x)
+        gens = []
+        for p in parts_like:
+            g = torch.Generator(device=p.device)
+            g.manual_seed(base + 1000003 * (n0 + 1))
+            gens.append(g)
+        ones = x.new_ones([x.shape[0]])
+        before = before_h = None
+        used = x
+        for i in range(steps):
+            sigma, to = sigmas[i], sigmas[i + 1]
+            denoised = model(x, sigma * ones, **extra_args)
+            if callback is not None:
+                callback({"x": x, "i": i, "sigma": sigma, "sigma_hat": sigma, "denoised": denoised})
+            progress = (n0 + i) / max(1, total - 1)
+            eta = n_start + (n_end - n_start) * progress
+            h = float((sigma - to).abs())
+            d_parts = _parts(denoised)
+            used_parts = list(d_parts)
+            if order > 0.0 and before is not None and before_h > 1e-7 and h > 1e-7:
+                r = min(5.0, max(0.1, h / before_h))
+                used_parts[0] = d_parts[0] + order * progress * (d_parts[0] - before[0]) * (r / 2)   # picture only
+            before, before_h = d_parts, h
+            used = _join(denoised, used_parts)
+            if float(to) == 0:
+                break
+            nxt = []
+            for k, (xp, up) in enumerate(zip(_parts(x), used_parts)):
+                e = eta if k == 0 else 0.0                                    # the sound gets no fresh noise
+                nxt.append(step_to(xp, up, sigma, to, e, _noise(xp, gens[k]) if e > 0.0 else None))
+            x = _join(x, nxt)
+        else:
+            return x
+        if tail > 0.0 and last_stage:
+            noisy = []
+            for k, up in enumerate(_parts(used)):
+                t = torch.as_tensor(tail, device=up.device, dtype=up.dtype)
+                noisy.append((1 - t) * up + t * _noise(up, gens[k]))
+            t = torch.as_tensor(tail, device=parts_like[0].device, dtype=parts_like[0].dtype)
+            used = model(_join(used, noisy), t * ones, **extra_args)
+        return used
+    return sample
+
+
+def _settings(values):
+    return {name: float(values.get(name, spec["default"])) for name, spec in SETTINGS.items() if name != "enabled"}
+
+
+_CURRENT = {}                                   # the settings of the run in hand, for the registered sampler
+
+
+def registered(model, x, sigmas, extra_args=None, callback=None, disable=None, **options):
+    """`funui_turbo` in ComfyUI's sampler list: the settings when the module is on, else the defaults."""
+    return make(_CURRENT or _settings({}))(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, **options)
+
+
+def schedule(model_sampling, steps):
+    """`funui_turbo` in ComfyUI's scheduler list: the evenly spaced timesteps of `sgm_uniform`, the grid the turbo
+    merges are trained on. The step spacing per stage is applied by the sampler, which knows the stage."""
+    import comfy.samplers
+    return comfy.samplers.normal_scheduler(model_sampling, steps, sgm=True)
+
+
+def register():
+    """Add the sampler and the scheduler to ComfyUI's own lists (in place: its Sampler nodes share the lists)."""
+    import comfy.k_diffusion.sampling as ksampling
+    import comfy.samplers
+    if "funui_turbo" not in comfy.samplers.SAMPLER_NAMES:
+        ksampling.sample_funui_turbo = registered
+        comfy.samplers.SAMPLER_NAMES.append("funui_turbo")
+    if "funui_turbo" not in comfy.samplers.SCHEDULER_HANDLERS:
+        comfy.samplers.SCHEDULER_HANDLERS["funui_turbo"] = comfy.samplers.SchedulerHandler(schedule)
+        comfy.samplers.SCHEDULER_NAMES.append("funui_turbo")
+
+
+register()
+
+
+def install(patcher, values, key):
+    _CURRENT.clear()
+    if not values.get("enabled"):
+        return None
+    settings = _settings(values)
+    _CURRENT.update(settings)
+
+    def sampler_sample(executor, guider, sigmas, *args, **kwargs):
+        sampler = executor.class_obj
+        inner = getattr(sampler, "sampler_function", None)
+        if inner is None:
+            log.once("turbo:shape", log.INFO, TITLE, "off this run: the sampler is not a plain sampler function")
+            return executor(guider, sigmas, *args, **kwargs)
+        sampler.sampler_function = make(settings)       # swapped for this call and passed ON: later wrappers still run
+        try:
+            return executor(guider, sigmas, *args, **kwargs)
+        finally:
+            sampler.sampler_function = inner
+
+    patcher.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, key, sampler_sample)
+    return (f"fresh noise {settings['noise_start']:g} to {settings['noise_end']:g}, second order {settings['second_order']:g}"
+            + (f", polish at {settings['polish']:g}" if settings["polish"] > 0 else ""))
+
+
+PROVIDES = {"modifier": install}

modules/sampling/turbo/tests/test_turbo.py

@@ -0,0 +1,169 @@
+"""Turbo sampler: the step maths on toy models, then the real wrapper through ComfyUI's guider on the tiny H3 (CPU)."""
+
+import pytest
+import torch
+
+from modules.sampling.turbo import SETTINGS, install, make, step_to, warp
+from modules._core import dit_hooks
+from modules.sampling.two_stage.tests.test_two_stage import model  # noqa: F401 -- the module-scoped fixture
+
+
+def _settings(**over):
+    values = {name: spec["default"] for name, spec in SETTINGS.items() if name != "enabled"}
+    values.update(over)
+    return values
+
+
+class Behind:
+    """A guider stand-in: which stage Two stages says it is in, and the whole schedule."""
+    def __init__(self, stage=None, full=None):
+        setattr(self, dit_hooks.STAGE, stage)
+        setattr(self, dit_hooks.FULL, full)
+
+
+class Oracle:
+    """A model that always knows the clean picture `target` (a flow model with a perfect prediction)."""
+    def __init__(self, target, behind=None):
+        self.target, self.calls, self.inner_model = target, [], behind
+
+    def __call__(self, x, sigma, **kw):
+        self.calls.append(float(sigma[0]))
+        return self.target.to(x)
+
+
+def test_warp_keeps_the_ends_and_order():
+    s = torch.tensor([1.0, 0.8, 0.6, 0.3, 0.0])
+    assert torch.equal(warp(s, 1.0), s)
+    for k in (0.6, 1.5, 2.0):
+        w = warp(s, k)
+        assert w[0] == s[0] and w[-1] == s[-1] and len(w) == len(s)
+        assert (w[:-1] > w[1:]).all()
+    assert warp(s, 2.0)[1] < s[1]                    # k > 1 pulls the inner steps toward the end of the stage
+    assert warp(s, 0.6)[1] > s[1]
+
+
+def test_eta_zero_is_euler():
+    x, x0 = torch.randn(1, 4, 2, 3, 3), torch.randn(1, 4, 2, 3, 3)
+    sigma, to = torch.tensor(0.8), torch.tensor(0.5)
+    euler = x + ((x - x0) / sigma) * (to - sigma)
+    assert torch.allclose(step_to(x, x0, sigma, to, 0.0, None), euler, atol=1e-5)
+
+
+def test_perfect_model_lands_on_the_target_for_any_step_count_and_split():
+    target = torch.randn(1, 4, 2, 3, 3)
+    for sigmas in ([1.0, 0.5, 0.0], [1.0, 0.9, 0.7, 0.4, 0.2, 0.0]):
+        s = torch.tensor(sigmas)
+        for behind in (None, Behind(2, torch.tensor([1.0] * 4 + sigmas))):
+            oracle = Oracle(target, behind)
+            out = make(_settings(noise_start=0.8, polish=0.0))(oracle, torch.randn_like(target), s)
+            assert torch.allclose(out, target, atol=1e-5)
+            assert oracle.calls == pytest.approx([float(v) for v in sigmas[:-1]], rel=1e-6)
+
+
+def test_polish_adds_one_call_in_the_last_stage_only():
+    target = torch.randn(1, 4, 2, 3, 3)
+    s = torch.tensor([1.0, 0.5, 0.0])
+    for last, calls in ((True, 3), (False, 2)):
+        oracle = Oracle(target, Behind(None if last else 1, torch.tensor([1.0, 0.5, 0.3, 0.0])))
+        out = make(_settings(polish=0.05))(oracle, torch.randn_like(target), s)
+        assert len(oracle.calls) == calls and torch.allclose(out, target, atol=1e-5)
+        if last:
+            assert oracle.calls[-1] == pytest.approx(0.05)
+
+
+def test_same_seed_same_result_and_noise_actually_used():
+    target = torch.randn(1, 4, 2, 3, 3)
+    s = torch.tensor([1.0, 0.7, 0.4, 0.0])
+    noisy_start = torch.randn_like(target)
+
+    class Sloppy(Oracle):                              # answers differ with x, so injected noise shows in the result
+        def __call__(self, x, sigma, **kw):
+            return 0.5 * self.target.to(x) + 0.5 * x * (1 - sigma[0])
+    run = lambda eta, seed: make(_settings(noise_start=eta, noise_end=eta, second_order=0.0))(
+        Sloppy(target), noisy_start.clone(), s, extra_args={"seed": seed})
+    assert torch.equal(run(0.7, 1), run(0.7, 1))
+    assert not torch.allclose(run(0.7, 1), run(0.7, 2))
+    assert not torch.allclose(run(0.7, 1), run(0.0, 1))
+
+
+def test_follows_the_input_dtype_and_nested_sound_gets_no_noise():
+    from comfy.nested_tensor import NestedTensor
+    video, audio = torch.randn(1, 4, 2, 3, 3, dtype=torch.float64), torch.randn(1, 6, 2, 3, dtype=torch.float64)
+    seen = []
+
+    class Sound:
+        def __call__(self, x, sigma, **kw):
+            seen.append(x.unbind()[1].clone())
+            return NestedTensor((torch.zeros_like(x.unbind()[0]), torch.zeros_like(x.unbind()[1])))
+    s = torch.tensor([1.0, 0.6, 0.3, 0.0], dtype=torch.float32)       # sigmas in another dtype than the latent
+    out = make(_settings(noise_start=1.0, polish=0.1, second_order=1.0))(
+        Sound(), NestedTensor((video, audio)), s)
+    v, a = out.unbind()
+    assert v.dtype == torch.float64 and a.dtype == torch.float64
+    # the sound is stepped by plain Euler toward zero: x * (to / sigma) each time
+    assert torch.allclose(seen[1], audio * (0.6 / 1.0), atol=1e-9)
+
+
+def _guided(model, values, sigmas, with_two_stages=None):
+    import comfy.samplers
+    from comfy.nested_tensor import NestedTensor
+    patched = model.clone()
+    if with_two_stages:
+        from modules.sampling.two_stage import install as two
+        two(patched, with_two_stages, "funui.two_stage")
+    note = install(patched, values, "funui.turbo")
+    calls = []
+    hook = patched.model.diffusion_model.register_forward_pre_hook(
+        lambda m, a, kw: calls.append(float(kw["transformer_options"]["sigmas"].reshape(-1)[0])), with_kwargs=True)
+    guider = comfy.samplers.CFGGuider(patched)
+    cond = [[torch.randn(1, 5, 48), {}]]
+    guider.set_conds(cond, cond)
+    guider.set_cfg(1.0)
+    latent = NestedTensor((torch.zeros(1, 24, 2, 8, 8), torch.zeros(1, 32, 2, 3)))
+    noise = NestedTensor((torch.randn(1, 24, 2, 8, 8), torch.randn(1, 32, 2, 3)))
+    try:
+        out = guider.sample(noise, latent, comfy.samplers.sampler_object("euler"), torch.tensor(sigmas),
+                            seed=1, disable_pbar=True)
+    finally:
+        hook.remove()
+    return note, out, calls
+
+
+def test_off_installs_nothing(model):
+    assert install(model.clone(), {"enabled": False}, "k") is None
+
+
+@pytest.mark.parametrize("sigmas,cut", [((1.0, 0.8, 0.5, 0.25, 0.0), 2), ((1.0, 0.9, 0.7, 0.55, 0.4, 0.2, 0.0), 3),
+                                        ((1.0, 0.6, 0.3, 0.0), 1)])
+def test_runs_through_two_stages_with_any_split(model, sigmas, cut):
+    values = {"enabled": True, **_settings(spacing_first=1.3, spacing_last=0.8, polish=0.05)}
+    note, out, calls = _guided(model, values, sigmas, {"enabled": True, "size": 0.5, "steps": cut})
+    assert "fresh noise" in note
+    v, a = out.unbind()
+    assert torch.isfinite(v).all() and torch.isfinite(a).all()
+    assert len(calls) == len(sigmas) - 1 + 1                 # one call per step plus the polish call
+    assert calls[cut - 1] != calls[cut] and calls[-1] == pytest.approx(0.05)
+    assert calls[cut] == pytest.approx(sigmas[cut], rel=1e-4) # the hand-over sigma did not move
+
+
+def test_without_two_stages_and_without_polish(model):
+    values = {"enabled": True, **_settings()}
+    _, out, calls = _guided(model, values, (1.0, 0.7, 0.3, 0.0))
+    assert len(calls) == 3 and torch.isfinite(out.unbind()[0]).all()
+
+
+def test_registered_in_comfys_lists_and_runs_like_the_module(model):
+    import comfy.samplers
+    from comfy.nested_tensor import NestedTensor
+    assert "funui_turbo" in comfy.samplers.SAMPLER_NAMES and "funui_turbo" in comfy.samplers.SCHEDULER_NAMES
+    patched = model.clone()
+    guider = comfy.samplers.CFGGuider(patched)
+    cond = [[torch.randn(1, 5, 48), {}]]
+    guider.set_conds(cond, cond)
+    guider.set_cfg(1.0)
+    sigmas = comfy.samplers.calculate_sigmas(patched.get_model_object("model_sampling"), "funui_turbo", 4)
+    assert torch.allclose(sigmas, comfy.samplers.calculate_sigmas(patched.get_model_object("model_sampling"), "sgm_uniform", 4))
+    noise = NestedTensor((torch.randn(1, 24, 2, 8, 8), torch.randn(1, 32, 2, 3)))
+    latent = NestedTensor((torch.zeros(1, 24, 2, 8, 8), torch.zeros(1, 32, 2, 3)))
+    out = guider.sample(noise, latent, comfy.samplers.sampler_object("funui_turbo"), sigmas, seed=1, disable_pbar=True)
+    assert torch.isfinite(out.unbind()[0]).all()

modules/sampling/two_stage/stages.py

@@ -44,7 +44,7 @@ def _why_not(guider, latent_image, denoise_mask, sigmas, latent_shapes, steps):
     return None
 
 
-FULL = "funui_full_sigmas"
+FULL = dit_hooks.FULL
 
 
 def keep_whole_schedule(executor, guider, sigmas, extra_args, *args, **kwargs):