nodes/lookuters22-ComfyUI-ApplySpectralLock
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix blur: replace energy-normalization with detail-preserving HF soft-clip (clip_k)

Browse Source
This commit is contained in:
lookuters22
2026-05-29 22:47:05 +02:00
parent 151c4ef4ee
commit 2df084da20
2 changed files with 23 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+3 -3
View File
@@ -34,7 +34,8 @@ Restart ComfyUI.
| `kernel_size` | Odd kernel size for separable Gaussian low-pass (default 15). |
| `blur_sigma` | Gaussian sigma for the 1D kernels (default 3.0). |
| `alpha_lock` | Structure lock: blend weight for anchoring low frequencies to `latent_orig` (0–1, default 0.85). Active from step 0. |
| `gamma_max` | Detail gate: maximum strength of high-frequency variance matching toward the reference (0–1, default 0.70). Higher = texture energy pulled harder toward the original (less grain). |
| `gamma_max` | Detail gate: maximum blend toward the soft-clipped high band over the final steps (0–1, default 0.70). Higher = stronger grain suppression. |
| `clip_k` | Soft-clip knee in units of the high-band sigma (default 2.0). Spikes beyond `k*sigma` are tamed; ordinary detail (`|H| << knee`) passes through. Lower = more aggressive grain removal; higher = gentler / more detail kept. |
| `decay_type` | `cosine` (default, slow start then firm asymptote) or `linear` progress scaling for the detail gate. |
| `edge_taper` | Spatial edge-gating width in latent cells (default 48, `0` disables). The spectral adjustment is cosine-tapered to zero at the tile border so USDU seam blending sees no discontinuity. Auto-clamped to half the tile size. |
@@ -51,8 +52,7 @@ The hook is registered with **`disable_cfg1_optimization=True`** so the post-CFG
## Implementation notes
- **Depthwise separable** Gaussian: two `F.conv2d` passes with `groups=16`, **`F.pad(..., mode='reflect')`** before each pass.
- **HF variance matching**: per-channel spatial std of the reference high-band is cached once per tile; the generated high-band is rescaled by `std_orig / (std_t + eps)` and blended in via the cosine gate. A gain clamp prevents runaway amplification on flat patches.
- **Covariance-preserving (joint) epsilon**: the denominator floor is tied to the *mean energy across the 16 channels* (`eps_abs + eps_rel * mean_c(std_t)`), not to each channel in isolation. Near-flat channels inherit the joint floor instead of being independently amplified, keeping inter-channel scale ratios stable and preventing microscopic color/tonal drift (and division-by-zero on studio backdrops).
- **Detail-preserving grain suppression**: the high band is soft-clipped with a `tanh` knee at `clip_k * sigma`. Grain/"deep-fry" is sparse extreme spikes, so clipping tames them while ordinary structured detail (`|H| << knee`) passes through unchanged — it does NOT normalize energy toward the reference (which would blur detail down to a soft base).
- **Spatial edge-gating**: a cached cosine taper mask zeroes the spectral delta `(X_new - denoised)` at tile borders (`edge_taper`), so USDU seam blending never reconciles a discontinuity introduced here.
- **Decoupled envelopes**: structure lock (`alpha`) is applied every step; the variance gate (`gamma`) ramps with schedule progress.
- **Float32** for blur / variance / NCC internals; outputs are cast back to **`denoised.dtype`** and device.
spectral_lock.py
+20 -29
View File
@@ -6,11 +6,10 @@ high (H) frequency bands. Two decoupled controls are then applied:
* Structure lock (alpha): low frequencies are anchored to a reference latent.
Active from step 0 to preserve composition / color / lighting.
* Detail gate (gamma, cosine envelope): attenuate-only grain suppression. A single
per-tile scalar (reduced over all channels, so color/tone is preserved) scales the
high band DOWN toward the reference energy when the tile is grainier than the
reference, and never amplifies. Flat reference tiles are left untouched to avoid
per-tile gain differences (seams). Engages over the final steps via the cosine gate.
* Detail-preserving grain suppression (gamma, cosine envelope): the high band is
soft-clipped with a tanh knee at k*sigma. Sparse extreme spikes (grain / "deep-fry")
are tamed while ordinary structured detail passes through unchanged, so detail is
preserved rather than blurred. Engages over the final steps via the cosine gate.
Latent-space normalized cross-correlation aligns USDU sub-tiles to the reference,
and a cosine spatial edge-taper forces the adjustment to zero at tile borders so
@@ -159,6 +158,7 @@ class SpectralLockPatcher:
"blur_sigma": ("FLOAT", {"default": 3.0, "min": 0.1, "max": 10.0, "step": 0.1}),
"alpha_lock": ("FLOAT", {"default": 0.85, "min": 0.0, "max": 1.0, "step": 0.01}),
"gamma_max": ("FLOAT", {"default": 0.70, "min": 0.0, "max": 1.0, "step": 0.01}),
"clip_k": ("FLOAT", {"default": 2.0, "min": 0.5, "max": 6.0, "step": 0.1}),
"decay_type": (["cosine", "linear"],),
"edge_taper": ("INT", {"default": 8, "min": 0, "max": 256, "step": 1}),
}
@@ -179,6 +179,7 @@ class SpectralLockPatcher:
gamma_max: float,
decay_type: str,
edge_taper: int = 8,
clip_k: float = 2.0,
) -> Tuple[Any]:
m = model.clone()
X_orig_full = latent_orig["samples"]
@@ -293,8 +294,7 @@ class SpectralLockPatcher:
L_t = gaussian(denoised).to(dtype=dtype)
H_t = denoised - L_t
L_orig_tile = state["cached_L_orig"]
X_orig_cropped = state["cached_X_orig_cropped"]
if L_orig_tile is None or X_orig_cropped is None:
if L_orig_tile is None:
return denoised
# --- Structure lock (decoupled): fully active from step 0 ---
@@ -302,29 +302,20 @@ class SpectralLockPatcher:
alpha = float(alpha_lock)
L_new = (L_orig_tile * alpha) + (L_t * (1.0 - alpha))
# --- Detail gate: attenuate-only, color-safe grain suppression ---
# Goal: remove *excess* high-frequency energy (grain/overshoot) without
# manufacturing contrast. We therefore only ever scale H DOWN, never up,
# and use a single per-tile scalar (reduced over channels too) so the
# inter-channel balance -- i.e. color/tone -- is preserved.
H_orig = X_orig_cropped - L_orig_tile
std_orig = H_orig.to(torch.float32).std(dim=[1, 2, 3], keepdim=True)
std_t = H_t.to(torch.float32).std(dim=[1, 2, 3], keepdim=True)
# --- Detail-preserving grain suppression (soft-clip on HF outliers) ---
# KEY: do NOT normalize energy toward the (soft) reference -- that erases
# the detail the refiner adds and blurs the result. Grain / "deep-fry" is
# sparse EXTREME spikes in the high band, whereas genuine detail is moderate.
# A tanh soft-knee at k*sigma tames the spikes while passing ordinary detail
# through unchanged (tanh(x)~x for |x| << knee), so detail is preserved.
H_t_32 = H_t.to(torch.float32)
sigma_h = H_t_32.std(dim=[2, 3], keepdim=True) # per-channel spatial spread
knee = (float(clip_k) * sigma_h) + 1e-4
H_soft = knee * torch.tanh(H_t_32 / knee)
# Ratio < 1 only when the generated tile is grainier than the reference.
# Clamped to <= 1.0: we suppress excess energy but never amplify (no deep-fry).
gain = torch.clamp(std_orig / (std_t + 1e-3), max=1.0)
# Flat-region guard: if the reference tile is essentially textureless
# (studio backdrop), leave it untouched. This keeps every tile identical in
# flat areas, eliminating the per-tile gain differences that cause seams.
flat = std_orig < 1e-2
gain = torch.where(flat, torch.ones_like(gain), gain)
# Cosine gate: no change early (gamma_t~0), full suppression late.
# multiplier = 1 - gamma_t * (1 - gain) in [gain, 1].
multiplier = (1.0 - gamma_t * (1.0 - gain)).to(dtype)
H_new = H_t * multiplier
# Cosine gate: no change early (gamma_t~0), clip engages over the final steps.
# gamma_t already folds in gamma_max as the maximum blend toward the clipped HF.
H_new = (H_t_32 + gamma_t * (H_soft - H_t_32)).to(dtype)
# --- Spatial edge-gating (FIX 3) ---
# Force the entire spectral adjustment to zero at the absolute tile border