979 lines
44 KiB
Python
979 lines
44 KiB
Python
"""
|
||
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)",
|
||
}
|