"""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, sysinfo from ..._core import control as control_mod, 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) def _first_step_clock(model): """-> (model clone carrying the clock, report). One log line saying where the time before step 1 goes: model preparation (loading / patching onto the GPU), the setup after it up to the first forward pass, and the first forward pass itself. ComfyUI's own wrapper hooks, on a clone, so the model is untouched.""" import time from comfy.patcher_extension import WrappersMP t = {"start": time.perf_counter()} def stamp(name): t.setdefault(name, time.perf_counter()) def prepare(executor, *a, **k): try: return executor(*a, **k) finally: stamp("prepared") def apply_model(executor, *a, **k): if "first" in t: return executor(*a, **k) stamp("first") try: return executor(*a, **k) finally: if comfy.model_management.get_torch_device().type == "cuda": import torch torch.cuda.synchronize() stamp("first_done") try: clone = model.clone() clone.add_wrapper_with_key(WrappersMP.PREPARE_SAMPLING, "funui.clock", prepare) clone.add_wrapper_with_key(WrappersMP.APPLY_MODEL, "funui.clock", apply_model) except Exception: # noqa: BLE001 -- a clock must never stop a run return model, lambda: None def report(): if not {"prepared", "first", "first_done"} <= t.keys(): return log.info("FunUI Timing", f"before step 1: model prepared in {t['prepared'] - t['start']:.1f}s, " f"then {t['first'] - t['prepared']:.1f}s of setup before the first forward pass, " f"which took {t['first_done'] - t['first']:.1f}s") return clone, report class FunUISampler(io.ComfyNode): @classmethod def define_schema(cls) -> io.Schema: samplers, schedulers = sampler_names(), scheduler_names() return io.Schema( node_id="FunUISampler", display_name="FunUI Sampler", category="FunUI/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("FUNUI_SETTINGS").Input("settings", optional=True), io.Sigmas.Input("sigmas", optional=True, tooltip="A schedule to run exactly as written. When wired, " "steps, scheduler and denoise are not used. For a " "second pass its first value is the strength: 0.4 " "polishes, 0.8 reworks the shot."), ], 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, sigmas=None) -> io.NodeOutput: if settings is not None and not isinstance(settings, dict): raise RuntimeError( f"FunUI Sampler: settings must be an object keyed by module id, " f"got {type(settings).__name__}.") from comfy_extras.nodes_custom_sampler import Noise_RandomNoise # A generation begins here, every time. Taking the id only when none # exists left every run after the first reporting as the first, because # the node that used to start them is cached away by ComfyUI. run = run_mod.start() log.info("FunUI Sampler", f"{run}: starting. {sysinfo.memory_line()}") # The record travels on the model, which ComfyUI may hand us again from # its cache -- so it is emptied here rather than trusted to be fresh. # Otherwise a modifier dropped in one generation stays dropped in every # later one, with nothing said. dropped = getattr(model, "model_options", {}).get("funui_dropped") if dropped is None: dropped = getattr(model, "funui_dropped", None) if dropped is None: dropped = patching.Dropped() dropped.clear() 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 wired = sigmas is not None if wired: if sigmas.numel() < 2: raise RuntimeError("FunUI Sampler: the wired schedule needs at least two " "values (a start and an end).") if bool((sigmas[1:] > sigmas[:-1]).any()): raise RuntimeError("FunUI Sampler: the wired schedule must run high to low.") else: 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) model, report = _first_step_clock(model) 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)) try: 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, ) except comfy.model_management.InterruptProcessingException: # Said, then re-raised. Cancelling is a normal thing to do and the # log should show it happening rather than a run simply stopping -- # but it is not ours to absorb: ComfyUI ends the queue item. log.alert("FunUI Sampler", f"{run} was cancelled part way through") raise # 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. report() 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 schedule = f"wired schedule from {float(sigmas[0]):.3g}" if wired else scheduler headline = (f"{run}: {sampler_name} / {schedule}, {len(sigmas) - 1} step(s), " f"cfg {cfg}") log.info("FunUI 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 fingerprint_inputs(cls, **_): """A module going into or out of quarantine changes what this run does.""" return "quarantine=" + control_mod.fingerprint() @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) compatible, switched = control_mod.partition(compatible, settings) ordered, rejected = relations_mod.order(compatible) guards_off = any(answer(settings or {}) for _spec, answer in registry.providers("guards_off")) values, problems = {}, [] for spec in ordered: clean, said = schema_mod.check_values( spec, (settings or {}).get(spec.id), keep_bad=guards_off) values[spec.id] = clean problems.extend(said) schema_mod.refuse_or_warn(problems, "FunUI Sampler", guards_off=guards_off) 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 + switched @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("FunUI 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") + "_funui" return comfy.samplers.KSAMPLER( with_modifiers, extra_options=getattr(sampler, "extra_options", {}), inpaint_options=getattr(sampler, "inpaint_options", {}))