From 52e4e0aa6f8d0a96d63cbdc01b4e090d5a457136 Mon Sep 17 00:00:00 2001
From: akashzeno <77930996+akashzeno@users.noreply.github.com>
Date: Mon, 1 Jun 2026 01:39:04 +0530
Subject: [PATCH] feat(engine): color_tensor FFT shaping with renorm, ndim
 dispatch, cpu fallback

---
 src/engine/spectral.py | 50 ++++++++++++++++++++++++++++++++++++++
 tests/test_spectral.py | 55 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+)

diff --git a/src/engine/spectral.py b/src/engine/spectral.py
index 0e3ec30..48f062a 100644
--- a/src/engine/spectral.py
+++ b/src/engine/spectral.py
@@ -48,3 +48,53 @@ def parametric_scaling(num_bins, alpha, device=torch.device("cpu")):
     if alpha > 0:
         scaling[0] = 0.0
     return scaling
+
+
+def _renorm_unit(t, energy_scale):
+    t = t - t.mean()
+    std = t.std()
+    if float(std) > 1e-9:
+        t = t / std
+    return t * energy_scale
+
+
+def color_tensor(white, scaling_vector, energy_scale=1.0):
+    """Spectrally shape ``white`` by a per-bin ``scaling_vector``.
+
+    Returns zero-mean, unit-variance (times ``energy_scale``) noise of the same
+    shape/dtype. FFT axes are derived from the rank: >=4D -> 2D spatial FFT over
+    the last two dims (image [B,C,H,W] and video [B,C,T,H,W]); 3D -> 1D FFT over
+    the last dim (audio [B,C,L]); <3D -> returned unchanged. Nested-tensor safe.
+    The FFT runs in float32 (half precision is unsupported on most backends) with
+    a CPU fallback for backends (DirectML/MPS) that reject ``torch.fft``.
+    """
+    if white.is_nested:
+        import comfy.nested_tensor as nt
+        return nt.NestedTensor(
+            [color_tensor(t, scaling_vector, energy_scale) for t in white.unbind()]
+        )
+
+    num_bins = scaling_vector.shape[0]
+    ndim = white.dim()
+    if ndim >= 4:
+        h, w = white.shape[-2], white.shape[-1]
+        grid = scaling_vector.to(white.device)[radial_bin_map_2d(h, w, num_bins, white.device)]
+        dims = (-2, -1)
+    elif ndim == 3:
+        length = white.shape[-1]
+        grid = scaling_vector.to(white.device)[radial_bin_map_1d(length, num_bins, white.device)]
+        dims = (-1,)
+    else:
+        return white  # cannot meaningfully shape (<3d, e.g. flat audio not in [B,C,L] form)
+
+    orig_dtype = white.dtype
+    wf = white.float()
+    try:
+        spec = torch.fft.fftn(wf, dim=dims)
+        shaped = torch.fft.ifftn(spec * grid, dim=dims).real
+    except RuntimeError:
+        # DirectML / MPS reject torch.fft -> CPU fallback (mirrors comfy_extras FreeU/FreSca).
+        spec = torch.fft.fftn(wf.cpu(), dim=dims)
+        shaped = torch.fft.ifftn(spec * grid.detach().cpu(), dim=dims).real.to(white.device)
+
+    return _renorm_unit(shaped, energy_scale).to(orig_dtype)
diff --git a/tests/test_spectral.py b/tests/test_spectral.py
index 6e57df3..ab64cc6 100644
--- a/tests/test_spectral.py
+++ b/tests/test_spectral.py
@@ -32,3 +32,58 @@ def test_parametric_blue_emphasizes_high_freq():
     s = spectral.parametric_scaling(64, alpha=-2.0)
     assert s[-1] > s[1]          # high-freq amplitude > low-freq
     assert torch.isfinite(s).all()
+
+
+def _unit(x):
+    return abs(float(x.mean())) < 1e-4 and abs(float(x.std()) - 1.0) < 1e-3
+
+
+def test_color_tensor_white_is_unit_variance_4d():
+    w = torch.randn(2, 4, 32, 32)
+    s = spectral.parametric_scaling(64, 0.0)
+    out = spectral.color_tensor(w, s)
+    assert out.shape == w.shape and _unit(out)
+
+
+def test_color_tensor_5d_video():
+    w = torch.randn(1, 16, 3, 24, 24)
+    s = spectral.parametric_scaling(64, -1.0)
+    out = spectral.color_tensor(w, s)
+    assert out.shape == w.shape and _unit(out)
+
+
+def test_color_tensor_3d_audio_uses_1d():
+    w = torch.randn(1, 64, 256)
+    s = spectral.parametric_scaling(32, 1.0)
+    out = spectral.color_tensor(w, s)
+    assert out.shape == w.shape and _unit(out)
+
+
+def test_color_tensor_2d_returns_unchanged():
+    w = torch.randn(8, 8)
+    s = spectral.parametric_scaling(16, 1.0)
+    out = spectral.color_tensor(w, s)
+    assert torch.equal(out, w)  # <3d -> passthrough
+
+
+def test_color_tensor_dtype_roundtrip_fp16():
+    w = torch.randn(1, 4, 16, 16, dtype=torch.float16)
+    s = spectral.parametric_scaling(32, 2.0)
+    out = spectral.color_tensor(w, s)
+    assert out.dtype == torch.float16
+
+
+def test_color_tensor_red_has_more_lowfreq_energy_than_blue():
+    torch.manual_seed(0)
+    w = torch.randn(1, 4, 64, 64)
+    red = spectral.color_tensor(w, spectral.parametric_scaling(128, 2.0))
+    blue = spectral.color_tensor(w, spectral.parametric_scaling(128, -2.0))
+
+    def lowfreq_frac(t):
+        m = torch.fft.fftshift(torch.fft.fft2(t[0, 0].float()).abs())
+        h, w_ = m.shape
+        cy, cx = h // 2, w_ // 2
+        center = m[cy - 4:cy + 4, cx - 4:cx + 4].sum()
+        return float(center / m.sum())
+
+    assert lowfreq_frac(red) > lowfreq_frac(blue)