diff --git a/core/chain.py b/core/chain.py new file mode 100644 index 00000000..ced7a4fb --- /dev/null +++ b/core/chain.py @@ -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 diff --git a/core/contract.py b/core/contract.py index b9269c62..e0eb82e1 100644 --- a/core/contract.py +++ b/core/contract.py @@ -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 diff --git a/core/registry.py b/core/registry.py index 8ef39e56..e6b64a6d 100644 --- a/core/registry.py +++ b/core/registry.py @@ -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", } diff --git a/core/schema.py b/core/schema.py index 58cbf5b6..b5de7184 100644 --- a/core/schema.py +++ b/core/schema.py @@ -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, diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index c852c12a..00ec4ede 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -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"], diff --git a/modules/_core.py b/modules/_core.py index debdf00a..070f1273 100644 --- a/modules/_core.py +++ b/modules/_core.py @@ -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"] diff --git a/modules/sampling/sampler/__init__.py b/modules/sampling/sampler/__init__.py new file mode 100644 index 00000000..7bef766b --- /dev/null +++ b/modules/sampling/sampler/__init__.py @@ -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] diff --git a/modules/sampling/sampler/nodes.py b/modules/sampling/sampler/nodes.py new file mode 100644 index 00000000..e2be63fa --- /dev/null +++ b/modules/sampling/sampler/nodes.py @@ -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", {})) diff --git a/modules/sampling/tests/test_sampler.py b/modules/sampling/tests/test_sampler.py new file mode 100644 index 00000000..bb568f3f --- /dev/null +++ b/modules/sampling/tests/test_sampler.py @@ -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