# Wan 2.1 flow-matching training variant of Any4DModel. # Extracted from custom/any4d/a4d_model.py during the wan-subdirectory split. import torch from einops import rearrange from custom.any4d.a4d_model import Any4DModel, apply_cond_aug from custom.any4d.a4d_logistics import ( multiply_dicts, rgb_latent_span, sanitize_nan_inf_streams, unpack_entries_from_streams, ) class Any4DWanModel(Any4DModel): """Any4D with Wan 2.1 Fun 1.3B InP backbone, trained natively with flow matching. Diffs vs Any4DModel (which uses Cosmos EDM conventions): - σ sampling: delegates to Wan's scheduler (shift-transformed uniform s ∈ [0,1]). - Forward process: `yt = (1 − s)·x0 + s·ε` (flow matching), NOT `x0 + σ·ε` (EDM). - Loss weighting: uniform 1.0 per sigma, matching Wan's flow-match training. - Everything else (masks, streams, conditioner, visualization) inherited unchanged. Inference already uses Wan-native math via `Any4DWanPipeline.denoise`; this class makes training follow the same convention so the pretrained DiT stays in-distribution. """ def __init__(self, config): super().__init__(config) assert getattr(config, 'use_wan_backbone', False), \ 'Any4DWanModel requires any4d_config.use_wan_backbone=True' # Replaces Predict2Video2WorldModel.draw_training_sigma_and_epsilon for Wan. def draw_training_sigma_and_epsilon(self, x0_size, condition): """Draw (σ, ε) for a training batch — Wan-native, no multipliers or spikes. Returns σ (i.e. Wan `s`) of shape (B, 1) in [sigma_min, sigma_max] after Wan's shift transform, and ε of the same shape as x0_size. """ batch_size = x0_size[0] epsilon = torch.randn(x0_size, device='cuda') # WanFlowMatchScheduler.sample_sigma returns shift-transformed s ∈ [0,1]. s_B = self.pipe.scheduler.sample_sigma(batch_size, device='cuda').to(device='cuda') s_B_1 = rearrange(s_B, 'b -> b 1') return s_B_1, epsilon # Replaces Predict2Video2WorldModel.get_per_sigma_loss_weights (which returns # (σ²+σ_data²)/(σ·σ_data)² for EDM). Wan's flow-match training uses uniform weight. def get_per_sigma_loss_weights(self, sigma): return torch.ones_like(sigma) # Overrides Any4DModel.compute_loss with Wan flow-match forward process: # yt = (1 − s) · x0 + s · ε (instead of EDM's yt = x0 + σ · ε). # The rest of the method is identical to Any4DModel.compute_loss so we reuse its # machinery for loss masks, entries unpacking, Kendall reduction, etc. def compute_loss( self, x0_streams: dict[str, torch.Tensor], masks: dict[str, torch.Tensor], condition, epsilons: dict[str, torch.Tensor], sigmas: dict[str, torch.Tensor], train_step: int = -1, sample_info: list = None, ): streams = list(sigmas.keys()) # --- PPD-style structured noise (https://arxiv.org/abs/2512.05106) --- # Replace ``epsilons[v*]`` with frequency-selective structured noise so the # clean latent's low-frequency PHASE is preserved into the diffusion input, # while magnitude / high-freq phase still come from the standard Gaussian. # One cutoff_radius is sampled per iter (per the reference PPD train.py); # cutoff ~ Exp(scale) with scale=10 means most iters see weak structure # preservation (close to plain Gaussian) with a long tail of strong # preservation. Only the per-view ``v*`` streams are touched; lowdim # streams (sattn / xattn / adaln) get standard Gaussian. if getattr(self.config, 'use_structured_noise', False): import numpy as _np from einops import rearrange as _rearrange from custom.wan.structured_noise import generate_structured_noise_subspace cutoff_scale = float(getattr(self.config, 'structured_noise_cutoff_scale', 10.0)) # Full-frequency mode: cutoff_radius=None injects the phase at ALL frequencies # (freq_mask=ones), instead of the per-iter Exp(scale) low-pass draw. if bool(getattr(self.config, 'structured_noise_full_frequency', False)): cutoff_radius = None else: cutoff_radius = float(_np.random.exponential(scale=cutoff_scale)) K = int(getattr(self.config, 'structured_noise_channels', 16)) seed = int(getattr(self.config, 'structured_noise_seed', 0)) # Learnable orthonormal projection (always): re-orthonormalize the DiT's # sn_P_raw via QR each step (stays exactly orthonormal). full_tensor() gathers # the FSDP-sharded param (grad reduce-scatters back); the QR graph carries # gradients to it. raw is None only in full-channel mode (K >= C), where the # subspace projection is unused and P_learn stays None. def _qr_named(name): rw = getattr(self.pipe.dit, name, None) if rw is None: return None from torch.distributed.tensor import DTensor as _DTensor rw_full = rw.full_tensor() if isinstance(rw, _DTensor) else rw _Q, _ = torch.linalg.qr(rw_full.float()) # (C, K) orthonormal columns return _Q.t().contiguous() # (K, C) orthonormal rows P_learn = _qr_named('sn_P_raw') # noise projection (P_rgb) # PER-MODALITY P (variant 4): the noise is still projected/lifted by P_rgb (P_learn), # but this step's SOURCE modality phase is read via its own P_m. One source per forward # (the dataloader picks it and publishes it on the same rank/thread). rgb-source steps # keep P_phase=None (== P_rgb), matching the shared-P case. P_phase_permod = None if getattr(self.config, 'structured_noise_per_modality_p', False) and P_learn is not None: from custom.wan.control_maps import get_last_sn_pick _m = get_last_sn_pick() if _m and _m != 'rgb': P_phase_permod = _qr_named(f'sn_P_raw_{_m}') # Phase SOURCE: the clean RGB latent itself (PPD, variant 1). When a # ``structured_noise_source_mix`` is configured (variant 2, multimodal), the per-step # source was routed into rgb{v}_edge_control by the dataloader, so read the phase from # that slot instead — the noise then carries depth/edge/seg structure with no explicit # control input to the DiT. edge_phase = getattr(self.config, 'structured_noise_source_mix', None) is not None new_eps = dict(epsilons) for k in list(epsilons.keys()): if not k.startswith('v'): continue x0 = x0_streams[k] eps = epsilons[k] if x0.ndim != 5 or eps.ndim != 5 or x0.shape != eps.shape: continue B, C, T, H, W = x0.shape v = int(k[1:]) entries = self.config.video_entries[v] # Noise is injected into the generated RGB-latent channels (e.g. [0:16]); # the other stream channels (masks/ref/canny) are read clean during assemble, # so their noise is irrelevant and stays plain Gaussian. rs, re_ = rgb_latent_span(entries) # Phase source span: the edge-control latent if requested, else the rgb span. if edge_phase: ek = f'rgb{v}_edge_control' if ek not in entries: raise AssertionError( f'structured_noise_source_mix is set but view {v} has no ' f'{ek!r} entry in video_entries.') ps, pe = int(entries[ek][0]), int(entries[ek][1]) assert (pe - ps) == (re_ - rs), ( f'edge-phase channels {pe-ps} != rgb channels {re_-rs}; ' f'phase source and noise must share C for the projection.') else: ps, pe = rs, re_ x0_btchw = _rearrange(x0[:, ps:pe], 'b c t h w -> (b t) c h w') # phase source eps_btchw = _rearrange(eps[:, rs:re_], 'b c t h w -> (b t) c h w') # rgb noise struct = generate_structured_noise_subspace( x0_btchw.float(), eps_btchw.float(), cutoff_radius=cutoff_radius, num_channels=K, seed=seed, P=P_learn, P_phase=P_phase_permod, ) struct = _rearrange( struct, '(b t) c h w -> b c t h w', b=B, t=T, ).to(dtype=eps.dtype, device=eps.device) eps_new = eps.clone() eps_new[:, rs:re_] = struct new_eps[k] = eps_new epsilons = new_eps # --- Wan forward process (ONLY real change from Any4DModel.compute_loss) --- yt_streams = {} for k in streams: s = sigmas[k] # (B, 1/T, 1) or (B, 1, T, 1, 1) yt_streams[k] = (1.0 - s) * x0_streams[k] + s * epsilons[k] # Conditioning augmentation on x0 (network input side), unchanged. (x0_streams_noised, cond_aug_sigmas) = apply_cond_aug( x0_streams, self.config.train_cond_aug_sigma_range) sanitize_nan_inf_streams( [x0_streams, x0_streams_noised, yt_streams], masks, train_step, sample_info=sample_info) # Any4DWanPipeline.denoise implements Wan flow-match math; returns y0_pred etc. (xt_streams, y0_pred_streams, eps_pred_streams, logvar) = self.pipe.denoise( x0_streams_noised, yt_streams, masks, sigmas, condition, iteration=train_step, phase='train') # Uniform loss weight (overridden get_per_sigma_loss_weights returns ones). weights_per_sigma = {k: self.get_per_sigma_loss_weights(sigma=sigma) for (k, sigma) in sigmas.items()} # MSE on y0 (Wan model is a velocity predictor, but denoise has already converted # velocity → y0, so the loss is written on the clean-sample target as in Any4DModel). unmasked_mse_losses = {k: (y0_pred_streams[k] - x0_streams[k]) ** 2.0 for k in streams} mse_loss_streams = multiply_dicts(unmasked_mse_losses, masks['supervise']) edm_loss_streams = multiply_dicts(mse_loss_streams, weights_per_sigma) # Reduce to per-entry and aggregate (identical to Any4DModel.compute_loss). x0_entries = unpack_entries_from_streams(x0_streams, self.config, detach=True, ignore_masks=False) yt_entries = unpack_entries_from_streams(yt_streams, self.config, detach=True, ignore_masks=True) xt_entries = unpack_entries_from_streams(xt_streams, self.config, detach=True, ignore_masks=True) y0_pred_entries = unpack_entries_from_streams(y0_pred_streams, self.config, detach=True, ignore_masks=True) edm_loss_entries = unpack_entries_from_streams(edm_loss_streams, self.config, detach=False, ignore_masks=True) for (k, lw) in self.config.loss_weights.items(): if k in edm_loss_entries: edm_loss_entries[k] = edm_loss_entries[k] * lw (edm_loss_means, active_edm_loss_means) = self._reduce_entries_and_filter_active(edm_loss_entries, train_step) kendall_loss = torch.mean(torch.stack(list(active_edm_loss_means.values()))) return { 'x0_streams': x0_streams, 'x0_streams_noised': x0_streams_noised, 'x0_entries': x0_entries, 'masks': masks, 'yt_streams': yt_streams, 'yt_entries': yt_entries, 'xt_streams': xt_streams, 'xt_entries': xt_entries, 'y0_pred_streams': y0_pred_streams, 'y0_pred_entries': y0_pred_entries, 'eps_pred_streams': eps_pred_streams, 'sigmas': sigmas, 'cond_aug_sigmas': cond_aug_sigmas, 'weights_per_sigma': weights_per_sigma, 'condition': condition, 'logvar': logvar, 'mse_loss_streams': mse_loss_streams, 'edm_loss_streams': edm_loss_streams, 'edm_loss_entries': edm_loss_entries, 'edm_loss_means': edm_loss_means, 'active_edm_loss_means': active_edm_loss_means, 'loss': kendall_loss, }