diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..2f1cb60 --- /dev/null +++ b/LICENSE @@ -0,0 +1,27 @@ +MIT License + +Copyright (c) 2026 Akash (akashzeno) + +Portions of this software derive from "colored-noise-sampling" +(https://github.com/hadardavidson/colored-noise-sampling), Copyright (c) 2026 +Hadar Davidson, Roy Issachar, and Sagie Benaim, also under the MIT License: +- the bundled gamma matrices in gamma_matrices/ (verbatim), and +- the frequency-domain noise-shaping algorithm reimplemented in src/engine/. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..55a2ce0 --- /dev/null +++ b/README.md @@ -0,0 +1,122 @@ +# ComfyUI-ColoredNoiseSampling + +Frequency-shaped (**colored**) noise for ComfyUI's stochastic diffusion sampling — a clean, +native port of the *transferable mechanism* behind +[Colored Noise Diffusion Sampling (CNS)](https://github.com/hadardavidson/colored-noise-sampling) +(Davidson, Issachar & Benaim, 2026, arXiv:2605.30332). + +Instead of injecting plain white Gaussian noise at each stochastic step, this pack injects noise +whose **power spectrum is shaped** — and whose shape can **vary across the sampling trajectory** +(broadband early → high-frequency detail late). It plugs into ComfyUI's standard +`noise_sampler` seam, so it works with the existing custom-sampling graph and the all-in-one node. + +> **Honest scope.** The CNS *technique* ports cleanly and is genuinely useful as a quality/character +> knob. The paper's quantitative results (FID 6.27 vs 8.26) are specific to **SiT on ImageNet-256** +> and will **not** reproduce on SD/SDXL/Flux. The optional γ-matrix mode ships the paper's original +> matrices, but on non-SiT models they are a documented heuristic, not a calibrated schedule. For +> general use, prefer the **parametric** mode. + +--- + +## What it does (in one paragraph) + +At each stochastic step it draws white noise, FFTs it, multiplies the spectrum by a +radially-symmetric per-frequency profile, inverse-FFTs, and **renormalizes to unit variance** — +then hands that to the sampler, which applies its own per-step amplitude (`sigma_up · s_noise`). +The per-frequency profile is either a **power-law** `amplitude(f) ∝ f^(−α/2)` (α: 0 white, +1 pink, ++2 brown/red, −1 blue, −2 violet) with a time-varying exponent, or a **γ-matrix residual** +`(1 − γ(progress, f))` reproducing the paper's algorithm. + +## Nodes + +All three live under the **`sampling/colored_noise`** category. + +### 1. Colored Noise Sampler — `ColoredNoise_SamplerSelect` → `SAMPLER` +The core node. Outputs a `SAMPLER` for `SamplerCustom` / `SamplerCustomAdvanced`. Wraps a chosen +**stochastic** base sampler and colors its per-step noise. + +Key inputs: `base_sampler` (all stochastic ancestral/SDE samplers, incl. the `_RF` rectified-flow +variants for Flux/SD3), `eta`, `s_noise`, `mode` (`parametric` | `gamma_matrix`), +`alpha_start`/`alpha_end` + `interpolation` (the time-varying spectral tilt), the γ-matrix controls, +and `energy_scale`. + +### 2. Colored Noise (Initial) — `ColoredNoise_Noise` → `NOISE` +Outputs a `NOISE` object producing **colored initial latent noise** (a single constant color: +white/pink/brown/blue/violet or a custom α). Use it on the `noise` input of `SamplerCustomAdvanced`. +Honors `batch_index` determinism and nested (video) latents. + +### 3. Colored Noise KSampler — `ColoredNoise_KSampler` → `LATENT` +All-in-one (model + conditioning + latent → latent), for when you don't want to wire the custom +sampling graph. Picks a stochastic base sampler + scheduler + steps + cfg + the coloring controls, +plus an optional `color_initial_noise` toggle. + +## Two modes + +| Mode | What it does | When to use | +|---|---|---| +| **parametric** (default) | Power-law spectrum `f^(−α/2)`, α interpolated `alpha_start → alpha_end` across the trajectory (linear or exponential). Model-agnostic. | Everything. Start here. | +| **gamma_matrix** | Loads a `[steps, bins]` γ matrix and injects noise into the *unresolved* bands per the paper. Bundled SiT/ImageNet matrices included. | Reproducing/experimenting with the paper; SiT-like setups. | + +### Quick parametric recipes +- **White** (baseline / sanity): `alpha_start = 0`, `alpha_end = 0`. +- **CNS-like** (broadband → high-freq late): `alpha_start = 0`, `alpha_end = -1` … `-2`. +- **Constant pink/brown** (softer, low-freq emphasis): `alpha_start = alpha_end = 1` … `2`. + +`energy_scale` (default `1.0`) scales the noise std *after* unit-variance renorm — a deliberate +"heat" knob; values ≠ 1 intentionally change the effective noise level. + +## γ-matrix files + +A model folder **`colored_noise_gamma`** is registered automatically. The two bundled matrices +(`gamma_matrix_scaled.pt`, `gamma_matrix_scaled_cfg_1.5.pt`) appear in the dropdown. To add your +own, drop a raw `[steps, bins]` float32 tensor (`torch.save`, `.pt`/`.pth`) into +`models/colored_noise_gamma/`. *(New files may require a ComfyUI restart to appear.)* See +[`gamma_matrices/SOURCE.md`](gamma_matrices/SOURCE.md) for provenance. + +## Compatibility + +- **Attention backends** — fully compatible with SageAttention 2/3 and FlashAttention 2/3 (and + xformers / pytorch / split / sub-quad) **by construction**. This pack only shapes the stochastic + noise term; it never touches the model forward pass or `transformer_options`, where attention is + selected. Your attention choice is untouched. +- **Memory** — uses ComfyUI's native memory management only. Noise/FFT tensors are transient, + latent-sized, hold no model references, and do not impede offloading. +- **Precision** — the FFT runs in float32 with a CPU fallback (half-precision FFT is unsupported on + most backends), then casts back; output matches the latent dtype. +- **Latent ranks** — 4D image `[B,C,H,W]` and 5D video `[B,C,T,H,W]` get a 2D spatial FFT; 3D audio + `[B,C,L]` gets a 1D time FFT; nested (video) latents are colored per sub-tensor. +- **Samplers** — only **stochastic** samplers are offered (deterministic ones never inject per-step + noise, so coloring them would be a silent no-op). + +## Install + +Clone into `ComfyUI/custom_nodes/` and restart ComfyUI. No extra dependencies (torch comes with +ComfyUI). The nodes appear under `sampling/colored_noise`. + +## Usage + +**All-in-one:** `Load Checkpoint → Colored Noise KSampler → VAE Decode`. Set `base_sampler` +(e.g. `dpmpp_2m_sde`), `mode = parametric`, `alpha_start = 0`, `alpha_end = -1`. Render. + +**Custom-sampling graph:** +``` +Colored Noise (Initial) ─┐ +Colored Noise Sampler ───┤ +BasicScheduler ──────────┼─► SamplerCustomAdvanced ─► VAE Decode +BasicGuider ─────────────┘ +``` + +## Development / tests + +Pure-engine unit tests (no model weights), run in the `comfyenv` environment: +``` +cd custom_nodes/ComfyUI-ColoredNoiseSampling +python -m pytest tests/ -q # 26 tests: spectral shaping, schedule, batch determinism, ... +ruff check . +``` + +## Credits & license + +Method and bundled γ matrices: **Hadar Davidson, Roy Issachar, Sagie Benaim**, +[colored-noise-sampling](https://github.com/hadardavidson/colored-noise-sampling) (MIT), +arXiv:2605.30332. This pack is MIT-licensed; see [LICENSE](LICENSE). diff --git a/gamma_matrices/SOURCE.md b/gamma_matrices/SOURCE.md new file mode 100644 index 0000000..5e3897e --- /dev/null +++ b/gamma_matrices/SOURCE.md @@ -0,0 +1,19 @@ +# Bundled gamma matrices — provenance + +`gamma_matrix_scaled.pt` and `gamma_matrix_scaled_cfg_1.5.pt` are copied verbatim from +[hadardavidson/colored-noise-sampling](https://github.com/hadardavidson/colored-noise-sampling) +(MIT License), directory `gamma_matrix/`. + +Each is a float32 tensor of shape **[250 steps, 32 radial frequency bins]**, values in `[0, 1]`, +encoding empirical per-band *structural completion* `γ(step, freq)` for **SiT on ImageNet-256**. + +These values are **model- and dataset-specific**. The *shape* of the schedule (low frequencies +resolve early, high frequencies late) is a general property of diffusion models, so the matrices +behave sensibly as a heuristic on other models — but they are **not calibrated** for SD/SDXL/Flux +and will not reproduce the paper's FID figures there. Prefer the **parametric** mode for general +use; use these for experimentation or SiT-like setups. + +Paper: Davidson, Issachar & Benaim (2026), *Colored Noise Diffusion Sampling*, arXiv:2605.30332. + +To add your own matrices: drop a raw `[steps, bins]` float32 tensor saved with `torch.save` +(`.pt`/`.pth`) into `models/colored_noise_gamma/`; it will appear in the node's dropdown.