""" DCW – Differential Correction in Wavelet domain CWM – CFG Wavelet Mixing SMC – Sliding Mode Control CFG (integrated into CWM hook) ComfyUI Custom Node References: DCW: "Elucidating the SNR-t Bias of Diffusion Probabilistic Models" Yu et al., arXiv:2604.16044v1 (2026) SMC: "CFG-Ctrl: Control-Based Classifier-Free Diffusion Guidance" Wang et al., CVPR 2026 / arXiv:2603.03281 ───────────────────────────────────────────────────────────────────────────── Architecture overview ───────────────────────────────────────────────────────────────────────────── sampler_cfg_function (CWM hook): Standard path (SMC off, alpha_l = alpha_h = 0): output = uncond + w · (cond − uncond) CWM only (SMC off): e = cond − uncond DWT(e) → LL, LH, HL, HH output = uncond + IDWT(w_LL·LL, w_mid·LH, w_mid·HL, w_HH·HH) SMC only (alpha_l = alpha_h = 0): e = cond − uncond s = (e − e_prev) + λ · e_prev Δe = −k · s / ‖s‖₂ unit_2 switching (paper Table 4) e* = e + Δe output = uncond + w · e* SMC + CWM (full integration — the new combined path): e = cond − uncond s = (e − e_prev) + λ · e_prev Δe = −k · s / ‖s‖₂ e* = e + Δe e_prev ← e* DWT(e*) → LL*, LH*, HL*, HH* output = uncond + IDWT(w_LL·LL*, w_mid·LH*, w_mid·HL*, w_HH·HH*) Rationale for SMC → CWM order: SMC corrects the *quality* of the guidance error (reduces oscillation, improves semantic alignment). CWM then distributes the already-corrected error across frequency bands. Applying wavelet mixing to a noisy error first and then trying to correct it would mis-attribute the oscillation energy to the wrong frequency bands, so SMC must precede CWM. sampler_post_cfg_function (DCW hook): Always safe to chain; corrects x0_pred after CFG guidance is complete. ───────────────────────────────────────────────────────────────────────────── SMC hyperparameter presets (λ, k) — paper Sec 7.3 / Table 3 ───────────────────────────────────────────────────────────────────────────── SD1.5 / SD2 λ=5, k=0.10 SDXL λ=5, k=0.10 SD3 / SD3.5 λ=6, k=0.10 (grid-searched) Flux λ=6, k=0.70 (grid-searched) Qwen-Image λ=6, k=0.10 (grid-searched) Cosmos / Wan λ=6, k=0.20 (conservative) """ import math import torch import torch.nn.functional as F # ───────────────────────────────────────────────────────────────────────────── # SMC presets # ───────────────────────────────────────────────────────────────────────────── _SMC_PRESETS: dict[str, tuple[float, float]] = { "SD1.5 / SD2": (5.0, 0.10), "SDXL": (5.0, 0.10), "SD3 / SD3.5": (6.0, 0.10), "Flux": (6.0, 0.70), "Qwen-Image": (6.0, 0.10), "Cosmos / Wan": (6.0, 0.20), "Custom": (6.0, 0.10), } SMC_PRESET_NAMES: list[str] = ["Off", "Auto"] + list(_SMC_PRESETS.keys()) # State key for cross-step e_prev persistence inside model_options. # model_options is a fresh dict per KSampler run → auto-resets between runs. _SMC_STATE_KEY = "_dcw_smc_state" # SMC minimum scale guard (prevents division by zero in tanh normalisation) _SMC_CLAMP_MIN = 1e-6 # State key for cross-step RDC EMA persistence inside model_options. # Same pattern as _SMC_STATE_KEY — fresh dict per KSampler run → auto-resets. _RDC_STATE_KEY = "_dcw_rdc_state" # ───────────────────────────────────────────────────────────────────────────── # Model family auto-detection for SMC preset # ───────────────────────────────────────────────────────────────────────────── def _detect_smc_preset(model) -> str: """Infer the best SMC preset from the model object's class name / attributes.""" try: inner = model.model cls = type(inner).__name__.lower() if "flux" in cls: return "Flux" if "cosmos" in cls or "predict2" in cls: return "Cosmos / Wan" if "wan" in cls: return "Cosmos / Wan" if "anima" in cls: return "Cosmos / Wan" if "sd3" in cls or "mmdit" in cls: return "SD3 / SD3.5" if "sdxl" in cls: return "SDXL" if "qwen" in cls: return "Qwen-Image" mt = str(getattr(inner, "model_type", "")).lower() if "flow" in mt: return "SD3 / SD3.5" if "v_pred" in mt: return "SD1.5 / SD2" diff = getattr(inner, "diffusion_model", None) if diff is not None: dm = type(diff).__name__.lower() if "flux" in dm: return "Flux" if "cosmos" in dm or "wan" in dm: return "Cosmos / Wan" if "mmdit" in dm or "sd3" in dm: return "SD3 / SD3.5" if "sdxl" in dm: return "SDXL" except Exception: pass return "SD1.5 / SD2" # ───────────────────────────────────────────────────────────────────────────── # Haar Wavelet Utilities (pure PyTorch, zero extra dependencies) # ───────────────────────────────────────────────────────────────────────────── def _pad_even(x): """ Pad H and W to even numbers if necessary (Haar DWT requires even dims). Returns (padded_tensor, (original_H, original_W)). Cosmos VAE (used by Anima) compresses space by 8×, so e.g. resolution 1208 → latent H 151 (odd), which triggers this path on a 5-D video tensor. """ H, W = x.shape[-2], x.shape[-1] ph, pw = H % 2, W % 2 if ph or pw: extra = max(0, x.dim() - 4) pad = (0, pw, 0, ph) + (0, 0) * extra x = F.pad(x, pad, mode="reflect") return x, (H, W) def haar_dwt2d(x): """ 2-D Haar Discrete Wavelet Transform. Input : (B, C, H, W) — H and W must be even Output: four subbands, each (B, C, H//2, W//2) LL – low-low : coarse structure / global energy LH – low-high : horizontal edges HL – high-low : vertical edges HH – high-high : diagonal / fine texture """ L = (x[..., 0::2, :] + x[..., 1::2, :]) * 0.5 Hi = (x[..., 0::2, :] - x[..., 1::2, :]) * 0.5 LL = (L[..., 0::2] + L[..., 1::2]) * 0.5 LH = (L[..., 0::2] - L[..., 1::2]) * 0.5 HL = (Hi[..., 0::2] + Hi[..., 1::2]) * 0.5 HH = (Hi[..., 0::2] - Hi[..., 1::2]) * 0.5 return LL, LH, HL, HH def haar_idwt2d(LL, LH, HL, HH): """ 2-D Haar Inverse Discrete Wavelet Transform. Handles arbitrary leading dimensions: 4-D image : (B, C, h, w) → (B, C, 2h, 2w) 5-D video : (B, C, T, h, w) → (B, C, T, 2h, 2w) """ *leading, h, w = LL.shape dev, dt = LL.device, LL.dtype L = torch.empty(*leading, h, w * 2, device=dev, dtype=dt) Hi = torch.empty(*leading, h, w * 2, device=dev, dtype=dt) L[..., 0::2] = LL + LH L[..., 1::2] = LL - LH Hi[..., 0::2] = HL + HH Hi[..., 1::2] = HL - HH out = torch.empty(*leading, h * 2, w * 2, device=dev, dtype=dt) out[..., 0::2, :] = L + Hi out[..., 1::2, :] = L - Hi return out # ───────────────────────────────────────────────────────────────────────────── # Shared utilities # ───────────────────────────────────────────────────────────────────────────── _FP8_DTYPES = { dt for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz") if (dt := getattr(torch, name, None)) is not None } def _safe_compute_dtype(dtype: torch.dtype) -> torch.dtype: """fp8 → bfloat16 for arithmetic safety; everything else unchanged.""" if dtype in _FP8_DTYPES: return torch.bfloat16 return dtype def _sigma_norm(sigma, tensor: torch.Tensor, compute_dtype: torch.dtype): """ σ_norm = σ / (σ + 1) → [0, 1) Consistent normalisation across EDM / DDPM / Flow schedulers. Always computed in fp32 to avoid precision loss at small sigma. """ if isinstance(sigma, torch.Tensor): s = sigma.float() / (sigma.float() + 1.0) if s.dim() == 0: pass elif s.dim() == 1: extra = tensor.dim() - 1 s = s.view(-1, *([1] * extra)) return s.to(dtype=compute_dtype, device=tensor.device) else: sigma_f = float(sigma) return sigma_f / (sigma_f + 1.0) def _ch_energy_weight(x: torch.Tensor, clamp_lo: float = 0.25, clamp_hi: float = 4.0) -> torch.Tensor: """ Per-channel energy weight for DCW correction. Motivation ---------- In multi-channel VAEs (e.g. Cosmos 16-ch used by Anima), each channel encodes different semantic content — pose, accessory presence, lighting, etc. Applying a uniform correction lambda across all channels treats a semantically-rich channel and a near-empty background channel identically. This causes small lambda changes (~0.01) to push high-energy semantic channels past decision boundaries, producing the discrete jumps (body angle, accessory appearing/disappearing) observed in practice. Formula ------- For a subband tensor x of shape (B, C, ...) : energy_c = mean_{spatial}( x_c² ) per channel weight_c = energy_c / mean_c( energy_c ) normalise → mean=1 weight_c = clamp( weight_c, lo, hi ) stability guard Properties ---------- • Mean weight is always 1.0 → total correction energy is unchanged. • High-energy (semantically active) channels get stronger correction; low-energy channels get weaker correction. • In early steps HH_x is dominated by noise (uniform across channels) → weights ≈ 1 → effectively no per-channel bias when reference is meaningless, which is the correct behaviour. • Clamping [0.25, 4.0] limits the range to 1/4× – 4×; prevents a single loud channel from monopolising or a quiet channel from being corrected into instability. Parameters ---------- x : reference subband tensor (B, C, ...) clamp_lo : minimum weight (default 0.25) clamp_hi : maximum weight (default 4.0) Returns tensor of same shape, broadcastable with x. """ spatial = tuple(range(2, x.dim())) # all dims after C energy = x.float().pow(2).mean(dim=spatial, keepdim=True) # (B, C, 1...) mean_e = energy.mean(dim=1, keepdim=True).clamp(min=1e-8) # (B, 1, 1...) return (energy / mean_e).clamp(clamp_lo, clamp_hi).to(x.dtype) # SMC minimum norm/scale guard (prevents division-by-zero in unit_2 normalisation) _SMC_CLAMP_MIN = 1e-8 # ───────────────────────────────────────────────────────────────────────────── # DCW Core (post-CFG: corrects x0_pred) # ───────────────────────────────────────────────────────────────────────────── def apply_dcw(denoised: torch.Tensor, x_t: torch.Tensor, sigma, lambda_l: float, lambda_h: float, rdc_alpha_ll: float = 0.03, rdc_alpha_hh: float = 0.0, rdc_tau: float = 0.0, rdc_state: dict | None = None) -> torch.Tensor: """ Apply Differential Correction in Wavelet domain to x0_pred. Correction formula per frequency band f: corrected_f = denoised_f + λ_f(t) · w_ch_f · (x_t_f − denoised_f) Band timing ----------- LL (global structure) λ_l(t) = lambda_l · σ_norm early steps LH (horiz. edges) ┐ λ_m(t) = (λ_l + λ_h) / 2 mid steps HL (vert. edges) ┘ (linear interpolation of LL and HH schedules) HH (fine texture) λ_h(t) = lambda_h · (1 − σ_norm) late steps LH and HL represent directional edge information that forms between global structure (LL) and fine texture (HH) in the coarse-to-fine denoising process. Giving them the same timing as HH (late-only) was incorrect; the interpolated schedule is more physically accurate. Per-channel energy weighting ---------------------------- w_ch_f = _ch_energy_weight(x_t_f) Channels with higher signal energy receive proportionally stronger correction. This prevents small lambda changes from causing discrete semantic jumps (body pose, accessory presence) in multi-channel VAEs such as Cosmos 16-ch (Anima), where different channels encode different semantic concepts at different energy levels. When x_t_f is dominated by noise (early steps, HH band), the energy is approximately uniform across channels → weights ≈ 1 → no artificial bias is introduced where the reference signal is not yet meaningful. """ rdc_on = rdc_tau > 0.0 if lambda_l == 0.0 and lambda_h == 0.0 and not rdc_on: return denoised orig_dtype = denoised.dtype compute_dtype = _safe_compute_dtype(orig_dtype) need_cast = (compute_dtype != orig_dtype) if need_cast: denoised = denoised.to(dtype=compute_dtype) s = _sigma_norm(sigma, denoised, compute_dtype) lam_l = lambda_l * s lam_h = lambda_h * (1.0 - s) lam_mid = (lam_l + lam_h) * 0.5 # LH / HL: mid-step timing dn_p, (H, W) = _pad_even(denoised) xt_p, _ = _pad_even(x_t.to(dtype=compute_dtype, device=denoised.device)) LL_d, LH_d, HL_d, HH_d = haar_dwt2d(dn_p) LL_x, LH_x, HL_x, HH_x = haar_dwt2d(xt_p) # Per-channel energy weights derived from x_t subbands cw_LL = _ch_energy_weight(LL_x) cw_LH = _ch_energy_weight(LH_x) cw_HL = _ch_energy_weight(HL_x) cw_HH = _ch_energy_weight(HH_x) LL_c = LL_d + lam_l * cw_LL * (LL_x - LL_d) LH_c = LH_d + lam_mid * cw_LH * (LH_x - LH_d) HL_c = HL_d + lam_mid * cw_HL * (HL_x - HL_d) HH_c = HH_d + lam_h * cw_HH * (HH_x - HH_d) # ── RDC: band-wise EMA drift compensation (trajectory correction) ────── # Unlike the DCW terms above (instantaneous x_t vs denoised comparison), # RDC corrects cross-step drift by pulling each band back toward its own # running EMA. See DCW_RDC_확장_개발문서.md §3.2–3.3 for derivation. if rdc_on and rdc_state is not None: s_val = float(s.mean()) if isinstance(s, torch.Tensor) else float(s) s_prev = rdc_state.get("_s_prev", s_val) delta_s = abs(s_prev - s_val) beta = 1.0 - math.exp(-delta_s / max(rdc_tau, 1e-6)) rdc_state["_s_prev"] = s_val alpha_lh_hl = (rdc_alpha_ll + rdc_alpha_hh) * 0.5 bands = (("LL", LL_c, rdc_alpha_ll), ("LH", LH_c, alpha_lh_hl), ("HL", HL_c, alpha_lh_hl), ("HH", HH_c, rdc_alpha_hh)) updated = {} for name, band, alpha in bands: if alpha == 0.0: continue band_detached = band.detach() prev_ema = rdc_state.get(name) if (prev_ema is None or prev_ema.shape != band_detached.shape or prev_ema.device != band_detached.device): # First step (or shape/device change, e.g. resolution switch): # no baseline yet, seed the EMA and skip correction this step. rdc_state[name] = band_detached.to(dtype=compute_dtype).clone() continue new_ema = (1.0 - beta) * prev_ema + beta * band_detached rdc_state[name] = new_ema.to(dtype=compute_dtype) drift = band - new_ema updated[name] = band - alpha * drift LL_c = updated.get("LL", LL_c) LH_c = updated.get("LH", LH_c) HL_c = updated.get("HL", HL_c) HH_c = updated.get("HH", HH_c) out = haar_idwt2d(LL_c, LH_c, HL_c, HH_c) out = out[..., :H, :W] if need_cast: out = out.to(dtype=orig_dtype) return out # ───────────────────────────────────────────────────────────────────────────── # Combined CFG hook: SMC + CWM # ───────────────────────────────────────────────────────────────────────────── def apply_cfg_combined( cond: torch.Tensor, uncond: torch.Tensor, sigma, cfg_scale: float, # CWM params alpha_l: float, alpha_h: float, # SMC params (lambda_val=0 / k=0 → SMC disabled) smc_lambda: float, smc_k: float, # SMC cross-step state container (mutated in-place) smc_state: dict, # fp32 upcast for SMC (prevents NaN with fp16/bf16 in edge cases) fp32_smc: bool = True, ) -> torch.Tensor: """ Unified CFG computation. Execution order: 1. Compute raw guidance error e = cond − uncond 2. [Optional] SMC correction: e* = SMC(e, e_prev, λ, k) 3. [Optional] CWM scaling: apply per-band adaptive weights 4. Reconstruct: uncond + guided_error When both SMC and CWM are inactive, reduces exactly to standard CFG. When only SMC is active (alpha_l = alpha_h = 0), no wavelet ops run. When only CWM is active (smc_lambda = 0 or smc_k = 0), no SMC state ops run. """ smc_active = (smc_lambda != 0.0 and smc_k != 0.0) cwm_active = (alpha_l != 0.0 or alpha_h != 0.0) # ── Fast path: pure standard CFG ────────────────────────────────────────── if not smc_active and not cwm_active: return uncond + cfg_scale * (cond - uncond) orig_dtype = cond.dtype compute_dtype = _safe_compute_dtype(orig_dtype) need_cast = (compute_dtype != orig_dtype) # SMC upcast target smc_dtype = torch.float32 if fp32_smc else compute_dtype if need_cast: cond_w = cond.to(dtype=compute_dtype) uncond_w = uncond.to(dtype=compute_dtype) else: cond_w, uncond_w = cond, uncond # ── Step 1: raw guidance error ──────────────────────────────────────────── e = torch.nan_to_num(cond_w - uncond_w, nan=0.0, posinf=0.0, neginf=0.0) # ── Step 2: SMC correction ──────────────────────────────────────────────── if smc_active: e_smc = e.to(dtype=smc_dtype) e_prev_raw = smc_state.get("e_prev") if e_prev_raw is None or e_prev_raw.shape != e_smc.shape: # First step OR resolution changed (e.g. SPEED spectral expansion): # reset e_prev = e so SMC restarts cleanly at the new resolution. # This prevents the tensor size mismatch that occurs when # spectral_expand changes spatial dims between sampler segments. if e_prev_raw is not None and e_prev_raw.shape != e_smc.shape: print( f"[CWM/SMC] Resolution change detected " f"({list(e_prev_raw.shape)} → {list(e_smc.shape)}): " f"resetting SMC state." ) # First step: initialise e_prev = e (Algorithm 1, line 7-9) e_prev = e_smc.detach().clone() else: e_prev = torch.nan_to_num( e_prev_raw.to(device=e_smc.device, dtype=smc_dtype), nan=0.0, posinf=0.0, neginf=0.0, ) # Sliding surface: s = (e - e_prev) + λ · e_prev s = torch.nan_to_num( (e_smc - e_prev) + smc_lambda * e_prev, nan=0.0, posinf=0.0, neginf=0.0, ) # Switching control: Δe = −k · s / ‖s‖₂ (unit_2 normalisation) # # Paper Table 4 / Notation Table explicitly defines: # sign(s_t) ≡ s_t / ‖s_t‖₂ (unit_2, NOT element-wise sign) # # Rationale vs element-wise sign(s): # Element-wise ±1 discards the relative magnitude across spatial # positions — in a high-dim latent every position votes equally # regardless of how strong the surface signal is there, which # produces the checkerboard ±k pattern and colour corruption seen # in Cosmos/Anima 16-ch latents. # unit_2 preserves the full directional structure of s while # normalising its global scale; each position's contribution is # proportional to its local |s| relative to the tensor's L2 norm. # # Norm is reduced over all non-batch dims so each sample in the # batch gets its own normalisation (handles video / variable res). reduce_dims = tuple(range(1, s.ndim)) s_norm = torch.linalg.vector_norm(s, dim=reduce_dims, keepdim=True).clamp(min=_SMC_CLAMP_MIN) s_unit = s / s_norm delta_e_raw = -smc_k * s_unit # Adaptive clamp |Δe| ≤ 0.5 × mean|e| (paper Supp. Eq. 45) # Prevents the switching term from dominating in early steps when # the surface norm is large. max_delta = ( 0.5 * e_smc.abs().mean(dim=reduce_dims, keepdim=True) ).clamp(min=_SMC_CLAMP_MIN) delta_e = delta_e_raw.clamp(-max_delta, max_delta) e_star = torch.nan_to_num(e_smc + delta_e, nan=0.0, posinf=0.0, neginf=0.0) # Update cross-step state with corrected error smc_state["e_prev"] = e_star.detach().clone() # Cast corrected error back to compute_dtype for subsequent ops e = e_star.to(dtype=compute_dtype) # ── Step 3: CWM per-band scaling ────────────────────────────────────────── if cwm_active: w = float(cfg_scale) s_n = _sigma_norm(sigma, e, compute_dtype) w_LL = w * (1.0 + float(alpha_l) * s_n) w_HH = w * (1.0 + float(alpha_h) * (1.0 - s_n)) # Geometric mean for LH/HL: physically correct for a band that lies # between LL and HH in the coarse-to-fine hierarchy. # AM ≥ GM always, so arithmetic mean systematically over-weights the # mid-band; geometric mean keeps energy scaling consistent with the # log-linear interpolation implicit in wavelet decomposition. w_mid = (w_LL * w_HH) ** 0.5 e_p, (H, W) = _pad_even(e) LL_e, LH_e, HL_e, HH_e = haar_dwt2d(e_p) LL_g = w_LL * LL_e LH_g = w_mid * LH_e HL_g = w_mid * HL_e HH_g = w_HH * HH_e e_guided = haar_idwt2d(LL_g, LH_g, HL_g, HH_g) e_guided = e_guided[..., :H, :W] out = uncond_w + e_guided else: # SMC only path — no wavelet overhead out = uncond_w + float(cfg_scale) * e out = torch.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0) if need_cast: out = out.to(dtype=orig_dtype) return out # ───────────────────────────────────────────────────────────────────────────── # ComfyUI Node Definition # ───────────────────────────────────────────────────────────────────────────── class DCWModelPatch: """ DCW + CWM + SMC Model Patch node. DCW (sampler_post_cfg_function): Corrects x0_pred in the wavelet domain to mitigate SNR-t bias. Applied after CFG combination, every denoising step. Always safe to chain with other post-cfg hooks. CWM (sampler_cfg_function): Replaces scalar CFG with per-band adaptive guidance scaling. Low-freq bands (structure/layout) are boosted in early steps; high-freq bands (edges/texture) are boosted in late steps. alpha_l = alpha_h = 0 → identical to standard CFG. SMC (integrated into CWM hook): Sliding Mode Control correction of the guidance error prior to wavelet decomposition. Reduces oscillation in the guidance trajectory and improves semantic alignment. smc_preset = Off → SMC disabled entirely. When CWM is also enabled, SMC runs first (corrects error quality), then CWM distributes the corrected error across frequency bands. When CWM is disabled, SMC runs as a pure CFG replacement with no wavelet overhead. Compatibility: DCW can co-exist with any other sampler_cfg_function node because it uses a different hook (sampler_post_cfg_function). CWM/SMC use sampler_cfg_function. If another node has already registered that hook before this node runs, CWM/SMC are skipped automatically with a console warning. """ RETURN_TYPES = ("MODEL",) RETURN_NAMES = ("model",) FUNCTION = "patch" CATEGORY = "model_patches" DESCRIPTION = ( "DCW – Differential Correction in Wavelet domain (SNR-t bias correction).\n" "CWM – CFG Wavelet Mixing (adaptive per-band guidance scale).\n" "SMC – Sliding Mode Control CFG (integrated into CWM, corrects guidance error).\n" "All three are training-free and plug-and-play." ) @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("MODEL",), # ── DCW parameters ──────────────────────────────────────────── "lambda_l": ("FLOAT", { "default": 0.05, "min": -0.5, "max": 0.5, "step": 0.005, "round": 0.001, "tooltip": ( "[DCW] Low-frequency correction strength (x0_pred).\n" "Positive: push x0_pred toward x_t (more signal).\n" "Negative: push away from x_t (less signal).\n" "Active mainly in EARLY denoising steps.\n" "Start: 0.04–0.07 for DDPM/EDM; ~2× for Flux/Flow.\n" "Set 0.0 to disable DCW low-freq correction." ), }), "lambda_h": ("FLOAT", { "default": 0.01, "min": -0.3, "max": 0.3, "step": 0.001, "round": 0.001, "tooltip": ( "[DCW] High-frequency correction strength (x0_pred).\n" "Positive: push x0_pred toward x_t (more detail).\n" "Negative: push away from x_t (less detail).\n" "Active mainly in LATE denoising steps.\n" "Start: 0.008–0.015 for DDPM/EDM; ~2× for Flux/Flow.\n" "Set 0.0 to disable DCW high-freq correction." ), }), "dcw_enabled": ("BOOLEAN", { "default": True, "tooltip": "[DCW] Toggle SNR-t bias correction on/off.", }), # ── CWM parameters ──────────────────────────────────────────── "alpha_l": ("FLOAT", { "default": 0.0, "min": -1.0, "max": 2.0, "step": 0.01, "round": 0.001, "tooltip": ( "[CWM] Low-frequency CFG boost strength.\n" "Scales the LL guidance band relative to base CFG scale.\n" "Positive: amplify coarse structure/composition guidance in early steps.\n" "Negative: suppress low-freq guidance.\n" "0.0 = standard CFG for this band.\n" "Suggested start: 0.1–0.3. Flow models (Flux): try 0.2–0.5." ), }), "alpha_h": ("FLOAT", { "default": 0.0, "min": -1.0, "max": 2.0, "step": 0.01, "round": 0.001, "tooltip": ( "[CWM] High-frequency CFG boost strength.\n" "Scales the HH guidance band relative to base CFG scale.\n" "Positive: amplify edge/texture/detail guidance in late steps.\n" "Negative: suppress high-freq guidance.\n" "0.0 = standard CFG for this band.\n" "Suggested start: 0.1–0.2. Excessive values cause over-sharpening." ), }), "cwm_enabled": ("BOOLEAN", { "default": True, "tooltip": "[CWM] Toggle CFG Wavelet Mixing on/off.", }), # ── SMC parameters ──────────────────────────────────────────── "smc_preset": ( SMC_PRESET_NAMES, { "default": "Off", "tooltip": ( "[SMC] Sliding Mode Control CFG preset.\n" "'Off' — SMC disabled entirely.\n" "'Auto' — auto-detect model family, apply paper defaults.\n" "Named presets — paper-verified (λ, k) per model family.\n" "'Custom' — use smc_lambda / smc_k sliders below.\n\n" "SMC runs inside the CWM hook (sampler_cfg_function).\n" "It corrects guidance error quality BEFORE CWM distributes\n" "it across frequency bands.\n" "If cwm_enabled=False, SMC still activates the cfg hook\n" "and runs without wavelet overhead (pure SMC-CFG mode)." ), }, ), "smc_lambda": ("FLOAT", { "default": 6.0, "min": 0.5, "max": 30.0, "step": 0.1, "tooltip": ( "[SMC] Shape parameter λ of the sliding mode surface.\n" "Only used when smc_preset = Custom.\n" "Paper recommends 2–8. Extreme values distort the manifold." ), }), "smc_k": ("FLOAT", { "default": 0.1, "min": 0.0, "max": 5.0, "step": 0.01, "tooltip": ( "[SMC] Switching gain k.\n" "Only used when smc_preset = Custom.\n" "Low k → better realism, weaker text alignment.\n" "High k → stronger text alignment, risk of chattering." ), }), # ── RDC parameters ───────────────────────────────────────────── # No separate on/off toggle: rdc_tau = 0.0 fully disables RDC # (β always evaluates to 0 → drift always 0 → no-op, see apply_dcw), # rdc_tau > 0.0 enables it. Keeps the parameter count at 3. "rdc_tau": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 0.5, "step": 0.01, "round": 0.001, "tooltip": ( "[RDC] Band-wise EMA drift compensation — on/off AND time-\n" "constant in one value.\n" "0.0 = RDC fully OFF (default; identical to not having RDC).\n" "> 0.0 = RDC ON. Value = memory span in sigma_norm units\n" "([0,1) range, scheduler/step-count independent): after the\n" "sampler has moved this far in sigma_norm, the EMA's memory\n" "of older steps has decayed to ~37%.\n" "Smaller (e.g. 0.05–0.1) → reacts fast, short memory.\n" "Larger (e.g. 0.2–0.3) → reacts slowly, long memory, smoother.\n" "Corrects cross-step trajectory drift (pose/face/composition\n" "slowly wandering over steps) — different from DCW, which only\n" "compares each step to itself, not the trajectory.\n" "Reuses DCW's existing wavelet decomposition (no extra cost)." ), }), "rdc_alpha_ll": ("FLOAT", { "default": 0.03, "min": 0.0, "max": 0.3, "step": 0.005, "round": 0.001, "tooltip": ( "[RDC] Structure-band (LL) drift correction strength.\n" "Pulls the low-frequency band back toward its running EMA\n" "each step, damping slow composition/pose drift.\n" "Has no effect while rdc_tau = 0.0 (RDC off).\n" "Suggested start: 0.02–0.05.\n" "Too high → composition gets 'stuck', resists intended change." ), }), "rdc_alpha_hh": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 0.1, "step": 0.001, "round": 0.001, "tooltip": ( "[RDC] Texture-band (HH) drift correction strength.\n" "HH is expected to regenerate every step (grain/detail);\n" "correcting it pulls detail toward a blurred running average.\n" "0.0 (default) = off — recommended unless you see texture\n" "flicker you specifically want to dampen.\n" "Suggested ceiling if used: ~0.01. Higher values blur detail.\n" "Has no effect while rdc_tau = 0.0 (RDC off)." ), }), } } # ── Patch method ────────────────────────────────────────────────────────── def patch( self, model, lambda_l: float, lambda_h: float, dcw_enabled: bool, alpha_l: float, alpha_h: float, cwm_enabled: bool, smc_preset: str, smc_lambda: float, smc_k: float, rdc_tau: float = 0.0, rdc_alpha_ll: float = 0.03, rdc_alpha_hh: float = 0.0, ): # ── Resolve SMC hyperparameters ─────────────────────────────────────── smc_on = (smc_preset != "Off") if smc_on: if smc_preset == "Auto": detected = _detect_smc_preset(model) resolved_lambda, resolved_k = _SMC_PRESETS[detected] print(f"[DCW/SMC] Auto → '{detected}' λ={resolved_lambda} k={resolved_k}") elif smc_preset == "Custom": resolved_lambda, resolved_k = smc_lambda, smc_k print(f"[DCW/SMC] Custom λ={resolved_lambda} k={resolved_k}") else: resolved_lambda, resolved_k = _SMC_PRESETS[smc_preset] print(f"[DCW/SMC] '{smc_preset}' λ={resolved_lambda} k={resolved_k}") else: resolved_lambda = resolved_k = 0.0 # ── Determine active features ───────────────────────────────────────── rdc_on = rdc_tau > 0.0 dcw_active = dcw_enabled and (lambda_l != 0.0 or lambda_h != 0.0 or rdc_on) # cfg hook is needed if CWM alpha is non-zero OR SMC is on cwm_alpha_active = cwm_enabled and (alpha_l != 0.0 or alpha_h != 0.0) cfg_hook_needed = cwm_alpha_active or smc_on if not dcw_active and not cfg_hook_needed: return (model,) m = model.clone() m.model_options = {**m.model_options} # ── CFG hook: CWM + (optional) SMC ─────────────────────────────────── if cfg_hook_needed: if "sampler_cfg_function" in m.model_options: print( "[DCW/CWM/SMC] Warning: sampler_cfg_function already registered " "by another node. CWM and SMC skipped to avoid conflict.\n" " → Place this node earlier in the pipeline, or disable the " "conflicting node's CFG hook." ) else: # Capture parameters and mutable SMC state in closure _al = alpha_l if cwm_alpha_active else 0.0 _ah = alpha_h if cwm_alpha_active else 0.0 _smc_lambda = resolved_lambda _smc_k = resolved_k # smc_state dict is mutated per-step to carry e_prev. # model_options is a fresh dict per KSampler run, but the # closure itself persists across the node's lifetime; we # embed the state inside model_options so it resets correctly. _model_options_ref = m.model_options def cfg_hook(args: dict) -> torch.Tensor: cond_t = args.get("cond") uncond_t = args.get("uncond") cond_scale = args.get("cond_scale", 7.0) sigma = args.get("sigma") if cond_t is None or uncond_t is None: return uncond_t + float(cond_scale) * (cond_t - uncond_t) # Retrieve or create per-run SMC state opts = args.get("model_options", _model_options_ref) smc_state = opts.setdefault(_SMC_STATE_KEY, {}) try: return apply_cfg_combined( cond_t, uncond_t, sigma, float(cond_scale), _al, _ah, _smc_lambda, _smc_k, smc_state, fp32_smc=True, ) except Exception as exc: print(f"[CWM/SMC] Warning: skipped at this step – {exc}") return uncond_t + float(cond_scale) * (cond_t - uncond_t) m.model_options["sampler_cfg_function"] = cfg_hook m.model_options["disable_cfg1_optimization"] = True # ── Post-CFG hook: DCW ──────────────────────────────────────────────── if dcw_active: existing_fns = list(m.model_options.get("sampler_post_cfg_function", [])) _ll = lambda_l _lh = lambda_h _rdc = rdc_on _rdc_a_ll = rdc_alpha_ll if rdc_on else 0.0 _rdc_a_hh = rdc_alpha_hh if rdc_on else 0.0 _rdc_tau = rdc_tau if rdc_on else 0.0 # model_options is a fresh dict per KSampler run, but the closure # persists across the node's lifetime — same pattern as the SMC # cfg_hook above (_model_options_ref fallback). _model_options_ref_dcw = m.model_options if rdc_on: print( f"[DCW/RDC] Band-wise drift compensation ON " f"alpha_LL={_rdc_a_ll} alpha_HH={_rdc_a_hh} tau={_rdc_tau}" ) def dcw_post_cfg(args: dict) -> torch.Tensor: denoised = args.get("denoised") x_t = args.get("input") sigma = args.get("sigma") if denoised is None or x_t is None or sigma is None: return denoised try: rdc_state = None if _rdc: opts = args.get("model_options", _model_options_ref_dcw) rdc_state = opts.setdefault(_RDC_STATE_KEY, {}) return apply_dcw( denoised, x_t, sigma, _ll, _lh, rdc_alpha_ll=_rdc_a_ll, rdc_alpha_hh=_rdc_a_hh, rdc_tau=_rdc_tau, rdc_state=rdc_state, ) except Exception as exc: print(f"[DCW] Warning: correction skipped at this step – {exc}") return denoised existing_fns.append(dcw_post_cfg) m.model_options["sampler_post_cfg_function"] = existing_fns return (m,) # ───────────────────────────────────────────────────────────────────────────── # ComfyUI Registration # ───────────────────────────────────────────────────────────────────────────── NODE_CLASS_MAPPINGS = { "DCWModelPatch": DCWModelPatch, } NODE_DISPLAY_NAME_MAPPINGS = { "DCWModelPatch": "DCW(+a)", }