A ComfyUI custom node for Colored Noise Sampling (CNS), based on Colored Noise Diffusion Sampling by Hadar Davidson, Noam Issachar, and Sagie Benaim.

This plugin adapts the public CNS noise-shaping code from HadarDavidson/colored-noise-sampling to ComfyUI's SAMPLER interface.

What's New

Aligned gamma matrix handling with the official implementation: bundled matrices are read as [step, freq_bin].
Moved energy_scale after FFT filtering and unit-std normalization, matching the official CNS code path.
Updated alpha tilting to follow the official residual-energy behavior.
Kept compatibility with ComfyUI latent tensors that have extra leading dimensions; the last two dimensions are treated as spatial H/W.
Cleaned up README/code comments and kept the node API stable for existing workflows.

Install

Place this folder in your ComfyUI custom_nodes directory:

ComfyUI/
+-- custom_nodes/
    +-- ComfyUI-CNS-Sampler-CHENGOU/
        +-- __init__.py
        +-- nodes.py
        +-- gamma_matrix_scaled.pt
        +-- README.md

Restart ComfyUI. The node appears under:

sampling / custom_sampling / samplers / CNS Sampler (Colored Noise) | CHENGOU

Workflow

Use this node with SamplerCustomAdvanced. It outputs a normal ComfyUI SAMPLER.

BasicScheduler
ModelPatcher
Conditioning (positive/negative)
RandomNoise / DisableNoise
CNS Sampler (Colored Noise)
        |
        +-- SamplerCustomAdvanced -- VAEDecode -- Image

KSampler does not expose the custom sampler slot in the same way, so SamplerCustomAdvanced is the intended path.

Parameters

Parameter	Default	Notes
`s_churn`	`0.5`	Stochastic strength for the ComfyUI Euler adaptation. `0` disables CNS noise injection.
`power_gamma`	`0.75`	Power on residual energy. Official unguided setting: `0.75`.
`gamma_divider`	`1.73`	Divides gamma before residual energy is computed. Official unguided setting: `1.73`; guided setting: `25.0`.
`energy_scale`	`0.98`	Applied after FFT filtering and unit-std normalization. Official unguided setting: `0.98`; guided setting: `0.998`.
`alpha_tilt_start`	`0.15`	Frequency tilt at the beginning of sampling.
`alpha_tilt_end`	`-0.5`	Frequency tilt at the end of sampling.
`alpha_use_fnorm`	`True`	Uses normalized radial frequency for alpha tilting. This matches the published settings.
`alpha_exp_interp`	`True`	Uses exponential interpolation between start/end alpha values.
`alpha_exp_sharpness`	`0.75`	Official unguided setting: `0.75`.
`num_freq_bins`	`32`	Official matrices use 32 radial frequency bins.
`use_bundled_gamma_matrix`	`True`	Loads `gamma_matrix_scaled.pt` from this folder. Disable to use the sigma-schedule fallback.

Suggested Settings

Unguided / low CFG:

s_churn          = 0.5
power_gamma      = 0.75
gamma_divider    = 1.73
energy_scale     = 0.98
alpha_tilt_start = 0.15
alpha_tilt_end   = -0.5
alpha_use_fnorm  = True
alpha_exp_interp = True

Guided / higher CFG:

s_churn          = 0.5
power_gamma      = 0.5
gamma_divider    = 25.0
energy_scale     = 0.998
alpha_tilt_start = -0.1
alpha_tilt_end   = 0.03
alpha_use_fnorm  = True
alpha_exp_interp = False

The bundled gamma_matrix_scaled.pt is the official unguided matrix from gamma_matrix/gamma_matrix_scaled.pt. For guided experiments, you can replace it with the official guided matrix, but keep the filename as gamma_matrix_scaled.pt unless you also edit the loader path.

Implementation Notes

The official repository implements CNS inside its SiT transport SDE integrator. This plugin keeps the same CNS frequency-shaping path while fitting into ComfyUI's Euler-style sampler call:

gamma matrix layout is [step, freq_bin], matching the official code;
residual energy is computed from 1 - gamma / gamma_divider;
alpha tilting uses exp(alpha * f_norm) * residual when alpha_use_fnorm=True;
FFT-filtered noise is normalized to unit std, then energy_scale is applied;
the final colored noise is injected through ComfyUI's sigma schedule.

If gamma_matrix_scaled.pt is missing or disabled, the plugin falls back to a simple sigma-schedule approximation. That fallback is useful for booting, but it is not a replacement for an ODE-derived gamma matrix.

中文说明

这是一个 ComfyUI 自定义采样器节点，基于 CNS 论文和官方代码实现。它不是官方 SiT 训练/评估仓库的完整移植，而是把官方 CNS 的频域噪声调度接到 ComfyUI 的 SAMPLER 接口上。

本版重点修正：

gamma_matrix_scaled.pt 按官方 [step, freq_bin] 方向读取；
energy_scale 放在频域滤波和标准差归一化之后，参数会真正生效；
alpha tilt 逻辑按官方 residual energy 写法实现；
保留 SamplerCustomAdvanced 工作流和原节点参数。

推荐先使用默认参数。如果你使用 guided / 高 CFG，可以参考上面的 guided 参数，但最好同时替换为官方 guided gamma matrix。

Citation

@misc{davidson2026colorednoisediffusionsampling,
      title={Colored Noise Diffusion Sampling},
      author={Hadar Davidson and Noam Issachar and Sagie Benaim},
      year={2026},
      eprint={2605.30332},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2605.30332},
}