diff --git a/README.md b/README.md index 29b9e6a7..dd01f509 100644 --- a/README.md +++ b/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. diff --git a/core/dit_hooks.py b/core/dit_hooks.py index 83e31e6c..38aae8e0 100644 --- a/core/dit_hooks.py +++ b/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) diff --git a/modules/sampling/turbo/__init__.py b/modules/sampling/turbo/__init__.py new file mode 100644 index 00000000..563d9d87 --- /dev/null +++ b/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} diff --git a/modules/sampling/turbo/tests/test_turbo.py b/modules/sampling/turbo/tests/test_turbo.py new file mode 100644 index 00000000..1261bfb8 --- /dev/null +++ b/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() diff --git a/modules/sampling/two_stage/stages.py b/modules/sampling/two_stage/stages.py index 6a4e9b10..b96be83d 100644 --- a/modules/sampling/two_stage/stages.py +++ b/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):