# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Wan 2.1 flow-matching scheduler, adapted from DiffSynth-Studio's FlowMatchScheduler # (diffsynth/schedulers/flow_match.py, Apache-2.0). # # Design: # - Everything lives in Wan s-space natively: s ∈ [0, 1], where s=0 is clean and s=1 # is pure noise. Forward: xt = (1 - s) * x0 + s * noise. Target: v = noise - x0. # Euler step: xt_{s-Δs} = xt_s + (s_next - s_cur) * v. # - No EDM↔s conversion and no preconditioning scaling layer. The companion # Any4DWanPipeline.denoise override in custom/wan/a4d_pipe_wan.py uses the same # s-native formulas, so the pipeline's `sample` tensors ARE the Wan xt tensors. # - API surface matches the other schedulers in cosmos_predict2/schedulers/ # (sigmas, timesteps, set_timesteps, step_streams) so pipe.generate's loop # drop-in-substitutes without change. from types import SimpleNamespace from typing import Dict, Optional import torch class WanFlowMatchScheduler: """Wan 2.1 flow-matching scheduler. Operates natively in s ∈ [0, 1].""" def __init__( self, shift: float = 5.0, sigma_min: float = 0.003 / 1.002, sigma_max: float = 1.0, num_train_timesteps: int = 1000, extra_one_step: bool = True, ): # Match the `.config` style of diffusers-based schedulers used in Any4D. self.config = SimpleNamespace( shift=shift, sigma_min=sigma_min, sigma_max=sigma_max, num_train_timesteps=num_train_timesteps, extra_one_step=extra_one_step, ) self.sigmas: Optional[torch.Tensor] = None # s values ∈ [0, 1] self.timesteps: Optional[torch.Tensor] = None # s * num_train_timesteps (DiT input) self.num_inference_steps: Optional[int] = None # ------------------------------------------------------------------ sampling schedule def set_timesteps( self, num_inference_steps: int, device: Optional[torch.device] = None, num_train_timesteps: Optional[int] = None, denoising_strength: float = 1.0, **_unused, ): """Build Wan's shift-transformed flow-match schedule in s-space. Shape: `sigmas` has `num_inference_steps + 1` entries if extra_one_step=True, else `num_inference_steps` entries.""" device = device or torch.device("cpu") shift = self.config.shift n_train = num_train_timesteps or self.config.num_train_timesteps s_min = self.config.sigma_min s_max = self.config.sigma_max s_start = s_min + (s_max - s_min) * denoising_strength # linear-in-s schedule before shift. if self.config.extra_one_step: sigmas = torch.linspace(s_start, s_min, num_inference_steps + 1, device=device, dtype=torch.float64) else: sigmas = torch.linspace(s_start, s_min, num_inference_steps, device=device, dtype=torch.float64) # Wan shift transform: concentrates steps near s=1 (the noisy end). sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas) self.sigmas = sigmas.to(dtype=torch.float32) self.timesteps = (self.sigmas[: num_inference_steps] * n_train).to(dtype=torch.float32) self.num_inference_steps = num_inference_steps return self.timesteps # ------------------------------------------------------------------ training sampling def sample_sigma(self, batch_size: int, device=None) -> torch.Tensor: """Sample s values for a training batch. Wan draws from a shift-transformed uniform distribution over the full [sigma_min, sigma_max] range.""" device = device or torch.device("cuda") s_lin = torch.rand(batch_size, device=device, dtype=torch.float64) s_lin = self.config.sigma_min + (self.config.sigma_max - self.config.sigma_min) * s_lin s = self.config.shift * s_lin / (1.0 + (self.config.shift - 1.0) * s_lin) return s.to(dtype=torch.float32) # ------------------------------------------------------------------ forward / target def add_noise(self, original_samples: torch.Tensor, noise: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: """Wan forward process: xt = (1 - s) * x0 + s * noise. `sigma` here is s.""" s = sigma.to(dtype=original_samples.dtype, device=original_samples.device) while s.ndim < original_samples.ndim: s = s.unsqueeze(-1) return (1.0 - s) * original_samples + s * noise def training_target(self, original_samples: torch.Tensor, noise: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: """Wan's velocity target: v = noise - x0. Independent of s.""" return noise - original_samples # ------------------------------------------------------------------ reverse (sampling) step def step( self, x0_pred: torch.Tensor, i: int, sample: torch.Tensor, x0_prev: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None, ): """One Euler step on the x0 prediction, in Wan s-space. Derivation: given xt_s and y0_pred, recover noise via noise_pred = (xt - (1-s) * y0_pred) / s Rebuild at s_next: xt_{s_next} = (1 - s_next) * y0_pred + s_next * noise_pred Substituting and simplifying gives the one-liner below: xt_{s_next} = y0_pred + (s_next / s) * (xt - y0_pred) """ dtype_target = sample.dtype s_cur = self.sigmas[i].to(device=sample.device, dtype=sample.dtype) s_next = (torch.zeros_like(s_cur) if i + 1 >= len(self.sigmas) else self.sigmas[i + 1].to(device=sample.device, dtype=sample.dtype)) ratio = s_next / s_cur.clamp(min=1e-8) next_sample = x0_pred + ratio * (sample - x0_pred) return next_sample.to(dtype=dtype_target), x0_pred.to(dtype=dtype_target) def step_streams( self, x0_pred_streams: Dict[str, torch.Tensor], i: int, sample_streams: Dict[str, torch.Tensor], x0_prev_streams: Optional[Dict[str, torch.Tensor]] = None, ): """Per-stream Euler step. Mirrors RectifiedFlowAB2Scheduler.step_streams API.""" next_samples, next_x0 = {}, {} for k in sample_streams.keys(): if k not in x0_pred_streams: next_samples[k] = sample_streams[k] next_x0[k] = None continue ns, nx = self.step(x0_pred=x0_pred_streams[k], i=i, sample=sample_streams[k], x0_prev=None if x0_prev_streams is None else x0_prev_streams.get(k)) next_samples[k] = ns next_x0[k] = nx return next_samples, next_x0