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, }