624 lines
28 KiB
Python
624 lines
28 KiB
Python
"""
|
||
CNS Sampler Patch — Colored Noise Sampling for ComfyUI
|
||
=======================================================
|
||
|
||
Patches any ancestral-type SAMPLER to replace uniform white noise injection
|
||
with frequency-aware colored noise via Haar wavelet subband energy weighting.
|
||
|
||
Paper: "Colored Noise Diffusion Sampling"
|
||
Davidson et al., arXiv:2605.30332v1 (2026)
|
||
|
||
Wavelet utilities: adapted from dcw_node.py (Haar DWT, identical API)
|
||
|
||
─────────────────────────────────────────────────────────────────────────────
|
||
Architecture overview
|
||
─────────────────────────────────────────────────────────────────────────────
|
||
|
||
SAMPLER → CNSSamplerPatch → SAMPLER (patched)
|
||
|
||
The patch intercepts the sampler's noise_sampler callable and replaces it
|
||
with a CNS colored-noise sampler. A callback hook captures x_t each step
|
||
so the wavelet-based γ proxy stays synchronized with the current state.
|
||
|
||
CNS core (per step):
|
||
1. Haar DWT on x_t → LL, LH, HL, HH subband energies
|
||
2. γ proxy per band: γ_f ≈ e_f / Σe (energy fraction = "built-ness")
|
||
3. deficit per band: d_f = 1 − γ_f (how unresolved that band still is)
|
||
4. β_f = d_f^gamma_power (unnormalized noise scale weight)
|
||
5. RMS-normalize β so that mean(β²) = 1
|
||
6. Apply β to noise subbands, IDWT back to spatial domain
|
||
7. Crop-energy correction: rescale std after [..,:H,:W] crop (odd-dim guard)
|
||
8. Lerp with white noise by `strength` (0=off, 1=full CNS)
|
||
|
||
Wavelet γ proxy vs. CNS original:
|
||
CNS uses FFT + radial frequency bins + precomputed γ(f,t) ODE trajectory.
|
||
This implementation approximates γ from the live x_t subband energies,
|
||
requiring NO separate ODE prepass and no precomputed matrices.
|
||
|
||
Early steps (x_t ≈ pure noise):
|
||
All subband energies roughly equal → γ_f ≈ 0.25 for all f
|
||
→ all β ≈ 1 → white noise (correct: CNS starts as white noise)
|
||
|
||
Late steps (x_t ≈ clean image, natural 1/f² spectrum):
|
||
LL energy >> HH energy
|
||
→ γ_LL high → d_LL low → β_LL low (suppress resolved low-freq noise)
|
||
→ γ_HH low → d_HH high → β_HH high (boost unresolved high-freq noise)
|
||
|
||
─────────────────────────────────────────────────────────────────────────────
|
||
Compatibility
|
||
─────────────────────────────────────────────────────────────────────────────
|
||
Tested: euler_ancestral, euler_ancestral_cfg_pp,
|
||
dpm_2_ancestral, dpmpp_2s_ancestral, dpmpp_sde
|
||
Deterministic ODEs (euler, dpm++2m, ddim …) do not inject noise →
|
||
CNS patch is silently inert (noise_sampler never called).
|
||
"""
|
||
|
||
import inspect
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
import comfy.samplers
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# Debug Logging Switch
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# True로 바꾸면 매 스텝마다 [CNS DEBUG] x_t.shape / parity 정보를 콘솔에 출력합니다.
|
||
# (5D 텐서 발생 여부, H/W 홀짝 전환 시점 등을 진단할 때 사용)
|
||
# 평소에는 False로 두세요 — 매 스텝 print는 콘솔이 매우 시끄러워집니다.
|
||
DEBUG_LOGGING = False
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# Haar Wavelet Utilities (API-identical to dcw_node.py)
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
|
||
def _pad_even(x: torch.Tensor):
|
||
"""
|
||
Pad H and W to even numbers (Haar DWT requires even spatial dims).
|
||
Returns (padded_tensor, (original_H, original_W)).
|
||
|
||
Works for any number of leading dimensions (4-D image, 5-D video, …).
|
||
|
||
IMPORTANT: F.pad's reflect/replicate modes require the pad tuple length
|
||
to match the input rank in a strict way:
|
||
3-D / 4-D input → pad length 4 (pads last 2 dims)
|
||
4-D / 5-D input → pad length 6 (pads last 3 dims)
|
||
A 5-D video latent (B, C, T, H, W) therefore needs a 6-length pad tuple
|
||
(0, pw, 0, ph, 0, 0) — the trailing (0, 0) leaves T untouched — or
|
||
PyTorch raises "Padding size N is not supported for 5D input tensor".
|
||
We build the pad tuple dynamically so both 4-D and 5-D inputs work.
|
||
"""
|
||
H, W = x.shape[-2], x.shape[-1]
|
||
ph, pw = H % 2, W % 2
|
||
if ph or pw:
|
||
ndim = x.dim()
|
||
if ndim >= 5:
|
||
# Pad only the last 2 dims (H, W); leave dim -3 (e.g. T) alone.
|
||
pad = (0, pw, 0, ph, 0, 0)
|
||
else:
|
||
pad = (0, pw, 0, ph)
|
||
|
||
mode = "reflect" if (H >= 2 and W >= 2) else "replicate"
|
||
try:
|
||
x = F.pad(x, pad, mode=mode)
|
||
except RuntimeError:
|
||
# Final safety net for any unsupported rank/size combination
|
||
# (e.g. H or W == 1, or ndim outside the reflect/replicate range).
|
||
x = F.pad(x, pad, mode="constant", value=0.0)
|
||
return x, (H, W)
|
||
|
||
|
||
def haar_dwt2d(x: torch.Tensor):
|
||
"""
|
||
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 – coarse structure / global energy
|
||
LH – horizontal edges
|
||
HL – vertical edges
|
||
HH – 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 dims (4-D image, 5-D video).
|
||
"""
|
||
*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
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# CNS Core
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
|
||
def _subband_energy(sub: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Mean squared energy of a wavelet subband, averaged over C×H×W per batch item.
|
||
Returns shape (B, 1, 1, 1) — broadcastable with (B, C, H, W) subband tensors.
|
||
"""
|
||
dims = tuple(range(1, sub.dim())) # all dims after B
|
||
return sub.float().pow(2).mean(dim=dims, keepdim=True).clamp(min=1e-8)
|
||
|
||
|
||
def color_noise_wavelet(
|
||
noise: torch.Tensor,
|
||
x_t: torch.Tensor,
|
||
strength: float = 1.0,
|
||
gamma_power: float = 0.5,
|
||
gamma_scale: float = 2.0,
|
||
) -> torch.Tensor:
|
||
"""
|
||
Apply CNS wavelet coloring to a white noise tensor.
|
||
|
||
Parameters
|
||
----------
|
||
noise White noise to color, shape (B, C, H, W).
|
||
x_t Current noisy latent at this sampling step, same shape.
|
||
strength Lerp weight: 0.0 = white noise unchanged, 1.0 = full CNS.
|
||
gamma_power Exponent applied to the deficit before computing β.
|
||
0.5 (√) matches CNS paper default.
|
||
gamma_scale Divides γ proxy values to soften the routing effect.
|
||
Equivalent to "γ Divider" in CNS paper (default 1.73–25.0).
|
||
Returns
|
||
-------
|
||
Colored noise tensor with identical std to input (variance-preserving).
|
||
"""
|
||
if strength == 0.0:
|
||
return noise
|
||
|
||
# ── DEBUG LOGGING ──────────────────────────────────────────────────────
|
||
# 상단의 DEBUG_LOGGING = True 로 바꾸면 활성화됩니다 (기본값: False).
|
||
# 매 스텝 x_t의 shape / 차원 / H,W 패리티를 출력하고, 직전 스텝과 모양이
|
||
# 달라진 경우 "← 변경됨" 표시를 붙여 줍니다 (progressive upscale 전환
|
||
# 시점이나 5D 텐서 발생 시점을 한눈에 찾을 수 있습니다).
|
||
if DEBUG_LOGGING:
|
||
shape_now = tuple(x_t.shape)
|
||
prev_shape = getattr(color_noise_wavelet, "_prev_shape", None)
|
||
step_i = getattr(color_noise_wavelet, "_step_i", 0)
|
||
|
||
changed_mark = ""
|
||
if prev_shape is not None and prev_shape != shape_now:
|
||
changed_mark = " ← 변경됨! (이전 단계와 shape 다름)"
|
||
|
||
print(
|
||
f"[CNS DEBUG] step={step_i:04d} | x_t.shape={shape_now} "
|
||
f"(dim={x_t.dim()}) | H,W parity=({x_t.shape[-2] % 2},{x_t.shape[-1] % 2})"
|
||
f"{changed_mark}"
|
||
)
|
||
|
||
color_noise_wavelet._prev_shape = shape_now
|
||
color_noise_wavelet._step_i = step_i + 1
|
||
|
||
orig_dtype = noise.dtype
|
||
noise_f = noise.float()
|
||
x_f = x_t.float()
|
||
|
||
# ── Step 1: Haar DWT on x_t → subband energies ───────────────────────────
|
||
x_p, (H, W) = _pad_even(x_f)
|
||
LL_x, LH_x, HL_x, HH_x = haar_dwt2d(x_p)
|
||
|
||
e_LL = _subband_energy(LL_x) # (B,1,1,1)
|
||
e_LH = _subband_energy(LH_x)
|
||
e_HL = _subband_energy(HL_x)
|
||
e_HH = _subband_energy(HH_x)
|
||
e_sum = e_LL + e_LH + e_HL + e_HH # total energy per batch item
|
||
|
||
# ── Step 2: γ proxy — energy fraction = "how built is this band" ─────────
|
||
# Divided by gamma_scale to prevent premature over-routing (CNS Appendix C.2)
|
||
g_LL = (e_LL / e_sum / gamma_scale).clamp(0.0, 1.0)
|
||
g_LH = (e_LH / e_sum / gamma_scale).clamp(0.0, 1.0)
|
||
g_HL = (e_HL / e_sum / gamma_scale).clamp(0.0, 1.0)
|
||
g_HH = (e_HH / e_sum / gamma_scale).clamp(0.0, 1.0)
|
||
|
||
# ── Step 3: deficit = 1 − γ ("how unresolved") ──────────────────────────
|
||
d_LL = (1.0 - g_LL).clamp(min=1e-8)
|
||
d_LH = (1.0 - g_LH).clamp(min=1e-8)
|
||
d_HL = (1.0 - g_HL).clamp(min=1e-8)
|
||
d_HH = (1.0 - g_HH).clamp(min=1e-8)
|
||
|
||
# ── Step 4: β weights — proportional to deficit, shaped by gamma_power ───
|
||
b_LL = d_LL.pow(gamma_power)
|
||
b_LH = d_LH.pow(gamma_power)
|
||
b_HL = d_HL.pow(gamma_power)
|
||
b_HH = d_HH.pow(gamma_power)
|
||
|
||
# ── Step 5: RMS normalization — enforces mean(β²) = 1 (CNS Eq. 11) ─────
|
||
# Shapes the relative noise distribution across subbands while keeping
|
||
# the expected total injected energy constant. β ratios are preserved;
|
||
# only the overall scale is set. This alone is exact when H and W are
|
||
# both even (no padding needed → no crop → no energy leak).
|
||
rms = ((b_LL**2 + b_LH**2 + b_HL**2 + b_HH**2) / 4.0).sqrt().clamp(min=1e-8)
|
||
b_LL = b_LL / rms
|
||
b_LH = b_LH / rms
|
||
b_HL = b_HL / rms
|
||
b_HH = b_HH / rms
|
||
|
||
# ── Step 6: Apply β to noise subbands ────────────────────────────────────
|
||
n_p, _ = _pad_even(noise_f)
|
||
nLL, nLH, nHL, nHH = haar_dwt2d(n_p)
|
||
|
||
# β shape (B,1,1,1) broadcasts over (B,C,h,w)
|
||
nLL = nLL * b_LL
|
||
nLH = nLH * b_LH
|
||
nHL = nHL * b_HL
|
||
nHH = nHH * b_HH
|
||
|
||
n_colored = haar_idwt2d(nLL, nLH, nHL, nHH)[..., :H, :W]
|
||
|
||
# ── Step 7: Crop-energy correction ───────────────────────────────────────
|
||
# When H or W is odd, _pad_even added a reflect row/col before DWT, and
|
||
# the [..,:H,:W] crop above discards the corresponding IDWT output row/col.
|
||
# That crop introduces a small energy leak that Step 5 cannot anticipate.
|
||
# A single scalar rescale restores the correct std without disturbing the
|
||
# subband β ratios established in Step 5.
|
||
# For even-sized inputs there is no crop, so this is a near-identity op.
|
||
orig_std = noise_f.std().clamp(min=1e-8)
|
||
n_colored = n_colored * (orig_std / n_colored.std().clamp(min=1e-8))
|
||
|
||
# ── Step 8: Strength lerp ─────────────────────────────────────────────────
|
||
if strength < 1.0:
|
||
n_colored = torch.lerp(noise_f, n_colored, strength)
|
||
|
||
return n_colored.to(dtype=orig_dtype)
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# Noise Sampler Factory
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
|
||
def make_cns_noise_sampler(
|
||
x_initial: torch.Tensor,
|
||
strength: float,
|
||
gamma_power: float,
|
||
gamma_scale: float,
|
||
original_noise_sampler=None,
|
||
):
|
||
"""
|
||
Build a CNS noise sampler callable and a mutable state dict.
|
||
|
||
The state dict holds ``state["x_current"]``, which must be updated each
|
||
step via a k-diffusion callback so the γ proxy reflects the live x_t.
|
||
|
||
Returns
|
||
-------
|
||
(cns_noise_sampler_fn, state_dict)
|
||
|
||
cns_noise_sampler_fn has signature fn(sigma, sigma_next) → noise tensor
|
||
state_dict caller should update ["x_current"] each step
|
||
"""
|
||
state = {"x_current": x_initial.detach().clone()}
|
||
|
||
# 새 샘플링 런이 시작될 때 디버그 카운터/이전 shape 기록을 리셋합니다.
|
||
# (리셋하지 않으면 이전 생성의 마지막 shape이 남아 있어 새 런의 첫 스텝에
|
||
# 불필요한 "← 변경됨" 표시가 잘못 뜰 수 있습니다.)
|
||
if DEBUG_LOGGING:
|
||
color_noise_wavelet._prev_shape = None
|
||
color_noise_wavelet._step_i = 0
|
||
|
||
def cns_noise_sampler(sigma, sigma_next):
|
||
x_t = state["x_current"]
|
||
|
||
# Generate base white noise
|
||
if original_noise_sampler is not None:
|
||
try:
|
||
noise = original_noise_sampler(sigma, sigma_next)
|
||
except Exception:
|
||
noise = torch.randn_like(x_t)
|
||
else:
|
||
noise = torch.randn_like(x_t)
|
||
|
||
try:
|
||
return color_noise_wavelet(
|
||
noise, x_t,
|
||
strength=strength,
|
||
gamma_power=gamma_power,
|
||
gamma_scale=gamma_scale,
|
||
)
|
||
except Exception as exc:
|
||
print(f"[CNS] Warning: coloring skipped at this step – {exc}")
|
||
return noise
|
||
|
||
return cns_noise_sampler, state
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# ComfyUI Node
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
|
||
class CNSSamplerPatch:
|
||
"""
|
||
CNS Sampler Patch node.
|
||
|
||
Wraps any ancestral SAMPLER and replaces its white noise injection with
|
||
frequency-aware colored noise (CNS).
|
||
|
||
Usage in SamplerCustomAdvanced
|
||
──────────────────────────────
|
||
SAMPLER (e.g. euler_ancestral_cfg_pp)
|
||
│
|
||
▼
|
||
CNSSamplerPatch
|
||
│
|
||
▼
|
||
SamplerCustomAdvanced ← MODEL (with DCW patch, optional)
|
||
|
||
The DCWModelPatch and CNSSamplerPatch are orthogonal:
|
||
• DCW corrects x0_pred in the post-cfg hook (model-level).
|
||
• CNS colors the ancestral noise in the sampler step (sampler-level).
|
||
Both can be active simultaneously with no hook conflicts.
|
||
"""
|
||
|
||
RETURN_TYPES = ("SAMPLER",)
|
||
RETURN_NAMES = ("sampler",)
|
||
FUNCTION = "patch"
|
||
CATEGORY = "sampling/custom_sampling/samplers"
|
||
DESCRIPTION = (
|
||
"CNS Sampler Patch: replaces white ancestral noise with wavelet-colored noise.\n"
|
||
"Low-frequency noise is suppressed once those structures are built; "
|
||
"high-frequency noise is boosted while fine details remain unresolved.\n"
|
||
"Variance is strictly preserved per step (CNS energy budget constraint).\n"
|
||
"Compatible with euler_ancestral, euler_ancestral_cfg_pp, dpm_2_ancestral, "
|
||
"dpmpp_2s_ancestral, dpmpp_sde. Inert on deterministic (ODE) samplers."
|
||
)
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"sampler": ("SAMPLER",),
|
||
|
||
"strength": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.05,
|
||
"round": 0.01,
|
||
"tooltip": (
|
||
"[CNS] Effect strength. Lerp between white noise (0.0) and "
|
||
"fully colored noise (1.0).\n"
|
||
"Start at 1.0 and reduce only if artifacts appear."
|
||
),
|
||
}),
|
||
|
||
"gamma_power": ("FLOAT", {
|
||
"default": 0.5,
|
||
"min": 0.1,
|
||
"max": 2.0,
|
||
"step": 0.05,
|
||
"round": 0.01,
|
||
"tooltip": (
|
||
"[CNS] Exponent applied to the subband deficit before computing β.\n"
|
||
"0.5 (√) = CNS paper default — smooth, proportional routing.\n"
|
||
"Higher (0.75–1.0) → sharper frequency routing, more aggressive.\n"
|
||
"Lower (0.25–0.4) → softer routing, closer to white noise."
|
||
),
|
||
}),
|
||
|
||
"gamma_scale": ("FLOAT", {
|
||
"default": 2.0,
|
||
"min": 0.1,
|
||
"max": 25.0,
|
||
"step": 0.1,
|
||
"round": 0.01,
|
||
"tooltip": (
|
||
"[CNS] γ proxy divider. Scales down the subband energy ratios "
|
||
"before computing deficits, preventing premature over-routing.\n"
|
||
"Equivalent to 'γ Divider' in CNS paper (range 1.73–25.0).\n"
|
||
"1.0 = most aggressive (full energy ratio).\n"
|
||
"2.0–5.0 = recommended starting range for most models.\n"
|
||
"Increase if early steps show structural instability."
|
||
),
|
||
}),
|
||
}
|
||
}
|
||
|
||
def patch(
|
||
self,
|
||
sampler,
|
||
strength: float,
|
||
gamma_power: float,
|
||
gamma_scale: float,
|
||
):
|
||
original_fn = sampler.sampler_function
|
||
|
||
# Capture params in closure
|
||
_strength = strength
|
||
_gamma_power = gamma_power
|
||
_gamma_scale = gamma_scale
|
||
def _wrap_kd_fn(kd_fn):
|
||
"""
|
||
Wrap a single k-diffusion ancestral sampler function with CNS.
|
||
Deterministic samplers never call noise_sampler so CNS is inert.
|
||
"""
|
||
import functools
|
||
|
||
@functools.wraps(kd_fn)
|
||
def wrapped(model, x, sigmas, **kwargs):
|
||
orig_noise_sampler = kwargs.pop("noise_sampler", None)
|
||
cns_ns, state = make_cns_noise_sampler(
|
||
x_initial=x,
|
||
strength=_strength,
|
||
gamma_power=_gamma_power,
|
||
gamma_scale=_gamma_scale,
|
||
original_noise_sampler=orig_noise_sampler,
|
||
)
|
||
user_callback = kwargs.pop("callback", None)
|
||
|
||
def cns_callback(info: dict):
|
||
x_current = info.get("x")
|
||
if x_current is not None:
|
||
state["x_current"] = x_current.detach()
|
||
if user_callback is not None:
|
||
user_callback(info)
|
||
|
||
kwargs["noise_sampler"] = cns_ns
|
||
kwargs["callback"] = cns_callback
|
||
return kd_fn(model, x, sigmas, **kwargs)
|
||
|
||
return wrapped
|
||
|
||
def patched_sampler_fn(model, x, sigmas, **kwargs):
|
||
"""
|
||
Top-level sampler function registered in KSAMPLER.
|
||
|
||
Two mutually exclusive paths:
|
||
|
||
A) Direct path — original_fn IS a kd sample_* function.
|
||
CNS is injected once via kwargs["noise_sampler"].
|
||
No monkey-patching needed; original_fn is called directly.
|
||
|
||
B) Wrapper path — original_fn is an outer wrapper (e.g. SPEED)
|
||
that internally looks up and calls kds.sample_* by name.
|
||
CNS is injected by temporarily monkey-patching all ancestral
|
||
sample_* functions in comfy.k_diffusion.sampling just before
|
||
the call, then restoring them in a finally block.
|
||
kwargs["noise_sampler"] is NOT touched here so the wrapper
|
||
can forward it unmodified (or ignore it).
|
||
|
||
The two paths are detected by checking whether original_fn's
|
||
unwrapped identity matches any function already in kds. If it
|
||
does, we are in path A and skip monkey-patching entirely.
|
||
"""
|
||
import comfy.k_diffusion.sampling as kds
|
||
|
||
# ── Identify kd sample_* functions and their noise_sampler support ─
|
||
kds_ancestral = {} # attr_name → (original_fn_obj, accepts_ns)
|
||
for attr in dir(kds):
|
||
if not attr.startswith("sample_"):
|
||
continue
|
||
fn = getattr(kds, attr, None)
|
||
if fn is None or not callable(fn):
|
||
continue
|
||
try:
|
||
unwrapped = inspect.unwrap(fn)
|
||
sig = inspect.signature(unwrapped)
|
||
has_var_kw = any(
|
||
p.kind == inspect.Parameter.VAR_KEYWORD
|
||
for p in sig.parameters.values()
|
||
)
|
||
accepts_ns = has_var_kw or "noise_sampler" in sig.parameters
|
||
except (ValueError, TypeError, StopIteration):
|
||
accepts_ns = True
|
||
kds_ancestral[attr] = (fn, accepts_ns)
|
||
|
||
# ── Detect which path applies ─────────────────────────────────────
|
||
try:
|
||
unwrapped_orig = inspect.unwrap(original_fn)
|
||
except Exception:
|
||
unwrapped_orig = original_fn
|
||
|
||
is_direct_kd = any(
|
||
inspect.unwrap(fn) is unwrapped_orig
|
||
for fn, _ in kds_ancestral.values()
|
||
if callable(fn)
|
||
)
|
||
|
||
orig_noise_sampler = kwargs.pop("noise_sampler", None)
|
||
|
||
if is_direct_kd:
|
||
# ── Path A: original_fn is itself a kd sampler ────────────────
|
||
# Check if it accepts noise_sampler (should, but be safe).
|
||
try:
|
||
sig_orig = inspect.signature(unwrapped_orig)
|
||
has_var_kw_orig = any(
|
||
p.kind == inspect.Parameter.VAR_KEYWORD
|
||
for p in sig_orig.parameters.values()
|
||
)
|
||
orig_accepts_ns = has_var_kw_orig or "noise_sampler" in sig_orig.parameters
|
||
except (ValueError, TypeError, StopIteration):
|
||
orig_accepts_ns = True
|
||
|
||
if orig_accepts_ns:
|
||
cns_ns, state = make_cns_noise_sampler(
|
||
x_initial=x,
|
||
strength=_strength,
|
||
gamma_power=_gamma_power,
|
||
gamma_scale=_gamma_scale,
|
||
original_noise_sampler=orig_noise_sampler,
|
||
)
|
||
user_callback = kwargs.pop("callback", None)
|
||
|
||
def cns_callback(info: dict):
|
||
x_current = info.get("x")
|
||
if x_current is not None:
|
||
state["x_current"] = x_current.detach()
|
||
if user_callback is not None:
|
||
user_callback(info)
|
||
|
||
kwargs["noise_sampler"] = cns_ns
|
||
kwargs["callback"] = cns_callback
|
||
else:
|
||
if orig_noise_sampler is not None:
|
||
kwargs["noise_sampler"] = orig_noise_sampler
|
||
|
||
return original_fn(model, x, sigmas, **kwargs)
|
||
|
||
else:
|
||
# ── Path B: wrapper — monkey-patch kds, then call ─────────────
|
||
patched_attrs = {
|
||
attr: (fn, _wrap_kd_fn(fn))
|
||
for attr, (fn, accepts_ns) in kds_ancestral.items()
|
||
if accepts_ns
|
||
}
|
||
for attr, (_, wrapped) in patched_attrs.items():
|
||
setattr(kds, attr, wrapped)
|
||
|
||
# Restore original noise_sampler into kwargs for the wrapper
|
||
# to forward (it is NOT wrapped again here — CNS lives in the
|
||
# monkey-patched kds functions only).
|
||
if orig_noise_sampler is not None:
|
||
kwargs["noise_sampler"] = orig_noise_sampler
|
||
|
||
try:
|
||
return original_fn(model, x, sigmas, **kwargs)
|
||
finally:
|
||
for attr, (original, _) in patched_attrs.items():
|
||
setattr(kds, attr, original)
|
||
|
||
new_sampler = comfy.samplers.KSAMPLER(
|
||
patched_sampler_fn,
|
||
extra_options=sampler.extra_options,
|
||
inpaint_options=sampler.inpaint_options,
|
||
)
|
||
return (new_sampler,)
|
||
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# Registration
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"CNSSamplerPatch": CNSSamplerPatch,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"CNSSamplerPatch": "CNS Sampler Patch (Colored Noise)",
|
||
}
|