From 2df084da20e894ef4aacc0679de9e068de8656a9 Mon Sep 17 00:00:00 2001 From: lookuters22 Date: Fri, 29 May 2026 22:47:05 +0200 Subject: [PATCH] Fix blur: replace energy-normalization with detail-preserving HF soft-clip (clip_k) --- README.md | 6 +++--- spectral_lock.py | 49 ++++++++++++++++++++---------------------------- 2 files changed, 23 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index 184eff7..ff71156 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/spectral_lock.py b/spectral_lock.py index 76007f1..53301d5 100644 --- a/spectral_lock.py +++ b/spectral_lock.py @@ -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