371 lines
12 KiB
Python
371 lines
12 KiB
Python
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,
|
|
}
|