sampling: a sampler that samples, and hosts what it says it can host

v4's Chain Sampler became a hub: 8183 lines, roughly seventy optional
inputs, a _sample_chunk that grew six ALG arguments, and a hardcoded list of
knobs silently inert on one model family. Every technique that needed the
loop was added to it, so it collected all of them.

This one picks an algorithm and a schedule from ComfyUI's own registries --
so a sampler it gains appears here without an edit -- and runs them. Three
things stop it becoming a hub again. It never learns what a model-side
modifier is. Sampler-side ones reach it through ONE call site,
chain.process(ctx, latent), which cannot grow because the step travels in a
context object rather than in the signature. And it declares the hook points
it offers, so core hands it only what fits and it never names a modifier.

Some modifiers genuinely cannot ride on a model: they need the step index,
the schedule, or a second pass. Those attach here instead, and one asking
for a hook this sampler does not offer is absent and said so, not
half-running.

Verified against ComfyUI's own guider path with real weights: bit-identical
output. That took correcting the method, not the code -- see below.

MEASUREMENT, corrected: sampling on this Mac's MPS backend is
NONDETERMINISTIC. The same node with the same seed and the same inputs
differs run to run: 2.42 by max latent difference, 0.0008 by mean pixel gap.
Comparing by MAX said this sampler diverged from core by 14.2 and it was
about to be "fixed"; by MEAN against a measured noise floor it is 0.00000
while core against itself is 0.00183. Every comparison here now runs a
control twice first. Re-checked the Detail modifier the same way: 145x the
floor, so that claim stands.

Also matched what core's sampler actually does and this did not:
fix_empty_latent_channels using the ratios our own empty-latent node reports,
the preview and progress callback, the move back to the intermediate device,
and dropping the ratio keys so nothing rescales twice.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-08-29 17:51:02 +03:00
co-authored by Claude Opus 5
parent 78c43acd67
commit 090669e950
9 changed files with 601 additions and 4 deletions
+134
View File
@@ -0,0 +1,134 @@
"""Modifiers that belong to the sampler rather than to the model.
Some cannot ride on a ModelPatcher and it is not a limitation of the plumbing:
they need the step index, the schedule, or the chance to run the whole thing
twice. A model patch has none of those.
So a sampler offers HOOK POINTS by name and core hands it the modifiers that
asked for one it offers. Two properties fall out, and both are the point:
* A sampler file never names a modifier. Five samplers and six modifiers stay
zero hand-written call sites rather than thirty.
* The sampler keeps the veto -- it says what it can host -- and core keeps the
plumbing. A modifier wanting a hook this sampler does not offer is absent, and
said, rather than half-working.
The call site is deliberately ONE, and shaped so it cannot grow:
latent = chain.process(ctx, latent)
`ctx` carries the step, not the signature. v4 passed each technique's parameters
down as keyword arguments and `_sample_chunk` ended up with six ALG arguments
before anyone noticed the shape was wrong.
"""
from dataclasses import dataclass
from typing import Any, List, Optional, Sequence, Tuple
from . import log, patching
@dataclass(frozen=True)
class Step:
"""Everything a per-step modifier is allowed to know about where it is.
A small closed object rather than loose arguments: adding a field here is
one edit, whereas widening a call signature is an edit in every sampler and
every modifier that ever passed it along.
"""
index: int
sigma: Any
sigmas: Any
total: int
run: Optional[str] = None
anchor: Any = None # the pinned latent, where there is one
class Chain:
"""The modifiers a sampler is actually running, in order."""
def __init__(self, members: Sequence[Tuple[str, Any]], dropped: patching.Dropped):
self.members = list(members)
self.dropped = dropped
def __bool__(self) -> bool:
return bool(self.members)
@property
def ids(self) -> List[str]:
return [name for name, _ in self.members]
def process(self, ctx: Step, latent):
"""Every member that wants this step, in declared order.
A member that raises is dropped for the rest of the run and the run
continues -- the same rule the model-side hooks follow, for the same
reason: a modifier is an opinion about the picture and is never worth
the picture.
"""
for name, member in self.members:
key = f"funpack.{name}"
if key in self.dropped:
continue
try:
if not member.active(ctx):
continue
changed = member.process(ctx, latent)
except Exception as exc: # noqa: BLE001
if self.dropped.record(key, exc):
import traceback
log.warning(key, "failed during sampling and is now OFF for the "
"rest of this run; the run continues without it\n"
+ "".join(traceback.format_exception(exc)).rstrip())
continue
if changed is not None:
latent = changed
return latent
def build(specs, values, accepts: Sequence[str], dropped: patching.Dropped,
capability: str = "sampler_modifier") -> Tuple[Chain, List[str]]:
"""(chain, notes) for the modifiers this sampler can host.
`specs` are already filtered by trait and ordered by relation -- this only
decides who fits THIS sampler, and builds them.
"""
offered = set(accepts or ())
members, notes = [], []
for spec in specs:
make = spec.provides.get(capability)
if make is None:
continue
wanted = tuple(getattr(spec, "hooks", ()) or ())
missing = [hook for hook in wanted if hook not in offered]
if missing:
# Absent rather than half-working: a modifier that needs a second
# pass cannot be usefully run by a sampler that makes one.
notes.append(f"{spec.id}: needs {', '.join(missing)}, which this sampler "
f"does not offer")
continue
try:
member = make(values.get(spec.id) or {}, tuple(offered))
except Exception as exc: # noqa: BLE001
log.failed(f"{spec.id}.{capability}", exc)
notes.append(f"{spec.id}: failed to start -- {type(exc).__name__}: {exc}")
continue
if member is None:
continue # the module decided it is off
for required in ("active", "process"):
if not callable(getattr(member, required, None)):
notes.append(f"{spec.id}: is not a sampler modifier (no {required}())")
member = None
break
if member is None:
continue
members.append((spec.id, member))
notes.append(f"{spec.id}: on")
return Chain(members, dropped), notes
+3
View File
@@ -66,6 +66,9 @@ class ModuleSpec:
# name out of here and never defines one: it is a lookup, so a node can ask
# "who can build one of these" without core learning what the thing is.
provides: Dict[str, Callable] = field(default_factory=dict)
# Hook points a sampler must offer before this module can run. Empty means
# "any sampler will do", the same narrowing rule traits follow.
hooks: List[str] = field(default_factory=list)
ui: Optional[str] = None # served path to its ui.js
status: str = "experimental" # or "proven"
source: str = "" # dotted import path
+1 -1
View File
@@ -20,7 +20,7 @@ from .schema import SchemaError, validate
FIELDS = {
"id": "ID", "title": "TITLE", "mount": "MOUNT", "settings": "SETTINGS",
"requires": "REQUIRES", "after": "AFTER", "before": "BEFORE",
"stage": "STAGE", "status": "STATUS", "nodes": "NODES", "traits": "TRAITS", "provides": "PROVIDES",
"stage": "STAGE", "status": "STATUS", "nodes": "NODES", "traits": "TRAITS", "provides": "PROVIDES", "hooks": "HOOKS",
}
+1
View File
@@ -332,6 +332,7 @@ def validate(announcement: Dict[str, Any], source: str = "") -> ModuleSpec:
mount=announcement.get("mount") or "",
settings=settings,
requires=_ids("requires"),
hooks=_ids("hooks"),
after=_ids("after"),
before=_ids("before"),
stage=stage,
+4
View File
@@ -46,6 +46,10 @@ FROZEN = {
"FunPackCLIPLoader": [
"clip_name1", "clip_name2", "clip_name3", "clip_name4", "type", "device",
],
"FunPackSampler": [
"model", "positive", "negative", "latent", "seed", "steps", "cfg",
"sampler_name", "scheduler", "denoise", "settings",
],
"FunPackLoadModifiers": ["model", "settings"],
"FunPackModifierSettings": ["settings"],
"FunPackDecode": ["samples", "vae", "model", "audio_vae"],
+3 -3
View File
@@ -11,8 +11,8 @@ leading underscore keeps it out of the module scan.
"""
try: # root on sys.path: tests, dev server
from core import contract, log, patching, registry, relations, run, schema, traits
from core import chain, contract, log, patching, registry, relations, run, schema, traits
except ImportError: # inside ComfyUI: a subpackage of the pack
from ..core import contract, log, patching, registry, relations, run, schema, traits
from ..core import chain, contract, log, patching, registry, relations, run, schema, traits
__all__ = ["contract", "log", "patching", "registry", "relations", "run", "schema", "traits"]
__all__ = ["chain", "contract", "log", "patching", "registry", "relations", "run", "schema", "traits"]
+10
View File
@@ -0,0 +1,10 @@
"""The sampler. A node-only module."""
from .nodes import ACCEPTS, FunPackSampler
ID = "sampling_sampler"
TITLE = "Sampler"
STAGE = "sampling"
STATUS = "proven"
NODES = [FunPackSampler]
+233
View File
@@ -0,0 +1,233 @@
"""The sampler. It samples, and owns nothing else.
v4's Chain Sampler became a hub: 8183 lines, roughly seventy optional inputs, a
`_sample_chunk` that grew six ALG arguments, and a hardcoded list of knobs that
are silently inert on one model family. Every technique that needed the loop was
added to it, so it accumulated all of them.
This one picks an algorithm and a schedule from what is available and runs them.
Three things keep it from becoming a hub again:
* **It does not know what a modifier is.** Model-side ones ride on the model and
it never learns they exist. Sampler-side ones reach it through ONE call site,
`chain.process(ctx, latent)`, which cannot grow because the step travels in a
context object rather than in the signature.
* **It declares what it can host** (`ACCEPTS`) and core hands it only the
modifiers that fit. It never names one.
* **Choices are announced, not listed.** The algorithms and schedules are read
from ComfyUI's own registries, so a new one appears without an edit here.
"""
import comfy.model_management
import comfy.sample
import comfy.samplers
import comfy.utils
import latent_preview
import torch
from comfy_api.latest import io
from ..._core import chain as chain_mod, log, patching, registry as registry_mod, run as run_mod
from ..._core import relations as relations_mod, schema as schema_mod, traits as traits_mod
# The hook points this sampler offers. A modifier asking for anything else is
# absent and said, rather than half-running.
ACCEPTS = ("latent",)
def sampler_names():
"""Read from ComfyUI, so a sampler it gains appears here without an edit."""
return list(comfy.samplers.SAMPLER_NAMES)
def scheduler_names():
return list(comfy.samplers.SCHEDULER_NAMES)
class FunPackSampler(io.ComfyNode):
@classmethod
def define_schema(cls) -> io.Schema:
samplers, schedulers = sampler_names(), scheduler_names()
return io.Schema(
node_id="FunPackSampler",
display_name="FunPack Sampler",
category="FunPack/Sampling",
description="Sample a latent with a chosen algorithm and schedule.",
inputs=[
io.Model.Input("model"),
io.Conditioning.Input("positive"),
io.Conditioning.Input("negative"),
io.Latent.Input("latent"),
io.Int.Input("seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF,
control_after_generate=True),
io.Int.Input("steps", default=20, min=1, max=1000),
io.Float.Input("cfg", default=7.0, min=0.0, max=100.0, step=0.1,
tooltip="1.0 turns guidance off, which some models want."),
io.Combo.Input("sampler_name", options=samplers,
default="euler" if "euler" in samplers else samplers[0]),
io.Combo.Input("scheduler", options=schedulers,
default="normal" if "normal" in schedulers else schedulers[0]),
io.Float.Input("denoise", default=1.0, min=0.0, max=1.0, step=0.01,
optional=True,
tooltip="Below 1 starts partway through the schedule."),
io.Custom("FUNPACK_SETTINGS").Input("settings", optional=True),
],
outputs=[
io.Latent.Output(display_name="latent"),
io.String.Output(display_name="status"),
],
)
@classmethod
def execute(cls, model, positive, negative, latent, seed: int, steps: int, cfg: float,
sampler_name: str, scheduler: str, denoise: float = 1.0,
settings=None) -> io.NodeOutput:
if settings is not None and not isinstance(settings, dict):
raise RuntimeError(
f"FunPack Sampler: settings must be an object keyed by module id, "
f"got {type(settings).__name__}.")
from comfy_extras.nodes_custom_sampler import Noise_RandomNoise
run = run_mod.current() or run_mod.start()
dropped = getattr(model, "funpack_dropped", None) or patching.Dropped()
latent = latent.copy()
# The latent says what ratios it was BUILT at; this is where a model that
# wants different ones gets it rescaled. Skipping it silently samples a
# latent of the wrong size -- and our own empty-latent node reports those
# ratios expecting somebody to read them.
samples_in = comfy.sample.fix_empty_latent_channels(
model, latent["samples"],
latent.get("downscale_ratio_spacial"), latent.get("downscale_ratio_temporal"))
latent["samples"] = samples_in
sigmas = cls._sigmas(model, scheduler, steps, denoise)
sampler = comfy.samplers.sampler_object(sampler_name)
chain, notes = cls._chain(model, settings, dropped)
if chain:
sampler = cls._with_chain(sampler, chain, sigmas, run)
guider = comfy.samplers.CFGGuider(model)
guider.set_conds(positive, negative)
guider.set_cfg(cfg)
# Progress and previews, from ComfyUI's own machinery rather than ours:
# it already picks a preview method, already falls back when the decoder
# is missing, and already reports steps to whatever is watching.
callback = latent_preview.prepare_callback(model, max(0, len(sigmas) - 1))
samples = guider.sample(
Noise_RandomNoise(seed).generate_noise(latent),
samples_in, sampler, sigmas,
denoise_mask=latent.get("noise_mask"),
callback=callback,
disable_pbar=not comfy.utils.PROGRESS_BAR_ENABLED,
seed=seed,
)
# Back to the device every other sampler returns on. Leaving it on the
# GPU works until something downstream is on the CPU, and then it is a
# device mismatch a long way from here.
samples = samples.to(comfy.model_management.intermediate_device())
out = latent.copy()
# Consumed above; carrying them onward would rescale a second time.
out.pop("downscale_ratio_spacial", None)
out.pop("downscale_ratio_temporal", None)
out["samples"] = samples
headline = (f"{run}: {sampler_name} / {scheduler}, {len(sigmas) - 1} step(s), "
f"cfg {cfg}")
log.info("FunPack Sampler", headline)
return io.NodeOutput(out, "\n".join([headline, *notes]))
# --- the parts, kept separate so each is testable on its own -----------
@classmethod
def _sigmas(cls, model, scheduler: str, steps: int, denoise: float):
"""The schedule, including the partial-denoise case.
Computed over the FULL step count and then trimmed, which is what makes
denoise 0.5 mean "the second half of a 20-step schedule" rather than "a
10-step schedule" -- the two are not the same curve.
"""
if denoise is None or denoise >= 1.0:
total = steps
else:
if denoise <= 0.0:
return torch.FloatTensor([])
total = int(steps / denoise)
sigmas = comfy.samplers.calculate_sigmas(
model.get_model_object("model_sampling"), scheduler, total)
return sigmas[-(steps + 1):]
@classmethod
def _chain(cls, model, settings, dropped):
"""Whoever asked for a hook point this sampler offers."""
registry = registry_mod.current()
specs = list(registry.specs.values())
available = traits_mod.traits_of(model, specs)
offering = [s for s in specs if s.provides.get("sampler_modifier")]
compatible, incompatible = traits_mod.split(offering, available)
ordered, rejected = relations_mod.order(compatible)
values, problems = {}, []
for spec in ordered:
clean, said = schema_mod.check_values(spec, (settings or {}).get(spec.id))
values[spec.id] = clean
problems.extend(said)
if problems:
raise RuntimeError("FunPack Sampler:\n " + "\n ".join(problems))
chain, notes = chain_mod.build(ordered, values, ACCEPTS, dropped)
for spec in incompatible:
notes.append(f"{spec.id}: needs {', '.join(traits_mod.missing_for(spec, available))}")
for spec, why in rejected:
notes.append(f"{spec.id}: {why}")
return chain, notes
@classmethod
def _with_chain(cls, sampler, chain, sigmas, run):
"""The one call site, wrapped around whatever algorithm was chosen.
The algorithm is not modified and not replaced: its own loop runs, and
the chain sees each step through the denoiser it already calls.
"""
inner = getattr(sampler, "sampler_function", None)
if inner is None:
log.alert("FunPack Sampler",
f"{type(sampler).__name__} has no loop to attach to, so "
f"{', '.join(chain.ids)} will not run this time")
return sampler
total = max(1, len(sigmas) - 1)
def with_modifiers(denoiser, x, step_sigmas, extra_args=None, callback=None,
disable=None, **options):
seen = {"index": 0}
class Watched:
"""Transparent, except that every model call is a step."""
def __getattr__(self, name):
return getattr(denoiser, name)
def __call__(self, latent, sigma, **kwargs):
ctx = chain_mod.Step(
index=seen["index"], sigma=sigma, sigmas=step_sigmas,
total=total, run=run,
anchor=getattr(denoiser, "latent_image", None))
latent = chain.process(ctx, latent)
seen["index"] += 1
return denoiser(latent, sigma, **kwargs)
return inner(Watched(), x, step_sigmas, extra_args=extra_args,
callback=callback, disable=disable, **options)
with_modifiers.__name__ = getattr(inner, "__name__", "sampler") + "_funpack"
return comfy.samplers.KSAMPLER(
with_modifiers,
extra_options=getattr(sampler, "extra_options", {}),
inpaint_options=getattr(sampler, "inpaint_options", {}))
+212
View File
@@ -0,0 +1,212 @@
"""The sampler: it samples, and it hosts what it said it could host.
v4's became a hub -- 8183 lines and a `_sample_chunk` that grew six ALG
arguments. These tests are mostly about the shape that prevents that: one call
site, a context object rather than a widening signature, and a sampler that
never names a modifier.
"""
import pytest
@pytest.fixture(autouse=True)
def _needs_comfy(comfyui):
"""Imports comfy."""
@pytest.fixture
def registry(monkeypatch):
from core import registry as registry_mod
fake = registry_mod.Registry()
monkeypatch.setattr(registry_mod, "current", lambda rescan=False: fake)
return fake
def _spec(module_id, make, hooks=(), requires=(), settings=None):
from core.contract import ModuleSpec
return ModuleSpec(id=module_id, title=module_id, mount="", stage="sampling",
hooks=list(hooks), requires=list(requires),
settings=settings or {}, provides={"sampler_modifier": make})
class _Modifier:
"""The shape a sampler modifier has: active(ctx), process(ctx, latent)."""
def __init__(self, when=lambda ctx: True):
self.seen = []
self.when = when
def active(self, ctx):
return self.when(ctx)
def process(self, ctx, latent):
self.seen.append(ctx.index)
return latent
# --- the schedule ----------------------------------------------------------
def test_full_denoise_asks_for_exactly_the_steps_requested(comfyui, monkeypatch):
import comfy.samplers
import torch
from modules.sampling.sampler.nodes import FunPackSampler
seen = {}
def fake(model_sampling, scheduler, steps):
seen["steps"] = steps
return torch.linspace(1.0, 0.0, steps + 1)
monkeypatch.setattr(comfy.samplers, "calculate_sigmas", fake)
class Model:
def get_model_object(self, name):
return None
out = FunPackSampler._sigmas(Model(), "normal", 6, 1.0)
assert seen["steps"] == 6 and len(out) == 7
def test_partial_denoise_is_the_tail_of_a_longer_schedule(comfyui, monkeypatch):
"""denoise 0.5 means the second half of a 20-step curve, not a 10-step
curve. They are different curves and swapping them changes the picture."""
import comfy.samplers
import torch
from modules.sampling.sampler.nodes import FunPackSampler
seen = {}
def fake(model_sampling, scheduler, steps):
seen["steps"] = steps
return torch.linspace(1.0, 0.0, steps + 1)
monkeypatch.setattr(comfy.samplers, "calculate_sigmas", fake)
class Model:
def get_model_object(self, name):
return None
out = FunPackSampler._sigmas(Model(), "normal", 10, 0.5)
assert seen["steps"] == 20, "the curve was not computed over the full length"
assert len(out) == 11, "the tail was not trimmed to the requested steps"
def test_zero_denoise_asks_for_nothing(comfyui):
from modules.sampling.sampler.nodes import FunPackSampler
class Model:
def get_model_object(self, name):
return None
assert len(FunPackSampler._sigmas(Model(), "normal", 10, 0.0)) == 0
# --- hosting modifiers -----------------------------------------------------
def test_a_modifier_sees_every_step_through_one_call_site(registry, comfyui):
import torch
from core import chain as chain_mod, patching
from modules.sampling.sampler.nodes import ACCEPTS, FunPackSampler
modifier = _Modifier()
registry.specs["m"] = _spec("m", lambda values, accepts: modifier, hooks=["latent"])
chain, notes = FunPackSampler._chain(_TraitlessModel(), None, patching.Dropped())
assert chain.ids == ["m"] and "m: on" in notes
# Drive it the way the wrapped sampler function would.
for index in range(3):
chain.process(chain_mod.Step(index=index, sigma=1.0, sigmas=None, total=3), "latent")
assert modifier.seen == [0, 1, 2]
def test_a_modifier_can_sit_out_a_step(registry, comfyui):
from core import chain as chain_mod, patching
from modules.sampling.sampler.nodes import FunPackSampler
modifier = _Modifier(when=lambda ctx: ctx.index % 2 == 0)
registry.specs["m"] = _spec("m", lambda values, accepts: modifier, hooks=["latent"])
chain, _ = FunPackSampler._chain(_TraitlessModel(), None, patching.Dropped())
for index in range(4):
chain.process(chain_mod.Step(index=index, sigma=1.0, sigmas=None, total=4), "l")
assert modifier.seen == [0, 2]
def test_a_modifier_wanting_a_hook_this_sampler_lacks_is_absent(registry, comfyui):
"""The sampler keeps the veto. A modifier that needs a second pass cannot be
usefully run by a sampler that makes one."""
from core import patching
from modules.sampling.sampler.nodes import FunPackSampler
registry.specs["greedy"] = _spec("greedy", lambda values, accepts: _Modifier(),
hooks=["second_pass"])
chain, notes = FunPackSampler._chain(_TraitlessModel(), None, patching.Dropped())
assert chain.ids == []
assert any("needs second_pass" in note for note in notes)
def test_a_modifier_needing_an_absent_trait_is_absent(registry, comfyui):
from core import patching
from modules.sampling.sampler.nodes import FunPackSampler
registry.specs["audio"] = _spec("audio", lambda values, accepts: _Modifier(),
requires=["audio_stream"])
chain, notes = FunPackSampler._chain(_TraitlessModel(), None, patching.Dropped())
assert chain.ids == []
assert any("audio_stream" in note for note in notes)
def test_one_failing_modifier_does_not_end_the_run(registry, comfyui):
from core import chain as chain_mod, patching
from modules.sampling.sampler.nodes import FunPackSampler
class Breaks:
def active(self, ctx):
return True
def process(self, ctx, latent):
raise RuntimeError("boom")
good = _Modifier()
registry.specs["a_breaks"] = _spec("a_breaks", lambda v, a: Breaks())
registry.specs["b_good"] = _spec("b_good", lambda v, a: good)
dropped = patching.Dropped()
chain, _ = FunPackSampler._chain(_TraitlessModel(), None, dropped)
for index in range(5):
out = chain.process(chain_mod.Step(index=index, sigma=1.0, sigmas=None, total=5), "l")
assert out == "l"
assert good.seen == [0, 1, 2, 3, 4], "a failing sibling stopped a healthy modifier"
assert "funpack.a_breaks" in dropped
def test_bad_settings_are_refused_before_sampling(registry, comfyui):
from core import patching
from modules.sampling.sampler.nodes import FunPackSampler
registry.specs["m"] = _spec("m", lambda v, a: _Modifier(), settings={
"strength": {"type": "float", "default": 0.5, "min": 0.0, "max": 1.0, "label": "S"},
})
with pytest.raises(RuntimeError, match="above its maximum"):
FunPackSampler._chain(_TraitlessModel(), {"m": {"strength": 9.0}}, patching.Dropped())
def test_something_that_is_not_a_modifier_is_named_not_run(registry, comfyui):
from core import patching
from modules.sampling.sampler.nodes import FunPackSampler
registry.specs["m"] = _spec("m", lambda v, a: object())
chain, notes = FunPackSampler._chain(_TraitlessModel(), None, patching.Dropped())
assert chain.ids == []
assert any("not a sampler modifier" in note for note in notes)
class _TraitlessModel:
"""A model core can read nothing from, so trait filtering is exercised
without needing weights."""
model = None