Add AnyFlow algorithm (any-step video diffusion via flow maps)

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

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 = "<path-to-stage2-pretrain-ckpt>"
+
+    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

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

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

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