"""PPD-style frequency-selective structured noise for diffusion training. Vendored from https://github.com/zengxianyu/structured-noise (Apache 2.0) — see the paper https://arxiv.org/abs/2512.05106. We keep only ``generate_structured_noise_batch_vectorized`` (and its helpers) — the demo ``main()`` is dropped. Used by ``Any4DWanModel.compute_loss`` to replace the standard Gaussian epsilon with a phase-mixed noise whose **low-frequency phase matches the clean target latent**. The high-frequency phase is taken from the input Gaussian noise; the magnitude is the input Gaussian's magnitude. The cutoff radius controls how much of the spectrum is "phase-preserved" — large radius preserves more structure (closer to clean), small radius collapses back to plain Gaussian. PPD trains by sampling ``cutoff_radius ~ Exp(scale=10)`` so the model sees the full sparse↔dense structure spectrum. """ import os import torch def create_frequency_soft_cutoff_mask(height: int, width: int, cutoff_radius: float, transition_width: float = 5.0, device: torch.device = None) -> torch.Tensor: """Smooth low-pass cutoff mask of shape (H, W). 1 inside ``cutoff_radius``, Gaussian rolloff (``transition_width``) outside.""" if device is None: device = torch.device('cpu') u = torch.arange(height, device=device) v = torch.arange(width, device=device) u, v = torch.meshgrid(u, v, indexing='ij') center_u, center_v = height // 2, width // 2 frequency_radius = torch.sqrt((u - center_u) ** 2 + (v - center_v) ** 2) mask = torch.exp(-(frequency_radius - cutoff_radius) ** 2 / (2 * transition_width ** 2)) mask = torch.where(frequency_radius <= cutoff_radius, torch.ones_like(mask), mask) return mask def clip_frequency_magnitude(noise_magnitudes: torch.Tensor, clip_percentile: float = 0.95) -> torch.Tensor: """Clip frequency-domain magnitudes at a quantile to prevent extreme values.""" clip_threshold = torch.quantile(noise_magnitudes, clip_percentile) return torch.clamp(noise_magnitudes, max=clip_threshold) def generate_structured_noise_batch_vectorized( image_batch: torch.Tensor, noise_std: float = 1.0, pad_factor: float = 1.5, cutoff_radius: float = None, transition_width: float = 2.0, input_noise: torch.Tensor = None, sampling_method: str = 'fft', ) -> torch.Tensor: """Generate structured noise for a batch of images using frequency soft cutoff. :param image_batch: ``(B, C, H, W)`` clean tensor whose phase we want to keep in the low-frequency portion of the produced noise. :param noise_std: scale applied to noise magnitudes (default 1.0). :param pad_factor: reflection-pad factor to reduce boundary artifacts. :param cutoff_radius: frequency cutoff radius. ``None`` means full structure preservation; ``0`` means pure Gaussian. :param transition_width: smooth-transition width around the cutoff (smaller = sharper cutoff). :param input_noise: optional ``(B, C, H, W)`` Gaussian noise to seed the process (vs sampling fresh). When given we keep its magnitude/phase and only mix in the image phase at low frequencies. :param sampling_method: ``'fft'`` (default), ``'cdf'`` or ``'two-gaussian'``. :return: ``(B, C, H, W)`` structured-noise tensor in the same dtype as input. """ assert sampling_method in ['fft', 'cdf', 'two-gaussian'] batch_size, channels, height, width = image_batch.shape dtype = image_batch.dtype device = image_batch.device image_batch = image_batch.float() pad_h = int(height * (pad_factor - 1)) pad_h = pad_h // 2 * 2 pad_w = int(width * (pad_factor - 1)) pad_w = pad_w // 2 * 2 padded_images = torch.nn.functional.pad( image_batch, (pad_w // 2, pad_w // 2, pad_h // 2, pad_h // 2), mode='reflect', ) padded_height = height + pad_h padded_width = width + pad_w if cutoff_radius is not None: cutoff_radius = min(min(padded_height / 2, padded_width / 2), cutoff_radius) freq_mask = create_frequency_soft_cutoff_mask( padded_height, padded_width, cutoff_radius, transition_width, device, ) else: freq_mask = torch.ones(padded_height, padded_width, device=device) fft = torch.fft.fft2(padded_images, dim=(-2, -1)) fft_shifted = torch.fft.fftshift(fft, dim=(-2, -1)) image_phases = torch.angle(fft_shifted) image_phases = clip_frequency_magnitude(image_phases) image_magnitudes = torch.abs(fft_shifted) if input_noise is not None: noise_batch = torch.nn.functional.pad( input_noise, (pad_w // 2, pad_w // 2, pad_h // 2, pad_h // 2), mode='reflect', ).float() else: noise_batch = torch.randn_like(padded_images) if sampling_method == 'fft': noise_fft = torch.fft.fft2(noise_batch, dim=(-2, -1)) noise_fft_shifted = torch.fft.fftshift(noise_fft, dim=(-2, -1)) noise_magnitudes = torch.abs(noise_fft_shifted) noise_phases = torch.angle(noise_fft_shifted) elif sampling_method == 'cdf': N = padded_height * padded_width rayleigh_scale = (N / 2) ** 0.5 uu = torch.rand(size=image_magnitudes.shape, device=device) noise_magnitudes = rayleigh_scale * torch.sqrt(-2.0 * torch.log(uu)) if input_noise is not None: noise_fft = torch.fft.fft2(noise_batch, dim=(-2, -1)) noise_fft_shifted = torch.fft.fftshift(noise_fft, dim=(-2, -1)) noise_magnitudes = torch.abs(noise_fft_shifted) noise_phases = torch.angle(noise_fft_shifted) else: noise_phases = torch.rand(size=image_magnitudes.shape, device=device) * 2 * torch.pi - torch.pi elif sampling_method == 'two-gaussian': N = padded_height * padded_width rayleigh_scale = (N / 2) ** 0.5 u1 = torch.randn_like(image_magnitudes) u2 = torch.randn_like(image_magnitudes) noise_magnitudes = rayleigh_scale * torch.sqrt(u1 ** 2 + u2 ** 2) if input_noise is not None: noise_fft = torch.fft.fft2(noise_batch, dim=(-2, -1)) noise_fft_shifted = torch.fft.fftshift(noise_fft, dim=(-2, -1)) noise_magnitudes = torch.abs(noise_fft_shifted) noise_phases = torch.angle(noise_fft_shifted) else: noise_phases = torch.rand(size=image_magnitudes.shape, device=device) * 2 * torch.pi - torch.pi else: raise ValueError(f'Unknown sampling method: {sampling_method}') noise_magnitudes = clip_frequency_magnitude(noise_magnitudes) noise_magnitudes = noise_magnitudes * noise_std # Low frequencies (within cutoff) take image phase; high freqs take noise phase. mixed_phases = freq_mask.unsqueeze(0).unsqueeze(0) * image_phases + \ (1 - freq_mask.unsqueeze(0).unsqueeze(0)) * noise_phases fft_combined = noise_magnitudes * torch.exp(1j * mixed_phases) fft_unshifted = torch.fft.ifftshift(fft_combined, dim=(-2, -1)) structured_noise_padded = torch.fft.ifft2(fft_unshifted, dim=(-2, -1)) structured_noise_padded = torch.real(structured_noise_padded) # Fallback to plain Gaussian where output blew up. clamp_mask = (structured_noise_padded < -5) + (structured_noise_padded > 5) clamp_mask = (clamp_mask > 0).float() structured_noise_padded = structured_noise_padded * (1 - clamp_mask) + noise_batch * clamp_mask # Strip the reflection padding. structured_noise_batch = structured_noise_padded[ :, :, pad_h // 2:pad_h // 2 + height, pad_w // 2:pad_w // 2 + width, ] return structured_noise_batch.to(dtype) # Cache of fixed semi-orthogonal projection matrices P, keyed by (C, K, seed). _P_CACHE: dict = {} # P matrices are PERSISTED here as .pt files so the EXACT same projection is reused at # inference (independent of torch/hardware RNG drift). Override with A4D_STRUCT_NOISE_P_DIR. # A given (C, K, seed) -> P_C{C}_K{K}_seed{seed}.pt; once written it is authoritative. P_CACHE_DIR = os.environ.get( 'A4D_STRUCT_NOISE_P_DIR', os.path.join(os.path.dirname(os.path.abspath(__file__)), 'structnoise_P_cache')) def _semi_orthogonal_P(C: int, K: int, seed: int, device, dtype): """Fixed random semi-orthogonal P of shape (K, C) with orthonormal ROWS (P Pᵀ = I_K). Deterministic per (C, K, seed). Persisted to ``P_CACHE_DIR`` on first build and loaded thereafter, so training and later inference use byte-identical P. Held in-memory too. """ key = (C, K, int(seed)) P = _P_CACHE.get(key) if P is None: path = os.path.join(P_CACHE_DIR, f'P_C{C}_K{K}_seed{int(seed)}.pt') if os.path.exists(path): P = torch.load(path, map_location='cpu') else: g = torch.Generator(device='cpu').manual_seed(int(seed)) M = torch.randn(C, K, generator=g) # (C, K) Q, _ = torch.linalg.qr(M) # (C, K), orthonormal columns P = Q.t().contiguous() # (K, C), orthonormal rows try: # atomic write so concurrent ranks never read a partial file os.makedirs(P_CACHE_DIR, exist_ok=True) tmp = f'{path}.tmp{os.getpid()}' torch.save(P, tmp) os.replace(tmp, path) except Exception: pass _P_CACHE[key] = P return P.to(device=device, dtype=dtype) def generate_structured_noise_subspace(image_batch, input_noise, cutoff_radius, num_channels=None, seed=0, P=None, P_phase=None): """Rank-K structured noise: inject the clean latent's low-freq phase only within a learned K-dimensional channel subspace (orthonormal P), leaving the orthogonal complement of the noise untouched. Keeps the result full-rank, ~unit-variance Gaussian (so it stays in-distribution) while capping the structural locking at "K channels' worth". With ``num_channels`` None or >= C this is exactly the full-channel ``generate_structured_noise_batch_vectorized`` (P unused). :param image_batch: (N, C, H, W) clean latent (the phase source). :param input_noise: (N, C, H, W) base Gaussian noise. :param num_channels: K, the subspace rank (e.g. 4 to mimic a 4-channel latent). :param seed: unused (kept for signature back-compat); P is always learned. :param P: (K, C) learned orthonormal projection — REQUIRED in K-subspace mode. :param P_phase: optional (K, C) projection used ONLY to read the PHASE SOURCE (``image_batch``). When given, the noise is still projected/lifted by ``P`` (= P_rgb, so the result stays in the RGB latent's unit-variance subspace) while the source modality's phase is read through its own ``P_phase`` (= P_m). This is the per-modality design: P_rgb and P_m are pretrained so channel k of each lands on the SAME common structural direction, so mixing n_k's magnitude with z_k's phase per-channel is meaningful. None => P_phase = P (shared-P behavior). The anchor (rgb) source always uses ``P`` (it lives in RGB space). """ N, C, H, W = image_batch.shape if num_channels is None or num_channels >= C or num_channels <= 0: return generate_structured_noise_batch_vectorized( image_batch, cutoff_radius=cutoff_radius, input_noise=input_noise) # P is the learnable orthonormal projection (already orthonormalized, shape (K, C)) # supplied by the caller — its graph carries gradients back to the DiT's sn_P_raw param. # There is no fixed-P path: P is always learned. assert P is not None, ( 'generate_structured_noise_subspace: K-subspace mode requires a learned P ' '(the DiT.sn_P_raw projection); the fixed seeded-P path has been removed.') P = P.to(device=image_batch.device, dtype=torch.float32) Pp = (P_phase.to(device=image_batch.device, dtype=torch.float32) if P_phase is not None else P) # phase-source projection (P_m or P) x = image_batch.float() n = input_noise.float() # Project the PHASE SOURCE into the K-D subspace via P_phase (1x1 conv == per-pixel matmul). z_k = torch.einsum('kc,nchw->nkhw', Pp, x) n_k = torch.einsum('kc,nchw->nkhw', P, n) # Structured noise inside the subspace. s_k = generate_structured_noise_batch_vectorized( z_k, cutoff_radius=cutoff_radius, input_noise=n_k) # Lift back and recombine with the untouched orthogonal complement of the noise: # s = Pᵀ s_k + (n - Pᵀ P n). Orthonormal P => result stays unit-variance Gaussian. s_lift = torch.einsum('ck,nkhw->nchw', P.t(), s_k.float()) n_proj = torch.einsum('ck,nkhw->nchw', P.t(), n_k) out = s_lift + (n - n_proj) return out.to(image_batch.dtype) def phase_sync_loss(rgb_latent, map_latent, P_rgb, P_map=None, eps=1e-6, exclude_dc=True): """RGB->map phase-SYNCHRONIZATION loss on the rank-k subspace (the recommended agreement objective). Returns -coherence in [-1, 1]; MINIMIZE to make RGB's projected Fourier phase synchronize with the map's, so P learns the cross-modal structural subspace and RGB's private (appearance) phase -- which synchronizes to no map -- is excluded. Magnitude-weighted complex coherence (= mag-weighted cos(dphase), no angle()/singularity): - sum Re(Z_rgb * conj(Z_map)) / sum |Z_rgb| |Z_map| over channels x freqs (DC dropped) On rgb-source steps (map_latent == rgb) coherence==1 with ~0 gradient, so it self-gates to map steps. :param rgb_latent: (N, C, H, W) clean RGB latent. :param map_latent: (N, C, H, W) clean map latent (depth/canny/seg) for this step. :param P_rgb: (k, C) projection applied to RGB. :param P_map: (k, C) projection applied to the map. None => use P_rgb (shared-P, variant 1). Pass a distinct P_m for per-modality projections (variant 2). """ import torch Pr = P_rgb.float() Pm = (P_map if P_map is not None else P_rgb).float() Zr = torch.fft.fft2(torch.einsum('kc,nchw->nkhw', Pr, rgb_latent.float()), dim=(-2, -1)) Zm = torch.fft.fft2(torch.einsum('kc,nchw->nkhw', Pm, map_latent.float()), dim=(-2, -1)) num = (Zr * Zm.conj()).real # |Zr||Zm| cos(phi_rgb - phi_map), per (n,k,h,w) den = Zr.abs() * Zm.abs() if exclude_dc: num = num.clone(); den = den.clone() num[..., 0, 0] = 0.0; den[..., 0, 0] = 0.0 return -(num.sum() / (den.sum() + eps)) # Canonical modality names + the source-mix picker. Shared by the dataloader (which fills the # phase slot) and the model/pipe (which must pick the matching per-modality projection P_m), so # the two can never drift. The canonical order also defines the per-modality P layout. SN_MODALITY_ALIAS = {'semantic': 'seg', 'edge': 'canny', 'image': 'rgb'} SN_MODALITIES = ['rgb', 'canny', 'depth', 'seg'] # rgb first = the noise-projection P_rgb def sn_canonical_mix(source_mix): """Map a config source_mix (may use aliases like 'edge'/'semantic'/'image') to canonical modality names, preserving order. e.g. ['depth','edge','semantic','image'] -> ['depth','canny','seg','rgb'].""" return [SN_MODALITY_ALIAS.get(s, s) for s in (source_mix or [])] def sn_pick_modality(candidates, train_step=-1, val_step=-1): """Deterministically pick ONE source modality from ``candidates`` (already canonical). TRAIN: uniform keyed on train_step. VAL: round-robin keyed on val_step. Matches the dataloader exactly so the model can re-derive / cross-check the chosen source.""" import random as _random if not candidates: return None if int(train_step) >= 0: return _random.Random(int(train_step)).choice(candidates) if int(val_step) >= 0: return candidates[int(val_step) % len(candidates)] return candidates[0] __all__ = [ 'create_frequency_soft_cutoff_mask', 'clip_frequency_magnitude', 'generate_structured_noise_batch_vectorized', 'generate_structured_noise_subspace', 'phase_sync_loss', 'SN_MODALITY_ALIAS', 'SN_MODALITIES', 'sn_canonical_mix', 'sn_pick_modality', ]