diff --git a/CHANGELOG.md b/CHANGELOG.md index 7ceefb9c..73a937b3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ## Unreleased -- **Distilled Flow samplers** (experimental, untested on a real model): plain names in every ComfyUI sampler list, no settings. `dflow_ord1` is Euler; `dflow_ord2` reuses the last prediction for second order (no extra model calls); `_ramp` fades that in across the run; `_heun1` / `_heun2` correct the last one or two steps with one extra model call each. Sound riding in the latent stays first order. Back from 4.x; its extras (velocity bias, ALG, momentum, sharpness) are separate modules now. +- **Distilled Flow samplers** (experimental, untested on a real model): plain names in every ComfyUI sampler list, no settings. `dflow_ord1` is Euler; `dflow_ord2` reuses the last prediction for second order (no extra model calls); `dflow_ord2_ab2` fades the second order in across the run. Sound riding in the latent stays first order. Back from 4.x; its extras (velocity bias, ALG, momentum, sharpness) are separate modules now. - **Sharpness** (experimental): wins back fine detail in the last part of the run with any sampler that calls the model once per step; no extra model calls. Settings > Engine > Sampling. Not the spatial latent sharpening that was tried and removed in 4.x. ## [5.0.0] "Cunning Cranberry" - 2026-10-16 diff --git a/modules/sampling/distilled_flow/__init__.py b/modules/sampling/distilled_flow/__init__.py index f2b0398d..c549af64 100644 --- a/modules/sampling/distilled_flow/__init__.py +++ b/modules/sampling/distilled_flow/__init__.py @@ -1,9 +1,10 @@ """Distilled Flow: samplers for few-step distilled flow models, listed in ComfyUI like any other. -Names read as parts: `dflow_ord1` / `dflow_ord2` (first order, or Adams-Bashforth 2 which reuses the -last prediction: no extra model calls), `_ramp` (the second order fades in across the run instead of -acting from step 2), `_heun1` / `_heun2` (the last one or two steps are corrected with a second model -call, so they cost one more call each). Plain `dflow_ord1` is Euler: same speed as `euler`. +`dflow_ord1` is Euler (same speed as `euler`). `dflow_ord2` is Adams-Bashforth 2: it reuses the last +prediction for a second-order step, with no extra model calls. `dflow_ord2_ab2` fades that second order +in across the run (early noisy steps stay first order, late detail steps get the full second order), +which is gentler on short schedules. A first-order version of the fade would be plain `dflow_ord1`, so it +has no name of its own. Deterministic, no noise injection. Flow models only (the model predicts a straight path to the picture). A sound stream riding in the latent is stepped first order, because second order corrupts it. @@ -28,13 +29,11 @@ def _picture_only(new, old): return new -def make(order, ramp, heun): +def make(order, ramp): @torch.no_grad() def sample(model, x, sigmas, extra_args=None, callback=None, disable=None): extra_args = {} if extra_args is None else extra_args steps = len(sigmas) - 1 - # The final step (sigma 0) cannot be corrected, so `heun` counts the steps that can. - first_heun = sum(1 for s in sigmas[1:] if float(s) > 0) - heun ones = x.new_ones([x.shape[0]]) before = before_h = None for i in range(steps): @@ -53,10 +52,6 @@ def make(order, ramp, heun): if float(to) == 0: return used d = ksampling.to_d(x, sigma, used) - if i >= first_heun: - guess = x + d * (to - sigma) - d = _picture_only((d + ksampling.to_d(guess, to, model(guess, to * ones, **extra_args))) / 2, d) - before = before_h = None # x moved by a corrected direction: no history to extrapolate from x = x + d * (to - sigma) return x return sample @@ -64,12 +59,10 @@ def make(order, ramp, heun): def register(): """Add every preset to ComfyUI's own sampler list (in place: its Sampler nodes share the list).""" - for order, ramp, tag in ((1, False, "ord1"), (2, False, "ord2"), (2, True, "ord2_ramp")): - for heun, suffix in ((0, ""), (1, "_heun1"), (2, "_heun2")): - name = f"dflow_{tag}{suffix}" - if name not in comfy.samplers.SAMPLER_NAMES: - setattr(ksampling, f"sample_{name}", make(order, ramp, heun)) - comfy.samplers.SAMPLER_NAMES.append(name) + for name, order, ramp in (("dflow_ord1", 1, False), ("dflow_ord2", 2, False), ("dflow_ord2_ab2", 2, True)): + if name not in comfy.samplers.SAMPLER_NAMES: + setattr(ksampling, f"sample_{name}", make(order, ramp)) + comfy.samplers.SAMPLER_NAMES.append(name) register() diff --git a/modules/sampling/distilled_flow/tests/test_distilled_flow.py b/modules/sampling/distilled_flow/tests/test_distilled_flow.py index 56f2545e..1741aa1f 100644 --- a/modules/sampling/distilled_flow/tests/test_distilled_flow.py +++ b/modules/sampling/distilled_flow/tests/test_distilled_flow.py @@ -19,7 +19,7 @@ def _model(x, sigma, **_): def test_the_presets_are_ordinary_comfy_sampler_names(): import comfy.samplers import modules.sampling.distilled_flow # noqa: F401 - for name in ("dflow_ord1", "dflow_ord2", "dflow_ord2_ramp", "dflow_ord1_heun2", "dflow_ord2_ramp_heun1"): + for name in ("dflow_ord1", "dflow_ord2", "dflow_ord2_ab2"): assert name in comfy.samplers.SAMPLER_NAMES assert comfy.samplers.sampler_object(name).sampler_function is not None @@ -34,24 +34,15 @@ def test_order_one_is_euler(): def test_order_two_differs_and_ramp_differs_again(): import modules.sampling.distilled_flow # noqa: F401 x, sigmas = torch.randn(1, 4, 2, 4, 4), torch.linspace(1.0, 0.0, 5) - one, two, ramp = (_line(_model, x, sigmas, n) for n in ("dflow_ord1", "dflow_ord2", "dflow_ord2_ramp")) + one, two, ramp = (_line(_model, x, sigmas, n) for n in ("dflow_ord1", "dflow_ord2", "dflow_ord2_ab2")) assert not torch.allclose(one, two) and not torch.allclose(two, ramp) -def test_heun_corrects_the_last_steps_with_one_more_call_each(): - import modules.sampling.distilled_flow # noqa: F401 - x, sigmas = torch.randn(1, 4, 2, 4, 4), torch.linspace(1.0, 0.0, 5) - for name, extra in (("dflow_ord1", 0), ("dflow_ord1_heun1", 1), ("dflow_ord1_heun2", 2)): - calls = [] - _line(lambda a, s, **k: calls.append(1) or _model(a, s), x, sigmas, name) - assert len(calls) == 4 + extra - - def test_the_sound_stream_stays_first_order(): from comfy.nested_tensor import NestedTensor import modules.sampling.distilled_flow # noqa: F401 video, audio = torch.randn(1, 4, 2, 4, 4), torch.randn(1, 8, 5) sigmas = torch.linspace(1.0, 0.0, 5) - both = _line(lambda x, s, **k: x * 0.5, NestedTensor((video, audio)), sigmas, "dflow_ord2_heun2") + both = _line(lambda x, s, **k: x * 0.5, NestedTensor((video, audio)), sigmas, "dflow_ord2_ab2") alone = _line(lambda x, s, **k: x * 0.5, audio, sigmas, "dflow_ord1") assert torch.allclose(both.unbind()[1], alone, atol=1e-5)