Add AnyFlow algorithm

AnyFlow is an any-step video diffusion method that trains a single model
u_theta(x_t, t, r) to predict the average velocity from t back to r, so
the same checkpoint supports arbitrary inference NFE.

Training has two stages, switched via config.loss_config.training_stage:

  * pretrain  — flow-map prediction with a central-difference target
                target = (eps - x0) - (t - r) * dF/dt
                with dF/dt estimated by central differences at (t ± delta).
                Per-batch sampling assigns r=t to a `diffusion_ratio`
                fraction (pure flow matching) and r=0 to a
                `consistency_ratio` fraction (consistency to clean data).

  * onpolicy  — distribution-matching distillation with r=0 conditioning
                on top of the pretrained flow-map weights. Inherits DMD2's
                alternating fake_score / teacher / discriminator updates.

The backbone requirement (a secondary timestep r) is already satisfied by
the Wan transformer with r_timestep=True, which MeanFlow also exercises;
no Wan-side changes are needed.

New files:
  fastgen/methods/distribution_matching/anyflow.py
  fastgen/methods/distribution_matching/anyflow_scheduler.py
  fastgen/configs/methods/config_anyflow.py
  fastgen/configs/experiments/WanT2V/config_anyflow.py
  tests/test_anyflowmodel.py

Modified:
  fastgen/methods/__init__.py                       (+1 import)
  fastgen/methods/distribution_matching/README.md   (+1 algorithm entry)

The multi-step rollout-with-gradient training (matching
self_forcing.py's rollout_with_gradient) is intentionally left for a
follow-up PR — the on-policy stage here uses single-step student
generation.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Enderfga <qq2639135175@gmail.com>

Author: Enderfga

fastgen/configs/experiments/WanT2V/config_anyflow.py

@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Reference AnyFlow experiment config on Wan-1.3B T2V.
+
+Mirrors the AnyFlow paper's pretrain configuration: 1.3B student initialised
+from a Wan2.1-T2V checkpoint, flow-matching shift=5, beta08 loss weighting,
+6k iterations with batch_size_global=32 and lr=5e-5.
+
+Switching to the on-policy stage:
+
+    config.model.loss_config.training_stage = "onpolicy"
+    config.model.pretrained_student_net_path = "<path-to-pretrain-ckpt>"
+
+and adjust ``student_update_freq`` / ``gan_loss_weight_gen`` to taste.
+"""
+
+import fastgen.configs.methods.config_anyflow as config_anyflow_default
+from fastgen.configs.data import VideoLoaderConfig
+from fastgen.configs.discriminator import Discriminator_Wan_1_3B_Config
+from fastgen.configs.net import Wan_1_3B_Config
+
+
+def create_config():
+    config = config_anyflow_default.create_config()
+
+    # Default to the pretrain stage; flip the switch to "onpolicy" once the
+    # flow-map pretrain checkpoint is available.
+    config.model.loss_config.training_stage = "pretrain"
+    config.model.loss_config.jvp_finite_diff_eps = 5e-3
+    config.model.loss_config.diffusion_ratio = 0.5
+    config.model.loss_config.consistency_ratio = 0.25
+    config.model.loss_config.weight_type = "beta08"
+    config.model.loss_config.shift = 5.0
+
+    config.model.net = Wan_1_3B_Config
+    config.model.net.r_timestep = True
+
+    # The on-policy stage uses these too, but they are harmless in pretrain.
+    config.model.discriminator = Discriminator_Wan_1_3B_Config
+    config.model.discriminator.disc_type = "multiscale_down_mlp_large"
+    config.model.discriminator.feature_indices = [15, 22, 29]
+    config.model.gan_loss_weight_gen = 0.0  # disabled by default in pretrain
+    config.model.guidance_scale = 5.0
+
+    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]
+
+    config.model.net_optimizer.lr = 5e-5
+    config.model.fake_score_optimizer.lr = 5e-5
+    config.model.discriminator_optimizer.lr = 5e-5
+
+    config.model.sample_t_cfg.time_dist_type = "shifted"
+    config.model.sample_t_cfg.min_t = 0.001
+    config.model.sample_t_cfg.max_t = 0.999
+
+    config.model.student_sample_type = "ode"
+    # Any-step model — multiple NFEs validated at inference time.
+    config.model.student_sample_steps = 4
+    config.model.sample_t_cfg.t_list = [0.999, 0.937, 0.833, 0.624, 0.0]
+
+    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/methods/config_anyflow.py

@@ -0,0 +1,92 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Config schema for the AnyFlow method.
+
+AnyFlow inherits the DMD2 model config (so the on-policy stage gets fake_score
+/ discriminator / alternating-step machinery for free) and adds a
+``LossConfig`` describing the flow-map pretrain hyperparameters.
+"""
+
+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.methods import AnyFlowModel
+from fastgen.utils import LazyCall as L
+
+
+@attrs.define(slots=False)
+class LossConfig:
+    """Hyperparameters for the AnyFlow flow-map loss and on-policy switch."""
+
+    # Which stage to train. "pretrain" runs the central-difference flow-map
+    # objective; "onpolicy" inherits DMD2's alternating distillation with
+    # dual-timestep r=0 conditioning.
+    training_stage: str = "pretrain"
+
+    # Central-difference step size for estimating dF/dt. Lives in the same
+    # units as the noise scheduler's timesteps. The default (5e-3) matches the
+    # AnyFlow paper's choice of epsilon=5 with num_train_timesteps=1000.
+    jvp_finite_diff_eps: float = 5e-3
+
+    # Per-batch fraction with r = t (recovers pure flow matching).
+    diffusion_ratio: float = 0.5
+    # Per-batch fraction with r = min_t (forces consistency to clean data).
+    consistency_ratio: float = 0.25
+
+    # Per-timestep loss weighting scheme — passed through to the flow-map
+    # scheduler. One of "gaussian", "beta08", "uniform".
+    weight_type: str = "beta08"
+    # Flow-matching schedule shift for the weighting / sampling scheduler.
+    # Wan video defaults use 5.0; image use 1.0.
+    shift: float = 1.0
+    # Resolution of the discrete weighting grid; matches the AnyFlow reference.
+    num_train_timesteps: int = 1000
+
+
+@attrs.define(slots=False)
+class ModelConfig(DMD2ModelConfig):
+    """AnyFlow model config — inherits DMD2 fields, adds the flow-map loss config."""
+
+    loss_config: LossConfig = attrs.field(factory=LossConfig)
+
+
+@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,
+        }
+    )
+
+    # Pretrain stage relies on a flow-matching net_pred_type and 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/methods/__init__.py

@@ -9,6 +9,7 @@
 
 from fastgen.methods.distribution_matching.causvid import CausVidModel as CausVidModel
 from fastgen.methods.distribution_matching.self_forcing import SelfForcingModel as SelfForcingModel
+from fastgen.methods.distribution_matching.anyflow import AnyFlowModel as AnyFlowModel
 
 from fastgen.methods.consistency_model.CM import CMModel as CMModel
 from fastgen.methods.consistency_model.TCM import TCMModel as TCMModel

fastgen/methods/distribution_matching/README.md

@@ -85,6 +85,31 @@ 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** — flow-map prediction with a central-difference target that reuses the network's own forward at `(t ± δ, r)` to estimate `dF/dt`. Per-batch sampling assigns `r = t` to a `diffusion_ratio` fraction (recovering plain flow matching) and `r = 0` to a `consistency_ratio` fraction (forcing consistency to clean data).
+2. **On-policy** — distribution-matching distillation with `r = 0` conditioning on top of the pretrained flow-map weights. Inherits DMD2's alternating fake_score / discriminator / VSD machinery.
+
+**Key Parameters:**
+- `loss_config.training_stage`: `"pretrain"` or `"onpolicy"`
+- `loss_config.jvp_finite_diff_eps`: central-difference step δ (in noise scheduler t-units)
+- `loss_config.diffusion_ratio` / `loss_config.consistency_ratio`: per-batch fraction with `r=t` / `r=0`
+- `loss_config.weight_type`: `gaussian` | `beta08` | `uniform` per-timestep loss weight
+- `loss_config.shift`: flow-matching schedule shift (5.0 for Wan video)
+- See also key parameters of DMD2 above (used by the on-policy stage)
+
+**Backbone requirement:** the student network must accept a secondary timestep `r` (Wan with `r_timestep=True`).
+
+**Note:** the on-policy stage in this PR uses single-step student generation. Multi-step rollout-with-gradient (matching `self_forcing.py`'s `rollout_with_gradient`) is intentionally deferred to a follow-up PR.
+
+**Configs:** [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py)
+
+---
+
 ## Self-Forcing
 
 **File:** [`self_forcing.py`](self_forcing.py) | **Reference:** [Huang et al., 2025](https://arxiv.org/abs/2506.08009)