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
co-authored by
Claude Opus 5
Claude Opus 5
parent
78c43acd67
commit
090669e950
@@ -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]
|
||||
@@ -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", {}))
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user