From 177f99f96afa7ff398e748af703bc07ec24db936 Mon Sep 17 00:00:00 2001 From: lookuters22 Date: Fri, 29 May 2026 22:02:11 +0200 Subject: [PATCH] Initial commit: ApplySpectralLock ComfyUI node (Flux/USDU spectral handoff) --- .gitignore | 16 +++ README.md | 63 +++++++++ __init__.py | 16 +++ pyproject.toml | 22 +++ spectral_lock.py | 358 +++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 475 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 pyproject.toml create mode 100644 spectral_lock.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f8ce888 --- /dev/null +++ b/.gitignore @@ -0,0 +1,16 @@ +__pycache__/ +*.py[cod] +*$py.class +.Python +*.so +.venv/ +venv/ +.env +*.egg-info/ +dist/ +build/ +.idea/ +.vscode/ +*.log +.DS_Store +Thumbs.db diff --git a/README.md b/README.md new file mode 100644 index 0000000..43d7d5c --- /dev/null +++ b/README.md @@ -0,0 +1,63 @@ +# ComfyUI-ApplySpectralLock + +**Apply Spectral Lock (Flux/USDU)** — a ComfyUI model patch that performs a **Latent Spectral Handoff** on **16-channel Flux** latents during sampling. It uses two **decoupled** controls: + +- **Structure lock** — low frequencies are anchored to a reference latent (`latent_orig`). Active from step 0, so composition / color / lighting are preserved even at 0.8 denoise. +- **Detail gate (variance matching)** — instead of blanket dampening, the high frequencies are **renormalized to match the reference's per-channel contrast/energy**. This keeps real, crisp detail while stopping the Rectified Flow solver from overshooting into digital grain ("Turbo paradox"). The gate follows a **schedule-aware cosine (or linear) envelope**, engaging mostly over the final steps. + +## Requirements + +- ComfyUI (recent master) +- PyTorch +- **Flux-family models** with **16 latent channels** (SD 4-channel latents are passed through unchanged) + +## Install + +Clone into `ComfyUI/custom_nodes/`: + +```text +ComfyUI/custom_nodes/ComfyUI-ApplySpectralLock/ +``` + +Restart ComfyUI. + +## Usage + +1. Load your **MODEL** as usual. +2. Provide **`latent_orig`**: the full-resolution (or reference) latent corresponding to the uncorrupted base structure you want to lock to (image2latent / VAE encode of the guide image). +3. Run **ApplySpectralLock** → connect **`patched_model`** into your sampler / Ultimate SD Upscale chain **before** sampling. + +### Parameters + +| Input | Role | +|--------|------| +| `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). | +| `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. | + +### Ultimate SD Upscale (USDU) + +For tiled runs, the patcher estimates the tile’s **(y, x)** offset in latent space via **normalized template matching** on **channel-mean maps**, then crops `latent_orig` once per tile and caches the blurred reference low-pass for the rest of that tile’s steps. + +**Note:** Very large Gaussian radii relative to USDU overlap can still stress seams; keep `kernel_size` coherent with your tile overlap. + +### Flux CFG = 1.0 + +The hook is registered with **`disable_cfg1_optimization=True`** so the post-CFG callback still runs when CFG is disabled (common for distilled / single-stream Flux runs). + +## 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). +- **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. +- **Sigma progress**: schedule-agnostic `progress = 1 - sigma / max_sigma_seen` for the current tile/run. + +## License + +MIT (replace if you use a different license for your repo). diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..6b0415e --- /dev/null +++ b/__init__.py @@ -0,0 +1,16 @@ +"""Advanced Latent Spectral Handoff node for ComfyUI. + +Optimized for Flux.1 Rectified Flow architectures and tiled refinement (USDU). +""" + +from .spectral_lock import SpectralLockPatcher + +NODE_CLASS_MAPPINGS = { + "ApplySpectralLock": SpectralLockPatcher, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ApplySpectralLock": "Apply Spectral Lock (Flux/USDU)", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..4c95b76 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,22 @@ +[build-system] +requires = ["setuptools>=61"] +build-backend = "setuptools.build_meta" + +[project] +name = "comfyui-applyspectrallock" +version = "1.0.0" +description = "ApplySpectralLock — latent spectral handoff for Flux.1 / Rectified Flow and USDU tiled upscaling." +readme = "README.md" +requires-python = ">=3.10" +license = { text = "MIT" } +authors = [{ name = "ComfyUI-ApplySpectralLock contributors" }] +dependencies = ["torch"] + +[project.urls] +Repository = "https://github.com/REPLACE_ME/ComfyUI-ApplySpectralLock" + +# Comfy Registry (optional): https://docs.comfy.org/registry/overview +[tool.comfy] +PublisherId = "REPLACE_ME" +DisplayName = "ApplySpectralLock" +Icon = "" diff --git a/spectral_lock.py b/spectral_lock.py new file mode 100644 index 0000000..0145049 --- /dev/null +++ b/spectral_lock.py @@ -0,0 +1,358 @@ +""" +Latent Spectral Handoff for Flux.1 / Rectified Flow. + +Depthwise separable Gaussian low-pass splits each denoised latent into low (L) and +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): instead of blanket dampening, the high + frequencies are *variance matched* to the reference's per-channel energy, so + generated texture keeps real detail while the solver is prevented from + overshooting into digital grain. Engages mostly over the final steps. A joint + (cross-channel) epsilon keeps per-channel scaling coherent to avoid color drift. + +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 +USDU seam blending stays clean. +""" + +from __future__ import annotations + +import math +from typing import Any, Dict, Tuple + +import torch +import torch.nn.functional as F + +import comfy.model_patcher + + +class SeparableGaussianSmoothing(torch.nn.Module): + """ + Depthwise separable 1D Gaussian smoothing for 16-channel Flux latents. + Each channel is filtered in isolation (groups=C). Two 1D passes approximate + a 2D Gaussian with lower cost than a full 2D conv. + """ + + def __init__( + self, + channels: int = 16, + kernel_size: int = 15, + sigma: float = 3.0, + dtype: torch.dtype = torch.float32, + ) -> None: + super().__init__() + if kernel_size % 2 == 0: + raise ValueError("kernel_size must be odd for symmetric padding.") + self.kernel_size = int(kernel_size) + self.padding = self.kernel_size // 2 + self.channels = int(channels) + + x_cord = torch.arange(self.kernel_size, dtype=torch.float32) + x_grid = x_cord - float(self.padding) + gaussian_1d = torch.exp(-(x_grid**2) / (2.0 * float(sigma) ** 2)) + gaussian_1d = gaussian_1d / torch.sum(gaussian_1d) + + kernel_x = gaussian_1d.view(1, 1, -1, 1).repeat(self.channels, 1, 1, 1) + kernel_y = gaussian_1d.view(1, 1, 1, -1).repeat(self.channels, 1, 1, 1) + + self.register_buffer("kernel_x", kernel_x.to(dtype=dtype)) + self.register_buffer("kernel_y", kernel_y.to(dtype=dtype)) + + @torch.no_grad() + def forward(self, x: torch.Tensor) -> torch.Tensor: + original_dtype = x.dtype + device = x.device + kx = self.kernel_x.to(device=device, dtype=torch.float32) + ky = self.kernel_y.to(device=device, dtype=torch.float32) + x32 = x.to(dtype=torch.float32) + + x_padded = F.pad( + x32, + (self.padding, self.padding, self.padding, self.padding), + mode="reflect", + ) + out = F.conv2d(x_padded, kx, groups=self.channels) + out = F.conv2d(out, ky, groups=self.channels) + return out.to(dtype=original_dtype) + + +def _find_tile_offset_ncc(full: torch.Tensor, tile: torch.Tensor) -> Tuple[int, int]: + """ + Normalized cross-correlation in latent space: channel-mean maps, ZNCC-style + template as conv weight, argmax of correlation map -> (y, x) top-left of tile in full. + """ + if full.ndim != 4 or tile.ndim != 4: + return 0, 0 + b_f, _, hf, wf = full.shape + b_t, _, ht, wt = tile.shape + if b_f < 1 or b_t < 1 or ht > hf or wt > wf or ht < 1 or wt < 1: + return 0, 0 + + full_32 = full[0:1].to(dtype=torch.float32).mean(dim=1, keepdim=True) + tile_32 = tile[0:1].to(dtype=torch.float32).mean(dim=1, keepdim=True) + + tile_32 = tile_32 - tile_32.mean() + norm = torch.sqrt(torch.sum(tile_32**2)) + if float(norm.item()) <= 1e-6: + return 0, 0 + tile_32 = tile_32 / norm + + weight = tile_32 # [1,1,ht,wt] + corr = F.conv2d(full_32, weight) # [1,1, hf-ht+1, wf-wt+1] + flat = corr.view(-1) + if flat.numel() == 0: + return 0, 0 + max_idx = int(torch.argmax(flat).item()) + ow = int(corr.size(3)) + y = max_idx // ow + x = max_idx % ow + y = max(0, min(y, hf - ht)) + x = max(0, min(x, wf - wt)) + return y, x + + +def _build_edge_taper_mask( + h: int, + w: int, + taper: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """ + 2D separable cosine edge-taper mask: 1.0 in the interior, smoothly cosine-faded + to exactly 0.0 at the absolute tile border over the outer `taper` rows/cols. + + Forces any spectral modification to vanish at tile edges so the USDU blend mask + never has to reconcile a discontinuity introduced by this node. Returns [1,1,H,W]. + """ + taper = int(max(0, min(taper, h // 2, w // 2))) + if taper <= 0: + return torch.ones((1, 1, h, w), device=device, dtype=dtype) + + def _ramp(n: int) -> torch.Tensor: + idx = torch.arange(n, device=device, dtype=torch.float32) + edge_dist = torch.minimum(idx, (n - 1) - idx) + r = torch.ones(n, device=device, dtype=torch.float32) + m = edge_dist < taper + # 0 at the edge (edge_dist=0) -> 1 at edge_dist=taper, smooth cosine in between. + r[m] = 0.5 * (1.0 - torch.cos(math.pi * edge_dist[m] / float(taper))) + return r + + ry = _ramp(h) + rx = _ramp(w) + mask = torch.outer(ry, rx) # [H, W] + return mask.view(1, 1, h, w).to(dtype=dtype) + + +class SpectralLockPatcher: + """Model patcher node: injects post-CFG hook for latent spectral handoff.""" + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + return { + "required": { + "model": ("MODEL",), + "latent_orig": ("LATENT",), + "kernel_size": ("INT", {"default": 15, "min": 3, "max": 63, "step": 2}), + "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}), + "decay_type": (["cosine", "linear"],), + "edge_taper": ("INT", {"default": 48, "min": 0, "max": 256, "step": 1}), + } + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("patched_model",) + FUNCTION = "patch" + CATEGORY = "model_patches/spectral" + + def patch( + self, + model: Any, + latent_orig: Dict[str, Any], + kernel_size: int, + blur_sigma: float, + alpha_lock: float, + gamma_max: float, + decay_type: str, + edge_taper: int = 48, + ) -> Tuple[Any]: + m = model.clone() + X_orig_full = latent_orig["samples"] + + state: Dict[str, Any] = { + "max_sigma": None, + "gaussian_filter": None, + "cached_L_orig": None, + "cached_H_orig_std": None, + "cached_edge_mask": None, + "cached_tile_shape": None, + "last_sigma": None, + } + + def post_cfg_function(args: Dict[str, Any]) -> torch.Tensor: + denoised = args["denoised"] + if not isinstance(denoised, torch.Tensor): + return denoised + if denoised.ndim != 4 or denoised.shape[1] != 16: + return denoised + if not isinstance(X_orig_full, torch.Tensor) or X_orig_full.ndim != 4: + return denoised + if int(X_orig_full.shape[1]) != 16: + return denoised + + sigma_tensor = args["sigma"] + if isinstance(sigma_tensor, torch.Tensor): + sigma_val = float(sigma_tensor.detach().float().max().item()) + else: + sigma_val = float(sigma_tensor) + + prev = state["last_sigma"] + is_new_run = prev is None or sigma_val > prev + 1e-4 + if is_new_run: + state["max_sigma"] = None + state["cached_L_orig"] = None + state["cached_H_orig_std"] = None + state["cached_edge_mask"] = None + state["cached_tile_shape"] = None + state["last_sigma"] = sigma_val + + if state["max_sigma"] is None or sigma_val > float(state["max_sigma"]): + state["max_sigma"] = sigma_val + + progress = 0.0 + max_sigma = state["max_sigma"] + if max_sigma is not None and float(max_sigma) > 0: + progress = 1.0 - (sigma_val / float(max_sigma)) + progress = max(0.0, min(1.0, progress)) + + if decay_type == "cosine": + gamma_t = float(gamma_max) * 0.5 * (1.0 - math.cos(math.pi * progress)) + else: + gamma_t = float(gamma_max) * progress + + if gamma_t < 0.01 and float(alpha_lock) < 0.01: + return denoised + + device = denoised.device + dtype = denoised.dtype + + if state["gaussian_filter"] is None: + state["gaussian_filter"] = SeparableGaussianSmoothing( + channels=16, + kernel_size=int(kernel_size), + sigma=float(blur_sigma), + dtype=torch.float32, + ).to(device=device) + + gaussian: SeparableGaussianSmoothing = state["gaussian_filter"] + if gaussian.kernel_x.device != device: + gaussian.to(device=device) + + _, _, h, w = denoised.shape + need_match = state["cached_L_orig"] is None or state["cached_tile_shape"] != (h, w) + + if need_match: + X_orig = X_orig_full.to(device=device, dtype=dtype) + oh, ow = int(X_orig.shape[2]), int(X_orig.shape[3]) + + if oh > h or ow > w: + input_noisy = args.get("input") + if isinstance(input_noisy, torch.Tensor) and input_noisy.shape == denoised.shape: + y_off, x_off = _find_tile_offset_ncc(X_orig, input_noisy) + else: + y_off, x_off = 0, 0 + y_end = min(y_off + h, oh) + x_end = min(x_off + w, ow) + y_off = max(0, y_end - h) + x_off = max(0, x_end - w) + X_orig_cropped = X_orig[:, :, y_off : y_off + h, x_off : x_off + w] + else: + X_orig_cropped = X_orig + + if X_orig_cropped.shape[2] != h or X_orig_cropped.shape[3] != w: + X_orig_cropped = F.interpolate( + X_orig_cropped.float(), + size=(h, w), + mode="bilinear", + align_corners=False, + ).to(dtype=dtype) + + L_orig = gaussian(X_orig_cropped) + # High-frequency residual of the reference tile and its per-channel + # spatial energy (std). Computed once per tile in fp32 for stability. + H_orig_32 = (X_orig_cropped.to(torch.float32) - L_orig.to(torch.float32)) + H_orig_std = H_orig_32.std(dim=(2, 3), keepdim=True, unbiased=False) + state["cached_L_orig"] = L_orig.to(dtype=dtype) + state["cached_H_orig_std"] = H_orig_std # keep in fp32 + # Cosine edge-taper mask for this tile geometry (fp32, broadcast over C). + state["cached_edge_mask"] = _build_edge_taper_mask( + h, w, int(edge_taper), device, torch.float32 + ) + state["cached_tile_shape"] = (h, w) + + L_t = gaussian(denoised).to(dtype=dtype) + H_t = denoised - L_t + L_orig_tile = state["cached_L_orig"] + H_orig_std = state["cached_H_orig_std"] + if L_orig_tile is None or H_orig_std is None: + return denoised + + # --- Structure lock (decoupled): fully active from step 0 --- + # Anchors low frequencies (composition/color/lighting) to the reference. + alpha = float(alpha_lock) + L_new = (L_orig_tile * alpha) + (L_t * (1.0 - alpha)) + + # --- Detail gate (decoupled): variance matching governed by gamma_t --- + # Instead of shrinking high-frequency energy, renormalize the generated + # texture so its per-channel contrast/energy matches the reference latent. + # This preserves real detail while preventing the Euler solver from + # overshooting into digital grain. The cosine envelope (gamma_t) controls + # how strongly this governance engages, ramping up over the final steps. + H_t_32 = H_t.to(torch.float32) + std_t = H_t_32.std(dim=(2, 3), keepdim=True, unbiased=False) + + # Covariance-preserving (joint) epsilon: tie the denominator floor to the + # mean energy ACROSS the 16 channels, not to each channel in isolation. A + # near-flat channel therefore inherits the joint floor instead of being + # independently amplified, which keeps the inter-channel scale ratios stable + # and prevents microscopic color/tonal drift in the Flux latent. + eps_abs = 1e-4 + eps_rel = 0.10 + joint_energy = std_t.mean(dim=1, keepdim=True) # [B,1,1,1] across channels + eps_joint = eps_abs + eps_rel * joint_energy + scale = H_orig_std / (std_t + eps_joint) + # Clamp the gain so a near-flat generated tile can't be blown up into noise. + scale = torch.clamp(scale, max=10.0) + + H_matched = H_t_32 * scale + # Blend raw texture -> variance-matched texture as governance engages. + H_new = H_t_32 * (1.0 - gamma_t) + H_matched * gamma_t + H_new = H_new.to(dtype=dtype) + + X_new = L_new + H_new + + # --- Spatial edge-gating --- + # Force the entire spectral adjustment to zero at the absolute tile border + # via a cosine taper, so USDU tile blending sees no seam discontinuity. + edge_mask = state.get("cached_edge_mask") + if edge_mask is not None: + if edge_mask.device != device: + edge_mask = edge_mask.to(device=device) + state["cached_edge_mask"] = edge_mask + delta = (X_new.to(torch.float32) - denoised.to(torch.float32)) * edge_mask + out = denoised.to(torch.float32) + delta + return out.to(dtype=dtype, device=device) + + return X_new.to(dtype=dtype, device=device) + + model_options = dict(m.model_options) + m.model_options = comfy.model_patcher.set_model_options_post_cfg_function( + model_options, + post_cfg_function, + disable_cfg1_optimization=True, + ) + return (m,)