nodes/SparknightLLC-ComfyUI-ModelSamplingSD3Advanced
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
main
SparknightLLC-ComfyUI-Model…/__init__.py
T
ThereforeGames fc63c23c69 Initial release
2025-12-22 12:05:10 -05:00

371 lines
12 KiB
Python
Raw Permalink Blame History

import torch
import comfy.model_sampling
# -----------------------------
# Math helpers
# -----------------------------
def time_snr_shift_tensor(alpha: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
"""
Tensor-safe version of comfy.model_sampling.time_snr_shift
Formula: alpha*t / (1 + (alpha-1)*t)
"""
return (alpha * t) / (1.0 + (alpha - 1.0) * t)
def curve_fn(p: torch.Tensor, curve: str) -> torch.Tensor:
# Curves are defined on [0,1]. We clamp for stability.
p = p.clamp(0.0, 1.0)
if curve == "linear":
return p
if curve == "smoothstep":
return p * p * (3.0 - 2.0 * p)
if curve == "cosine":
return 0.5 - 0.5 * torch.cos(torch.pi * p)
if curve == "ease_in":
return p * p
if curve == "ease_out":
return 1.0 - (1.0 - p) * (1.0 - p)
if curve == "sigmoid":
# centered sigmoid; stronger than others, use with care
k = 10.0
return torch.sigmoid(k * (p - 0.5))
raise ValueError(f"Unknown curve: {curve}")
def windowed_shift(
percent: torch.Tensor,
shift_start: torch.Tensor,
shift_end: torch.Tensor,
start_percent: float,
end_percent: float,
curve: str,
outside_window: str = "hold",
base_shift: torch.Tensor | None = None,
) -> torch.Tensor:
"""
percent: denoise progress in [0,1]
0 = start (high noise), 1 = end (low noise)
Inside the window [start_percent, end_percent]:
shift interpolates shift_start -> shift_end using 'curve'
Outside the window:
outside_window="hold":
percent <= start_percent -> shift_start
percent >= end_percent -> shift_end
outside_window="baseline":
outside the window uses base_shift (typically model's default shift)
outside_window="linear_extrapolate":
continues the (linear) lerp outside the window using u = (percent-sp)/(ep-sp)
(curve still applies inside the window)
"""
# sanitize window
sp = float(max(0.0, min(1.0, start_percent)))
ep = float(max(0.0, min(1.0, end_percent)))
if ep < sp:
sp, ep = ep, sp
denom = max(1e-8, ep - sp)
# unbounded interpolation coordinate
u = (percent - sp) / denom
# inside-window interpolation uses curve
u_in = u.clamp(0.0, 1.0)
w_in = curve_fn(u_in, curve)
shift_in = (1.0 - w_in) * shift_start + w_in * shift_end
if outside_window == "hold":
# exact endpoint behavior outside window
shift_out = torch.where(percent <= sp, shift_start, shift_in)
shift_out = torch.where(percent >= ep, shift_end, shift_out)
return shift_out
if outside_window == "baseline":
if base_shift is None:
# fallback to "hold" if caller didn't provide a baseline
shift_out = torch.where(percent <= sp, shift_start, shift_in)
shift_out = torch.where(percent >= ep, shift_end, shift_out)
return shift_out
shift_out = torch.where((percent <= sp) | (percent >= ep), base_shift, shift_in)
return shift_out
if outside_window == "linear_extrapolate":
# Outside window: linear w=u (unclamped), inside: use curved w_in.
# This keeps your chosen curve inside the window, but lets you extend trend beyond.
w = torch.where((percent < sp) | (percent > ep), u, w_in)
return (1.0 - w) * shift_start + w * shift_end
raise ValueError(f"Unknown outside_window mode: {outside_window}")
# -----------------------------
# ModelSampling implementation
# -----------------------------
class ModelSamplingDiscreteFlowShiftCurve(torch.nn.Module):
"""
Drop-in-ish replacement for comfy.model_sampling.ModelSamplingDiscreteFlow,
with a time-varying shift controlled by:
- shift_start, shift_end
- start_percent, end_percent (window for the transition)
- curve
- outside_window (what shift/sigmas do outside the window)
"""
def __init__(self, model_config=None):
super().__init__()
sampling_settings = model_config.sampling_settings if model_config is not None else {}
self._default_shift = float(sampling_settings.get("shift", 1.0))
self._default_timesteps = int(sampling_settings.get("timesteps", 1000))
self._default_multiplier = float(sampling_settings.get("multiplier", 1000))
# tiny cache so repeated parameter changes don't rebuild the base grid unnecessarily
self._cache_timesteps = None
self._cache_device = None
self._cache_dtype = None
self._cache_t = None
self._cache_percent = None
self.set_parameters(
shift_start=self._default_shift,
shift_end=self._default_shift,
start_percent=0.0,
end_percent=1.0,
curve="linear",
outside_window="hold",
timesteps=self._default_timesteps,
multiplier=self._default_multiplier,
)
def _get_base_grid(self, timesteps: int, device: torch.device, dtype: torch.dtype):
"""
Build (or reuse) base t and percent grids:
t in [1/timesteps, 1]
percent = 1 - t
"""
if (self._cache_t is not None and self._cache_timesteps == timesteps and self._cache_device == device and self._cache_dtype == dtype):
return self._cache_t, self._cache_percent
t = (torch.arange(1, timesteps + 1, device=device, dtype=dtype) / float(timesteps))
percent = 1.0 - t
self._cache_timesteps = timesteps
self._cache_device = device
self._cache_dtype = dtype
self._cache_t = t
self._cache_percent = percent
return t, percent
def set_parameters(
self,
shift_start=1.0,
shift_end=1.0,
start_percent=0.0,
end_percent=1.0,
curve="linear",
outside_window="hold",
timesteps=1000,
multiplier=1000.0,
enforce_monotonic=True,
):
self.shift_start = float(max(shift_start, 1e-8))
self.shift_end = float(max(shift_end, 1e-8))
self.start_percent = float(start_percent)
self.end_percent = float(end_percent)
self.curve = str(curve)
self.outside_window = str(outside_window)
self.multiplier = float(multiplier)
device = self.sigmas.device if hasattr(self, "sigmas") else torch.device("cpu")
dtype = torch.float32
# Use model-config timesteps unless user overrides. (Node passes model default by default.)
timesteps = int(max(1, timesteps))
t, percent = self._get_base_grid(timesteps, device=device, dtype=dtype)
shift_start_t = torch.tensor(self.shift_start, device=device, dtype=dtype)
shift_end_t = torch.tensor(self.shift_end, device=device, dtype=dtype)
base_shift_t = torch.tensor(self._default_shift, device=device, dtype=dtype)
shift_t = windowed_shift(
percent=percent,
shift_start=shift_start_t,
shift_end=shift_end_t,
start_percent=self.start_percent,
end_percent=self.end_percent,
curve=self.curve,
outside_window=self.outside_window,
base_shift=base_shift_t,
)
sigmas = time_snr_shift_tensor(shift_t, t) # in [0,1]
# Some downstream logic expects monotonic schedules; cheap safety.
if enforce_monotonic:
sigmas = torch.cummax(sigmas, dim=0).values
if hasattr(self, "sigmas"):
self.sigmas = sigmas
else:
self.register_buffer("sigmas", sigmas)
@property
def sigma_min(self):
return self.sigmas[0]
@property
def sigma_max(self):
return self.sigmas[-1]
def timestep(self, sigma):
# same scaling approach as flow-style samplers
return sigma * self.multiplier
def sigma(self, timestep):
# Allow arbitrary timestep tensors. Convert to t in [0,1].
# (This path is typically not the hot loop; schedule buffer is.)
t = (timestep / self.multiplier).clamp(0.0, 1.0).to(dtype=torch.float32)
percent = 1.0 - t
device = t.device
dtype = t.dtype
shift_start_t = torch.tensor(self.shift_start, device=device, dtype=dtype)
shift_end_t = torch.tensor(self.shift_end, device=device, dtype=dtype)
base_shift_t = torch.tensor(self._default_shift, device=device, dtype=dtype)
shift_t = windowed_shift(
percent=percent,
shift_start=shift_start_t,
shift_end=shift_end_t,
start_percent=self.start_percent,
end_percent=self.end_percent,
curve=self.curve,
outside_window=self.outside_window,
base_shift=base_shift_t,
)
return time_snr_shift_tensor(shift_t, t)
def percent_to_sigma(self, percent):
# Kept for compatibility with your earlier implementation.
# Note: For flow schedules, sigma isn't linear in percent; consider adding
# an optional LUT-based percent->sigma mapping as a future feature.
if percent <= 0.0:
return 1.0
if percent >= 1.0:
return 0.0
return 1.0 - percent
# -----------------------------
# ComfyUI node
# -----------------------------
class ModelSamplingSD3Advanced:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL", {
"tooltip": "Incoming model to patch. Output is a cloned MODEL with a modified flow (SD3-style) sampling schedule."
}),
"shift_start": ("FLOAT", {
"default": 6.0,
"min": 0.0,
"max": 100.0,
"step": 0.01,
"tooltip": "Shift value used at the beginning (early/high-noise steps). Higher values often increase prompt adherence early, but can reduce late-step quality if kept too high."
}),
"shift_end": ("FLOAT", {
"default": 3.0,
"min": 0.0,
"max": 100.0,
"step": 0.01,
"tooltip": "Shift value used at the end (late/low-noise steps). Lower values often improve late refinement/quality, at the cost of weaker early adherence if used too soon."
}),
"start_percent": ("FLOAT", {
"default": 0.0,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"round": False,
"tooltip": "Denoise progress where the shift transition window begins. 0.0 = very start (highest noise), 1.0 = very end (lowest noise)."
}),
"end_percent": ("FLOAT", {
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"round": False,
"tooltip": "Denoise progress where the shift transition window ends. Inside [start_percent, end_percent], shift moves from shift_start -> shift_end using the chosen curve."
}),
"curve": (["linear", "cosine", "smoothstep", "ease_in", "ease_out", "sigmoid"], {
"default": "ease_in",
"tooltip": "Controls how shift interpolates inside the window. 'ease_in' changes slowly at first then faster; 'ease_out' changes fast then slows; 'sigmoid' is more switch-like (strong)."
}),
"outside_window": (["hold", "baseline", "linear_extrapolate"], {
"default": "hold",
"tooltip": "Controls shift/sigma behavior outside the window:\n"
"- hold: before start_percent use shift_start; after end_percent use shift_end.\n"
"- baseline: outside the window use the model's default shift.\n"
"- linear_extrapolate: extend the shift trend beyond the window (advanced)."
}),
"multiplier": ("FLOAT", {
"default": 1000.0,
"min": 1.0,
"max": 100000.0,
"step": 1.0,
"tooltip": "Flow timestep scaling factor (sigma <-> timestep). Keep at 1000 unless your model/sampler expects a different flow scaling."
}),
}
}
RETURN_TYPES = ("MODEL", )
FUNCTION = "patch"
CATEGORY = "advanced/model"
def patch(self, model, shift_start, shift_end, start_percent, end_percent, curve, outside_window, multiplier):
m = model.clone()
# Use model-config defaults where possible (timesteps especially)
cfg = model.model.model_config
ss = getattr(cfg, "sampling_settings", {}) or {}
timesteps = int(ss.get("timesteps", 1000))
sampling_base = ModelSamplingDiscreteFlowShiftCurve
sampling_type = comfy.model_sampling.CONST
class ModelSamplingAdvanced(sampling_base, sampling_type):
pass
model_sampling = ModelSamplingAdvanced(cfg)
model_sampling.set_parameters(
shift_start=shift_start,
shift_end=shift_end,
start_percent=start_percent,
end_percent=end_percent,
curve=curve,
outside_window=outside_window,
timesteps=timesteps,
multiplier=multiplier,
)
m.add_object_patch("model_sampling", model_sampling)
return (m, )
NODE_CLASS_MAPPINGS = {
"ModelSamplingSD3Advanced": ModelSamplingSD3Advanced,
}
Reference in New Issue View Git Blame Copy Permalink