diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow.py b/fastgen/configs/experiments/WanT2V/config_anyflow.py new file mode 100644 index 0000000..65a9cc7 --- /dev/null +++ b/fastgen/configs/experiments/WanT2V/config_anyflow.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""AnyFlow flow-map pretrain config on Wan-1.3B T2V (paper Stage 2). + +AnyFlow's pretrain objective is the MeanFlow objective with a fixed +``beta08`` per-timestep loss weighting, a finite-difference JVP, shifted +timestep sampling, and a ``consistency_ratio`` fraction of the batch pinned +to ``r = min_t`` — so this config runs :class:`MeanFlowModel` directly. + +Mirrors the AnyFlow paper's pretrain recipe (``shift=5``, ``beta08`` +weighting, ``epsilon=5e-3``, ``diffusion_ratio=0.5``, +``consistency_ratio=0.25``, lr=5e-5, 6k iterations, batch_size_global=32, +text dropout 0.1 with velocity-fused guidance 3.0) on the gated dual-timestep +Wan architecture (``gate=0.25``, absolute-r conditioning, matching +``deltatime_type: r`` in the reference). + +The on-policy stage (paper Stage 3) lives in ``config_anyflow_onpolicy.py``. +""" + +import copy + +import fastgen.configs.methods.config_mean_flow as config_mean_flow +from fastgen.configs.data import VideoLoaderConfig +from fastgen.configs.net import Wan_1_3B_Config + + +def create_config(): + config = config_mean_flow.create_config() + + # ------ network: gated dual-timestep Wan (AnyFlow architecture) ------ + config.model.net = copy.deepcopy(Wan_1_3B_Config) + config.model.net.r_timestep = True + config.model.net.encoder_depth = None + # AnyFlow conditions the r-pathway on the absolute r (deltatime_type "r") + # and fuses the two time embeddings with a fixed convex gate of 0.25. + config.model.net.time_cond_type = "abs" + config.model.net.r_embedder_fusion = "gated" + config.model.net.r_embedder_gate_value = 0.25 + + config.model.precision = "bfloat16" + # VAE compress ratio: (1 + T/4) * H/8 * W/8. 81-frame, 480p clips. + config.model.input_shape = [16, 21, 60, 104] + + # ------ AnyFlow loss: MeanFlow l2 with fixed beta08 weighting ------ + config.model.loss_config.use_cd = False + config.model.loss_config.loss_type = "l2" + # Fixed beta08 per-timestep weighting on the per-element mean loss, no + # adaptive normalization — matching the reference train_bidirection. + config.model.loss_config.weight_type = "beta08" + config.model.loss_config.norm_method = None + config.model.loss_config.use_jvp_finite_diff = True + config.model.loss_config.jvp_finite_diff_eps = 5e-3 + # Rebalance the non-diffusion (flow-map / consistency) sample losses to + # the global diffusion-loss mean (reference scale_weight). + config.model.loss_config.rebalance_to_diffusion = True + config.model.precision_amp_jvp = "float32" + + # Prediction-side guidance fusion with text dropout (reference: + # drop_text_ratio=0.1, fuse_guidance_scale=3.0): the conditional output + # learns the guided flow directly. guidance_scale stays None — MeanFlow's + # target-side eq. 19 fusion is a different mechanism. + config.model.guidance_scale = None + config.model.loss_config.guidance_fuse_scale = 3.0 + config.model.cond_dropout_prob = 0.1 + + # ------ (t, r) sampling: shifted uniform pairs + AnyFlow buckets ------ + config.model.sample_t_cfg.time_dist_type = "shifted" + config.model.sample_t_cfg.shift = 5.0 + config.model.sample_t_cfg.min_t = 0.001 + config.model.sample_t_cfg.max_t = 0.999 + # diffusion_ratio=0.5 of the batch keeps r = t (pure flow matching); + # r_sample_ratio is the complementary fraction that keeps the sampled r. + config.model.sample_t_cfg.r_sample_ratio = 0.5 + # consistency_ratio=0.25 of the batch is pinned to r = min_t. + config.model.sample_t_cfg.consistency_ratio = 0.25 + + # ------ optimization (reference: AdamW lr=5e-5, wd=0, betas=(0.9, 0.95), + # max_grad_norm=1.0, 1000-step LR warmup, EMA decay 0.999) ------ + config.model.net_optimizer.optim_type = "adamw" + config.model.net_optimizer.lr = 5e-5 + config.model.net_optimizer.betas = (0.9, 0.95) + config.model.net_optimizer.weight_decay = 0.0 + config.model.net_scheduler.warm_up_steps = [1000] + config.trainer.callbacks.grad_clip.grad_norm = 1.0 + config.trainer.callbacks.ema.beta = 0.999 + + # ------ inference / validation ------ + config.model.student_sample_type = "ode" + config.model.student_sample_steps = 4 + # 4-step shifted schedule shift*s/(1+(shift-1)*s) on s=linspace(1,0,5), + # first entry clamped to max_t (the reference samples from t=1.0). + config.model.sample_t_cfg.t_list = [0.999, 0.9375, 0.83333333, 0.625, 0.0] + + # ------ data / trainer ------ + config.dataloader_train = VideoLoaderConfig + config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8) + config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 + config.dataloader_train.batch_size = 1 + + config.trainer.max_iter = 6000 + config.trainer.logging_iter = 100 + config.trainer.save_ckpt_iter = 500 + config.trainer.batch_size_global = 32 + + config.log_config.group = "wan_anyflow" + return config diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py new file mode 100644 index 0000000..6595dc6 --- /dev/null +++ b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""AnyFlow on-policy distillation config on Wan-1.3B T2V (paper Stage 3). + +DMD2-style distribution matching with a flow-map student that generates via +a compressed rollout (jump -> fine step -> jump, gradient through all +segments, NFE sampled per iteration from [2, 4, 8, 16, 50]) and co-trains the +Stage-2 flow-map loss at every student update — see +:class:`~fastgen.methods.distribution_matching.anyflow.AnyFlowModel`. + +Mirrors the reference recipe +(``train_wan1b_onpolicy_81f_480p_lr2e-6_1k_b32.yml``): 1.2k iterations at +lr=2e-6, AdamW betas=(0.0, 0.999), wd=0, grad clip 1.0, EMA 0.99, +real-score CFG strength 3, no adversarial loss (the reference +"discriminator" is the fake score network). Set +``config.model.pretrained_student_net_path`` to a Stage 2 checkpoint from +``config_anyflow.py`` before launching, and point the teacher +(``config.model.pretrained_model_path``) at a flow-map teacher checkpoint — +the reference initializes both real and fake score from a separately +fine-tuned flow-map teacher, not stock Wan2.1. + +Known deviations from the reference, both documented: +* full-rank fine-tuning instead of the paper's rank-256 LoRA (FastGen has no + PEFT path today); +* the fake-score noising time is shifted-uniform (FastGen's sampler) instead + of shifted-logit-normal(0, 1). +""" + +import copy + +import fastgen.configs.methods.config_anyflow as config_anyflow_default +from fastgen.configs.data import VideoLoaderConfig +from fastgen.configs.net import Wan_1_3B_Config + + +def create_config(): + config = config_anyflow_default.create_config() + + # ------ network: gated dual-timestep Wan, same as the pretrain stage ------ + config.model.net = copy.deepcopy(Wan_1_3B_Config) + config.model.net.r_timestep = True + config.model.net.encoder_depth = None + config.model.net.time_cond_type = "abs" + config.model.net.r_embedder_fusion = "gated" + config.model.net.r_embedder_gate_value = 0.25 + + config.model.pretrained_student_net_path = "" + + config.model.precision = "bfloat16" + # VAE compress ratio: (1 + T/4) * H/8 * W/8. 81-frame, 480p clips. + config.model.input_shape = [16, 21, 60, 104] + + # ------ DMD machinery ------ + # The reference Stage 3 has no adversarial loss: its "discriminator" is + # the fake score network, trained with denoising score matching only. + config.model.gan_loss_weight_gen = 0.0 + # The reference updates generator and fake score 1:1 (both every global + # step); in DMD2's alternating scheme that is student_update_freq=2. + config.model.student_update_freq = 2 + # Reference real_guidance_scale=3.0 applies cond + 3*(cond - uncond); + # FastGen's CFG formula is cond + (g-1)*(cond - uncond), so g=4. + config.model.guidance_scale = 4.0 + + config.model.sample_t_cfg.time_dist_type = "shifted" + config.model.sample_t_cfg.shift = 5.0 + config.model.sample_t_cfg.min_t = 0.001 + config.model.sample_t_cfg.max_t = 0.999 + + # ------ student rollout (reference rollout_cfg) ------ + config.model.student_sample_type = "ode" + config.model.student_sample_steps = 4 # used for validation sampling + config.model.student_sample_steps_list = [2, 4, 8, 16, 50] + config.model.sample_t_cfg.t_list = None # rollout schedules are computed per NFE + + # ------ co-trained Stage-2 flow-map loss (reference cotrain_forward_kl) ------ + # FastGen's VSD loss carries a 0.5 factor the reference's DMD loss does + # not; 0.5 here keeps the DMD : flow-map gradient ratio at the + # reference's 1 : 1 (the common 0.5 scale is absorbed by Adam). + config.model.cotrain_pretrain_weight = 0.5 + config.model.loss_config.use_cd = False + config.model.loss_config.loss_type = "l2" + config.model.loss_config.weight_type = "beta08" + config.model.loss_config.norm_method = None + config.model.loss_config.use_jvp_finite_diff = True + config.model.loss_config.jvp_finite_diff_eps = 5e-3 + config.model.loss_config.rebalance_to_diffusion = True + config.model.loss_config.guidance_fuse_scale = 3.0 + config.model.cond_dropout_prob = 0.1 + config.model.precision_amp_jvp = "float32" + config.model.sample_t_cfg.r_sample_ratio = 0.5 + config.model.sample_t_cfg.consistency_ratio = 0.25 + + # ------ optimization (reference: AdamW lr=2e-6, betas=(0.0, 0.999), wd=0, + # grad clip 1.0, EMA 0.99) ------ + config.model.net_optimizer.lr = 2e-6 + config.model.net_optimizer.betas = (0.0, 0.999) + config.model.net_optimizer.weight_decay = 0.0 + config.model.fake_score_optimizer.lr = 2e-6 + config.model.fake_score_optimizer.betas = (0.0, 0.999) + config.model.fake_score_optimizer.weight_decay = 0.0 + config.trainer.callbacks.grad_clip.grad_norm = 1.0 + config.trainer.callbacks.ema.beta = 0.99 + + # ------ data / trainer ------ + config.dataloader_train = VideoLoaderConfig + config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8) + config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 + config.dataloader_train.batch_size = 1 + + config.trainer.max_iter = 1200 + config.trainer.logging_iter = 100 + config.trainer.save_ckpt_iter = 200 + config.trainer.batch_size_global = 32 + + config.log_config.group = "wan_anyflow_onpolicy" + return config diff --git a/fastgen/configs/methods/config_anyflow.py b/fastgen/configs/methods/config_anyflow.py new file mode 100644 index 0000000..61671d3 --- /dev/null +++ b/fastgen/configs/methods/config_anyflow.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Config schema for the AnyFlow on-policy method (paper Stage 3). + +AnyFlow's on-policy stage is DMD2 with a flow-map student that generates via +a multi-step rollout-with-gradient, so the config is the DMD2 model config +unchanged. The flow-map pretrain stage (paper Stage 2) is MeanFlow with +AnyFlow's hyperparameters — see ``configs/experiments/WanT2V/config_anyflow.py``. +""" + +from typing import List, Optional + +import attrs +from omegaconf import DictConfig + +from fastgen.configs.callbacks import ( + EMA_CALLBACK, + GPUStats_CALLBACK, + GradClip_CALLBACK, + ParamCount_CALLBACK, + TrainProfiler_CALLBACK, + WANDB_CALLBACK, +) +from fastgen.configs.config import BaseConfig +from fastgen.configs.methods.config_dmd2 import ModelConfig as DMD2ModelConfig +from fastgen.configs.methods.config_mean_flow import ( + LossConfig as MeanFlowLossConfig, + SampleRConfig, + SampleTConfig as MeanFlowSampleTConfig, +) +from fastgen.methods import AnyFlowModel +from fastgen.utils import LazyCall as L + + +@attrs.define(slots=False) +class ModelConfig(DMD2ModelConfig): + """AnyFlow on-policy model config — DMD2 plus the rollout / cotrain knobs. + + The MeanFlow loss / sampling configs drive the co-trained Stage-2 + flow-map loss inside the student update (the reference's + ``cotrain_forward_kl``). + """ + + # MeanFlow-style (t, r) sampling for the co-trained flow-map loss; the + # extra fields are ignored by the DMD2 noising-time sampling. + sample_t_cfg: MeanFlowSampleTConfig = attrs.field(factory=MeanFlowSampleTConfig) + sample_r_cfg: SampleRConfig = attrs.field(factory=SampleRConfig) + loss_config: MeanFlowLossConfig = attrs.field(factory=MeanFlowLossConfig) + + # Weight of the co-trained Stage-2 flow-map loss in the student update. + # The reference runs it at weight 1 (cotrain_forward_kl: True); 0 disables. + cotrain_pretrain_weight: float = 1.0 + + # Rollout NFE list, sampled uniformly per iteration with rank-0 broadcast + # (reference rollout_cfg.num_inference_steps_list). None falls back to + # the fixed student_sample_steps. + student_sample_steps_list: Optional[List[int]] = None + + # Text dropout for the co-trained flow-map loss (reference drop_text_ratio). + cond_dropout_prob: Optional[float] = None + cond_keys_no_dropout: List[str] = attrs.field(factory=list) + + # Precision for autocast in the co-trained loss JVP (None = training precision). + precision_amp_jvp: str | None = None + + +@attrs.define(slots=False) +class Config(BaseConfig): + model: ModelConfig = attrs.field(factory=ModelConfig) + model_class: DictConfig = L(AnyFlowModel)( + config=None, + ) + + +def create_config(): + config = Config() + config.trainer.callbacks = DictConfig( + { + **GradClip_CALLBACK, + **EMA_CALLBACK, + **GPUStats_CALLBACK, + **TrainProfiler_CALLBACK, + **ParamCount_CALLBACK, + **WANDB_CALLBACK, + } + ) + + # The student is a flow-map network with a dual-timestep input. + config.model.use_ema = True + config.model.net.r_timestep = True + config.model.net_scheduler.warm_up_steps = [0] + config.model.fake_score_scheduler.warm_up_steps = [0] + config.model.discriminator_scheduler.warm_up_steps = [0] + + return config diff --git a/fastgen/configs/methods/config_mean_flow.py b/fastgen/configs/methods/config_mean_flow.py index a6dba34..6541f20 100644 --- a/fastgen/configs/methods/config_mean_flow.py +++ b/fastgen/configs/methods/config_mean_flow.py @@ -36,6 +36,10 @@ class SampleTConfig(BaseSampleTConfig): # ratio for randomly sampling r r_sample_ratio: float = 0.0 + # fraction of the batch forced to r = min_t (consistency-to-clean, + # used by AnyFlow; 0.0 keeps the original MeanFlow behavior) + consistency_ratio: float = 0.0 + @attrs.define(slots=False) class SampleRConfig(BaseSampleTConfig): @@ -61,8 +65,9 @@ class LossConfig: use_jvp_finite_diff: bool = False # epsilon for finite difference estimate of JVP jvp_finite_diff_eps: float = 1e-4 - # normalize JVP - norm_method: str = "poly_1.0" + # adaptive loss normalization as a function of the per-sample loss + # (None disables it; the l2 loss then reduces with a per-element mean) + norm_method: Optional[str] = "poly_1.0" # tangent warmup constant norm_const: float = 1e-1 # tangent warmup steps @@ -71,6 +76,19 @@ class LossConfig: tangent_spatial_invariance: bool = False # loss type (choice between l2 and opt_grad) loss_type: str = "opt_grad" + # optional fixed per-timestep loss weighting evaluated as a function of t + # ("beta08", "gaussian", "uniform"; used by AnyFlow). Multiplies the + # adaptive norm_method weight above; None disables it. + weight_type: Optional[str] = None + # prediction-side guidance fusion scale (AnyFlow guidance distillation): + # the conditional output is trained to be the guided flow directly via + # (u_cond + (g-1) * u_uncond) / g against the raw data velocity, with + # plain text dropout (cond_dropout_prob). None keeps MeanFlow's + # target-side guidance (guidance_scale, eq. 19). + guidance_fuse_scale: Optional[float] = None + # rebalance non-diffusion (r < t) sample losses to the global + # diffusion-loss mean via a detached per-sample factor (AnyFlow) + rebalance_to_diffusion: bool = False @attrs.define(slots=False) @@ -97,7 +115,7 @@ class ModelConfig(BaseModelConfig): cond_dropout_prob: Optional[float] = None # list of condition keys that do not drop - cond_keys_no_dropout: List[str] = [] + cond_keys_no_dropout: List[str] = attrs.field(factory=list) # guidance t start guidance_t_start: float = 0.0 diff --git a/fastgen/methods/__init__.py b/fastgen/methods/__init__.py index e902b82..0508a27 100644 --- a/fastgen/methods/__init__.py +++ b/fastgen/methods/__init__.py @@ -15,6 +15,9 @@ from fastgen.methods.consistency_model.sCM import SCMModel as SCMModel from fastgen.methods.consistency_model.mean_flow import MeanFlowModel as MeanFlowModel +# AnyFlow reuses MeanFlow's student sample loop, so it must come after. +from fastgen.methods.distribution_matching.anyflow import AnyFlowModel as AnyFlowModel + from fastgen.methods.fine_tuning.sft import SFTModel as SFTModel from fastgen.methods.fine_tuning.sft import CausalSFTModel as CausalSFTModel diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index 485e752..98cdd3f 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -11,6 +11,7 @@ from fastgen.methods import CMModel from fastgen.utils import basic_utils, expand_like from fastgen.utils.basic_utils import convert_cfg_to_dict +from fastgen.utils.distributed import get_rank, world_size import fastgen.utils.logging_utils as logger @@ -68,6 +69,10 @@ def __init__(self, config: ModelConfig): self.precision_amp_jvp = basic_utils.PRECISION_MAP[self.config.precision_amp_jvp] logger.critical(f"Using precision {self.precision_amp_jvp} for JVP") + # Normalization constant for the fixed per-timestep loss weighting + # (computed lazily on first use, see _get_timestep_weight). + self._timestep_weight_scale: Optional[float] = None + def _mix_condition( self, condition: Any, @@ -100,6 +105,25 @@ def _mix_condition( return condition, dxt_dt + def _drop_condition(self, condition: Any, neg_condition: Any) -> Any: + """Replace the condition with neg_condition for a random per-sample subset. + + Plain text dropout (used by the prediction-side guidance fusion); the + dropped samples then train the unconditional pathway. + """ + if self.config.cond_dropout_prob is None or neg_condition is None: + return condition + ref = neg_condition if isinstance(neg_condition, torch.Tensor) else next(iter(neg_condition.values())) + keep = torch.rand(ref.shape[0], device=ref.device) >= self.config.cond_dropout_prob + if isinstance(condition, torch.Tensor): + return torch.where(expand_like(keep, condition), condition, neg_condition) + if isinstance(condition, dict): + condition = condition.copy() + for k in condition.keys() - set(self.config.cond_keys_no_dropout): + condition[k] = torch.where(expand_like(keep, condition[k]), condition[k], neg_condition[k]) + return condition + raise TypeError(f"Unsupported condition type: {type(condition)}") + @torch.no_grad() def _get_velocity( self, @@ -251,20 +275,141 @@ def net_wrapper(x_t, t, r): return u_theta_jvp + def _sample_t_r_buckets(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Sample (t, r) with t >= r and assign the per-batch buckets. + + Returns ``(t, r, r_eq_t_mask)`` where ``r_eq_t_mask`` marks the + samples with ``r = t`` (pure flow matching). + + With ``consistency_ratio == 0`` (MeanFlow default) the flow-matching + head size is stochastic, as before. With ``consistency_ratio > 0`` + (AnyFlow) the buckets are assigned deterministically over the GLOBAL + batch by rank index, exactly like the AnyFlow reference + (``sample_timestep`` in ``trainer_wan_anyflow_pretrain.py``): the + first ``round((1 - r_sample_ratio) * global_bsz)`` samples get + ``r = t``, the next ``round(consistency_ratio * global_bsz)`` get + ``r = 0`` (consistency to clean data), and the rest keep the sampled + random pair. + """ + t_sample_kwargs = convert_cfg_to_dict(self.sample_t_cfg) + t = self.net.noise_scheduler.sample_t(batch_size, **t_sample_kwargs, device=self.device) + r_sample_kwargs = convert_cfg_to_dict(self.sample_r_cfg) if self.sample_r_cfg.enabled else t_sample_kwargs + r = self.net.noise_scheduler.sample_t(batch_size, **r_sample_kwargs, device=self.device) + t, r = torch.maximum(t, r), torch.minimum(t, r) + assert torch.all(t >= r), "r cannot be larger than t" + + if self.sample_t_cfg.consistency_ratio > 0: + global_bsz = world_size() * batch_size + n_flow_matching = round((1.0 - self.sample_t_cfg.r_sample_ratio) * global_bsz) + n_consistency = round(self.sample_t_cfg.consistency_ratio * global_bsz) + if (n_flow_matching == 0 or n_consistency == 0) and not getattr(self, "_bucket_warned", False): + self._bucket_warned = True + logger.warning( + f"The deterministic (t, r) bucket partition is degenerate: with " + f"world_size * batch_size = {global_bsz}, r_sample_ratio=" + f"{self.sample_t_cfg.r_sample_ratio} and consistency_ratio=" + f"{self.sample_t_cfg.consistency_ratio} yield bucket sizes " + f"({n_flow_matching}, {n_consistency}). The partition spans ranks but not " + f"gradient-accumulation rounds (as in the AnyFlow reference), so empty " + f"buckets stay empty every iteration." + ) + global_idx = get_rank() * batch_size + torch.arange(batch_size, device=self.device) + r_eq_t_mask = global_idx < n_flow_matching + is_consistency = (global_idx >= n_flow_matching) & (global_idx < n_flow_matching + n_consistency) + r = torch.where(r_eq_t_mask, t, r) + r = torch.where(is_consistency, torch.zeros_like(r), r) + else: + # set t=r (flow matching loss) for a subset of the batch + flow_matching_size = (torch.rand(batch_size, device=self.device) >= self.sample_t_cfg.r_sample_ratio).sum() + r_eq_t_mask = torch.arange(batch_size, device=self.device) < flow_matching_size + r = torch.where(r_eq_t_mask, t, r) + + return t, r, r_eq_t_mask + + def _reduce_mf_loss(self, mf_loss: torch.Tensor, r_eq_t_mask: torch.Tensor) -> torch.Tensor: + """Reduce the per-sample loss to a scalar, optionally rebalancing. + + With ``loss_config.rebalance_to_diffusion`` set (AnyFlow), every + flow-map / consistency (r < t) sample's loss is multiplied by the + detached factor ``mean(global flow-matching losses) / (own loss + 1e-5)``, + matching the AnyFlow reference (``train_bidirection``): flow-map / + consistency gradients are self-normalized and rescaled to the global + flow-matching-loss mean. + """ + # No rank-local condition here: the branch must be taken (or not) by + # every rank so the collective below cannot deadlock. Applying the + # scale to an empty ~r_eq_t_mask selection is a no-op. + if getattr(self.loss_config, "rebalance_to_diffusion", False): + with torch.no_grad(): + # The global flow-matching-loss mean only needs the global sum + # and count, so reduce two scalars instead of gathering the + # per-sample losses (equivalent to the reference's + # cat(all_gather(loss))[mask].mean(), and independent of the + # per-rank batch sizes). + fm_loss_sum = torch.where(r_eq_t_mask, mf_loss, torch.zeros_like(mf_loss)).sum() + fm_count = r_eq_t_mask.sum().to(mf_loss.dtype) + if world_size() > 1: + torch.distributed.all_reduce(fm_loss_sum) + torch.distributed.all_reduce(fm_count) + scale = torch.ones_like(mf_loss) + if fm_count > 0: + scale[~r_eq_t_mask] = (fm_loss_sum / fm_count) / (mf_loss[~r_eq_t_mask] + 1e-5) + mf_loss = mf_loss * scale + return mf_loss.mean() + + def _timestep_weight_raw(self, t: torch.Tensor) -> torch.Tensor: + """Unnormalized per-timestep weight as a direct function of t in [0, 1].""" + weight_type = self.loss_config.weight_type + if weight_type == "beta08": + return t * (1 - t).clamp(min=0).sqrt() + if weight_type == "gaussian": + # exp(-2 (t - 1/2)^2), shifted so the minimum over [0, 1] is zero. + return (torch.exp(-2 * (t - 0.5) ** 2) - float(torch.exp(torch.tensor(-0.5)))).clamp(min=0) + if weight_type == "uniform": + return torch.ones_like(t) + raise ValueError(f"Invalid weight_type: {weight_type!r}") + + @torch.no_grad() + def _get_timestep_weight(self, t: torch.Tensor) -> torch.Tensor: + """Fixed per-timestep loss weight w(t) (AnyFlow-style). + + The weight is evaluated directly as a function of t. The normalization + constant matches the AnyFlow reference scheduler, which normalizes the + weights over its shifted discrete timestep grid (``set_timesteps``: + 1000 points excluding t=0); we compute the same constant once and + cache it. + """ + if self._timestep_weight_scale is None: + shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type == "shifted" else 1.0 + grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64)[:-1] + grid = shift * grid / (1 + (shift - 1) * grid) + self._timestep_weight_scale = float(1000.0 / self._timestep_weight_raw(grid).sum()) + return self._timestep_weight_raw(t) * self._timestep_weight_scale + @torch.no_grad() - def _compute_weight(self, tensor: torch.Tensor) -> torch.Tensor: - norm_method, *norm_args = self.loss_config.norm_method.split("_") - - if norm_method == "poly": - power = float(norm_args[0]) - assert len(norm_args) == 1, "poly norm method requires 1 argument" - weight = 1 / (tensor + self.loss_config.norm_const).pow(power) - elif norm_method == "exp": - assert len(norm_args) == 2, "exp norm method requires 2 arguments" - const, scale = float(norm_args[0]), float(norm_args[1]) - weight = const * torch.exp(scale * tensor + self.loss_config.norm_const) + def _compute_weight(self, tensor: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """Per-sample loss weight: the adaptive normalization (``norm_method``, + a function of the per-sample loss; ``None`` disables it) times the + optional fixed per-timestep weight (``weight_type``, a function of t). + """ + if self.loss_config.norm_method is None: + weight = torch.ones_like(tensor) else: - raise ValueError(f"Invalid norm method: {self.loss_config.norm_method}") + norm_method, *norm_args = self.loss_config.norm_method.split("_") + + if norm_method == "poly": + power = float(norm_args[0]) + assert len(norm_args) == 1, "poly norm method requires 1 argument" + weight = 1 / (tensor + self.loss_config.norm_const).pow(power) + elif norm_method == "exp": + assert len(norm_args) == 2, "exp norm method requires 2 arguments" + const, scale = float(norm_args[0]), float(norm_args[1]) + weight = const * torch.exp(scale * tensor + self.loss_config.norm_const) + else: + raise ValueError(f"Invalid norm method: {self.loss_config.norm_method}") + + if self.loss_config.weight_type is not None: + weight = weight * self._get_timestep_weight(t) assert ( weight.shape == tensor.shape @@ -306,8 +451,13 @@ def _mf_pred_to_loss( if self.loss_config.loss_type == "l2": tangent = dxt_dt - warmup_weight * delta_t * u_theta_jvp loss = (u_theta - tangent).pow(2) - loss = torch.sum(loss, dim=list(range(1, loss.ndim))) - weight = self._compute_weight(loss) + if self.loss_config.norm_method is None: + # Without adaptive normalization the reduction matters: the + # per-element mean matches the AnyFlow reference loss scale. + loss = torch.mean(loss, dim=list(range(1, loss.ndim))) + else: + loss = torch.sum(loss, dim=list(range(1, loss.ndim))) + weight = self._compute_weight(loss, t) loss = loss * weight # use explicit gradient @@ -321,7 +471,7 @@ def _mf_pred_to_loss( tangent = tangent * sample_dim_inv opt_grad_norm = torch.linalg.vector_norm(tangent.flatten(1), dim=-1) - weight = self._compute_weight(opt_grad_norm) + weight = self._compute_weight(opt_grad_norm, t) weight = expand_like(weight, tangent) loss = (u_theta - (u_theta + tangent * weight).detach()).pow(2) loss = torch.sum(loss, dim=list(range(1, loss.ndim))) @@ -402,21 +552,51 @@ def _compute_mf_loss( z = torch.randn_like(real_data) x_t = self.net.noise_scheduler.forward_process(real_data, z, t) - condition, dxt_dt = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition) - # prevent JVP to use cached conversions (which can break the computational graph) that were created in the no_grad context of _get_velocity - torch.clear_autocast_cache() - u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) - assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients" - - # additional forward pass to get u_theta with gradient; see also https://github.com/Gsunshine/py-meanflow?tab=readme-ov-file#note-on-jvp - assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}" - u_theta = self.net( - x_t, - t, - r=r, - condition=condition, - fwd_pred_type="flow", - ) + guidance_fuse_scale = getattr(self.loss_config, "guidance_fuse_scale", None) + if guidance_fuse_scale is not None: + # AnyFlow-style guidance distillation: the guidance is fused on the + # PREDICTION side — the conditional output is trained to be the + # guided flow directly. Matches the reference ``train_bidirection``: + # the regression target stays the raw data velocity, the + # unconditional branch is queried at the SAME (t, r) flow-map + # slice, and the effective prediction is + # (u_cond + (g - 1) * u_uncond) / g. Text dropout replaces the + # condition with neg_condition for a random subset beforehand. + g = float(guidance_fuse_scale) + assert g > 0, f"guidance_fuse_scale must be > 0, got {g} (set it to None to disable guidance fusion)" + assert ( + neg_condition is not None + ), "guidance_fuse_scale requires neg_condition: the unconditional branch is queried at the same (t, r)" + condition = self._drop_condition(condition, neg_condition) + dxt_dt = self.net.noise_scheduler.cond_velocity(x=real_data, eps=z, t=t) + torch.clear_autocast_cache() + # dF/dt of the fused prediction: finite difference of the + # conditional output divided by g (the unconditional derivative is + # dropped, as in the reference's compute_central_difference). + u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) / g + assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients" + + assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}" + u_theta = self.net(x_t, t, r=r, condition=condition, fwd_pred_type="flow") + with torch.no_grad(): + u_uncond = self.net(x_t, t, r=r, condition=neg_condition, fwd_pred_type="flow") + u_theta = (u_theta + (g - 1.0) * u_uncond) / g + else: + condition, dxt_dt = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition) + # prevent JVP to use cached conversions (which can break the computational graph) that were created in the no_grad context of _get_velocity + torch.clear_autocast_cache() + u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) + assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients" + + # additional forward pass to get u_theta with gradient; see also https://github.com/Gsunshine/py-meanflow?tab=readme-ov-file#note-on-jvp + assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}" + u_theta = self.net( + x_t, + t, + r=r, + condition=condition, + fwd_pred_type="flow", + ) mf_loss, tangent, loss_weight, warmup_weight = self._mf_pred_to_loss( u_theta=u_theta, u_theta_jvp=u_theta_jvp, x_t=x_t, dxt_dt=dxt_dt, t=t, r=r, iteration=iteration ) @@ -454,19 +634,8 @@ def single_train_step( real_data, condition, neg_condition = self._prepare_training_data(data) batch_size = real_data.shape[0] - # sample t and r - t_sample_kwargs = convert_cfg_to_dict(self.sample_t_cfg) - t = self.net.noise_scheduler.sample_t(batch_size, **t_sample_kwargs, device=self.device) - r_sample_kwargs = convert_cfg_to_dict(self.sample_r_cfg) if self.sample_r_cfg.enabled else t_sample_kwargs - r = self.net.noise_scheduler.sample_t(batch_size, **r_sample_kwargs, device=self.device) - t, r = torch.maximum(t, r), torch.minimum(t, r) - assert torch.all(t >= r), "r cannot be larger than t" - - # set t=r (flow matching loss) for a subset of the batch - batch_size = real_data.shape[0] - flow_matching_size = (torch.rand(batch_size, device=self.device) >= self.sample_t_cfg.r_sample_ratio).sum() - zero_mask = torch.arange(batch_size, device=self.device) < flow_matching_size - r = torch.where(zero_mask, t, r) + # sample t and r, with per-batch buckets (flow matching / consistency) + t, r, r_eq_t_mask = self._sample_t_r_buckets(batch_size) ( mf_loss, @@ -486,7 +655,7 @@ def single_train_step( neg_condition=neg_condition, ) - loss = mf_loss.mean() + loss = self._reduce_mf_loss(mf_loss, r_eq_t_mask) loss_map = { "total_loss": loss, "mf_loss": loss, diff --git a/fastgen/methods/distribution_matching/README.md b/fastgen/methods/distribution_matching/README.md index 021e181..29abe84 100644 --- a/fastgen/methods/distribution_matching/README.md +++ b/fastgen/methods/distribution_matching/README.md @@ -85,6 +85,39 @@ DMD2 extended for causal video generation with autoregressive chunk-by-chunk pro --- +## AnyFlow + +**File:** [`anyflow.py`](anyflow.py) | **Reference:** AnyFlow — any-step video diffusion framework on flow maps + +Single model that supports arbitrary inference NFE by learning a flow map `u_θ(x_t, t, r)` (average velocity from `t` back to `r`). Trained in two stages: + +1. **Pretrain** — the MeanFlow objective with AnyFlow's hyperparameters, run directly on [`MeanFlowModel`](../consistency_model/mean_flow.py): finite-difference JVP, a fixed `beta08` per-timestep loss weight (`loss_config.weight_type`), prediction-side guidance fusion (`loss_config.guidance_fuse_scale` — the conditional output learns the guided flow directly), global rebalancing of the flow-map / consistency sample losses to the diffusion-loss mean (`loss_config.rebalance_to_diffusion`), shifted timestep sampling, and a `sample_t_cfg.consistency_ratio` fraction of the batch pinned to `r = 0`. There is no AnyFlow-specific pretrain code. +2. **On-policy** — distribution-matching distillation on top of the pretrained flow-map weights ([`anyflow.py`](anyflow.py)). DMD2 with the reference's three deviations: the student generates via a flow-map rollout compressed to at most three forwards — jump `t_0 → t_g`, fine step `t_g → t_{g+1}`, jump to 0 — with the NFE sampled per iteration from `student_sample_steps_list` and gradient through all segments (`gen_data_from_net`); the student always starts from pure noise at `max_t` (`_generate_noise_and_time`); and every student update co-trains the Stage-2 flow-map loss (`cotrain_pretrain_weight`). There is no adversarial loss (the reference "discriminator" is the fake score network). The teacher / fake_score are flow-map networks queried at the instantaneous velocity `r = t` (the network-level default for gated-fusion Wan when `r` is not passed). + +**Key Parameters (on-policy):** +- `student_sample_steps_list`: rollout NFEs sampled per iteration (reference: `[2, 4, 8, 16, 50]`) +- `cotrain_pretrain_weight`: weight of the co-trained Stage-2 flow-map loss (reference: 1) +- `guidance_scale`: real-score CFG; FastGen applies `cond + (g-1)·(cond - uncond)`, so the reference's strength 3 is `g=4` +- See also key parameters of DMD2 above + +**Key Parameters (pretrain, on MeanFlow):** +- `loss_config.weight_type`: `beta08` | `gaussian` | `uniform` fixed per-timestep loss weight +- `loss_config.guidance_fuse_scale` + `cond_dropout_prob`: prediction-side guidance distillation with text dropout +- `loss_config.rebalance_to_diffusion`: rescale non-diffusion losses to the global diffusion-loss mean +- `loss_config.use_jvp_finite_diff` / `loss_config.jvp_finite_diff_eps`: central-difference step δ +- `sample_t_cfg.r_sample_ratio` / `sample_t_cfg.consistency_ratio`: per-batch fraction with sampled `r` / `r = 0` +- `sample_t_cfg.time_dist_type="shifted"`, `sample_t_cfg.shift`: shifted timestep sampling (5.0 for Wan video) + +**Backbone requirement:** the student network must accept a secondary timestep `r` (Wan with `r_timestep=True`). When loading the published AnyFlow HF checkpoints, set `r_embedder_fusion="gated"` and `time_cond_type="abs"` on the Wan constructor — this routes the t/r mix through `Wan/network.py::_fuse_r_embedding`'s gated branch (shared with MeanFlow's additive default) so the released weights reproduce bit-for-bit. `Wan.load_state_dict` remaps the AnyFlow checkpoint layout (`condition_embedder.delta_embedder.*`) automatically; note that loading the HF folder via `model_id_or_local_path` goes through diffusers' `from_pretrained`, which silently drops the `delta_embedder` weights — load the transformer state dict through `Wan.load_state_dict` (or call `remap_anyflow_keys` yourself) instead. + +**Note:** correctness of the port is established via forward-parity and single-step training-step parity against the AnyFlow reference (see PR #25 discussion). End-to-end convergence-scale validation on the paper's training corpus is deferred to a follow-up. + +**Configs:** +- [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) — Stage 2 pretrain (6k iter, lr=5e-5, shift=5, beta08, paper-aligned) +- [`WanT2V/config_anyflow_onpolicy.py`](../../configs/experiments/WanT2V/config_anyflow_onpolicy.py) — Stage 3 on-policy distillation (1.2k iter, lr=2e-6, no adversarial loss; full-rank fine-tune of a Stage 2 pretrain ckpt — the paper's rank-256 LoRA variant awaits a PEFT path in FastGen) + +--- + ## Self-Forcing **File:** [`self_forcing.py`](self_forcing.py) | **Reference:** [Huang et al., 2025](https://arxiv.org/abs/2506.08009) diff --git a/fastgen/methods/distribution_matching/anyflow.py b/fastgen/methods/distribution_matching/anyflow.py new file mode 100644 index 0000000..271c7e1 --- /dev/null +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""AnyFlow — any-step video diffusion with flow maps and on-policy distillation. + +AnyFlow trains a single flow-map model :math:`u_\\theta(x_t, t, r)` that +predicts the average velocity from time ``t`` to ``r`` (with ``r \\le t``). +Once trained, the same model supports arbitrary inference step counts: each +Euler-like sampling step picks its own integration interval +``(t \\rightarrow r)``. + +Training has two stages: + +* **Flow-map pretrain** (paper Stage 2) is MeanFlow with AnyFlow's + hyperparameters — fixed ``beta08`` per-timestep loss weighting, + finite-difference JVP, prediction-side guidance fusion + (``guidance_fuse_scale``), global rebalancing of non-diffusion losses + (``rebalance_to_diffusion``), and a ``consistency_ratio`` fraction of the + batch pinned to ``r = 0``. It is configured directly on + :class:`~fastgen.methods.consistency_model.mean_flow.MeanFlowModel` + (see ``configs/experiments/WanT2V/config_anyflow.py``); there is no + AnyFlow-specific pretrain code. + +* **On-policy** (paper Stage 3, this module) is DMD2 on top of the pretrained + flow-map weights, with the reference's three deviations from stock DMD2: + + - the student generates via a flow-map rollout compressed into at most three + network forwards — one jump from ``t_0`` to ``t_g``, one fine step + ``t_g \\rightarrow t_{g+1}``, one jump to 0 — with the inference step + count sampled per iteration from ``student_sample_steps_list`` and the + fine-step position ``g`` sampled uniformly (both rank-0 broadcast). + Gradient flows through ALL segments + (:meth:`AnyFlowModel.gen_data_from_net`, mirroring + ``WanAnyFlowPipeline.training_rollout``); + - the student always starts from pure noise at ``max_t`` + (:meth:`AnyFlowModel._generate_noise_and_time`); + - every student update co-trains the Stage-2 flow-map loss on the real + batch (``cotrain_pretrain_weight``, mirroring the reference's + ``cotrain_forward_kl``), borrowing MeanFlow's loss machinery. + + The teacher and fake_score are flow-map networks queried at the + instantaneous velocity ``r = t`` — for gated-fusion Wan networks this is + the network-level default when ``r`` is not passed (see + ``fastgen/networks/Wan/network.py``), so DMD2's update steps run unchanged. +""" + +from __future__ import annotations + +from typing import Any, Optional, TYPE_CHECKING + +import torch + +from fastgen.methods.consistency_model.mean_flow import MeanFlowModel +from fastgen.methods.distribution_matching.dmd2 import DMD2Model +from fastgen.utils import basic_utils, expand_like +from fastgen.utils.basic_utils import convert_cfg_to_dict +from fastgen.utils.distributed import world_size +import fastgen.utils.logging_utils as logger + + +if TYPE_CHECKING: + from typing import Dict, Callable + + from fastgen.configs.methods.config_anyflow import ModelConfig + + +class AnyFlowModel(DMD2Model): + """AnyFlow on-policy stage: DMD2 with a flow-map rollout student.""" + + # Validation sampling integrates the flow map with r = t_next (ode) just + # like MeanFlow; FastGenModel's default x0-prediction loop never passes r. + _student_sample_loop = MeanFlowModel._student_sample_loop + + # MeanFlow loss machinery borrowed for the co-trained Stage-2 flow-map + # loss (the reference's cotrain_forward_kl runs the full bidirection loss + # at every generator step). + _sample_t_r_buckets = MeanFlowModel._sample_t_r_buckets + _reduce_mf_loss = MeanFlowModel._reduce_mf_loss + _compute_mf_loss = MeanFlowModel._compute_mf_loss + _drop_condition = MeanFlowModel._drop_condition + _jvp = MeanFlowModel._jvp + _estimate_jvp_finite_difference = MeanFlowModel._estimate_jvp_finite_difference + _mf_pred_to_loss = MeanFlowModel._mf_pred_to_loss + _compute_weight = MeanFlowModel._compute_weight + _timestep_weight_raw = MeanFlowModel._timestep_weight_raw + _get_timestep_weight = MeanFlowModel._get_timestep_weight + + @torch.no_grad() + def _get_velocity(self, x, z, t, condition=None, neg_condition=None): + """Raw data velocity for the co-trained flow-map loss. + + The DMD teacher CFG (``config.guidance_scale``) must not leak into the + co-trained loss — AnyFlow's guidance lives in + ``loss_config.guidance_fuse_scale`` (prediction-side fusion, handled + inside ``_compute_mf_loss``), so MeanFlow's target-side eq. 19 path is + deliberately not used here. + """ + del neg_condition + return condition, self.net.noise_scheduler.cond_velocity(x=x, eps=z, t=t) + + def __init__(self, config: ModelConfig): + super().__init__(config) + self.config = config + + # Attributes required by the borrowed MeanFlow methods. + self.sample_t_cfg = self.config.sample_t_cfg + self.sample_r_cfg = self.config.sample_r_cfg + self.loss_config = self.config.loss_config + if self.config.precision_amp_jvp is None or self.config.precision_amp_jvp == self.precision_amp: + self.precision_amp_jvp = None + else: + self.precision_amp_jvp = basic_utils.PRECISION_MAP[self.config.precision_amp_jvp] + self._timestep_weight_scale = None + + logger.info( + f"AnyFlow on-policy: student_sample_steps_list={self.config.student_sample_steps_list}, " + f"student_update_freq={self.config.student_update_freq}, " + f"cotrain_pretrain_weight={self.config.cotrain_pretrain_weight}, " + f"gan_loss_weight_gen={self.config.gan_loss_weight_gen}" + ) + + # ------------------------------------------------------------------ + # Rollout + # ------------------------------------------------------------------ + + def _broadcast_choice(self, high: int) -> int: + """Pick an index in [0, high) on rank 0 and broadcast it.""" + idx = torch.randint(0, high, (1,), device=self.device, dtype=torch.long) + if world_size() > 1: + torch.distributed.broadcast(idx, src=0) + return int(idx.item()) + + def _sample_rollout_steps(self) -> int: + """Sample the rollout NFE for this iteration. + + Matches the reference: ``random.choice(num_inference_steps_list)`` + broadcast from rank 0 in both the generator and fake-score updates. + Falls back to the fixed ``student_sample_steps`` when no list is set. + """ + steps_list = self.config.student_sample_steps_list + if not steps_list: + return int(self.config.student_sample_steps) + return int(steps_list[self._broadcast_choice(len(steps_list))]) + + def _rollout_t_list(self, num_steps: int) -> torch.Tensor: + """Shifted timestep schedule for ``num_steps`` rollout steps. + + Equivalent to the reference scheduler's ``set_timesteps``: a uniform + grid mapped through ``shift * s / (1 + (shift - 1) * s)``. The first + entry is clamped to the scheduler's ``max_t``. A config-provided + ``t_list`` takes precedence when its length matches. + """ + ns = self.net.noise_scheduler + cfg_t_list = self.config.sample_t_cfg.t_list + if cfg_t_list is not None and len(cfg_t_list) == num_steps + 1: + return torch.tensor(cfg_t_list, device=self.device, dtype=ns.t_precision) + + shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type == "shifted" else 1.0 + s = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64, device=self.device) + s = shift * s / (1 + (shift - 1) * s) + return s.clamp(max=float(ns.max_t)).to(ns.t_precision) + + def _generate_noise_and_time( + self, real_data: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The student rollout always starts from pure noise at ``max_t``.""" + batch_size = real_data.shape[0] + + eps_student = torch.randn(batch_size, *self.input_shape, device=self.device, dtype=real_data.dtype) + t_student = torch.full( + (batch_size,), + self.net.noise_scheduler.max_t, + device=self.device, + dtype=self.net.noise_scheduler.t_precision, + ) + input_student = self.net.noise_scheduler.latents(noise=eps_student) + + t = self.net.noise_scheduler.sample_t( + batch_size, **convert_cfg_to_dict(self.config.sample_t_cfg), device=self.device + ) + eps = torch.randn_like(real_data, device=self.device, dtype=real_data.dtype) + return input_student, t_student, t, eps + + def gen_data_from_net( + self, + input_student: torch.Tensor, + t_student: torch.Tensor, + condition: Optional[Any] = None, + ) -> torch.Tensor: + """Flow-map rollout compressed into at most three network forwards. + + Mirrors the reference ``WanAnyFlowPipeline.training_rollout``: with an + ``num_steps``-step schedule ``t_0 > ... > t_N = 0`` and a sampled fine + step position ``g``, the model jumps ``t_0 -> t_g`` in ONE flow-map + forward, takes the single fine step ``t_g -> t_{g+1}``, then jumps + ``t_{g+1} -> 0`` in one forward. Gradient flows through all segments + (the reference applies no ``no_grad`` inside ``training_rollout``); + the fake-score update wraps this call in ``no_grad`` at the caller. + Each step uses mean-velocity sampling ``u_theta(x_t, t, r=t_next)``. + """ + del t_student # the rollout schedule is built below + ns = self.net.noise_scheduler + batch_size = input_student.shape[0] + + num_steps = self._sample_rollout_steps() + if num_steps < 1: + raise ValueError(f"rollout steps must be >= 1, got {num_steps}") + grad_step = self._broadcast_choice(num_steps) + t_list = self._rollout_t_list(num_steps) + + segments = ( + (t_list[0], t_list[grad_step]), + (t_list[grad_step], t_list[grad_step + 1]), + (t_list[grad_step + 1], t_list[-1]), + ) + + x = input_student + for t_cur, t_next in segments: + if float(t_cur) == float(t_next): + continue + t_batch = t_cur.expand(batch_size).to(ns.t_precision) + r_batch = t_next.expand(batch_size).to(ns.t_precision) + flow_pred = self.net(x, t_batch, r=r_batch, condition=condition, fwd_pred_type="flow") + x = x - expand_like(t_batch - r_batch, x).to(x.dtype) * flow_pred + + return x + + # ------------------------------------------------------------------ + # Training step — DMD2 plus the co-trained Stage-2 flow-map loss + # ------------------------------------------------------------------ + + def single_train_step( + self, data: "Dict[str, Any]", iteration: int + ) -> tuple[dict[str, torch.Tensor], dict[str, "torch.Tensor | Callable"]]: + real_data, condition, neg_condition = self._prepare_training_data(data) + self._setup_grad_requirements(iteration) + input_student, t_student, t, eps = self._generate_noise_and_time(real_data) + + if iteration % self.config.student_update_freq == 0: + loss_map, outputs = self._student_update_step( + input_student, t_student, t, eps, data, condition=condition, neg_condition=neg_condition + ) + if self.config.cotrain_pretrain_weight > 0: + # Reference cotrain_forward_kl: every generator update also + # runs the full Stage-2 bidirection (flow-map) loss on the + # real batch. + t_mf, r_mf, r_eq_t_mask = self._sample_t_r_buckets(real_data.shape[0]) + mf_outputs = self._compute_mf_loss( + real_data=real_data, + t=t_mf, + r=r_mf, + iteration=iteration, + condition=condition, + neg_condition=neg_condition, + ) + bidirection_loss = self._reduce_mf_loss(mf_outputs[0], r_eq_t_mask) + loss_map["bidirection_loss"] = bidirection_loss + loss_map["total_loss"] = loss_map["total_loss"] + self.config.cotrain_pretrain_weight * bidirection_loss + return loss_map, outputs + + return self._fake_score_discriminator_update_step( + input_student, t_student, t, eps, real_data, condition=condition + ) diff --git a/fastgen/networks/EDM/network.py b/fastgen/networks/EDM/network.py index 0c9a550..331b481 100644 --- a/fastgen/networks/EDM/network.py +++ b/fastgen/networks/EDM/network.py @@ -924,6 +924,12 @@ def forward( else condition.reshape(-1, self.label_dim) ) + # A dual-timestep network always consumes the r-pathway (the embedding + # widths are sized for it), so r=None is not a valid input for it. + # Default to r=t, i.e. the instantaneous velocity u(x_t, t, t). + if r is None and getattr(self.model, "r_timestep", None) is not None: + r = t + # Preconditioning weights for input x_t_in, t_in = x_t, t if self.drop_precond not in ["input", "both"]: diff --git a/fastgen/networks/Wan/network.py b/fastgen/networks/Wan/network.py index f6c3668..2dadc87 100644 --- a/fastgen/networks/Wan/network.py +++ b/fastgen/networks/Wan/network.py @@ -274,6 +274,53 @@ def classify_forward( return out +def _fuse_r_embedding( + self, + temb: torch.Tensor, + timestep_proj: torch.Tensor, + remb: torch.Tensor, + rs_seq_len: Optional[torch.LongTensor], +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Combine the t- and r-time embeddings into the final timestep projection. + + Two fusion modes share the same ``r_embedder.time_proj`` / ``act_fn`` modules: + + * ``additive`` (MeanFlow default, ``r_embedder.fusion_mode`` absent or "additive"): + ``r_timestep_proj = time_proj(act_fn(remb))`` is added to ``timestep_proj`` + and ``remb`` is added to ``temb``. T2V dual-stream variants (``encoder_depth`` + set) instead replace ``temb`` with ``remb`` and leave ``timestep_proj`` alone. + + * ``gated`` (AnyFlow, opt-in): convex-combine the two embeddings *before* the + shared projection — ``rt_emb = (1-g)·temb + g·remb`` then + ``timestep_proj = time_proj(act_fn(rt_emb))``. This matches the + ``WanTwoTimeTextImageEmbedding.forward_timestep`` path in the AnyFlow + reference and is required to reproduce the published HF checkpoint + forward bit-for-bit. With ``encoder_depth`` set, the first + ``encoder_depth`` blocks keep the t-only ``timestep_proj`` and the + remaining blocks switch to the gated projection (mirroring additive). + + Returns ``(temb, timestep_proj, r_timestep_proj)`` where ``r_timestep_proj`` + is ``None`` when the r-information is already folded into ``timestep_proj``. + """ + fusion = getattr(self.r_embedder, "fusion_mode", "additive") + + if fusion == "gated": + gate = self.r_embedder.gate_value + rt_emb = (1 - gate) * temb + gate * remb + rt_ts_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(rt_emb)) + rt_ts_proj = unflatten_timestep_proj(rt_ts_proj, rs_seq_len) + if self.encoder_depth is None: + return rt_emb, rt_ts_proj, None + return rt_emb, timestep_proj, rt_ts_proj + + # additive — MeanFlow original path, bit-identical + r_ts_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(remb)) + r_ts_proj = unflatten_timestep_proj(r_ts_proj, rs_seq_len) + if self.encoder_depth is None: + return temb + remb, timestep_proj + r_ts_proj, r_ts_proj + return remb, timestep_proj, r_ts_proj + + def classify_forward_prepare( self, hidden_states: torch.Tensor, @@ -338,14 +385,8 @@ def classify_forward_prepare( r_timestep = r_timestep.to(time_embedder_dtype) remb = self.r_embedder.time_embedder(r_timestep).type_as(encoder_hidden_states) - r_timestep_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(remb)) - r_timestep_proj = unflatten_timestep_proj(r_timestep_proj, rs_seq_len) - if self.encoder_depth is None: - timestep_proj = timestep_proj + r_timestep_proj - temb = temb + remb - else: - temb = remb + temb, timestep_proj, r_timestep_proj = self._fuse_r_embedding(temb, timestep_proj, remb, rs_seq_len) elif r_timestep is not None: # Raise an error here, otherwise we silently ignore the r_timestep raise ValueError("r_timestep provided but no r_embedder is present") @@ -423,6 +464,50 @@ def classify_forward_block_forward( return hidden_states, features +def remap_anyflow_keys(state_dict: Mapping[str, Any]) -> Mapping[str, Any]: + """Remap an AnyFlow HF release state_dict to FastGen's Wan layout. + + AnyFlow's ``FAR_Wan_Transformer3DModel`` stores the r-pathway inside the + main ``condition_embedder`` as ``delta_embedder``, and uses ONE shared + ``time_proj`` for both t and (t, r). FastGen exposes a separate top-level + ``r_embedder`` with its own ``time_embedder`` + ``time_proj``. The two + layouts are functionally equivalent (FastGen's ``r_embedder.time_proj`` + starts as a deepcopy of ``condition_embedder.time_proj`` per + :meth:`Wan.init_embedder`), so we just rename / duplicate the tensors. + + The function is a no-op when no ``condition_embedder.delta_embedder.*`` + keys are present, so it's safe to call unconditionally. Keys with or + without the ``transformer.`` module prefix are both handled. + """ + delta_marker = "condition_embedder.delta_embedder." + delta_keys = [k for k in state_dict if delta_marker in k] + if not delta_keys: + return state_dict + + new_sd = dict(state_dict) + prefixes = set() + for k in delta_keys: + # [transformer.]condition_embedder.delta_embedder.linear_1.weight + # -> [transformer.]r_embedder.time_embedder.linear_1.weight + prefix, _, suffix = k.partition(delta_marker) + prefixes.add(prefix) + new_sd[f"{prefix}r_embedder.time_embedder.{suffix}"] = new_sd.pop(k) + # AnyFlow's gated fusion shares the final time_proj. FastGen has a + # separate r_embedder.time_proj that substitutes for the shared one when + # fusion="gated"; copy the weights so the two projections start identical. + for prefix in prefixes: + for sub in ("weight", "bias"): + src = f"{prefix}condition_embedder.time_proj.{sub}" + dst = f"{prefix}r_embedder.time_proj.{sub}" + if src in new_sd and dst not in new_sd: + new_sd[dst] = new_sd[src].clone() + logger.info( + f"remap_anyflow_keys: rewrote {len(delta_keys)} delta_embedder tensors " + "and duplicated time_proj weights into r_embedder." + ) + return new_sd + + class WanTextEncoder: def __init__(self, model_id_or_local_path): self.text_encoder = UMT5EncoderModel.from_pretrained( @@ -557,6 +642,8 @@ def __init__( load_pretrained: bool = True, use_fsdp_checkpoint: bool = True, use_wan_official_sinusoidal: bool = False, + r_embedder_fusion: str = "additive", + r_embedder_gate_value: float = 0.25, **model_kwargs, ): """Wan2.1/2.2 model constructor. @@ -610,6 +697,18 @@ def __init__( if r_timestep: logger.info(f"Initializing r embedder with {r_embedder_init}") self.transformer.r_embedder = self.init_embedder(r_embedder_init) + # Stash fusion config on the r_embedder module so the (method-bound) + # forward override can branch on it without changing its signature. + if r_embedder_fusion not in ("additive", "gated"): + raise ValueError(f"r_embedder_fusion must be 'additive' or 'gated', got {r_embedder_fusion!r}") + self.transformer.r_embedder.fusion_mode = r_embedder_fusion + # Plain float (not a buffer) so FSDP's reset_parameters does not + # need to re-materialize it after meta-device init. + self.transformer.r_embedder.gate_value = float(r_embedder_gate_value) + logger.info( + f"r_embedder fusion={r_embedder_fusion}" + + (f" gate_value={r_embedder_gate_value}" if r_embedder_fusion == "gated" else "") + ) else: self.transformer.r_embedder = None @@ -832,6 +931,7 @@ def override_transformer_forward(self, inner_dim: int) -> None: # Override transformer forward methods with custom implementations for block in self.transformer.blocks: block.forward = types.MethodType(block_forward, block) + self.transformer._fuse_r_embedding = types.MethodType(_fuse_r_embedding, self.transformer) self.transformer.classify_forward_prepare = types.MethodType(classify_forward_prepare, self.transformer) self.transformer.classify_forward_block_forward = types.MethodType( classify_forward_block_forward, self.transformer @@ -1067,6 +1167,10 @@ def rename_key(k: str) -> str: new_state[new_k] = v state_dict = new_state + # AnyFlow HF releases store the r-pathway as condition_embedder.delta_embedder; + # remap to FastGen's r_embedder layout (no-op for all other checkpoints). + state_dict = remap_anyflow_keys(state_dict) + return super().load_state_dict(state_dict, **kwargs) def forward( @@ -1116,6 +1220,16 @@ def forward( assert fwd_pred_type in NET_PRED_TYPES, f"{fwd_pred_type} is not supported as fwd_pred_type" condition = torch.stack(condition, dim=0) if isinstance(condition, list) else condition + + # A gated-fusion flow-map network always consumes the r-pathway (the + # fusion happens inside the shared time projection), so r=None has no + # in-distribution meaning for it. Default to r=t, i.e. the + # instantaneous velocity u(x_t, t, t) — this is how the AnyFlow + # reference queries its real/fake score models. Additive-fusion + # (MeanFlow) networks keep the skip-r-pathway behavior. + if r is None and getattr(self.transformer.r_embedder, "fusion_mode", None) == "gated": + r = t + timestep_mask = torch.ones_like(x_t[:, 0]) # shape: [batch_size, num_latent_frames, H, W] timestep = self._compute_timestep_inputs(t, timestep_mask) r_timestep = None if r is None else self._compute_timestep_inputs(r, timestep_mask) diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py new file mode 100644 index 0000000..fef00e4 --- /dev/null +++ b/tests/test_anyflowmodel.py @@ -0,0 +1,452 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for AnyFlow. + +The pretrain stage is MeanFlow with AnyFlow's hyperparameters (fixed +per-timestep loss weighting, finite-difference JVP, consistency bucket), so +the pretrain tests drive :class:`MeanFlowModel` directly. The on-policy tests +exercise :class:`AnyFlowModel` (DMD2 with a multi-step rollout-with-gradient +student). Both run on the tiny EDM backbone with ``r_timestep=True`` and +``schedule_type=rf`` so they execute on CPU without pretrained weights. +""" + +import gc +import types + +import pytest +import torch + +from fastgen.configs.config_utils import override_config_with_opts +from fastgen.configs.methods.config_anyflow import ModelConfig as AnyFlowModelConfig +from fastgen.configs.methods.config_mean_flow import ModelConfig as MeanFlowModelConfig +from fastgen.methods import AnyFlowModel, MeanFlowModel +from fastgen.utils.test_utils import check_grad_zero + + +def _build_pretrain_model(weight_type="beta08", consistency_ratio=0.25, r_sample_ratio=0.5): + """MeanFlow configured the AnyFlow way (paper Stage 2).""" + gc.collect() + instance = MeanFlowModelConfig() + + instance.loss_config.loss_type = "l2" + instance.loss_config.weight_type = weight_type + instance.loss_config.norm_method = None + instance.loss_config.use_jvp_finite_diff = True + instance.loss_config.jvp_finite_diff_eps = 1e-2 + + instance.sample_t_cfg.time_dist_type = "shifted" + instance.sample_t_cfg.shift = 5.0 + instance.sample_t_cfg.min_t = 0.001 + instance.sample_t_cfg.max_t = 0.999 + instance.sample_t_cfg.r_sample_ratio = r_sample_ratio + instance.sample_t_cfg.consistency_ratio = consistency_ratio + + opts = ["-", "img_resolution=2", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True", "+schedule_type=rf"] + instance.net = override_config_with_opts(instance.net, opts) + instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + instance.precision = "float32" if instance.device == torch.device("cpu") else "bfloat16" + instance.pretrained_model_path = "" + instance.input_shape = [3, 2, 2] + + model = MeanFlowModel(instance) + model.on_train_begin() + model.init_optimizers() + return model + + +def _build_onpolicy_model(): + """On-policy fixture mirrors test_dmd2model: img_resolution=8 so the + discriminator's 4x4 conv kernels can operate.""" + gc.collect() + instance = AnyFlowModelConfig() + + opts = ["-", "img_resolution=8", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True", "+schedule_type=rf"] + instance.net = override_config_with_opts(instance.net, opts) + opts_disc = ["-", "feature_indices=[0]", "all_res=[8]", "in_channels=128"] + instance.discriminator = override_config_with_opts(instance.discriminator, opts_disc) + + instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + instance.precision = "float32" if instance.device == torch.device("cpu") else "bfloat16" + instance.pretrained_model_path = "" + instance.student_update_freq = 2 + instance.student_sample_steps = 2 + instance.input_shape = [3, 8, 8] + + # Co-trained flow-map loss settings (MeanFlow machinery on the tiny net). + instance.sample_t_cfg.time_dist_type = "uniform" + instance.sample_t_cfg.min_t = 0.001 + instance.sample_t_cfg.max_t = 0.999 + instance.sample_t_cfg.r_sample_ratio = 0.5 + instance.sample_t_cfg.consistency_ratio = 0.25 + instance.loss_config.loss_type = "l2" + instance.loss_config.weight_type = "uniform" + instance.loss_config.norm_method = None + instance.loss_config.use_jvp_finite_diff = True + instance.loss_config.jvp_finite_diff_eps = 1e-2 + + model = AnyFlowModel(instance) + model.on_train_begin() + model.init_optimizers() + return model + + +def _make_data(model, img_resolution: int = 2, batch_size: int = 1): + labels = torch.nn.functional.one_hot(torch.randint(0, 10, (batch_size,)), num_classes=10) + neg_labels = torch.zeros(batch_size, 10) + return { + "real": torch.randn(batch_size, 3, img_resolution, img_resolution).to(model.device, model.precision), + "condition": labels.to(model.device, model.precision), + "neg_condition": neg_labels.to(model.device, model.precision), + } + + +# --------------------------------------------------------------------------- +# Pretrain stage — MeanFlow with AnyFlow options +# --------------------------------------------------------------------------- + + +def test_pretrain_single_train_step(): + model = _build_pretrain_model() + data = _make_data(model) + + loss_map, outputs = model.single_train_step(data, 0) + + assert "total_loss" in loss_map + assert "mf_loss" in loss_map + assert torch.isfinite(loss_map["total_loss"]).all() + assert "gen_rand" in outputs + + +def test_pretrain_optimizer_step(): + model = _build_pretrain_model() + data = _make_data(model) + for iteration in range(2): + model.optimizers_zero_grad(iteration) + loss_map, _ = model.single_train_step(data, iteration) + model.grad_scaler.scale(loss_map["total_loss"]).backward() + model.optimizers_schedulers_step(iteration) + # After one zero_grad with no backward in between, gradients should be cleared. + model.optimizers_zero_grad(2) + check_grad_zero(model.net) + + +def test_pretrain_finite_difference_falls_back_at_boundaries(): + """When (t ± eps) leaves [min_t, max_t], the one-sided fallback should + still produce a finite JVP estimate. We synthesise worst-case boundary t.""" + model = _build_pretrain_model() + ns = model.net.noise_scheduler + + real = torch.randn(2, 3, 2, 2, device=model.device, dtype=model.precision) + cond = torch.nn.functional.one_hot(torch.tensor([0, 1]), num_classes=10).to(model.device, model.precision) + + t = torch.tensor([float(ns.min_t), float(ns.max_t)], device=model.device, dtype=ns.t_precision) + r = torch.tensor([float(ns.min_t), float(ns.min_t)], device=model.device, dtype=ns.t_precision) + + eps_noise = torch.randn_like(real) + x_t = ns.forward_process(real, eps_noise, t) + dxt_dt = eps_noise - real + + u_theta_jvp = model._jvp(x_t, t, r, dxt_dt, condition=cond) + assert torch.isfinite(u_theta_jvp).all(), "boundary samples must yield finite JVP estimates" + + +def test_pretrain_consistency_bucket_pins_r_to_zero(): + """With consistency_ratio=1.0 (and no flow-matching head), every sample's + r must be pinned to 0 (consistency to clean data, as in the reference).""" + model = _build_pretrain_model(consistency_ratio=1.0, r_sample_ratio=1.0) + _t, r, r_eq_t_mask = model._sample_t_r_buckets(4) + assert not r_eq_t_mask.any() + assert torch.allclose(r.float(), torch.zeros_like(r.float())) + + +def test_pretrain_bucket_partition_is_deterministic(): + """With consistency_ratio > 0, buckets follow the reference's deterministic + global partition: diffusion head, consistency middle, random-pair tail.""" + model = _build_pretrain_model(consistency_ratio=0.25, r_sample_ratio=0.5) + batch_size = 8 + t, r, r_eq_t_mask = model._sample_t_r_buckets(batch_size) + + n_flow_matching = round(0.5 * batch_size) + n_consistency = round(0.25 * batch_size) + assert r_eq_t_mask.tolist() == [True] * n_flow_matching + [False] * (batch_size - n_flow_matching) + assert torch.equal(r[:n_flow_matching], t[:n_flow_matching]) + assert torch.allclose( + r[n_flow_matching : n_flow_matching + n_consistency].float(), + torch.zeros(n_consistency, device=r.device), + ) + + +def test_pretrain_rebalance_to_diffusion(): + """With rebalancing on, each non-diffusion loss is rescaled to the + diffusion-loss mean by a detached per-sample factor.""" + model = _build_pretrain_model() + model.loss_config.rebalance_to_diffusion = True + + mf_loss = torch.tensor([2.0, 4.0, 10.0, 100.0], dtype=torch.float64, requires_grad=True) + r_eq_t_mask = torch.tensor([True, True, False, False]) + loss = model._reduce_mf_loss(mf_loss, r_eq_t_mask) + + # diffusion mean = 3.0; each non-diffusion sample becomes ~3.0. + expected = (2.0 + 4.0 + 3.0 * (10.0 / 10.00001) + 3.0 * (100.0 / 100.00001)) / 4.0 + assert abs(loss.item() - expected) < 1e-3 + loss.backward() + assert mf_loss.grad is not None and torch.isfinite(mf_loss.grad).all() + + +def test_pretrain_rebalance_all_flow_matching_batch(): + """A rank whose batch is entirely flow-matching (r = t) must still take + the rebalance branch (the collective inside must run on every rank) and + reduce to the plain mean.""" + model = _build_pretrain_model() + model.loss_config.rebalance_to_diffusion = True + + mf_loss = torch.tensor([2.0, 4.0], dtype=torch.float64) + r_eq_t_mask = torch.tensor([True, True]) + loss = model._reduce_mf_loss(mf_loss, r_eq_t_mask) + assert abs(loss.item() - 3.0) < 1e-8 + + +def test_pretrain_prediction_side_guidance_fusion(): + """The AnyFlow guidance-distillation branch (guidance_fuse_scale) must run + end to end and keep gradients on the fused prediction.""" + model = _build_pretrain_model() + model.loss_config.guidance_fuse_scale = 3.0 + model.config.cond_dropout_prob = 0.5 + data = _make_data(model, batch_size=2) + + loss_map, _ = model.single_train_step(data, 0) + assert torch.isfinite(loss_map["total_loss"]).all() + loss_map["total_loss"].backward() + grad_seen = any(p.grad is not None for p in model.net.parameters()) + assert grad_seen + + +@pytest.mark.parametrize("weight_type", ["beta08", "gaussian", "uniform"]) +def test_timestep_weight_function(weight_type): + """The fixed per-timestep weight is a direct function of t: non-negative, + finite, and normalized like the reference scheduler.""" + model = _build_pretrain_model(weight_type=weight_type) + model._timestep_weight_scale = None # force re-derivation for this weight_type + + t = torch.linspace(0.0, 1.0, 101, dtype=torch.float64) + w = model._get_timestep_weight(t) + assert torch.all(w >= 0), f"{weight_type} weights must be non-negative" + assert torch.isfinite(w).all() + + if weight_type == "uniform": + # Uniform normalizes to exactly 1.0 over the reference grid. + assert torch.allclose(w, torch.ones_like(w)) + + # The normalization matches the reference set_timesteps grid (1000 points, + # t=0 excluded): sum over the shifted grid == 1000. + shift = model.sample_t_cfg.shift + grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64)[:-1] + grid = shift * grid / (1 + (shift - 1) * grid) + assert abs(model._get_timestep_weight(grid).sum().item() - 1000.0) < 1e-6 + + +# --------------------------------------------------------------------------- +# On-policy stage — DMD2 with the rollout-with-gradient student +# --------------------------------------------------------------------------- + + +def test_onpolicy_student_update_step(): + model = _build_onpolicy_model() + data = _make_data(model, img_resolution=8) + loss_map, outputs = model.single_train_step(data, 0) # iteration 0 -> student update + assert "total_loss" in loss_map + assert "vsd_loss" in loss_map + # The co-trained Stage-2 flow-map loss is part of every student update. + assert "bidirection_loss" in loss_map + assert torch.isfinite(loss_map["total_loss"]).all() + assert "gen_rand" in outputs + + +def test_onpolicy_cotrain_can_be_disabled(): + model = _build_onpolicy_model() + model.config.cotrain_pretrain_weight = 0.0 + data = _make_data(model, img_resolution=8) + loss_map, _ = model.single_train_step(data, 0) + assert "bidirection_loss" not in loss_map + + +def test_onpolicy_rollout_compresses_to_three_forwards(): + """Regardless of the sampled NFE, the rollout must run at most three + network forwards (jump -> fine step -> jump), with gradient through all.""" + model = _build_onpolicy_model() + model.config.student_sample_steps_list = [16] + real = torch.randn(1, 3, 8, 8, device=model.device, dtype=model.precision) + cond = torch.nn.functional.one_hot(torch.tensor([0]), num_classes=10).to(model.device, model.precision) + input_student, t_student, _, _ = model._generate_noise_and_time(real) + + calls = [] + orig_forward = model.net.forward + + def counting_forward(*args, **kwargs): + calls.append(1) + return orig_forward(*args, **kwargs) + + model.net.forward = counting_forward + gen = model.gen_data_from_net(input_student, t_student, condition=cond) + model.net.forward = orig_forward + + assert len(calls) <= 3, f"rollout must compress to <= 3 forwards, got {len(calls)}" + assert gen.requires_grad + + +def test_onpolicy_fake_score_discriminator_update_step(): + model = _build_onpolicy_model() + model.precision = torch.float32 + model.on_train_begin() + data = _make_data(model, img_resolution=8) + for k, v in data.items(): + if isinstance(v, torch.Tensor): + data[k] = v.to(model.precision) + loss_map, outputs = model.single_train_step(data, 1) # iteration 1 -> fake_score/disc update + assert "fake_score_loss" in loss_map + assert "gan_loss_disc" in loss_map + assert "gen_rand" in outputs + + +def test_onpolicy_rollout_propagates_gradient(): + """The rollout output must keep the autograd graph so the DMD generator + update has a valid gradient. + + Mirrors AnyFlow's ``training_rollout`` (pipeline_wan_anyflow.py): the + compressed jump -> fine step -> jump rollout runs with gradient through + all segments. + """ + model = _build_onpolicy_model() + real = torch.randn(1, 3, 8, 8, device=model.device, dtype=model.precision) + cond = torch.nn.functional.one_hot(torch.tensor([0]), num_classes=10).to(model.device, model.precision) + + input_student, t_student, _, _ = model._generate_noise_and_time(real) + gen = model.gen_data_from_net(input_student, t_student, condition=cond) + + assert tuple(gen.shape) == (1, 3, 8, 8), f"rollout output shape mismatch: {gen.shape}" + assert gen.requires_grad, "rollout output must keep autograd graph at the chosen step" + loss = gen.float().pow(2).mean() + loss.backward() + grad_seen = any(p.grad is not None and torch.isfinite(p.grad).all() for p in model.net.parameters()) + assert grad_seen, "no gradient reached the student network through the rollout" + + +def test_onpolicy_rollout_no_grad_under_no_grad(): + """Under torch.no_grad (the fake-score update path), the rollout must run + fully gradient-free.""" + model = _build_onpolicy_model() + real = torch.randn(1, 3, 8, 8, device=model.device, dtype=model.precision) + cond = torch.nn.functional.one_hot(torch.tensor([0]), num_classes=10).to(model.device, model.precision) + + input_student, t_student, _, _ = model._generate_noise_and_time(real) + with torch.no_grad(): + gen = model.gen_data_from_net(input_student, t_student, condition=cond) + assert not gen.requires_grad + + +def test_onpolicy_optimizer_step(): + model = _build_onpolicy_model() + data = _make_data(model, img_resolution=8) + for iteration in range(2): + model.optimizers_zero_grad(iteration) + loss_map, _ = model.single_train_step(data, iteration) + model.grad_scaler.scale(loss_map["total_loss"]).backward() + model.optimizers_schedulers_step(iteration) + + +# --------------------------------------------------------------------------- +# Wan r-embedder fusion + AnyFlow checkpoint remap (pure functions) +# --------------------------------------------------------------------------- + + +def _make_fake_fusion_self(fusion_mode, gate_value=0.25, encoder_depth=None, dim=4, proj_dim=12): + r_embedder = torch.nn.Module() + r_embedder.time_proj = torch.nn.Linear(dim, proj_dim) + r_embedder.act_fn = torch.nn.SiLU() + r_embedder.fusion_mode = fusion_mode + r_embedder.gate_value = gate_value + return types.SimpleNamespace(r_embedder=r_embedder, encoder_depth=encoder_depth) + + +def test_fuse_r_embedding_gated(): + from fastgen.networks.Wan.network import _fuse_r_embedding + + torch.manual_seed(0) + fake = _make_fake_fusion_self("gated") + temb = torch.randn(2, 4) + remb = torch.randn(2, 4) + timestep_proj = torch.randn(2, 6, 2) + + out_temb, out_proj, out_r_proj = _fuse_r_embedding(fake, temb, timestep_proj, remb, None) + + gate = fake.r_embedder.gate_value + rt_emb = (1 - gate) * temb + gate * remb + expected = fake.r_embedder.time_proj(fake.r_embedder.act_fn(rt_emb)).unflatten(1, (6, -1)) + assert torch.allclose(out_temb, rt_emb) + assert torch.allclose(out_proj, expected) + assert out_r_proj is None + + +def test_fuse_r_embedding_gated_respects_encoder_depth(): + from fastgen.networks.Wan.network import _fuse_r_embedding + + torch.manual_seed(0) + fake = _make_fake_fusion_self("gated", encoder_depth=2) + temb = torch.randn(2, 4) + remb = torch.randn(2, 4) + timestep_proj = torch.randn(2, 6, 2) + + out_temb, out_proj, out_r_proj = _fuse_r_embedding(fake, temb, timestep_proj, remb, None) + + # Encoder blocks keep the t-only projection; the gated projection is + # returned separately so the block loop switches at encoder_depth. + assert torch.allclose(out_proj, timestep_proj) + assert out_r_proj is not None and out_r_proj.shape == timestep_proj.shape + gate = fake.r_embedder.gate_value + assert torch.allclose(out_temb, (1 - gate) * temb + gate * remb) + + +def test_fuse_r_embedding_additive_unchanged(): + from fastgen.networks.Wan.network import _fuse_r_embedding + + torch.manual_seed(0) + fake = _make_fake_fusion_self("additive") + temb = torch.randn(2, 4) + remb = torch.randn(2, 4) + timestep_proj = torch.randn(2, 6, 2) + + out_temb, out_proj, out_r_proj = _fuse_r_embedding(fake, temb, timestep_proj, remb, None) + + r_proj = fake.r_embedder.time_proj(fake.r_embedder.act_fn(remb)).unflatten(1, (6, -1)) + assert torch.allclose(out_temb, temb + remb) + assert torch.allclose(out_proj, timestep_proj + r_proj) + assert torch.allclose(out_r_proj, r_proj) + + +@pytest.mark.parametrize("prefix", ["", "transformer."]) +def test_remap_anyflow_keys(prefix): + from fastgen.networks.Wan.network import remap_anyflow_keys + + sd = { + f"{prefix}condition_embedder.delta_embedder.linear_1.weight": torch.randn(2, 2), + f"{prefix}condition_embedder.time_proj.weight": torch.randn(2, 2), + f"{prefix}condition_embedder.time_proj.bias": torch.randn(2), + f"{prefix}blocks.0.attn1.to_q.weight": torch.randn(2, 2), + } + out = remap_anyflow_keys(sd) + + assert f"{prefix}r_embedder.time_embedder.linear_1.weight" in out + assert f"{prefix}condition_embedder.delta_embedder.linear_1.weight" not in out + # The shared time_proj is duplicated into the r_embedder. + assert torch.equal(out[f"{prefix}r_embedder.time_proj.weight"], sd[f"{prefix}condition_embedder.time_proj.weight"]) + assert torch.equal(out[f"{prefix}r_embedder.time_proj.bias"], sd[f"{prefix}condition_embedder.time_proj.bias"]) + # Unrelated keys untouched. + assert torch.equal(out[f"{prefix}blocks.0.attn1.to_q.weight"], sd[f"{prefix}blocks.0.attn1.to_q.weight"]) + + +def test_remap_anyflow_keys_noop_without_delta_keys(): + from fastgen.networks.Wan.network import remap_anyflow_keys + + sd = {"transformer.blocks.0.attn1.to_q.weight": torch.randn(2, 2)} + assert remap_anyflow_keys(sd) is sd