nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Distilled Flow: three named presets (ord1, ord2, ord2_ab2), no Heun variants

Browse Source 
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Dy3QqNuYH9NSJgob83xcTX
This commit is contained in:
Claude
2026-10-20 16:11:19 +00:00
parent 73adf3e65d
commit 0dc605f013
3 changed files with 14 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+1 -1
View File
@@ -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
modules/sampling/distilled_flow/__init__.py
+10 -17
View File
@@ -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()
modules/sampling/distilled_flow/tests/test_distilled_flow.py
+3 -12
View File
@@ -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)