""" SigmaCurves -- step-locked sigma editor. The frontend plots one point per sampling step (steps + 1 points total), with the y-axis representing the normalized sigma at that step (1.0 = sigma_max, 0.0 = sigma_min). Picking a base scheduler seeds the plot with that scheduler's natural shape (fetched from the preview endpoint below); the user then drags individual y-values up or down, or selects a range of steps and applies a curve type to re-fill that range. The node accepts: model, scheduler, steps, denoise, curve_data (JSON of {values, ...}). When ``curve_data`` is empty / malformed the node falls back to the real scheduler output, so a fresh node behaves like a regular scheduler until the user starts editing. """ import json import logging import os from typing import Any, Dict, List import torch import comfy.samplers as comfy_samplers def _get_basic_scheduler(): """Return ComfyUI's stock BasicScheduler instance, or None if the expected module isn't available. Same node KSampler-side users would drop on the canvas; calling its ``get_sigmas(model, scheduler, steps, denoise)`` returns the canonical sigmas tensor for the chosen model + scheduler.""" try: from comfy_extras.nodes_custom_sampler import BasicScheduler return BasicScheduler() except Exception: # noqa: BLE001 return None _BASIC_SCHEDULER = _get_basic_scheduler() # --------------------------------------------------------------------- # On-demand model loading via the connected loader's own node class. # # Lets the preview endpoint resolve "the model connected to SigmaCurves' # input" without waiting for workflow execution. The frontend walks the # graph back to the model loader, sends its node type + widget values, # and the backend instantiates that node class, calls its FUNCTION, and # extracts the resulting MODEL output. Result is session-cached so we # don't re-load on every preview request. # --------------------------------------------------------------------- # (loader_type, widgets_tuple) -> ModelPatcher _LOADER_PATCHER_CACHE: Dict[Any, Any] = {} def _resolve_node_class(loader_type: str): try: import nodes as comfy_nodes ncm = getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", None) if isinstance(ncm, dict) and loader_type in ncm: return ncm[loader_type] except Exception: # noqa: BLE001 return None return None def _load_model_via_loader(loader_type: str, widgets_values: list): """Instantiate the named loader node class, call its FUNCTION with widget values mapped by INPUT_TYPES order, and return the MODEL output. Caches the resulting patcher so subsequent calls are instant for the same (loader_type, widgets) tuple.""" key = (loader_type, tuple(widgets_values or [])) if key in _LOADER_PATCHER_CACHE: return _LOADER_PATCHER_CACHE[key] cls = _resolve_node_class(loader_type) if cls is None: return None try: instance = cls() func_name = getattr(cls, "FUNCTION", None) if not func_name: return None func = getattr(instance, func_name, None) if not callable(func): return None # Map widgets_values to kwargs in INPUT_TYPES order. Required # inputs first, then optional. Skips socket-typed inputs. try: input_types = cls.INPUT_TYPES() except Exception: # noqa: BLE001 input_types = {"required": {}} kwargs = {} wi = 0 for section in ("required", "optional"): for name, spec in (input_types.get(section, {}) or {}).items(): # Only widget-style inputs map to widgets_values; socket # inputs (MODEL, CLIP, ...) are uppercase tuple types. if isinstance(spec, tuple) and len(spec) >= 1: t = spec[0] is_widget = isinstance(t, list) or t in ( "INT", "FLOAT", "STRING", "BOOLEAN", "BOOL") if not is_widget: continue if wi < len(widgets_values): kwargs[name] = widgets_values[wi] wi += 1 result = func(**kwargs) # Extract MODEL from the return tuple by RETURN_TYPES position. return_types = getattr(cls, "RETURN_TYPES", ()) if isinstance(result, tuple): for i, t in enumerate(return_types): if t == "MODEL" and i < len(result): _LOADER_PATCHER_CACHE[key] = result[i] return result[i] return None except Exception as e: # noqa: BLE001 LOGGER.warning("Could not load model via %s: %s", loader_type, e) return None LOGGER = logging.getLogger("SigmaCurves") # --------------------------------------------------------------------- # Runtime sigma cache. # # The preview endpoint defaults to a synthetic ModelSamplingDiscrete so # the canvas widget can show baseline shapes without a loaded model. # That works for model-agnostic schedulers (karras, exponential), but # schedulers whose shape depends on the model's timestep table -- such # as RES4LYF's bong_tangent / beta57 / FlowMatch variants -- look # different against the synthetic model than against the user's real # model. # # To fix the visual mismatch, ``SigmaCurves.build`` writes the real # sigmas it just computed into this cache; the preview endpoint then # prefers cached real-model values when available. Effect: after the # user runs the workflow once, the preview snaps to the true shape on # the next refresh. # --------------------------------------------------------------------- # Map (scheduler, steps) -> {"values": [...], "trailing_zero": bool, # "sigma_min": float, "sigma_max": float}. _REAL_SIGMA_CACHE: Dict[Any, Dict[str, Any]] = {} _CACHE_FILE = os.path.join( os.path.dirname(__file__), "presets", "real_sigma_cache.json") def _persist_cache_to_disk() -> None: try: os.makedirs(os.path.dirname(_CACHE_FILE), exist_ok=True) # Convert tuple keys to "scheduler:steps" strings for JSON. serializable = {f"{k[0]}|{k[1]}": v for k, v in _REAL_SIGMA_CACHE.items()} tmp = _CACHE_FILE + ".tmp" with open(tmp, "w", encoding="utf-8") as fh: json.dump(serializable, fh) os.replace(tmp, _CACHE_FILE) except Exception: # noqa: BLE001 LOGGER.debug("Could not persist sigma cache.", exc_info=True) def _load_cache_from_disk() -> None: if not os.path.exists(_CACHE_FILE): return try: with open(_CACHE_FILE, "r", encoding="utf-8") as fh: data = json.load(fh) if not isinstance(data, dict): return for k, v in data.items(): if "|" not in k: continue sched, steps_s = k.rsplit("|", 1) try: steps = int(steps_s) except ValueError: continue _REAL_SIGMA_CACHE[(sched, steps)] = v if _REAL_SIGMA_CACHE: LOGGER.info( "SigmaCurves: loaded %d cached real-model sigmas from disk.", len(_REAL_SIGMA_CACHE)) except Exception: # noqa: BLE001 LOGGER.debug("Could not load sigma cache from disk.", exc_info=True) def _notify_frontend(scheduler: str, steps: int) -> None: """WebSocket-push a 'sigmas updated' event so any open SigmaCurves node can re-fetch its preview and snap to the real-model shape.""" try: from server import PromptServer PromptServer.instance.send_sync( "res4sho.sigmas_updated", {"scheduler": scheduler, "steps": int(steps)}, ) except Exception: # noqa: BLE001 LOGGER.debug("Could not send sigmas_updated event.", exc_info=True) def _cache_real_sigmas(scheduler: str, steps: int, sigmas, sigma_min: float, sigma_max: float) -> None: """Stash the actual real-model sigmas the SigmaCurves node just computed so the preview endpoint can match what the workflow actually runs.""" try: sigmas_list = sigmas.cpu().tolist() if hasattr(sigmas, "cpu") \ else list(sigmas) except Exception: # noqa: BLE001 return if not sigmas_list: return trailing = (len(sigmas_list) >= 2 and abs(sigmas_list[-1]) <= 1e-6) non_term = sigmas_list[:-1] if trailing else sigmas_list if not non_term or max(non_term) - min(non_term) < 1e-9: return hi = max(non_term) lo = min(non_term) denom = max(hi - lo, 1e-9) values = [] for i, s in enumerate(sigmas_list): if trailing and i == len(sigmas_list) - 1 and abs(s) <= 1e-6: values.append(0.0) else: v = (s - lo) / denom values.append(max(0.0, min(1.0, v))) _REAL_SIGMA_CACHE[(scheduler, int(steps))] = { "values": values, "raw_sigmas": sigmas_list, "trailing_zero": trailing, "sigma_min": float(sigma_min), "sigma_max": float(sigma_max), "source": "real_model", } _persist_cache_to_disk() _notify_frontend(scheduler, int(steps)) _load_cache_from_disk() # --------------------------------------------------------------------- # Server endpoint -- supplies the frontend with normalized previews # of any scheduler's shape so the canvas widget can populate the # per-step y values without needing the user's actual model. # --------------------------------------------------------------------- def _list_schedulers() -> List[str]: names = getattr(comfy_samplers, "SCHEDULER_NAMES", None) if isinstance(names, (list, tuple)) and len(names) > 0: return list(names) KSampler = getattr(comfy_samplers, "KSampler", None) if KSampler is not None and hasattr(KSampler, "SCHEDULERS"): return list(getattr(KSampler, "SCHEDULERS")) return ["normal", "karras", "exponential", "simple", "sgm_uniform", "ddim_uniform", "beta"] # Cache the RES4SHO scheduler dict at import time so we have a guaranteed # direct-dispatch path even if comfy's SCHEDULER_HANDLERS missed our # entries (on some versions the registration in sampling.py is a no-op # because the handlers dict is empty at the moment we try to register). _RES4SHO_SCHEDULERS = {} try: from . import sampling as _sampling_mod if hasattr(_sampling_mod, "_SCHEDULERS"): _RES4SHO_SCHEDULERS = dict(_sampling_mod._SCHEDULERS) LOGGER.info("SigmaCurves: cached %d RES4SHO schedulers for direct " "dispatch (%s).", len(_RES4SHO_SCHEDULERS), ", ".join(sorted(_RES4SHO_SCHEDULERS.keys()))) else: LOGGER.warning("SigmaCurves: RES4SHO sampling module has no " "_SCHEDULERS dict; preview will rely on comfy's " "calculate_sigmas only.") except Exception as _e: # noqa: BLE001 LOGGER.warning("SigmaCurves: could not cache RES4SHO schedulers: %s " "(preview will use comfy's calculate_sigmas only).", _e) def _res4lyf_special_sigmas(scheduler: str, model_sampling, steps: int, denoise: float = 1.0): """Compute sigmas for RES4LYF schedulers that are *not* registered in ``comfy.samplers.SCHEDULER_HANDLERS``. RES4LYF lists ``beta57`` in its custom dropdown (and uses it as the default for many of its nodes) but never registers a handler -- it short-circuits inside its own ``get_sigmas`` to ``beta_scheduler(model_sampling, steps, alpha=0.5, beta=0.7)``. Calling ``calculate_sigmas`` for ``beta57`` therefore raises, which is what made our preview drop to a linear-ramp stub. Returns a sigmas tensor (with denoise crop applied like BasicScheduler does) or ``None`` if *scheduler* is not a special. """ if scheduler != "beta57": return None beta_fn = getattr(comfy_samplers, "beta_scheduler", None) if beta_fn is None: return None if denoise >= 1.0 or denoise <= 0.0: return beta_fn(model_sampling, steps, alpha=0.5, beta=0.7) total_steps = int(steps / max(denoise, 1e-4)) sigmas = beta_fn(model_sampling, total_steps, alpha=0.5, beta=0.7) return sigmas[-(steps + 1):] def _compute_scheduler_sigmas(model_sampling, scheduler: str, steps: int): """Get a sigmas list for *scheduler* + *steps*. Tries the cached RES4SHO ``_SCHEDULERS`` dict first so our schedulers work regardless of whether comfy's ``SCHEDULER_HANDLERS`` picked them up. Then handles RES4LYF specials (``beta57``). Falls back to comfy's ``calculate_sigmas`` otherwise. Returns ``(sigmas_list, dispatch_label)`` -- a tuple so the endpoint can surface which path was used. ``sigmas_list`` is ``None`` only when every path failed. """ if scheduler in _RES4SHO_SCHEDULERS: try: sigmas = _RES4SHO_SCHEDULERS[scheduler](model_sampling, steps) if hasattr(sigmas, "cpu"): return sigmas.cpu().tolist(), "res4sho_direct" return list(sigmas), "res4sho_direct" except Exception as e: # noqa: BLE001 LOGGER.exception( "SigmaCurves: RES4SHO scheduler '%s' raised: %s", scheduler, e) return None, f"res4sho_error: {e!r}" try: special = _res4lyf_special_sigmas(scheduler, model_sampling, steps) if special is not None: return special.cpu().tolist(), "res4lyf_special" except Exception as e: # noqa: BLE001 LOGGER.warning( "SigmaCurves: RES4LYF special '%s' raised: %s", scheduler, e) return None, f"res4lyf_special_error: {e!r}" try: sigmas = comfy_samplers.calculate_sigmas( model_sampling, scheduler, steps) return sigmas.cpu().tolist(), "comfy_calculate_sigmas" except Exception as e: # noqa: BLE001 LOGGER.warning( "SigmaCurves: calculate_sigmas failed for '%s': %s", scheduler, e) return None, f"comfy_error: {e!r}" def _get_loaded_model_patcher(): """Return any ModelPatcher already loaded in this ComfyUI session. The preview endpoint hands this directly to ``BasicScheduler``.""" try: import comfy.model_management as mm except ImportError: return None candidates = getattr(mm, "current_loaded_models", None) or [] for loaded in candidates: for attr in ("model", "real_model", "model_patcher"): patcher = getattr(loaded, attr, None) if patcher is None: continue getter = getattr(patcher, "get_model_object", None) if not callable(getter): continue try: ms = getter("model_sampling") except Exception: # noqa: BLE001 continue if ms is not None and hasattr(ms, "sigma_min") \ and hasattr(ms, "sigma_max"): return patcher return None class _BareModelSampling: """Minimal fallback when ``ModelSamplingDiscrete`` cannot be instantiated. Covers the only attributes the RES4SHO schedulers actually consult. """ sigma_max = 14.61 sigma_min = 0.0292 sigma_data = 1.0 _SYNTH_MS = None def _synthetic_model_sampling(): """Lazy-create a default model_sampling object that the preview endpoint can hand to schedulers without a loaded model. Tries comfy's ``ModelSamplingDiscrete`` first (gives sigmas table for schedulers that need it), falls back to a hand-rolled stub with SDXL-typical sigma_min / sigma_max scalars. """ global _SYNTH_MS if _SYNTH_MS is not None: return _SYNTH_MS try: from comfy.model_sampling import ModelSamplingDiscrete ms = ModelSamplingDiscrete(model_config=None) if not (hasattr(ms, "sigma_max") and hasattr(ms, "sigma_min")): raise AttributeError("missing sigma_min/sigma_max") _SYNTH_MS = ms LOGGER.info( "SigmaCurves: synthetic ModelSamplingDiscrete ready " "(sigma_max=%.4f, sigma_min=%.4f)", float(ms.sigma_max), float(ms.sigma_min)) except Exception as e: # noqa: BLE001 LOGGER.warning( "SigmaCurves: ModelSamplingDiscrete unavailable (%s); " "using bare fallback (sigma_max=%.2f, sigma_min=%.4f).", e, _BareModelSampling.sigma_max, _BareModelSampling.sigma_min) _SYNTH_MS = _BareModelSampling() return _SYNTH_MS def _register_routes(): try: from server import PromptServer from aiohttp import web except ImportError: LOGGER.info("PromptServer / aiohttp unavailable; SigmaCurves preview " "endpoint disabled.") return if getattr(PromptServer, "_res4sho_sigma_curves_route", False): return # idempotent @PromptServer.instance.routes.post("/RES4SHO/sigma_curves/preview_for_loader") async def preview_for_loader(request): """Frontend POSTs: { loader_type: "CheckpointLoaderSimple", widgets_values: ["model.safetensors"], scheduler: "beta57", steps: 20 } We instantiate the loader, get the MODEL, run BasicScheduler, return normalized sigmas. """ try: body = await request.json() except Exception: # noqa: BLE001 return web.json_response( {"error": "invalid JSON body"}, status=400) loader_type = body.get("loader_type", "") widgets_values = body.get("widgets_values") or [] scheduler = body.get("scheduler", "normal") try: steps = max(1, min(1000, int(body.get("steps", 20)))) except (ValueError, TypeError): steps = 20 if not loader_type: return web.json_response( {"error": "loader_type required"}, status=400) patcher = _load_model_via_loader(loader_type, widgets_values) if patcher is None: return web.json_response( {"error": f"could not resolve a MODEL via " f"{loader_type!r} -- the loader class is " f"unknown to ComfyUI or its FUNCTION did " f"not return a MODEL output."}, status=404) bs = _BASIC_SCHEDULER or _get_basic_scheduler() ms = patcher.get_model_object("model_sampling") try: # RES4LYF specials (beta57) aren't registered in # SCHEDULER_HANDLERS, so BasicScheduler would raise. special = _res4lyf_special_sigmas(scheduler, ms, steps) if special is not None: sigmas_tensor = special.cpu() elif bs is None: return web.json_response( {"error": "BasicScheduler unavailable"}, status=500) else: sigmas_tensor = bs.get_sigmas( patcher, scheduler, steps, 1.0)[0].cpu() except Exception as e: # noqa: BLE001 LOGGER.error("BasicScheduler call failed: %s", e, exc_info=True) return web.json_response( {"error": f"BasicScheduler failed: {e!r}"}, status=500) sigma_min = float(ms.sigma_min) sigma_max = float(ms.sigma_max) _cache_real_sigmas(scheduler, steps, sigmas_tensor, sigma_min, sigma_max) cached = _REAL_SIGMA_CACHE.get((scheduler, int(steps))) if cached is None: return web.json_response( {"error": "cache write failed"}, status=500) return web.json_response({ "values": cached["values"], "raw_sigmas": cached.get("raw_sigmas"), "trailing_zero": cached.get("trailing_zero", True), "dispatch": "real_model_live", "sigma_min": sigma_min, "sigma_max": sigma_max, }) @PromptServer.instance.routes.get("/RES4SHO/sigma_curves/preview") async def get_preview(request): try: scheduler = request.query.get("scheduler", "normal") try: steps = int(request.query.get("steps", 20)) except (ValueError, TypeError): steps = 20 steps = max(1, min(1000, steps)) # Priority order for sourcing sigmas: # 1. Cache populated by a prior SigmaCurves.build() at this # exact (scheduler, steps) -- pixel-perfect match for # what the workflow actually produces. # 2. ANY currently-loaded model in ComfyUI's session -- # computes sigmas the same way BasicScheduler does. # This is the right path whenever the user has loaded a # checkpoint (which they have, since they're working in # a graph that uses one). # 3. Synthetic ModelSamplingDiscrete fallback (only used # when no model has ever been loaded -- e.g. fresh # ComfyUI start before the user touches any node). cached = _REAL_SIGMA_CACHE.get((scheduler, int(steps))) if cached is not None: values = cached["values"] target_n = steps + 1 if len(values) != target_n: values = _resample_linear(values, target_n) if cached.get("trailing_zero") and len(values) >= 1: values[-1] = 0.0 return web.json_response({ "values": values, "raw_sigmas": cached.get("raw_sigmas"), "trailing_zero": cached.get("trailing_zero", True), "dispatch": "real_model_cache", "sigma_min": cached.get("sigma_min"), "sigma_max": cached.get("sigma_max"), }) patcher = _get_loaded_model_patcher() bs = _BASIC_SCHEDULER or _get_basic_scheduler() if patcher is not None: try: ms = patcher.get_model_object("model_sampling") # RES4LYF specials (beta57) bypass BasicScheduler -- # they aren't registered in SCHEDULER_HANDLERS, so # bs.get_sigmas would raise. special = _res4lyf_special_sigmas(scheduler, ms, int(steps)) if special is not None: sigmas_tensor = special.cpu() elif bs is not None: sigmas_tensor = bs.get_sigmas( patcher, scheduler, int(steps), 1.0)[0].cpu() else: sigmas_tensor = None if sigmas_tensor is not None: sigma_min_real = float(ms.sigma_min) sigma_max_real = float(ms.sigma_max) _cache_real_sigmas( scheduler, int(steps), sigmas_tensor, sigma_min_real, sigma_max_real, ) cached = _REAL_SIGMA_CACHE.get( (scheduler, int(steps))) if cached is not None: return web.json_response({ "values": cached["values"], "raw_sigmas": cached.get("raw_sigmas"), "trailing_zero": cached.get("trailing_zero", True), "dispatch": "real_model_live", "sigma_min": sigma_min_real, "sigma_max": sigma_max_real, }) except Exception as e: # noqa: BLE001 LOGGER.debug( "BasicScheduler live preview failed for %s: %s", scheduler, e) ms = _synthetic_model_sampling() if ms is None: vals = [1.0 - i / max(steps, 1) for i in range(steps + 1)] return web.json_response({ "values": vals, "trailing_zero": True, "fallback": True, }) sigmas_list, dispatch = _compute_scheduler_sigmas( ms, scheduler, steps) if sigmas_list is None: vals = [1.0 - i / max(steps, 1) for i in range(steps + 1)] return web.json_response({ "values": vals, "trailing_zero": True, "fallback": True, "dispatch": dispatch, "error": f"could not compute '{scheduler}': {dispatch}", }) # Normalize using the *actual* output range of this scheduler so # FlowMatch (sigmas in [0,1]) and EPS (sigmas in [σmin, σmax]) # both fill the plot vertically. The user's edited curve is # always denormalized against the real model's bounds at run # time -- the preview just shows shape. trailing = (len(sigmas_list) >= 2 and abs(sigmas_list[-1]) <= 1e-6) non_term = sigmas_list[:-1] if trailing else sigmas_list if (not non_term) or (max(non_term) - min(non_term)) < 1e-9: values = [1.0 - i / max(steps, 1) for i in range(len(sigmas_list))] if trailing and len(values) >= 1: values[-1] = 0.0 else: hi = max(non_term) lo = min(non_term) denom = max(hi - lo, 1e-9) values = [] for i, s in enumerate(sigmas_list): if trailing and i == len(sigmas_list) - 1 and abs(s) <= 1e-6: values.append(0.0) else: v = (s - lo) / denom values.append(max(0.0, min(1.0, v))) # Some schedulers (e.g. RES4SHO's bong_tangent-derived ones) # return ``steps + 2`` sigmas instead of ``steps + 1``. The # frontend strictly expects ``steps + 1``, so resample here # while preserving the trailing zero terminator. target_n = steps + 1 if len(values) != target_n: resampled = _resample_linear(values, target_n) if trailing and len(resampled) >= 1: resampled[-1] = 0.0 values = resampled return web.json_response({ "values": values, "raw_sigmas": sigmas_list, "trailing_zero": trailing, "dispatch": dispatch, }) except Exception as e: # noqa: BLE001 LOGGER.error("SigmaCurves preview error: %s", e, exc_info=True) return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.get("/RES4SHO/sigma_curves/presets") async def list_presets_endpoint(request): from . import presets as _presets out = {} for n in _presets.list_names(): p = _presets.get(n) if p: out[n] = p return web.json_response({ "presets": out, "prefix": _presets.SCHEDULER_PREFIX, }) @PromptServer.instance.routes.post("/RES4SHO/sigma_curves/preset") async def save_preset_endpoint(request): from . import presets as _presets try: body = await request.json() except Exception: # noqa: BLE001 return web.json_response( {"error": "invalid JSON body"}, status=400) name = body.get("name", "") values = body.get("values") if not _presets.is_valid_name(name): return web.json_response( {"error": "invalid preset name (use letters, digits, " "spaces, dashes, underscores; max 64 chars)"}, status=400) if not isinstance(values, list) or len(values) < 2: return web.json_response( {"error": "values must be a list of >= 2 numbers"}, status=400) name = name.strip() try: cleaned = [float(v) for v in values] _presets.save( name, cleaned, scheduler=body.get("scheduler"), steps=body.get("steps"), trailing_zero=body.get("trailing_zero", True), ) except (ValueError, TypeError) as e: return web.json_response({"error": str(e)}, status=400) # Dynamically register so the new scheduler shows up in dropdowns # (after the frontend triggers a node-defs refresh). sched_name = _presets.SCHEDULER_PREFIX + name handler = _make_preset_scheduler_handler( cleaned, bool(body.get("trailing_zero", True))) _register_one_scheduler(sched_name, handler) return web.json_response({"ok": True, "scheduler": sched_name}) @PromptServer.instance.routes.delete("/RES4SHO/sigma_curves/preset") async def delete_preset_endpoint(request): from . import presets as _presets name = request.query.get("name", "").strip() ok = _presets.delete(name) if ok: _unregister_one_scheduler(_presets.SCHEDULER_PREFIX + name) return web.json_response({"ok": ok}) PromptServer._res4sho_sigma_curves_route = True # --------------------------------------------------------------------- # Dynamic scheduler registration for saved presets. # --------------------------------------------------------------------- def _make_preset_scheduler_handler(values_snapshot, trailing_zero): """Closure that produces sigmas for a saved preset given the active model_sampling and the requested step count. Resamples the saved normalized values to ``steps + 1`` length and denormalizes against the actual model's sigma_min / sigma_max. """ snap = list(values_snapshot) tz = bool(trailing_zero) def handler(model_sampling, steps): sigma_min = float(model_sampling.sigma_min) sigma_max = float(model_sampling.sigma_max) target_n = max(2, steps + 1) v = _resample_linear(snap, target_n) denorm = [x * (sigma_max - sigma_min) + sigma_min for x in v] # Force trailing zero AFTER denormalization so KSampler sees an # actual 0 (the sampling terminator), not sigma_min. if tz and len(denorm) >= 1 and v[-1] <= 1e-3: denorm[-1] = 0.0 return torch.tensor(denorm, dtype=torch.float32) return handler def _scheduler_handler_type(): """Find the type ComfyUI uses for SCHEDULER_HANDLERS values so we can construct compatible entries.""" handlers = getattr(comfy_samplers, "SCHEDULER_HANDLERS", None) if isinstance(handlers, dict) and len(handlers) > 0: return type(next(iter(handlers.values()))) try: from comfy.samplers import SchedulerHandler # type: ignore return SchedulerHandler except (ImportError, AttributeError): return None def _register_one_scheduler(name: str, handler_fn) -> None: """Add (or replace) *name* in comfy's scheduler registries.""" handlers = getattr(comfy_samplers, "SCHEDULER_HANDLERS", None) HandlerType = _scheduler_handler_type() if isinstance(handlers, dict) and HandlerType is not None: try: handlers[name] = HandlerType(handler=handler_fn, use_ms=True) except TypeError: try: handlers[name] = HandlerType(handler_fn, True) except Exception: # noqa: BLE001 LOGGER.warning( "Could not register %r in SCHEDULER_HANDLERS.", name) sched_names = getattr(comfy_samplers, "SCHEDULER_NAMES", None) if isinstance(sched_names, list): if name not in sched_names: sched_names.append(name) elif isinstance(sched_names, tuple): sched_names = list(sched_names) if name not in sched_names: sched_names.append(name) comfy_samplers.SCHEDULER_NAMES = sched_names KSampler = getattr(comfy_samplers, "KSampler", None) if KSampler is not None and hasattr(KSampler, "SCHEDULERS"): sl = list(getattr(KSampler, "SCHEDULERS")) if name not in sl: sl.append(name) KSampler.SCHEDULERS = sl def _unregister_one_scheduler(name: str) -> None: handlers = getattr(comfy_samplers, "SCHEDULER_HANDLERS", None) if isinstance(handlers, dict): handlers.pop(name, None) for attr in ("SCHEDULER_NAMES",): names = getattr(comfy_samplers, attr, None) if isinstance(names, list) and name in names: names.remove(name) setattr(comfy_samplers, attr, names) KSampler = getattr(comfy_samplers, "KSampler", None) if KSampler is not None and hasattr(KSampler, "SCHEDULERS"): sl = list(getattr(KSampler, "SCHEDULERS")) if name in sl: sl.remove(name) KSampler.SCHEDULERS = sl def _register_preset_schedulers_on_load() -> None: """Read every saved preset and register it as a comfy scheduler.""" try: from . import presets as _presets except ImportError: LOGGER.warning("SigmaCurves: presets module not found.") return count = 0 for name in _presets.list_names(): preset = _presets.get(name) if not preset or not isinstance(preset.get("values"), list): continue handler = _make_preset_scheduler_handler( preset["values"], bool(preset.get("trailing_zero", True)), ) _register_one_scheduler(_presets.SCHEDULER_PREFIX + name, handler) count += 1 if count: LOGGER.info("SigmaCurves: registered %d saved preset(s) as " "scheduler(s).", count) _register_routes() _register_preset_schedulers_on_load() # --------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------- def _resample_linear(values: List[float], target_n: int) -> List[float]: """Stretch / shrink a list of values to *target_n* via linear interp.""" n = len(values) if n == target_n: return list(values) if n <= 1 or target_n <= 1: if n == 0: return [0.0] * target_n return [values[0]] * target_n out = [] for i in range(target_n): t = i / (target_n - 1) * (n - 1) lo = int(t) hi = min(lo + 1, n - 1) frac = t - lo out.append(values[lo] * (1.0 - frac) + values[hi] * frac) return out # --------------------------------------------------------------------- # Node # --------------------------------------------------------------------- class SigmaCurves: @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("MODEL",), "scheduler": (_list_schedulers(), {"default": "normal"}), "steps": ("INT", {"default": 20, "min": 1, "max": 1000}), "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), "curve_data": ("STRING", { "default": "", "multiline": False, "tooltip": "Hand-edited via the canvas widget; not meant " "for direct entry. JSON of " "{values: [...], scheduler, steps, ...}.", }), }, } RETURN_TYPES = ("SIGMAS",) FUNCTION = "build" CATEGORY = "sampling/custom_sampling/schedulers" def build(self, model, scheduler: str, steps: int, denoise: float, curve_data: str): if denoise <= 0.0: return (torch.FloatTensor([]),) # Used by the user-curve resample path below. total_steps = steps if denoise >= 1.0 else int(steps / max(denoise, 1e-4)) # Use the stock BasicScheduler node directly. It does the right # thing for every scheduler against any connected model and is # what the user expects "scheduler with model + denoise" to mean. # RES4LYF's `beta57` is the one exception -- it isn't registered # in SCHEDULER_HANDLERS, so we mirror their special case # (beta_scheduler with alpha=0.5, beta=0.7) ourselves. model_sampling = model.get_model_object("model_sampling") special = _res4lyf_special_sigmas(scheduler, model_sampling, steps, denoise) if special is not None: base_used = special.cpu() else: bs = _BASIC_SCHEDULER or _get_basic_scheduler() if bs is None: raise RuntimeError("BasicScheduler unavailable -- " "comfy_extras.nodes_custom_sampler not " "importable.") base_used = bs.get_sigmas(model, scheduler, steps, denoise)[0].cpu() sigma_min_real = float(model_sampling.sigma_min) sigma_max_real = float(model_sampling.sigma_max) # Stash the real-model sigmas so the preview endpoint can show # the actual shape next time the frontend asks. _cache_real_sigmas(scheduler, steps, base_used, sigma_min_real, sigma_max_real) # Parse user values values = None if curve_data: try: data = json.loads(curve_data) raw = data.get("values") if isinstance(raw, list) and len(raw) >= 2: values = [max(0.0, min(1.0, float(v))) for v in raw] except (ValueError, TypeError): LOGGER.warning("SigmaCurves: invalid curve_data; falling " "back to scheduler.") values = None # No user edits -> use the scheduler verbatim. if values is None: return (base_used.float(),) # Treat the user's curve as the FULL intended schedule (sigma_max # -> 0). When denoise < 1.0, mimic ComfyUI's BasicScheduler / # KSampler denoise math: resample the curve to total_steps + 1 # then take the last (steps + 1) entries. That starts the schedule # from a partial-noise point, which is what img2img needs. target_full = total_steps + 1 if len(values) != target_full: values = _resample_linear(values, target_full) if denoise < 1.0: values = values[-(steps + 1):] # Denormalize: 0.0 -> sigma_min, 1.0 -> sigma_max. out = (torch.tensor(values, dtype=torch.float32) * (sigma_max_real - sigma_min_real) + sigma_min_real) # Preserve trailing zero termination if the real scheduler ends # at zero AND the user's last value is near zero. Ensures the # KSampler denoises fully when the user wants it. if base_used.shape[0] >= 2 and float(base_used[-1]) <= 1e-6: if values[-1] <= 1e-3: out[-1] = 0.0 return (out,) NODE_CLASS_MAPPINGS = {"SigmaCurves": SigmaCurves} NODE_DISPLAY_NAME_MAPPINGS = {"SigmaCurves": "Sigma Curves"}