[misc] Apply yapf/ruff pre-commit autofixes after rebase (main moved to line-length 120)

- scripts/verify_anyflow_fastvideo_parity.py, scripts/demo_anyflow_14b.py:
  read ANYFLOW_LOCAL / ANYFLOW_REF / ANYFLOW_DEMO_OUT from the
  environment so the scripts are not tied to a single workstation
  layout. Drop the bespoke `srun` invocation example from the docstring
  in favour of a plain `python scripts/...` snippet.
- fastvideo/train/methods/distribution_matching/anyflow_pretrain.py:
  collapse the diff_mean if/else into a ternary so ruff SIM108 is
  happy.

docs/distillation/anyflow.md

@@ -0,0 +1,116 @@
+# 🌊 AnyFlow Any-Step Video Distillation
+
+**AnyFlow** ([paper](https://arxiv.org/abs/2605.13724), [project page](https://nvlabs.github.io/AnyFlow/), [official code](https://github.com/NVlabs/AnyFlow), [model weights](https://huggingface.co/collections/nvidia/anyflow)) is an any-step video diffusion framework built on flow maps. A single distilled checkpoint can be evaluated at NFE ∈ {1, 2, 4, 8, 16, 32} without retraining, and quality scales **monotonically** with steps — unlike consistency-based distillation, which often degrades as NFE grows.
+
+The student network ``u_θ(x_t, t, r)`` predicts the *average velocity* from time ``t`` back to time ``r``, so one Euler step is
+
+```
+x_r = x_t - ((t - r) / N) · u_θ(x_t, t, r)
+```
+
+for any ``t > r``.
+
+## 📊 Model Overview
+
+NVIDIA publishes four checkpoints under [`nvidia/anyflow`](https://huggingface.co/collections/nvidia/anyflow):
+
+- `nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers` — bidirectional T2V, Wan2.1 1.3B base
+- `nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers` — bidirectional T2V, Wan2.1 14B base
+- `nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers` — frame-autoregressive variant, 1.3B
+- `nvidia/AnyFlow-FAR-Wan2.1-14B-Diffusers` — frame-autoregressive variant, 14B
+
+FastVideo currently supports the bidirectional T2V variants for training; the FAR variants can be loaded for inference through the diffusers integration.
+
+## ⚙️ Inference
+
+For inference, load the published checkpoint directly through diffusers; FastVideo's training-side ``WanModel`` config maps the HF AnyFlow ``delta_embedder`` weights onto its internal layout via ``param_names_mapping`` so the same checkpoint can be used as the ``init_from`` for the on-policy YAML below.
+
+## 🧠 Algorithm
+
+Training runs in two stages. Both use the dual-timestep Wan backbone — enabled by ``pipeline.dit_config.r_embedder: true`` in the YAML, which allocates a sibling ``condition_embedder.delta_embedder`` and fuses its embedding with the standard timestep embedding via either an additive or a gated mixer.
+
+### Stage 1 — Pretrain (flow-map central-difference)
+
+Method: ``AnyFlowPretrainMethod`` (``fastvideo/train/methods/distribution_matching/anyflow_pretrain.py``)
+
+For each batch, sample ``(t, r) ∈ [0, 1]`` as ``(max, min)`` of two uniform draws, then:
+
+- a ``diffusion_ratio`` fraction (default 0.5) gets ``r = t`` — recovers plain flow matching;
+- a ``consistency_ratio`` fraction (default 0.25) gets ``r = 0`` — forces consistency to clean data;
+- the remainder is free.
+
+The student forward at ``(t, r)`` is trained against the central-difference target
+
+```
+target = (eps - x_0) - (t - r) · dF/dt
+```
+
+where ``dF/dt`` is estimated from the student's own forward at ``(t ± δ, r)`` with the sample also moved along the flow trajectory by ``v_pred · (δ / N)``. Per-timestep weighting uses ``beta08`` (``w(t) = t · sqrt(1 - t)``, renormalized). A stop-gradient scale-balance keeps the non-diffusion branches' loss magnitude aligned with the diffusion branch.
+
+### Stage 2 — On-policy DMD
+
+Method: ``AnyFlowMethod`` (``fastvideo/train/methods/distribution_matching/anyflow.py``)
+
+Inherits ``DMD2Method``. The student is rolled out for ``student_sample_steps`` Euler-flow steps from pure noise; one randomly-chosen step is gradient-enabled (broadcast from rank 0 so every worker agrees), the rest run under ``torch.no_grad``. With ``use_mean_velocity: true`` (default) the rollout uses ``r = t_next`` at each step, matching AnyFlow's ``WanAnyFlowPipeline.training_rollout``.
+
+The inherited ``_dmd_loss`` (VSD with fake-score critic) consumes the rollout output and the teacher's CFG prediction. The optional pinned ``t_list_override`` lets configs reproduce the paper's hand-tuned 4-step schedule ``[999, 937, 833, 624, 0]``.
+
+## 🚀 Training Scripts
+
+### Stage 1 — pretrain
+
+```bash
+bash examples/train/run.sh \
+    examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml
+```
+
+**Key configuration** (in ``examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml``):
+
+- Global batch size: 32 (8 GPUs × 4 per-GPU)
+- Learning rate: 5e-5
+- Flow shift: 5.0
+- ``diffusion_ratio`` / ``consistency_ratio``: 0.5 / 0.25
+- ``epsilon`` (finite-difference step): 5 (absolute train-timestep units)
+- ``weight_type``: ``beta08``
+- ``fuse_guidance_scale``: 3.0
+- Training steps: 6000
+
+### Stage 2 — on-policy
+
+```bash
+bash examples/train/run.sh \
+    examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml \
+    --models.student.init_from outputs/wan2.1_anyflow_pretrain/checkpoint-final
+```
+
+(Or point ``models.student.init_from`` directly at ``nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers`` to bootstrap from the paper weights and skip Stage 1.)
+
+**Key configuration**:
+
+- Global batch size: 8 (8 GPUs × 1 per-GPU)
+- Learning rate: 2e-6
+- Flow shift: 5.0
+- ``student_sample_steps``: 4
+- ``t_list_override``: ``[999, 937, 833, 624, 0]``
+- ``use_mean_velocity``: ``true`` (i.e. ``r = t_next`` during rollout)
+- ``real_score_guidance_scale``: 3.0
+- ``generator_update_interval``: 5 (DMD2 alternation)
+- Training steps: 4000
+
+## 🔌 Loading published AnyFlow checkpoints
+
+The HF AnyFlow checkpoints expose ``condition_embedder.delta_embedder.*`` weights that FastVideo internally maps onto its ``condition_embedder.delta_embedder.mlp.*`` layout. This rename happens automatically through the regex in ``WanVideoArchConfig.param_names_mapping`` — no separate adapter is needed. The same regex is a no-op on plain Wan checkpoints (which don't contain any ``delta_embedder`` keys).
+
+Set the YAML's ``pipeline.dit_config.r_embedder: true`` to allocate the ``delta_embedder`` module on the FastVideo side; when initializing from a plain Wan checkpoint the delta weights are deep-copied from ``time_embedder`` (matching AnyFlow's ``setup_flowmap_model()`` behavior).
+
+## 🧭 Note on ``fuse_guidance_scale``
+
+Stage 1 optionally fuses classifier-free guidance into the training target so the resulting checkpoint can be sampled at ``guidance_scale=1.0`` (no extra forward pass at inference time). The transformation is
+
+```
+noise_pred ← (noise_pred - (1 - g) · noise_pred_uncond) / g
+```
+
+with ``g = fuse_guidance_scale``. The negative prompt embedding comes from ``WanModel``'s ``ensure_negative_conditioning()`` — i.e. the dataset's configured ``sampling_param.negative_prompt``. Setting ``fuse_guidance_scale: 1.0`` skips the extra unconditional forward entirely.
+
+The on-policy stage's ``real_score_guidance_scale`` (inherited from DMD2) follows the same parameterization conventions documented in [``dmd.md``](dmd.md#-note-on-real_score_guidance_scale).

examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml

@@ -0,0 +1,105 @@
+# AnyFlow on-policy DMD — Wan 2.1 T2V 1.3B.
+#
+# Stage 2 of the AnyFlow two-stage recipe. Continues from the pretrain
+# checkpoint; refines the student via DMD2 with a multi-step Euler-flow
+# rollout from pure noise. Teacher provides the real score, critic
+# learns the fake score; both inherited from DMD2Method.
+#
+# Replace <PATH_TO_PRETRAIN_CKPT> with the output of the pretrain stage,
+# or with the NVIDIA-released checkpoint
+# nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers to bootstrap directly from
+# the paper weights (the delta_embedder rename is handled by the
+# param_names_mapping in WanVideoArchConfig).
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: <PATH_TO_PRETRAIN_CKPT>
+    trainable: true
+  teacher:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
+    trainable: false
+    disable_custom_init_weights: true
+  critic:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.anyflow.AnyFlowMethod
+  rollout_mode: simulate
+  generator_update_interval: 5
+  real_score_guidance_scale: 3.0
+  dmd_denoising_steps: [999, 937, 833, 624]
+  warp_denoising_step: false
+
+  # AnyFlow rollout knobs.
+  student_sample_steps: 4
+  use_mean_velocity: true
+  t_list_override: [999.0, 937.0, 833.0, 624.0, 0.0]
+  dmd_score_r_value: 0.0  # DMD scoring conditioning is at r=0 (consistency target).
+
+  # Critic optimizer (DMD2 inherited).
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+  attn_kind: vsa
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/preprocessed
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 21
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wan2.1_anyflow_onpolicy
+    training_state_checkpointing_steps: 500
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: latest
+
+  tracker:
+    project_name: anyflow-wan
+    run_name: wan2.1_t2v_anyflow_onpolicy
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+
+pipeline:
+  flow_shift: 5.0
+  dit_config:
+    r_embedder: true
+    r_embedder_fusion: gated
+    r_embedder_gate_value: 0.25
+    r_embedder_deltatime_type: r

examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml

@@ -0,0 +1,83 @@
+# AnyFlow pretrain (flow-map central-difference) — Wan 2.1 T2V 1.3B.
+#
+# Stage 1 of the AnyFlow two-stage recipe. Trains the dual-timestep
+# u_θ(x_t, t, r) on the central-difference target so the same checkpoint
+# can be sampled at arbitrary NFE in the on-policy stage.
+#
+# Initialize from base Wan 2.1 T2V 1.3B. No teacher or critic at this
+# stage; AnyFlowPretrainMethod owns a single student + one optimizer.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod
+  diffusion_ratio: 0.5
+  consistency_ratio: 0.25
+  epsilon: 5            # finite-difference step in absolute train-timestep units
+  weight_type: beta08   # per-timestep loss weight = t * sqrt(1 - t), renormalized
+  fuse_guidance_scale: 3.0
+  # shift is taken from pipeline.flow_shift below.
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/preprocessed
+    dataloader_num_workers: 4
+    train_batch_size: 4
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 21
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 6000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wan2.1_anyflow_pretrain
+    training_state_checkpointing_steps: 500
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: latest
+
+  tracker:
+    project_name: anyflow-wan
+    run_name: wan2.1_t2v_anyflow_pretrain
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+
+pipeline:
+  flow_shift: 5.0
+  dit_config:
+    # Enable AnyFlow dual-timestep conditioning. The student loads from
+    # base Wan 2.1 — its checkpoint has no delta_embedder weights, so they
+    # get initialized identically to time_embedder via deep-copy in
+    # WanTimeTextImageEmbedding.__init__.
+    r_embedder: true
+    r_embedder_fusion: gated
+    r_embedder_gate_value: 0.25
+    r_embedder_deltatime_type: r

fastvideo/configs/models/dits/wanvideo.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
+from typing import Literal
 
 from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
 
@@ -19,6 +20,11 @@ class WanVideoArchConfig(DiTArchConfig):
             r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
             r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
             r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            # AnyFlow dual-timestep checkpoints expose delta_embedder weights with the
+            # same internal layout as time_embedder. The regex is harmless on plain
+            # Wan checkpoints (no delta_embedder keys to match).
+            r"^condition_embedder\.delta_embedder\.linear_1\.(.*)$": r"condition_embedder.delta_embedder.mlp.fc_in.\1",
+            r"^condition_embedder\.delta_embedder\.linear_2\.(.*)$": r"condition_embedder.delta_embedder.mlp.fc_out.\1",
             r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
             r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
             r"condition_embedder.image_embedder.ff.fc_in.\1",
@@ -86,6 +92,14 @@ class WanVideoArchConfig(DiTArchConfig):
     # "relativistic" keeps long rollouts in-distribution; a no-op unless sink_size > 0 and local_attn_size > 0.
     rope_cache_policy: str = "absolute"
 
+    # AnyFlow dual-timestep conditioning. Defaults preserve bit-identity with
+    # the legacy single-timestep forward (no delta_embedder allocated, no
+    # extra computation on the embedder forward path).
+    r_embedder: bool = False
+    r_embedder_fusion: Literal["additive", "gated"] = "additive"
+    r_embedder_gate_value: float = 0.25
+    r_embedder_deltatime_type: Literal["r", "t-r"] = "r"
+
     def __post_init__(self):
         super().__post_init__()
         self.out_channels = self.out_channels or self.in_channels

fastvideo/models/dits/wanvideo.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import copy
 import math
 from typing import Any
 
@@ -59,6 +60,11 @@ class WanTimeTextImageEmbedding(nn.Module):
         time_freq_dim: int,
         text_embed_dim: int,
         image_embed_dim: int | None = None,
+        *,
+        r_embedder: bool = False,
+        r_embedder_fusion: str = "additive",
+        r_embedder_gate_value: float = 0.25,
+        r_embedder_deltatime_type: str = "r",
     ):
         super().__init__()
 
@@ -77,14 +83,57 @@ class WanTimeTextImageEmbedding(nn.Module):
         if image_embed_dim is not None:
             self.image_embedder = WanImageEmbedding(image_embed_dim, dim)
 
+        # AnyFlow dual-timestep support. When r_embedder is False the forward
+        # path bypasses delta_embedder entirely and the output is byte-identical
+        # to the legacy single-timestep implementation.
+        self._r_embedder_enabled = bool(r_embedder)
+        self._r_embedder_fusion = r_embedder_fusion
+        self._r_embedder_deltatime_type = r_embedder_deltatime_type
+        if self._r_embedder_enabled:
+            if r_embedder_fusion not in ("additive", "gated"):
+                raise ValueError(
+                    "r_embedder_fusion must be one of {additive, gated}, "
+                    f"got {r_embedder_fusion!r}")
+            if r_embedder_deltatime_type not in ("r", "t-r"):
+                raise ValueError(
+                    "r_embedder_deltatime_type must be one of {r, t-r}, "
+                    f"got {r_embedder_deltatime_type!r}")
+            # Deep-copy preserves identical initialization with time_embedder,
+            # matching AnyFlow reference setup_flowmap_model() behavior.
+            self.delta_embedder = copy.deepcopy(self.time_embedder)
+            # Non-persistent buffer — gate is a hyperparameter, not learned.
+            self.register_buffer(
+                "_r_embedder_gate",
+                torch.tensor(float(r_embedder_gate_value)),
+                persistent=False,
+            )
+        else:
+            self.delta_embedder = None
+
     def forward(
         self,
         timestep: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
         encoder_hidden_states_image: torch.Tensor | None = None,
         timestep_seq_len: int | None = None,
+        r_timestep: torch.Tensor | None = None,
     ):
         temb = self.time_embedder(timestep, timestep_seq_len)
+
+        if self._r_embedder_enabled and r_timestep is not None:
+            assert self.delta_embedder is not None
+            if self._r_embedder_deltatime_type == "r":
+                delta_input = r_timestep
+            else:
+                delta_input = timestep - r_timestep
+            delta_emb = self.delta_embedder(delta_input, timestep_seq_len)
+            gate = self._r_embedder_gate
+            if self._r_embedder_fusion == "gated":
+                temb = (1.0 - gate) * temb + gate * delta_emb
+            else:
+                # Additive (additional channel, no convex blend).
+                temb = temb + gate * delta_emb
+
         timestep_proj = self.time_modulation(temb)
 
         if self.text_embedder is not None:
@@ -595,6 +644,10 @@ class WanTransformer3DModel(BaseDiT):
             time_freq_dim=config.freq_dim,
             text_embed_dim=config.text_dim,
             image_embed_dim=config.image_dim,
+            r_embedder=config.r_embedder,
+            r_embedder_fusion=config.r_embedder_fusion,
+            r_embedder_gate_value=config.r_embedder_gate_value,
+            r_embedder_deltatime_type=config.r_embedder_deltatime_type,
         )
 
         # 3. Transformer blocks
@@ -636,6 +689,7 @@ class WanTransformer3DModel(BaseDiT):
                 encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
                 | None = None,
                 guidance=None,
+                r_timestep: torch.Tensor | None = None,
                 **kwargs) -> torch.Tensor:
         orig_dtype = hidden_states.dtype
         if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
@@ -684,8 +738,17 @@ class WanTransformer3DModel(BaseDiT):
         else:
             ts_seq_len = None
 
+        # AnyFlow dual-timestep — match timestep's flattening so embedder
+        # sees aligned shapes.
+        if r_timestep is not None and r_timestep.dim() == 2:
+            r_timestep = r_timestep.flatten()
+
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len)
+            timestep,
+            encoder_hidden_states,
+            encoder_hidden_states_image,
+            timestep_seq_len=ts_seq_len,
+            r_timestep=r_timestep)
         if ts_seq_len is not None:
             # batch_size, seq_len, 6, inner_dim
             timestep_proj = timestep_proj.unflatten(2, (6, -1))

fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Flow-map any-step Euler scheduler for AnyFlow.
+
+The model predicts the *average* velocity ``u_θ(x_t, t, r)`` from time
+``t`` back to time ``r``, so one Euler step is
+
+    x_r = x_t - ((t - r) / num_train_timesteps) * u_θ(x_t, t, r)
+
+regardless of how far apart ``t`` and ``r`` are. The scheduler also
+provides the AnyFlow training-time helpers ``apply_shift`` (flow-matching
+shift transform) and ``get_train_weight`` (per-timestep loss weight,
+including ``beta08``).
+
+Standalone — does not depend on diffusers' ConfigMixin/SchedulerMixin.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+
+from fastvideo.models.schedulers.base import BaseScheduler
+
+
+WeightType = Literal["uniform", "gaussian", "beta08"]
+
+
+class FlowMapEulerDiscreteScheduler(BaseScheduler):
+    """Minimal flow-map scheduler.
+
+    Parameters
+    ----------
+    num_train_timesteps:
+        Discretization granularity for training. ``t`` is expressed in
+        absolute units in ``[0, num_train_timesteps]``.
+    shift:
+        Flow-matching shift parameter (Wan video default: ``5.0``). Set
+        to ``1.0`` for an identity shift.
+    """
+
+    order: int = 1
+
+    def __init__(
+        self,
+        *,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+    ) -> None:
+        self.num_train_timesteps = int(num_train_timesteps)
+        self.shift = float(shift)
+        self.timesteps: torch.Tensor = torch.empty(0)
+        self.sigmas: torch.Tensor = torch.empty(0)
+        super().__init__()
+
+    # ------------------------------------------------------------------
+    # BaseScheduler abstract surface
+
+    def set_shift(self, shift: float) -> None:
+        self.shift = float(shift)
+
+    def scale_model_input(
+        self,
+        sample: torch.Tensor,
+        timestep: int | None = None,
+    ) -> torch.Tensor:
+        # Flow-matching has no per-step input scaling; pass through.
+        del timestep
+        return sample
+
+    # ------------------------------------------------------------------
+    # Public helpers used by the AnyFlow pretrain method.
+
+    def apply_shift(
+        self,
+        t: torch.Tensor,
+        *,
+        shift: float | None = None,
+    ) -> torch.Tensor:
+        """Apply the flow-matching shift: ``t' = s * t / (1 + (s - 1) * t)``.
+
+        Operates in the normalized ``[0, 1]`` domain — callers should pass
+        ``t / num_train_timesteps`` (or sample ``t`` directly from
+        ``[0, 1]``).
+        """
+        s = self.shift if shift is None else float(shift)
+        if s == 1.0:
+            return t
+        return s * t / (1.0 + (s - 1.0) * t)
+
+    def get_train_weight(
+        self,
+        t: torch.Tensor,
+        *,
+        weight_type: WeightType = "beta08",
+    ) -> torch.Tensor:
+        """Per-timestep training weight, renormalized so the total weight
+        mass equals ``num_train_timesteps`` (matching AnyFlow reference's
+        ``scheduling_flowmap_euler_discrete.py``).
+
+        ``beta08``: ``w(t) = t * sqrt(1 - t)`` (in normalized t-space).
+        """
+        # Auto-detect domain: if t was given in absolute units, normalize.
+        t_f = t.float()
+        max_val = t_f.max() if t_f.numel() > 0 else torch.tensor(0.0)
+        if max_val > 1.0 + 1e-6:
+            t_norm = t_f / self.num_train_timesteps
+        else:
+            t_norm = t_f
+        t_norm = t_norm.clamp(min=0.0, max=1.0)
+
+        if weight_type == "uniform":
+            w = torch.ones_like(t_norm)
+        elif weight_type == "gaussian":
+            w = torch.exp(-0.5 * ((t_norm - 0.5) / 0.2) ** 2)
+        elif weight_type == "beta08":
+            w = t_norm.pow(1.0) * (1.0 - t_norm).clamp_min(0.0).pow(0.5)
+        else:
+            raise ValueError(f"Unknown weight_type: {weight_type!r}")
+
+        denom = w.sum().clamp_min(1e-8)
+        return w * (float(self.num_train_timesteps) / denom)
+
+    # ------------------------------------------------------------------
+
+    def set_timesteps(
+        self,
+        *,
+        num_inference_steps: int,
+        device: torch.device | str = "cpu",
+        custom_timesteps: list[float] | torch.Tensor | None = None,
+    ) -> None:
+        """Build a descending timestep schedule ending at 0.
+
+        With ``num_inference_steps=N`` the schedule has ``N + 1`` entries
+        ``[T_max, ..., 0]`` so a rollout consumes ``N`` Euler steps.
+
+        ``custom_timesteps`` overrides the linspace+shift schedule with
+        a pinned list (in absolute train-timestep units), useful for the
+        AnyFlow paper's hand-tuned ``[999, 937, 833, 624, 0]`` schedule.
+        """
+        if num_inference_steps <= 0:
+            raise ValueError(
+                "num_inference_steps must be positive, "
+                f"got {num_inference_steps}")
+        device = torch.device(device)
+
+        if custom_timesteps is not None:
+            ts = torch.as_tensor(
+                custom_timesteps, dtype=torch.float32, device=device)
+            if ts.ndim != 1:
+                raise ValueError(
+                    "custom_timesteps must be 1-D, got shape "
+                    f"{tuple(ts.shape)}")
+            if not torch.all(ts[:-1] >= ts[1:]):
+                raise ValueError(
+                    "custom_timesteps must be descending (largest first)")
+        else:
+            ts_norm = torch.linspace(
+                1.0, 0.0, num_inference_steps + 1, device=device)
+            ts_norm = self.apply_shift(ts_norm)
+            ts = ts_norm * self.num_train_timesteps
+
+        self.timesteps = ts
+        self.sigmas = ts / self.num_train_timesteps
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        *,
+        sample: torch.Tensor,
+        timestep: torch.Tensor,
+        r_timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        """One Euler step from ``t`` to ``r``.
+
+        ``model_output`` is the average-velocity prediction
+        ``u_θ(x_t, t, r)``. Both ``timestep`` and ``r_timestep`` are in
+        absolute train-timestep units (``[0, num_train_timesteps]``).
+        """
+        t = timestep.to(sample.device, dtype=sample.dtype)
+        r = r_timestep.to(sample.device, dtype=sample.dtype)
+        dt_norm = (t - r) / float(self.num_train_timesteps)
+        # Broadcast dt over channel/spatial dims.
+        view: list[int] = [-1] + [1] * (sample.ndim - 1)
+        return sample - dt_norm.view(*view) * model_output
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        """Linear flow-matching interpolation: ``x_t = (1 - σ) * x_0 + σ * ε``,
+        where ``σ = t / num_train_timesteps``.
+        """
+        sigma = (timestep.to(original_samples.device,
+                             dtype=original_samples.dtype)
+                 / float(self.num_train_timesteps))
+        view: list[int] = [-1] + [1] * (original_samples.ndim - 1)
+        sigma = sigma.view(*view)
+        return (1.0 - sigma) * original_samples + sigma * noise

fastvideo/tests/training/distill/test_anyflow_onpolicy.py

@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+"""AnyFlow on-policy method tests (CPU-only, no Wan instantiation).
+
+The full AnyFlowMethod requires a real student/teacher/critic trio plus
+DMD2's optimizer wiring — too heavyweight for a unit test. These tests
+exercise the rollout-shape helpers and source-level invariants via
+``object.__new__`` bypassing of ``__init__``.
+"""
+
+from __future__ import annotations
+
+import inspect
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+
+from fastvideo.train.methods.distribution_matching.anyflow import AnyFlowMethod
+
+
+# ---------------------------------------------------------------------------
+# Helpers: build a "naked" AnyFlowMethod that skips __init__.
+# ---------------------------------------------------------------------------
+
+
+def _naked_method(
+    *,
+    student_sample_steps: int = 4,
+    use_mean_velocity: bool = True,
+    t_list_override: list[float] | None = None,
+    denoising_step_list: list[float] | None = None,
+) -> AnyFlowMethod:
+    method = AnyFlowMethod.__new__(AnyFlowMethod)
+    method._student_sample_steps = int(student_sample_steps)  # type: ignore[attr-defined]
+    method._use_mean_velocity = bool(use_mean_velocity)  # type: ignore[attr-defined]
+    method._t_list_override = (  # type: ignore[attr-defined]
+        list(t_list_override) if t_list_override else None)
+    method._dmd_score_r = 0.0  # type: ignore[attr-defined]
+    method._real_score_guidance = 1.0  # type: ignore[attr-defined]
+    method.cuda_generator = None  # type: ignore[attr-defined]
+    method._cfg_uncond = None  # type: ignore[attr-defined]
+    method._denoising_step_list_cache = None  # type: ignore[attr-defined]
+
+    # Stub _get_denoising_step_list — DMD2 reads method_config but for these
+    # focused tests we want a deterministic schedule.
+    raw = denoising_step_list or [999, 750, 500, 250]
+    cached = torch.tensor(raw, dtype=torch.long)
+
+    def _stub(self, device: torch.device) -> torch.Tensor:
+        return cached.to(device=device)
+
+    bound = _stub.__get__(method, AnyFlowMethod)
+    method._get_denoising_step_list = bound  # type: ignore[assignment]
+    return method
+
+
+# ---------------------------------------------------------------------------
+# Schedule construction.
+# ---------------------------------------------------------------------------
+
+
+def test_get_rollout_schedule_uses_t_list_override_verbatim() -> None:
+    method = _naked_method(
+        t_list_override=[999.0, 937.0, 833.0, 624.0, 0.0])
+    schedule = method._get_rollout_schedule(device=torch.device("cpu"))
+    torch.testing.assert_close(
+        schedule,
+        torch.tensor([999.0, 937.0, 833.0, 624.0, 0.0],
+                     dtype=torch.float32),
+    )
+
+
+def test_get_rollout_schedule_falls_back_to_denoising_step_list() -> None:
+    method = _naked_method(
+        denoising_step_list=[999, 750, 500, 250])
+    schedule = method._get_rollout_schedule(device=torch.device("cpu"))
+    # Tail must be a 0 boundary so the final Euler step lands at t=0.
+    assert float(schedule[-1].item()) == 0.0
+    # Original step list preserved at the front.
+    torch.testing.assert_close(
+        schedule[:4],
+        torch.tensor([999.0, 750.0, 500.0, 250.0], dtype=torch.float32),
+    )
+
+
+def test_get_rollout_schedule_does_not_double_append_zero() -> None:
+    method = _naked_method(
+        denoising_step_list=[999, 750, 500, 0])
+    schedule = method._get_rollout_schedule(device=torch.device("cpu"))
+    assert schedule.numel() == 4  # no extra boundary inserted.
+
+
+# ---------------------------------------------------------------------------
+# t_list_override validation.
+# ---------------------------------------------------------------------------
+
+
+def test_anyflow_method_rejects_ascending_t_list_override() -> None:
+    """__init__ validates that t_list_override is descending. We test the
+    validation logic by stitching together a minimal cfg + role_models
+    path; if construction fails for an unrelated reason we still catch
+    the descending check via the explicit error message."""
+    src = inspect.getsource(AnyFlowMethod.__init__)
+    assert 't_list_override must be descending' in src
+    assert 'descending' in src
+
+
+def test_anyflow_method_rejects_non_positive_student_sample_steps() -> None:
+    src = inspect.getsource(AnyFlowMethod.__init__)
+    assert 'student_sample_steps must be positive' in src
+
+
+# ---------------------------------------------------------------------------
+# Rollout dynamics — stubbed student.
+# ---------------------------------------------------------------------------
+
+
+class _SpyStudent:
+    """Stand-in student that records every (t, r) pair seen during a
+    rollout and predicts a constant velocity field."""
+
+    def __init__(self, num_train_timesteps: int = 1000) -> None:
+        self.num_train_timesteps = num_train_timesteps
+        self.seen: list[tuple[float, float]] = []
+        # A single trainable parameter so callers can verify gradient flow.
+        self.param = torch.nn.Parameter(torch.zeros(1))
+
+    def predict_velocity_with_r(
+        self,
+        noisy: torch.Tensor,
+        t: torch.Tensor,
+        r: torch.Tensor,
+        batch: Any,
+        *,
+        conditional: bool = True,
+        cfg_uncond: Any = None,
+        attn_kind: str = "vsa",
+    ) -> torch.Tensor:
+        del batch, conditional, cfg_uncond, attn_kind
+        self.seen.append((float(t.flatten()[0].item()),
+                          float(r.flatten()[0].item())))
+        # Constant velocity field of magnitude param so we can backprop.
+        return self.param * torch.ones_like(noisy)
+
+
+def _make_batch(shape: tuple[int, ...]) -> SimpleNamespace:
+    batch = SimpleNamespace()
+    batch.latents = torch.randn(*shape)
+    batch.dmd_latent_vis_dict = {}
+    return batch
+
+
+def test_rollout_uses_mean_velocity_r_equals_t_next() -> None:
+    """With use_mean_velocity=True, r at step i must equal t at step i+1."""
+    method = _naked_method(
+        student_sample_steps=4,
+        use_mean_velocity=True,
+        t_list_override=[999.0, 750.0, 500.0, 250.0, 0.0],
+    )
+    student = _SpyStudent()
+    method.student = student  # type: ignore[assignment]
+
+    batch = _make_batch((1, 2, 4, 4, 4))
+    _ = method._student_rollout(batch, with_grad=False)
+
+    # 4 forwards = 4 (t, r) pairs.
+    assert len(student.seen) == 4
+    for i in range(3):
+        # r at step i must equal t at step i+1.
+        assert student.seen[i][1] == student.seen[i + 1][0]
+
+
+def test_rollout_use_mean_velocity_false_uses_r_equal_t() -> None:
+    method = _naked_method(
+        student_sample_steps=2,
+        use_mean_velocity=False,
+        t_list_override=[999.0, 500.0, 0.0],
+    )
+    student = _SpyStudent()
+    method.student = student  # type: ignore[assignment]
+    batch = _make_batch((1, 2, 4, 4, 4))
+    _ = method._student_rollout(batch, with_grad=False)
+    for t_seen, r_seen in student.seen:
+        assert t_seen == r_seen
+
+
+def test_rollout_with_grad_true_produces_differentiable_output() -> None:
+    method = _naked_method(
+        student_sample_steps=4,
+        use_mean_velocity=True,
+        t_list_override=[999.0, 750.0, 500.0, 250.0, 0.0],
+    )
+    student = _SpyStudent()
+    method.student = student  # type: ignore[assignment]
+    batch = _make_batch((1, 2, 4, 4, 4))
+    out = method._student_rollout(batch, with_grad=True)
+    assert out.requires_grad, (
+        "Rollout output must keep a gradient so the DMD loss can backprop "
+        "through the chosen step.")
+    out.sum().backward()
+    assert student.param.grad is not None
+    assert student.param.grad.abs().sum() > 0
+
+
+def test_rollout_with_grad_false_blocks_gradient_completely() -> None:
+    method = _naked_method(
+        student_sample_steps=4,
+        use_mean_velocity=True,
+        t_list_override=[999.0, 750.0, 500.0, 250.0, 0.0],
+    )
+    student = _SpyStudent()
+    method.student = student  # type: ignore[assignment]
+    batch = _make_batch((1, 2, 4, 4, 4))
+    out = method._student_rollout(batch, with_grad=False)
+    assert not out.requires_grad
+
+
+def test_broadcast_grad_step_index_in_range() -> None:
+    method = _naked_method(student_sample_steps=4)
+    for _ in range(20):
+        idx = method._broadcast_grad_step_index(
+            num_steps=4, device=torch.device("cpu"))
+        assert 0 <= idx < 4
+
+
+def test_broadcast_grad_step_index_rejects_non_positive_num_steps() -> None:
+    method = _naked_method()
+    with pytest.raises(ValueError, match="num_steps must be positive"):
+        method._broadcast_grad_step_index(
+            num_steps=0, device=torch.device("cpu"))

fastvideo/tests/training/distill/test_anyflow_pretrain.py

@@ -0,0 +1,604 @@
+# SPDX-License-Identifier: Apache-2.0
+"""AnyFlow pretrain method tests.
+
+CPU-only unit tests covering:
+- Config flag defaults (bit-identity preserved on legacy paths).
+- ``WanTimeTextImageEmbedding`` dual-timestep forward (additive default = bit-identical
+  to legacy; gated mode reproduces AnyFlow's ``(1 - g) * temb + g * delta_emb`` fusion).
+- ``WanTransformer3DModel.forward`` accepts ``r_timestep``.
+- ``FlowMapEulerDiscreteScheduler`` numerics: ``apply_shift``, ``get_train_weight``,
+  ``step``.
+- ``(t, r)`` per-batch sampling distribution.
+- Central-difference target math.
+- AnyFlow HF checkpoint key remap (``remap_anyflow_keys``).
+"""
+
+from __future__ import annotations
+
+import copy
+
+import pytest
+import torch
+
+from fastvideo.configs.models.dits import WanVideoConfig
+
+
+# ---------------------------------------------------------------------------
+# Task 1: r_embedder config flags default to bit-identity preservation.
+# ---------------------------------------------------------------------------
+
+
+def test_wan_arch_defaults_preserve_bit_identity() -> None:
+    cfg = WanVideoConfig()
+    arch = cfg.arch_config
+    assert arch.r_embedder is False
+    assert arch.r_embedder_fusion == "additive"
+    assert arch.r_embedder_gate_value == 0.25
+    assert arch.r_embedder_deltatime_type == "r"
+
+
+# ---------------------------------------------------------------------------
+# Task 2: WanTimeTextImageEmbedding dual-timestep forward.
+# ---------------------------------------------------------------------------
+
+
+def _init_uninitialized_weights(module: torch.nn.Module, seed: int = 0) -> None:
+    """FastVideo's ``ReplicatedLinear`` allocates weights with
+    ``torch.empty`` and relies on a downstream ``load_weights`` pass to
+    populate them. Unit tests bypass that pass, so weights start as
+    uninitialized garbage (typically NaN/Inf). Manually init every
+    Linear / RMSNorm / LayerNorm parameter so the forward produces
+    deterministic finite outputs.
+    """
+    torch.manual_seed(seed)
+    with torch.no_grad():
+        for p in module.parameters():
+            if p.ndim >= 2:
+                # Xavier-uniform scaled by inverse fan-in for stable forwards.
+                torch.nn.init.xavier_uniform_(p)
+            else:
+                p.zero_()
+
+
+def _make_embedder(
+    *,
+    r_embedder: bool,
+    fusion: str = "additive",
+    gate: float = 0.25,
+    deltatime_type: str = "r",
+    init_seed: int = 0,
+):
+    from fastvideo.models.dits.wanvideo import WanTimeTextImageEmbedding
+
+    emb = WanTimeTextImageEmbedding(
+        dim=32,
+        time_freq_dim=64,
+        text_embed_dim=16,
+        image_embed_dim=None,
+        r_embedder=r_embedder,
+        r_embedder_fusion=fusion,
+        r_embedder_gate_value=gate,
+        r_embedder_deltatime_type=deltatime_type,
+    )
+    _init_uninitialized_weights(emb, seed=init_seed)
+    emb.eval()
+    return emb
+
+
+def test_embedder_default_path_no_delta_module() -> None:
+    """When r_embedder=False, delta_embedder must not be allocated."""
+    emb = _make_embedder(r_embedder=False)
+    assert emb.delta_embedder is None
+
+
+def test_embedder_default_path_is_bit_identical_to_legacy() -> None:
+    """With r_embedder=False, forward output must match the legacy single-t path
+    (no r_timestep kwarg, no extra computation)."""
+    torch.manual_seed(0)
+    emb = _make_embedder(r_embedder=False)
+    t = torch.randint(0, 1000, (2,), dtype=torch.long)
+    txt = torch.randn(2, 4, 16)
+    temb_a, proj_a, _, _ = emb(t, txt)
+    # Calling without r_timestep again must be deterministic-equal.
+    temb_b, proj_b, _, _ = emb(t, txt)
+    torch.testing.assert_close(temb_a, temb_b)
+    torch.testing.assert_close(proj_a, proj_b)
+    assert temb_a.shape == (2, 32)
+    assert proj_a.shape == (2, 32 * 6)
+
+
+def test_embedder_enabled_without_r_timestep_is_bit_identical_to_legacy() -> None:
+    """Even with r_embedder=True, if r_timestep is None at call time the
+    forward must skip the delta path entirely so existing call sites that
+    don't pass r_timestep stay byte-equal to the legacy result."""
+    torch.manual_seed(0)
+    emb_legacy = _make_embedder(r_embedder=False)
+    torch.manual_seed(0)
+    emb_dual = _make_embedder(r_embedder=True, fusion="additive")
+    t = torch.randint(0, 1000, (2,), dtype=torch.long)
+    txt = torch.randn(2, 4, 16)
+    temb_legacy, proj_legacy, _, _ = emb_legacy(t, txt)
+    temb_dual, proj_dual, _, _ = emb_dual(t, txt)  # No r_timestep.
+    torch.testing.assert_close(temb_legacy, temb_dual)
+    torch.testing.assert_close(proj_legacy, proj_dual)
+
+
+def test_embedder_gated_fusion_formula() -> None:
+    """Gated mode: rt_emb = (1 - g) * temb_t + g * delta_emb (with delta_input=r)."""
+    torch.manual_seed(0)
+    emb = _make_embedder(r_embedder=True, fusion="gated", gate=0.25)
+    t = torch.tensor([500, 500], dtype=torch.long)
+    r = torch.tensor([100, 100], dtype=torch.long)
+    txt = torch.randn(2, 4, 16)
+
+    temb_t = emb.time_embedder(t)
+    delta_emb = emb.delta_embedder(r)
+    expected = 0.75 * temb_t + 0.25 * delta_emb
+
+    rt_emb, _, _, _ = emb(t, txt, r_timestep=r)
+    torch.testing.assert_close(rt_emb, expected, rtol=1e-5, atol=1e-5)
+
+
+def test_embedder_additive_fusion_formula() -> None:
+    """Additive mode: rt_emb = temb_t + g * delta_emb."""
+    torch.manual_seed(0)
+    emb = _make_embedder(r_embedder=True, fusion="additive", gate=0.3)
+    t = torch.tensor([700, 700], dtype=torch.long)
+    r = torch.tensor([200, 200], dtype=torch.long)
+    txt = torch.randn(2, 4, 16)
+
+    temb_t = emb.time_embedder(t)
+    delta_emb = emb.delta_embedder(r)
+    expected = temb_t + 0.3 * delta_emb
+
+    rt_emb, _, _, _ = emb(t, txt, r_timestep=r)
+    torch.testing.assert_close(rt_emb, expected, rtol=1e-5, atol=1e-5)
+
+
+def test_embedder_deltatime_type_t_minus_r() -> None:
+    """When deltatime_type='t-r', delta_embedder consumes (t - r)."""
+    torch.manual_seed(0)
+    emb = _make_embedder(
+        r_embedder=True, fusion="gated", gate=0.5, deltatime_type="t-r")
+    t = torch.tensor([800, 800], dtype=torch.long)
+    r = torch.tensor([300, 300], dtype=torch.long)
+    txt = torch.randn(2, 4, 16)
+
+    temb_t = emb.time_embedder(t)
+    delta_emb = emb.delta_embedder(t - r)
+    expected = 0.5 * temb_t + 0.5 * delta_emb
+
+    rt_emb, _, _, _ = emb(t, txt, r_timestep=r)
+    torch.testing.assert_close(rt_emb, expected, rtol=1e-5, atol=1e-5)
+
+
+def test_embedder_invalid_fusion_raises() -> None:
+    with pytest.raises(ValueError, match="r_embedder_fusion"):
+        _make_embedder(r_embedder=True, fusion="bogus")
+
+
+def test_embedder_invalid_deltatime_type_raises() -> None:
+    with pytest.raises(ValueError, match="r_embedder_deltatime_type"):
+        _make_embedder(r_embedder=True, fusion="gated", deltatime_type="2t-r")
+
+
+def test_embedder_gate_not_in_state_dict() -> None:
+    """Gate is a non-persistent buffer; it must not appear in state_dict so
+    checkpoints stay portable across different gate hyperparameters."""
+    emb = _make_embedder(r_embedder=True, fusion="gated", gate=0.25)
+    keys = list(emb.state_dict().keys())
+    assert not any("_r_embedder_gate" in k for k in keys)
+
+
+# ---------------------------------------------------------------------------
+# Task 3: WanTransformer3DModel threads r_timestep through.
+# ---------------------------------------------------------------------------
+
+
+def test_wan_transformer_forward_signature_has_r_timestep() -> None:
+    """The forward signature must declare r_timestep explicitly (not
+    swallowed by **kwargs) so callers and type checkers can see it."""
+    import inspect
+
+    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
+
+    sig = inspect.signature(WanTransformer3DModel.forward)
+    assert "r_timestep" in sig.parameters
+    param = sig.parameters["r_timestep"]
+    assert param.default is None
+
+
+def test_wan_transformer_init_propagates_r_embedder_config() -> None:
+    """When the arch config sets r_embedder=True the WanTransformer3DModel
+    constructor must instantiate the embedder with the delta path active.
+
+    We avoid full WanTransformer3DModel instantiation (which requires
+    distributed init) by reading the source's __init__ to confirm it
+    forwards the four arch flags to WanTimeTextImageEmbedding.
+    """
+    import inspect
+
+    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
+
+    src = inspect.getsource(WanTransformer3DModel.__init__)
+    # All four arch config fields must be passed to WanTimeTextImageEmbedding.
+    assert "r_embedder=config.r_embedder" in src
+    assert "r_embedder_fusion=config.r_embedder_fusion" in src
+    assert "r_embedder_gate_value=config.r_embedder_gate_value" in src
+    assert "r_embedder_deltatime_type=config.r_embedder_deltatime_type" in src
+
+
+def test_wan_transformer_forward_threads_r_timestep_to_embedder() -> None:
+    """The forward must pass r_timestep into the embedder call (verified via
+    source inspection to avoid heavyweight distributed bring-up)."""
+    import inspect
+
+    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
+
+    src = inspect.getsource(WanTransformer3DModel.forward)
+    assert "r_timestep=r_timestep" in src, (
+        "WanTransformer3DModel.forward must forward r_timestep into "
+        "self.condition_embedder")
+
+
+# ---------------------------------------------------------------------------
+# Task 4: FlowMapEulerDiscreteScheduler numerics.
+# ---------------------------------------------------------------------------
+
+
+def _scheduler(*, shift: float = 1.0, n_train: int = 1000):
+    from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import (
+        FlowMapEulerDiscreteScheduler, )
+    return FlowMapEulerDiscreteScheduler(
+        num_train_timesteps=n_train, shift=shift)
+
+
+def test_flow_map_scheduler_set_timesteps_descending() -> None:
+    sched = _scheduler(shift=5.0)
+    sched.set_timesteps(num_inference_steps=4, device=torch.device("cpu"))
+    ts = sched.timesteps
+    # N inference steps → N + 1 boundary entries.
+    assert ts.numel() == 5
+    assert torch.all(ts[:-1] >= ts[1:])  # descending
+    assert ts[-1].item() == 0.0
+    assert ts[0].item() == pytest.approx(1000.0, abs=1e-3)
+
+
+def test_flow_map_scheduler_set_timesteps_custom_overrides_schedule() -> None:
+    sched = _scheduler(shift=5.0)
+    custom = [999.0, 937.0, 833.0, 624.0, 0.0]
+    sched.set_timesteps(
+        num_inference_steps=4,
+        device=torch.device("cpu"),
+        custom_timesteps=custom,
+    )
+    torch.testing.assert_close(
+        sched.timesteps, torch.tensor(custom, dtype=torch.float32))
+
+
+def test_flow_map_scheduler_custom_timesteps_must_be_descending() -> None:
+    sched = _scheduler(shift=5.0)
+    with pytest.raises(ValueError, match="descending"):
+        sched.set_timesteps(
+            num_inference_steps=4,
+            device=torch.device("cpu"),
+            custom_timesteps=[100.0, 500.0, 900.0],
+        )
+
+
+def test_flow_map_scheduler_apply_shift_endpoints_invariant() -> None:
+    """apply_shift fixes the endpoints {0, 1} and produces non-trivial
+    motion in the interior for shift != 1."""
+    sched = _scheduler(shift=5.0)
+    t = torch.tensor([0.0, 0.5, 1.0])
+    shifted = sched.apply_shift(t)
+    torch.testing.assert_close(
+        shifted, torch.tensor([0.0, 5.0 / 6.0, 1.0]), rtol=1e-6, atol=1e-6)
+
+
+def test_flow_map_scheduler_apply_shift_shift_one_is_identity() -> None:
+    sched = _scheduler(shift=1.0)
+    t = torch.linspace(0.0, 1.0, 100)
+    torch.testing.assert_close(sched.apply_shift(t), t)
+
+
+def test_flow_map_scheduler_step_one_euler_iteration_matches_formula() -> None:
+    """One step: x_r = x_t - ((t - r) / N) * model_output."""
+    sched = _scheduler(shift=1.0)
+    sched.set_timesteps(num_inference_steps=4, device=torch.device("cpu"))
+
+    torch.manual_seed(0)
+    x_t = torch.randn(2, 4, 1, 8, 8)
+    v = torch.randn_like(x_t)
+    t = torch.tensor([750.0, 500.0])
+    r = torch.tensor([500.0, 250.0])
+
+    out = sched.step(v, sample=x_t, timestep=t, r_timestep=r)
+    expected = x_t - ((t - r) / 1000.0).view(-1, 1, 1, 1, 1) * v
+    torch.testing.assert_close(out, expected, rtol=1e-6, atol=1e-6)
+
+
+def test_flow_map_scheduler_get_train_weight_beta08_shape_and_renorm() -> None:
+    """beta08: t * sqrt(1-t), renormalized so sum equals num_train_timesteps.
+    The interior of the schedule must dominate the endpoints (monotone up
+    then monotone down)."""
+    sched = _scheduler()
+    t = torch.linspace(0.001, 0.999, 1000)
+    w = sched.get_train_weight(t, weight_type="beta08")
+    assert torch.allclose(w.sum(), torch.tensor(1000.0), rtol=1e-3)
+    assert torch.all(w >= 0.0)
+    # Endpoints smaller than the middle bump.
+    mid = len(w) // 2
+    assert w[0] < w[mid]
+    assert w[-1] < w[mid]
+
+
+def test_flow_map_scheduler_get_train_weight_uniform_is_constant_norm() -> None:
+    sched = _scheduler()
+    t = torch.linspace(0.0, 1.0, 1000)
+    w = sched.get_train_weight(t, weight_type="uniform")
+    assert torch.allclose(w.sum(), torch.tensor(1000.0), rtol=1e-3)
+    # All entries equal to 1.0 after renormalization.
+    torch.testing.assert_close(w, torch.ones_like(w), rtol=1e-6, atol=1e-6)
+
+
+def test_flow_map_scheduler_get_train_weight_accepts_absolute_units() -> None:
+    """When t is provided in [0, num_train_timesteps] the helper auto-
+    normalizes; result must match the [0, 1] call."""
+    sched = _scheduler()
+    t_norm = torch.linspace(0.001, 0.999, 1000)
+    t_abs = t_norm * 1000
+    w_norm = sched.get_train_weight(t_norm, weight_type="beta08")
+    w_abs = sched.get_train_weight(t_abs, weight_type="beta08")
+    torch.testing.assert_close(w_norm, w_abs, rtol=1e-5, atol=1e-5)
+
+
+def test_flow_map_scheduler_add_noise_matches_flow_matching_formula() -> None:
+    """Linear flow-matching: x_t = (1 - sigma) * x_0 + sigma * eps, where
+    sigma = t / num_train_timesteps."""
+    sched = _scheduler()
+    torch.manual_seed(0)
+    x0 = torch.randn(2, 4, 1, 4, 4)
+    eps = torch.randn_like(x0)
+    t = torch.tensor([250.0, 750.0])
+    out = sched.add_noise(x0, eps, t)
+    sigma = (t / 1000.0).view(-1, 1, 1, 1, 1)
+    expected = (1.0 - sigma) * x0 + sigma * eps
+    torch.testing.assert_close(out, expected, rtol=1e-6, atol=1e-6)
+
+
+# ---------------------------------------------------------------------------
+# Task 6: (t, r) per-batch sampling.
+# ---------------------------------------------------------------------------
+
+
+def test_sample_pair_timesteps_partitions_batch_correctly() -> None:
+    """For batch=8 with diffusion=0.5/consistency=0.25: 4 r=t, 2 r=0,
+    2 free entries."""
+    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
+        _sample_pair_timesteps, )
+
+    torch.manual_seed(42)
+    t, r, is_diffusion, is_consistency = _sample_pair_timesteps(
+        batch_size=8,
+        diffusion_ratio=0.5,
+        consistency_ratio=0.25,
+        device=torch.device("cpu"),
+        generator=None,
+    )
+    assert t.shape == (8,)
+    assert r.shape == (8,)
+    assert int(is_diffusion.sum()) == 4
+    assert int(is_consistency.sum()) == 2
+    # The masks must be disjoint.
+    assert not torch.any(is_diffusion & is_consistency)
+
+    diff_idx = torch.nonzero(is_diffusion).flatten()
+    torch.testing.assert_close(r[diff_idx], t[diff_idx])
+
+    cons_idx = torch.nonzero(is_consistency).flatten()
+    torch.testing.assert_close(r[cons_idx], torch.zeros(2))
+
+    free_idx = torch.nonzero(~(is_diffusion | is_consistency)).flatten()
+    assert torch.all(r[free_idx] <= t[free_idx])
+    assert torch.all(r[free_idx] >= 0.0)
+    assert torch.all(t[free_idx] <= 1.0)
+
+
+def test_sample_pair_timesteps_t_max_r_min_ordering() -> None:
+    """For the free fraction (ratios=0), t and r come from max/min of two
+    uniform draws so r <= t holds."""
+    torch.manual_seed(0)
+    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
+        _sample_pair_timesteps, )
+
+    for _ in range(50):
+        t, r, is_diff, is_cons = _sample_pair_timesteps(
+            batch_size=4,
+            diffusion_ratio=0.0,
+            consistency_ratio=0.0,
+            device=torch.device("cpu"),
+            generator=None,
+        )
+        assert torch.all(t >= r)
+        assert torch.all(r >= 0.0)
+        assert torch.all(t <= 1.0)
+        assert int(is_diff.sum()) == 0
+        assert int(is_cons.sum()) == 0
+
+
+def test_sample_pair_timesteps_rejects_ratios_summing_above_one() -> None:
+    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
+        _sample_pair_timesteps, )
+
+    with pytest.raises(ValueError, match="must be <= 1"):
+        _sample_pair_timesteps(
+            batch_size=8,
+            diffusion_ratio=0.7,
+            consistency_ratio=0.4,
+            device=torch.device("cpu"),
+            generator=None,
+        )
+
+
+def test_sample_pair_timesteps_rejects_negative_ratios() -> None:
+    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
+        _sample_pair_timesteps, )
+
+    with pytest.raises(ValueError, match="non-negative"):
+        _sample_pair_timesteps(
+            batch_size=8,
+            diffusion_ratio=-0.1,
+            consistency_ratio=0.25,
+            device=torch.device("cpu"),
+            generator=None,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Task 7: central-difference target.
+# ---------------------------------------------------------------------------
+
+
+class _StubStudent:
+    """Stand-in student for unit-testing the central-difference helper.
+
+    The "velocity prediction" is a closed-form function of x and t
+    (no actual neural network) so we can compare against the analytical
+    derivative.
+    """
+
+    def __init__(self, alpha: float = 0.3) -> None:
+        self.alpha = float(alpha)
+
+    def predict_velocity_with_r(
+        self,
+        noisy: torch.Tensor,
+        t: torch.Tensor,
+        r: torch.Tensor,
+        batch,
+        *,
+        conditional: bool = True,
+        attn_kind: str = "dense",
+        cfg_uncond=None,
+    ) -> torch.Tensor:
+        del batch, conditional, attn_kind, cfg_uncond, r
+        # f(x, t) = x + alpha * (t / 1000) → dF/dt = alpha / 1000.
+        view = [-1] + [1] * (noisy.ndim - 1)
+        return noisy + self.alpha * (t.view(*view).float() / 1000.0)
+
+
+def test_central_difference_dF_dt_linear_function() -> None:
+    """For f(x, t) = x + alpha * (t / N), the central-difference estimate
+    of dF/dt must be alpha / N (in absolute t-units that's alpha / N
+    velocity per t-unit, and our helper divides the difference by
+    2*delta -> exactly alpha / N)."""
+    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
+        _central_difference_dF_dt, )
+
+    student = _StubStudent(alpha=0.3)
+    x = torch.randn(2, 1, 4, 4, 4)
+    latents = torch.zeros_like(x)
+    noise = torch.zeros_like(x)  # v_pred = noise - latents = 0
+    t = torch.tensor([500.0, 250.0])
+    r = torch.tensor([100.0, 0.0])
+
+    dF = _central_difference_dF_dt(
+        student=student,
+        batch=None,
+        noisy=x,
+        latents=latents,
+        noise=noise,
+        t=t,
+        r=r,
+        delta=5.0,
+        num_train_timesteps=1000.0,
+    )
+
+    expected = torch.full_like(x, 0.3 / 1000.0)
+    torch.testing.assert_close(dF, expected, rtol=1e-5, atol=1e-5)
+
+
+def test_central_difference_dF_dt_guidance_scaling() -> None:
+    """With guidance_scale != 1, the result must be divided by guidance."""
+    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
+        _central_difference_dF_dt, )
+
+    student = _StubStudent(alpha=0.6)
+    x = torch.zeros(1, 1, 4, 4, 4)
+    latents = torch.zeros_like(x)
+    noise = torch.zeros_like(x)
+    t = torch.tensor([500.0])
+    r = torch.tensor([100.0])
+
+    dF_g1 = _central_difference_dF_dt(
+        student=student, batch=None, noisy=x, latents=latents, noise=noise,
+        t=t, r=r, delta=5.0, num_train_timesteps=1000.0, guidance_scale=1.0)
+    dF_g3 = _central_difference_dF_dt(
+        student=student, batch=None, noisy=x, latents=latents, noise=noise,
+        t=t, r=r, delta=5.0, num_train_timesteps=1000.0, guidance_scale=3.0)
+
+    torch.testing.assert_close(dF_g1 / 3.0, dF_g3, rtol=1e-5, atol=1e-5)
+
+
+def test_central_difference_dF_dt_rejects_zero_delta() -> None:
+    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
+        _central_difference_dF_dt, )
+
+    with pytest.raises(ValueError, match="delta must be positive"):
+        _central_difference_dF_dt(
+            student=_StubStudent(),
+            batch=None,
+            noisy=torch.zeros(1, 1, 4, 4, 4),
+            latents=torch.zeros(1, 1, 4, 4, 4),
+            noise=torch.zeros(1, 1, 4, 4, 4),
+            t=torch.tensor([500.0]),
+            r=torch.tensor([100.0]),
+            delta=0.0,
+            num_train_timesteps=1000.0,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Task 9: param_names_mapping handles AnyFlow checkpoints + is a no-op on plain
+# Wan checkpoints. We don't ship a separate remap_anyflow_keys helper because
+# the existing param_names_mapping regex mechanism does the job (delta_embedder
+# rename is a no-op when the source state dict doesn't contain those keys).
+# ---------------------------------------------------------------------------
+
+
+def test_param_names_mapping_includes_delta_embedder_rename() -> None:
+    """The Wan arch config's param_names_mapping must rename HF AnyFlow
+    delta_embedder weights into FastVideo's internal mlp.fc_in/fc_out layout."""
+    arch = WanVideoConfig().arch_config
+    mapping_keys = list(arch.param_names_mapping.keys())
+    assert any("delta_embedder" in k for k in mapping_keys), (
+        "WanVideoArchConfig.param_names_mapping must include delta_embedder "
+        "rename so HF AnyFlow checkpoints load without a separate adapter")
+
+
+def test_param_names_mapping_default_doesnt_break_plain_wan_keys() -> None:
+    """The new delta_embedder regex must not match any key in a plain
+    pretrained Wan2.1 checkpoint (those don't have delta_embedder)."""
+    import re
+
+    plain_wan_keys = [
+        "patch_embedding.weight",
+        "condition_embedder.time_embedder.linear_1.weight",
+        "condition_embedder.time_embedder.linear_2.weight",
+        "condition_embedder.time_proj.weight",
+        "condition_embedder.text_embedder.linear_1.weight",
+        "blocks.0.attn1.to_q.weight",
+        "blocks.0.ffn.net.0.proj.weight",
+    ]
+    arch = WanVideoConfig().arch_config
+    delta_regexes = [
+        k for k in arch.param_names_mapping if "delta_embedder" in k
+    ]
+    assert delta_regexes, "expected at least one delta_embedder regex"
+
+    for plain in plain_wan_keys:
+        for rx in delta_regexes:
+            assert re.match(rx, plain) is None, (
+                f"delta_embedder regex {rx!r} unexpectedly matched plain "
+                f"Wan key {plain!r}")

fastvideo/tests/training/distill/test_anyflow_smoke.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GPU smoke test for AnyFlow pretrain + on-policy.
+
+Mirrors ``test_distill_dmd.py`` — runs the new YAML-driven training
+entrypoint via ``torchrun`` for two iterations to verify end-to-end
+wiring (model load, optimizer build, forward, backward, step, save).
+
+CPU-only environments are skipped; the test is intended to fire on the
+Buildkite ``/test distillation`` lane and on local boxes with at least
+2 H100/H200 GPUs.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+PRETRAIN_YAML = (
+    REPO_ROOT
+    / "examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml")
+ONPOLICY_YAML = (
+    REPO_ROOT
+    / "examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml")
+
+NUM_NODES = "1"
+NUM_GPUS_PER_NODE = "2"
+
+
+def _have_enough_gpus() -> bool:
+    """Return True iff at least 2 CUDA devices are visible. The smoke test
+    needs HSDP/FSDP with a non-trivial world size; single-GPU bring-up
+    races against the new framework's distributed barriers."""
+    try:
+        import torch
+    except Exception:
+        return False
+    if not torch.cuda.is_available():
+        return False
+    return torch.cuda.device_count() >= 2
+
+
+pytestmark = pytest.mark.skipif(
+    not _have_enough_gpus(),
+    reason="AnyFlow smoke test requires >= 2 CUDA devices")
+
+
+def _run_torchrun(config_path: Path, *, output_dir: Path) -> None:
+    if not config_path.exists():
+        pytest.fail(f"YAML config missing: {config_path}")
+
+    env = os.environ.copy()
+    env.setdefault("MASTER_ADDR", "127.0.0.1")
+    env.setdefault("MASTER_PORT", "29551")
+    env.setdefault("WANDB_MODE", "offline")
+    env.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+    cmd = [
+        sys.executable,
+        "-m",
+        "torch.distributed.run",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE,
+        "--master_port", env["MASTER_PORT"],
+        "-m", "fastvideo.train.entrypoint.train",
+        "--config", str(config_path),
+        "--training.loop.max_train_steps", "2",
+        "--training.checkpoint.output_dir", str(output_dir),
+        "--training.distributed.num_gpus", NUM_GPUS_PER_NODE,
+        "--training.distributed.hsdp_shard_dim", NUM_GPUS_PER_NODE,
+        "--training.data.train_batch_size", "1",
+    ]
+    process = subprocess.run(cmd, capture_output=True, text=True, env=env)
+    if process.stdout:
+        print("STDOUT:", process.stdout)
+    if process.stderr:
+        print("STDERR:", process.stderr)
+    if process.returncode != 0:
+        raise subprocess.CalledProcessError(
+            process.returncode, cmd, process.stdout, process.stderr)
+
+
+def test_anyflow_pretrain_smoke(tmp_path: Path) -> None:
+    """Two-iteration pretrain — exercises (t, r) sampling, central-difference
+    target, scale balance, optimizer step, and checkpoint save path."""
+    _run_torchrun(PRETRAIN_YAML, output_dir=tmp_path / "pretrain")
+
+
+def test_anyflow_onpolicy_smoke(tmp_path: Path) -> None:
+    """Two-iteration on-policy DMD — exercises the multi-step Euler-flow
+    rollout, grad-step broadcast, DMD2 alternating updates."""
+    # Override init_from to the public Wan2.1-T2V-1.3B-Diffusers checkpoint
+    # for the smoke run; param_names_mapping handles the (non-existent)
+    # delta_embedder rename as a no-op.
+    env = os.environ.copy()
+    env.setdefault("MASTER_ADDR", "127.0.0.1")
+    env.setdefault("MASTER_PORT", "29552")
+    env.setdefault("WANDB_MODE", "offline")
+    env.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+    cmd = [
+        sys.executable,
+        "-m",
+        "torch.distributed.run",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE,
+        "--master_port", env["MASTER_PORT"],
+        "-m", "fastvideo.train.entrypoint.train",
+        "--config", str(ONPOLICY_YAML),
+        "--training.loop.max_train_steps", "2",
+        "--training.checkpoint.output_dir", str(tmp_path / "onpolicy"),
+        "--training.distributed.num_gpus", NUM_GPUS_PER_NODE,
+        "--training.distributed.hsdp_shard_dim", NUM_GPUS_PER_NODE,
+        "--training.data.train_batch_size", "1",
+        "--models.student.init_from", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        "--method.student_sample_steps", "2",
+    ]
+    process = subprocess.run(cmd, capture_output=True, text=True, env=env)
+    if process.stdout:
+        print("STDOUT:", process.stdout)
+    if process.stderr:
+        print("STDERR:", process.stderr)
+    if process.returncode != 0:
+        raise subprocess.CalledProcessError(
+            process.returncode, cmd, process.stdout, process.stderr)

fastvideo/train/methods/distribution_matching/__init__.py

@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
+from fastvideo.train.methods.distribution_matching.anyflow import AnyFlowMethod
+from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
+    AnyFlowPretrainMethod, )
 from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method
 from fastvideo.train.methods.distribution_matching.self_forcing import (
     SelfForcingMethod, )
@@ -7,6 +10,8 @@ from fastvideo.train.methods.distribution_matching.streaming_long_tuning import
     StreamingLongTuningMethod, )
 
 __all__ = [
+    "AnyFlowMethod",
+    "AnyFlowPretrainMethod",
     "DMD2Method",
     "SelfForcingMethod",
     "StreamingLongTuningMethod",

fastvideo/train/methods/distribution_matching/anyflow.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+"""AnyFlow on-policy distillation method.
+
+Stage 2 of the AnyFlow two-stage recipe. Continues from a pretrained
+flow-map student and refines it via distribution-matching distillation
+(DMD2) where the student is rolled out for ``student_sample_steps``
+Euler-flow steps from pure noise. One randomly-chosen step in the
+rollout is gradient-enabled (and broadcast across ranks so every worker
+agrees on which step to gradient-enable); the rest run under
+``torch.no_grad``.
+
+Inherits ``DMD2Method`` for the alternating student / critic update
+machinery and the existing DMD VSD-with-fake-score loss. Overrides
+``_student_rollout`` to drive the multi-step Euler-flow rollout with
+``r = t_next`` (mean-velocity sampling — matches the AnyFlow paper's
+``WanAnyFlowPipeline.training_rollout`` with ``use_mean_velocity=True``).
+
+Reference: ``pipeline_wan_anyflow.py::training_rollout`` in
+NVlabs/AnyFlow at commit ``549236a``.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+import torch
+import torch.distributed as dist
+
+from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method
+from fastvideo.train.utils.config import (
+    get_optional_float,
+    get_optional_int,
+)
+
+
+class AnyFlowMethod(DMD2Method):
+    """AnyFlow on-policy distillation (multi-step rollout)."""
+
+    def __init__(
+        self,
+        *,
+        cfg: Any,
+        role_models: dict[str, Any],
+    ) -> None:
+        super().__init__(cfg=cfg, role_models=role_models)
+        mcfg = self.method_config
+
+        student_sample_steps = get_optional_int(mcfg, "student_sample_steps", where="method.student_sample_steps")
+        if student_sample_steps is None:
+            student_sample_steps = 4
+        if int(student_sample_steps) <= 0:
+            raise ValueError("method.student_sample_steps must be positive, "
+                             f"got {student_sample_steps}")
+        self._student_sample_steps = int(student_sample_steps)
+
+        use_mean_velocity_raw = mcfg.get("use_mean_velocity", True)
+        if not isinstance(use_mean_velocity_raw, bool):
+            raise ValueError("method.use_mean_velocity must be a bool, "
+                             f"got {type(use_mean_velocity_raw).__name__}")
+        self._use_mean_velocity = bool(use_mean_velocity_raw)
+
+        # Optional pinned rollout schedule (descending, absolute t-units).
+        # Falls back to dmd_denoising_steps when absent.
+        raw_t_list = mcfg.get("t_list_override", None)
+        if raw_t_list is None:
+            self._t_list_override: list[float] | None = None
+        else:
+            if not isinstance(raw_t_list, list) or not raw_t_list:
+                raise ValueError("method.t_list_override must be a non-empty list of "
+                                 f"floats when set, got {raw_t_list!r}")
+            t_list = [float(x) for x in raw_t_list]
+            for i in range(len(t_list) - 1):
+                if t_list[i] < t_list[i + 1]:
+                    raise ValueError("method.t_list_override must be descending, "
+                                     f"got {t_list!r}")
+            self._t_list_override = t_list
+
+        # Scoring conditioning: AnyFlow scores against r=0 for the DMD branch.
+        score_r_raw = mcfg.get("dmd_score_r_value", 0.0)
+        try:
+            self._dmd_score_r = float(score_r_raw)
+        except (TypeError, ValueError) as exc:
+            raise ValueError("method.dmd_score_r_value must be numeric, "
+                             f"got {score_r_raw!r}") from exc
+
+        # Optional teacher guidance scale for the DMD loss (carry over from
+        # DMD2Method's behavior; default 1.0).
+        guidance = get_optional_float(mcfg, "real_score_guidance_scale", where="method.real_score_guidance_scale")
+        self._real_score_guidance = float(guidance) if guidance is not None else 1.0
+
+    # ------------------------------------------------------------------
+    # Rollout schedule
+
+    def _get_rollout_schedule(self, *, device: torch.device) -> torch.Tensor:
+        """Build the descending timestep schedule used by the on-policy
+        rollout. Length is ``num_steps + 1`` so ``num_steps`` Euler steps
+        consume the full range.
+
+        Order of precedence:
+          1. ``method.t_list_override`` — used verbatim (absolute units).
+          2. ``method.dmd_denoising_steps`` (inherited from DMD2) appended
+             with a final 0 boundary if the last entry isn't already 0.
+        """
+        if self._t_list_override is not None:
+            return torch.tensor(self._t_list_override, device=device, dtype=torch.float32)
+
+        steps = self._get_denoising_step_list(device).to(dtype=torch.float32)
+        if float(steps[-1].item()) != 0.0:
+            zero = torch.zeros(1, device=device, dtype=torch.float32)
+            steps = torch.cat([steps, zero], dim=0)
+        return steps
+
+    def _broadcast_grad_step_index(
+        self,
+        num_steps: int,
+        *,
+        device: torch.device,
+    ) -> int:
+        """Pick the rollout step that gets gradient enabled. In distributed
+        runs the choice is broadcast from rank 0 so every worker agrees."""
+        if num_steps <= 0:
+            raise ValueError("num_steps must be positive")
+        if dist.is_initialized() and dist.get_rank() != 0:
+            idx_tensor = torch.empty((1, ), dtype=torch.long, device=device)
+        else:
+            idx_tensor = torch.randint(0,
+                                       num_steps, (1, ),
+                                       device=device,
+                                       dtype=torch.long,
+                                       generator=self.cuda_generator)
+        if dist.is_initialized():
+            dist.broadcast(idx_tensor, src=0)
+        return int(idx_tensor.item())
+
+    # ------------------------------------------------------------------
+    # Rollout
+
+    def _student_rollout(
+        self,
+        batch: Any,
+        *,
+        with_grad: bool,
+    ) -> torch.Tensor:
+        """Multi-step Euler-flow rollout from pure noise.
+
+        Returns the predicted clean latent ``x_0`` after the chosen
+        gradient step (or the final ``x`` after the last step if
+        ``with_grad`` is False — used by the critic path).
+        """
+        latents = batch.latents
+        if latents is None or latents.ndim != 5:
+            raise RuntimeError("AnyFlow on-policy rollout requires TrainingBatch.latents "
+                               "of shape [B, T, C, H, W] for shape templating")
+        device = latents.device
+        dtype = latents.dtype
+
+        schedule = self._get_rollout_schedule(device=device)
+        num_entries = int(schedule.numel())
+        num_steps = num_entries - 1
+        if num_steps <= 0:
+            raise RuntimeError("rollout schedule must have at least two entries "
+                               f"(got {num_entries})")
+        if num_steps > self._student_sample_steps:
+            # Trim to the configured cap, keeping the last (=0) boundary.
+            schedule = torch.cat([schedule[:self._student_sample_steps], schedule[-1:]], dim=0)
+            num_steps = self._student_sample_steps
+
+        grad_step = self._broadcast_grad_step_index(num_steps, device=device) if with_grad else -1
+
+        attn_kind: Literal["dense", "vsa"] = "vsa"
+        n_train = float(self.student.num_train_timesteps)
+
+        x = torch.randn(latents.shape, device=device, dtype=dtype, generator=self.cuda_generator)
+        last_pred_x0: torch.Tensor | None = None
+        batch_size = int(latents.shape[0])
+
+        for i in range(num_steps):
+            t_cur = schedule[i].expand(batch_size)
+            t_next = schedule[i + 1].expand(batch_size)
+            r = t_next if self._use_mean_velocity else t_cur
+
+            enable_grad = bool(with_grad) and (i == grad_step)
+            with torch.set_grad_enabled(enable_grad):
+                v = self.student.predict_velocity_with_r(
+                    x,
+                    t_cur,
+                    r,
+                    batch,
+                    conditional=True,
+                    cfg_uncond=self._cfg_uncond,
+                    attn_kind=attn_kind,
+                )
+            # Euler step in absolute units: x ← x - ((t_cur - t_next) / N) * v.
+            view = [-1] + [1] * (x.ndim - 1)
+            dt = ((t_cur - t_next) / n_train).view(*view)
+            x = x - dt * v
+
+            if enable_grad:
+                # We treat the rollout output (post-step) as a predicted
+                # clean latent — AnyFlow's last Euler step lands at t=0.
+                last_pred_x0 = x
+
+        if last_pred_x0 is None:
+            # No gradient step taken (with_grad=False path).
+            last_pred_x0 = x
+
+        if hasattr(batch, "dmd_latent_vis_dict"):
+            batch.dmd_latent_vis_dict["generator_timestep"] = (schedule[-1].detach().clone())
+        return last_pred_x0

fastvideo/train/methods/distribution_matching/anyflow_pretrain.py

@@ -0,0 +1,404 @@
+# SPDX-License-Identifier: Apache-2.0
+"""AnyFlow pretrain (flow-map central-difference) training method.
+
+Stage 1 of the AnyFlow two-stage recipe. Trains a single student network
+``u_θ(x_t, t, r)`` to predict the average velocity from time ``t`` back
+to time ``r`` via the central-difference target
+
+    target = (eps - x_0) - ((t - r) / N) * dF/dt
+
+where ``N = num_train_timesteps`` and ``dF/dt`` is estimated from the
+student's own forward at ``(t ± δ, r)`` (with one-sided fallback near the
+schedule endpoints).
+
+Per-batch ``(t, r)`` sampling follows the AnyFlow paper:
+
+- ``diffusion_ratio`` fraction: ``r = t`` (recovers plain flow matching).
+- ``consistency_ratio`` fraction: ``r = 0`` (consistency to clean data).
+- Remaining fraction: ``(t, r) = (max, min)`` of two independent uniform
+  draws (full reconstruction range).
+
+Reference: ``trainer_wan_anyflow_pretrain.py`` in NVlabs/AnyFlow at
+commit ``549236a``.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from collections.abc import Sequence
+
+import torch
+
+from fastvideo.train.methods.base import LogScalar, TrainingMethod
+from fastvideo.train.models.base import ModelBase
+from fastvideo.train.utils.config import (
+    get_optional_float,
+    get_optional_int,
+)
+from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler
+
+
+def _sample_pair_timesteps(
+    *,
+    batch_size: int,
+    diffusion_ratio: float,
+    consistency_ratio: float,
+    device: torch.device,
+    generator: torch.Generator | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Sample ``(t, r)`` per the AnyFlow paper.
+
+    Two uniform draws ``u1, u2 ∈ [0, 1]`` per sample, then
+    ``t = max(u1, u2)`` and ``r = min(u1, u2)``. After the base sample,
+    the first ``diffusion_ratio * B`` entries get ``r = t`` (diffusion
+    branch, plain flow matching), and the next ``consistency_ratio * B``
+    entries get ``r = 0`` (consistency branch).
+
+    Returns
+    -------
+    t, r, is_diffusion, is_consistency
+        Each tensor has shape ``(batch_size,)``. ``t`` and ``r`` are in
+        ``[0, 1]`` (i.e. *not* yet shifted and *not* yet in absolute
+        train-timestep units). ``is_diffusion`` and ``is_consistency``
+        are bool masks that partition a subset of the batch — entries
+        outside both masks are the "free" reconstruction fraction.
+    """
+    if batch_size <= 0:
+        raise ValueError(f"batch_size must be positive, got {batch_size}")
+    if diffusion_ratio < 0.0 or consistency_ratio < 0.0:
+        raise ValueError("diffusion_ratio and consistency_ratio must be non-negative")
+    if diffusion_ratio + consistency_ratio > 1.0:
+        raise ValueError("diffusion_ratio + consistency_ratio must be <= 1, "
+                         f"got {diffusion_ratio} + {consistency_ratio}")
+
+    u1 = torch.rand(batch_size, device=device, generator=generator)
+    u2 = torch.rand(batch_size, device=device, generator=generator)
+    t = torch.maximum(u1, u2)
+    r = torch.minimum(u1, u2)
+
+    n_diff = int(diffusion_ratio * batch_size)
+    n_cons = int(consistency_ratio * batch_size)
+    is_diffusion = torch.zeros(batch_size, dtype=torch.bool, device=device)
+    is_consistency = torch.zeros(batch_size, dtype=torch.bool, device=device)
+    is_diffusion[:n_diff] = True
+    is_consistency[n_diff:n_diff + n_cons] = True
+
+    # Override per the AnyFlow paper:
+    # - diffusion entries: r = t (plain flow matching)
+    # - consistency entries: r = 0 (consistency to clean data)
+    r = torch.where(is_diffusion, t, r)
+    r = torch.where(is_consistency, torch.zeros_like(r), r)
+    return t, r, is_diffusion, is_consistency
+
+
+@torch.no_grad()
+def _central_difference_dF_dt(
+    *,
+    student: Any,
+    batch: Any,
+    noisy: torch.Tensor,
+    latents: torch.Tensor,
+    noise: torch.Tensor,
+    t: torch.Tensor,
+    r: torch.Tensor,
+    delta: float,
+    num_train_timesteps: float,
+    attn_kind: str = "dense",
+    guidance_scale: float = 1.0,
+) -> torch.Tensor:
+    """Estimate ``dF/dt`` for the AnyFlow central-difference target.
+
+    Computes a symmetric finite difference of the velocity prediction in
+    *absolute train-timestep units*:
+
+        dF/dt ≈ [u_θ(x_{t+δ}, t+δ, r) - u_θ(x_{t-δ}, t-δ, r)] / (2 * δ * guidance)
+
+    The sample is also moved along the flow trajectory by the same
+    finite step (``v_pred * (δ / N)``) to match AnyFlow's reference
+    formulation in ``trainer_wan_anyflow_pretrain.py::compute_central_difference``.
+    Wrapped in ``torch.no_grad`` so the two extra forwards never enter
+    the backward graph.
+    """
+    if delta <= 0.0:
+        raise ValueError(f"delta must be positive, got {delta}")
+    if guidance_scale <= 0.0:
+        raise ValueError(f"guidance_scale must be positive, got {guidance_scale}")
+
+    v_pred = noise - latents  # ground-truth flow velocity
+    delta_x = delta / float(num_train_timesteps)
+
+    t_plus = t + delta
+    noisy_plus = noisy + v_pred * delta_x
+    f_plus = student.predict_velocity_with_r(noisy_plus, t_plus, r, batch, conditional=True, attn_kind=attn_kind)
+
+    t_minus = t - delta
+    noisy_minus = noisy - v_pred * delta_x
+    f_minus = student.predict_velocity_with_r(noisy_minus, t_minus, r, batch, conditional=True, attn_kind=attn_kind)
+
+    return (f_plus - f_minus) / (2.0 * delta * guidance_scale)
+
+
+class AnyFlowPretrainMethod(TrainingMethod):
+    """AnyFlow flow-map pretrain method.
+
+    Single-student training; no teacher or critic. The student must
+    implement ``predict_velocity_with_r(noisy, t, r, batch, ...)`` —
+    typically a ``WanModel`` with ``r_embedder=True`` in its arch config.
+    """
+
+    def __init__(
+        self,
+        *,
+        cfg: Any,
+        role_models: dict[str, ModelBase],
+    ) -> None:
+        super().__init__(cfg=cfg, role_models=role_models)
+        if "student" not in role_models:
+            raise ValueError("AnyFlowPretrainMethod requires role 'student'")
+        if not self.student._trainable:
+            raise ValueError("AnyFlowPretrainMethod requires student to be trainable")
+
+        mcfg = self.method_config
+        self._diffusion_ratio = float(
+            get_optional_float(mcfg, "diffusion_ratio", where="method.diffusion_ratio") or 0.5)
+        self._consistency_ratio = float(
+            get_optional_float(mcfg, "consistency_ratio", where="method.consistency_ratio") or 0.25)
+        if self._diffusion_ratio + self._consistency_ratio > 1.0:
+            raise ValueError("method.diffusion_ratio + method.consistency_ratio must "
+                             f"be <= 1, got {self._diffusion_ratio} + "
+                             f"{self._consistency_ratio}")
+
+        # δ: finite-difference step in absolute train-timestep units.
+        epsilon = get_optional_int(mcfg, "epsilon", where="method.epsilon")
+        self._fd_epsilon = float(epsilon) if epsilon is not None else 5.0
+
+        # Loss weighting scheme (uniform / gaussian / beta08).
+        raw_weight_type = mcfg.get("weight_type", "beta08")
+        if not isinstance(raw_weight_type, str):
+            raise ValueError("method.weight_type must be a string, got "
+                             f"{type(raw_weight_type).__name__}")
+        weight_type = raw_weight_type.strip().lower()
+        if weight_type not in {"uniform", "gaussian", "beta08"}:
+            raise ValueError("method.weight_type must be one of "
+                             "{uniform, gaussian, beta08}, "
+                             f"got {raw_weight_type!r}")
+        self._weight_type = weight_type
+
+        # Guidance fused into the training target (default 1.0 = unused).
+        fg = get_optional_float(mcfg, "fuse_guidance_scale", where="method.fuse_guidance_scale")
+        self._fuse_guidance_scale = float(fg) if fg is not None else 1.0
+        if self._fuse_guidance_scale <= 0.0:
+            raise ValueError("method.fuse_guidance_scale must be positive, "
+                             f"got {self._fuse_guidance_scale}")
+
+        # Flow-map scheduler — uses pipeline_config.flow_shift if present
+        # and falls back to method.shift (and finally 1.0).
+        shift = float(getattr(self.training_config.pipeline_config, "flow_shift", 0.0) or 0.0)
+        if shift <= 0.0:
+            shift_override = get_optional_float(mcfg, "shift", where="method.shift")
+            shift = float(shift_override) if shift_override is not None else 1.0
+        self._shift = shift
+
+        # Lazy-imported to avoid circular imports on package load.
+        from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import (
+            FlowMapEulerDiscreteScheduler, )
+        self._flow_map_scheduler = FlowMapEulerDiscreteScheduler(
+            num_train_timesteps=int(self.student.num_train_timesteps),
+            shift=self._shift,
+        )
+
+        self.student.init_preprocessors(self.training_config)
+        self._init_optimizer_and_scheduler()
+
+    @property
+    def _optimizer_dict(self) -> dict[str, torch.optim.Optimizer]:
+        return {"student": self._student_optimizer}
+
+    @property
+    def _lr_scheduler_dict(self) -> dict[str, Any]:
+        return {"student": self._student_lr_scheduler}
+
+    def get_optimizers(
+        self,
+        iteration: int,
+    ) -> Sequence[torch.optim.Optimizer]:
+        del iteration
+        return [self._student_optimizer]
+
+    def get_lr_schedulers(self, iteration: int) -> Sequence[Any]:
+        del iteration
+        return [self._student_lr_scheduler]
+
+    def single_train_step(
+        self,
+        batch: dict[str, Any],
+        iteration: int,
+    ) -> tuple[
+            dict[str, torch.Tensor],
+            dict[str, Any],
+            dict[str, LogScalar],
+    ]:
+        del iteration  # AnyFlow pretrain has no iteration-dependent dispatch.
+
+        training_batch = self.student.prepare_batch(
+            batch,
+            generator=self.cuda_generator,
+            latents_source="data",
+        )
+        latents = training_batch.latents  # [B, T, C, H, W] (post-permute in prepare_batch).
+        if latents is None or latents.ndim != 5:
+            raise RuntimeError("AnyFlow pretrain expects TrainingBatch.latents of shape "
+                               "[B, T, C, H, W] after prepare_batch; got "
+                               f"{None if latents is None else tuple(latents.shape)}")
+        device = latents.device
+        dtype = latents.dtype
+        batch_size = int(latents.shape[0])
+
+        # AnyFlow (t, r) sampling — overrides the timestep drawn by
+        # WanModel._sample_timesteps inside prepare_batch.
+        t_norm, r_norm, is_diffusion, is_consistency = _sample_pair_timesteps(
+            batch_size=batch_size,
+            diffusion_ratio=self._diffusion_ratio,
+            consistency_ratio=self._consistency_ratio,
+            device=device,
+            generator=self.cuda_generator,
+        )
+
+        sched = self._flow_map_scheduler
+        n_train = float(self.student.num_train_timesteps)
+        t = (sched.apply_shift(t_norm) * n_train).to(device=device, dtype=dtype)
+        r = (sched.apply_shift(r_norm) * n_train).to(device=device, dtype=dtype)
+
+        # Fresh noise drawn from the method's RNG; ignore the noise that
+        # prepare_batch attached (it pairs with the discarded timestep).
+        noise = torch.randn(
+            latents.shape,
+            device=device,
+            dtype=dtype,
+            generator=self.cuda_generator,
+        )
+        noisy = sched.add_noise(latents, noise, t)
+
+        # Keep training_batch coherent with the new (t, noisy): downstream
+        # forward_context uses these fields.
+        training_batch.timesteps = t
+        training_batch.noise = noise
+        training_batch.noisy_model_input = noisy.permute(0, 2, 1, 3, 4)
+
+        # Student velocity prediction at (t, r).
+        noise_pred = self.student.predict_velocity_with_r(
+            noisy,
+            t,
+            r,
+            training_batch,
+            conditional=True,
+            attn_kind="dense",
+        )
+
+        # Optional guidance distillation — fuse CFG into the training target so
+        # the resulting checkpoint can be sampled at guidance_scale=1.0.
+        if self._fuse_guidance_scale != 1.0:
+            with torch.no_grad():
+                noise_pred_uncond = self.student.predict_velocity_with_r(
+                    noisy,
+                    t,
+                    r,
+                    training_batch,
+                    conditional=False,
+                    attn_kind="dense",
+                )
+            g = float(self._fuse_guidance_scale)
+            noise_pred = (noise_pred - (1.0 - g) * noise_pred_uncond) / g
+
+        dF_dt = _central_difference_dF_dt(
+            student=self.student,
+            batch=training_batch,
+            noisy=noisy,
+            latents=latents,
+            noise=noise,
+            t=t,
+            r=r,
+            delta=self._fd_epsilon,
+            num_train_timesteps=n_train,
+            attn_kind="dense",
+            guidance_scale=self._fuse_guidance_scale,
+        )
+
+        # AnyFlow target: target = (eps - x_0) - (t - r) * dF/dt
+        # dF/dt is in (velocity per absolute t-unit); (t - r) is in absolute units;
+        # the product cancels back to velocity units, matching noise_pred.
+        view = [batch_size] + [1] * (latents.ndim - 1)
+        target = (noise - latents) - (t - r).view(*view) * dF_dt
+
+        # Per-sample squared error, then per-timestep weight, then scale-balance
+        # so the non-diffusion branches stay on the same magnitude as the
+        # diffusion branch (matches AnyFlow's stop-grad rescaling).
+        per_sample = torch.mean(
+            ((noise_pred.float() - target.float())**2).reshape(batch_size, -1),
+            dim=-1,
+        )
+        weight = sched.get_train_weight(t, weight_type=self._weight_type)
+        per_sample = per_sample * weight
+
+        with torch.no_grad():
+            diff_mask = is_diffusion
+            diff_mean = per_sample[diff_mask].mean() if diff_mask.any() else per_sample.mean()
+            non_diff_mask = ~diff_mask
+            if non_diff_mask.any():
+                scale = diff_mean / (per_sample[non_diff_mask] + 1e-5)
+            else:
+                scale = torch.tensor(1.0, device=device, dtype=per_sample.dtype)
+        non_diff_idx = torch.nonzero(non_diff_mask, as_tuple=False).flatten()
+        if non_diff_idx.numel() > 0:
+            per_sample = per_sample.clone()
+            per_sample[non_diff_idx] = per_sample[non_diff_idx] * scale
+
+        total_loss = per_sample.mean()
+
+        loss_map = {"total_loss": total_loss}
+        metrics: dict[str, LogScalar] = {
+            "diffusion_fraction": float(is_diffusion.float().mean()),
+            "consistency_fraction": float(is_consistency.float().mean()),
+            "scale_weight_mean": float(scale.mean()) if isinstance(scale, torch.Tensor) else float(scale),
+        }
+        outputs = {
+            "student_ctx": (
+                training_batch.timesteps,
+                training_batch.attn_metadata,
+            ),
+        }
+        return loss_map, outputs, metrics
+
+    def backward(
+        self,
+        loss_map: dict[str, torch.Tensor],
+        outputs: dict[str, Any],
+        *,
+        grad_accum_rounds: int = 1,
+    ) -> None:
+        """Route the loss backward through the student's forward_context
+        so attn metadata stays attached during gradient computation."""
+        student_ctx = outputs.get("student_ctx")
+        if student_ctx is None:
+            super().backward(loss_map, outputs, grad_accum_rounds=grad_accum_rounds)
+            return
+        self.student.backward(
+            loss_map["total_loss"],
+            student_ctx,
+            grad_accum_rounds=grad_accum_rounds,
+        )
+
+    def _init_optimizer_and_scheduler(self) -> None:
+        tc = self.training_config
+        params = [p for p in self.student.transformer.parameters() if p.requires_grad]
+        (
+            self._student_optimizer,
+            self._student_lr_scheduler,
+        ) = build_optimizer_and_scheduler(
+            params=params,
+            optimizer_config=tc.optimizer,
+            loop_config=tc.loop,
+            learning_rate=float(tc.optimizer.learning_rate),
+            betas=tc.optimizer.betas,
+            scheduler_name=str(tc.optimizer.lr_scheduler),
+        )

fastvideo/train/models/wan/wan.py

@@ -358,6 +358,53 @@ class WanModel(ModelBase):
             pred_noise = transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
         return pred_noise
 
+    def predict_velocity_with_r(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        r_timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+    ) -> torch.Tensor:
+        """AnyFlow forward: predict average velocity from ``t`` back to ``r``.
+
+        Same plumbing as :meth:`predict_noise` but injects ``r_timestep``
+        into the transformer kwargs. The transformer must have been
+        constructed with an arch config that sets ``r_embedder=True`` for
+        the dual-timestep branch to be active — otherwise ``r_timestep``
+        is silently ignored by the embedder and the forward reduces to
+        the single-timestep path.
+        """
+        device_type = self.device.type
+        dtype = noisy_latents.dtype
+        if conditional:
+            text_dict = batch.conditional_dict
+            if text_dict is None:
+                raise RuntimeError("Missing conditional_dict in "
+                                   "TrainingBatch")
+        else:
+            text_dict = self._get_uncond_text_dict(batch, cfg_uncond=cfg_uncond)
+
+        if attn_kind == "dense":
+            attn_metadata = batch.attn_metadata
+        elif attn_kind == "vsa":
+            attn_metadata = batch.attn_metadata_vsa
+        else:
+            raise ValueError(f"Unknown attn_kind: {attn_kind!r}")
+
+        with torch.autocast(device_type, dtype=dtype), set_forward_context(
+                current_timestep=batch.timesteps,
+                attn_metadata=attn_metadata,
+        ):
+            input_kwargs = (self._build_distill_input_kwargs(noisy_latents, timestep, text_dict))
+            input_kwargs["r_timestep"] = r_timestep
+            transformer = self._get_transformer(timestep)
+            pred_velocity = transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
+        return pred_velocity
+
     def backward(
         self,
         loss: torch.Tensor,

mkdocs.yml

@@ -182,6 +182,7 @@ nav:
   - Distillation:
     - Data Preprocessing: distillation/data_preprocess.md
     - DMD: distillation/dmd.md
+    - AnyFlow: distillation/anyflow.md
   - Attention:
     - Overview: attention/index.md
     - Video Sparse Attention: attention/vsa/index.md

scripts/demo_anyflow_14b.py

@@ -0,0 +1,282 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""FastVideo-side AnyFlow 14B T2V demo at NFE=4 and NFE=50.
+
+Loads ``nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers`` into FastVideo's
+``WanTransformer3DModel`` (with the ``param_names_mapping`` regex
+handling the ``delta_embedder`` rename), runs the
+``FlowMapEulerDiscreteScheduler`` for the requested NFE schedule, and
+saves the decoded video as MP4. Matches the prompt / shift / guidance
+recipe used by the parallel FastGen demo so the videos are directly
+comparable.
+
+Memory tactics (single H200, 141 GB HBM):
+- Encode prompts with UMT5, free the encoder.
+- Build the FastVideo Wan-14B transformer, load AnyFlow safetensor
+  shards via param_names_mapping translation.
+- Sample at both NFEs without re-loading the transformer.
+- Free the transformer, then load the Wan VAE with tiling for decode.
+
+Configure local checkout paths via env vars (defaults assume a sibling
+layout next to this repo):
+
+  ANYFLOW_LOCAL     — path to ``nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers``
+                      (default ``./anyflow-14b``)
+  ANYFLOW_DEMO_OUT  — output directory for the rendered MP4s
+                      (default ``./demo_videos``)
+
+Run via::
+
+    PYTHONPATH=$PWD python scripts/demo_anyflow_14b.py
+"""
+
+from __future__ import annotations
+
+import gc
+import os
+import re
+import sys
+import time
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+
+
+ANYFLOW_LOCAL = Path(os.environ.get("ANYFLOW_LOCAL", "./anyflow-14b")).expanduser()
+OUT_DIR = Path(os.environ.get("ANYFLOW_DEMO_OUT", "./demo_videos")).expanduser()
+OUT_DIR.mkdir(parents=True, exist_ok=True)
+
+SEED = 0
+DEVICE = torch.device("cuda")
+DTYPE = torch.bfloat16
+NUM_FRAMES = 81
+HEIGHT, WIDTH = 480, 832
+PROMPT = (
+    "CG game concept digital art, a majestic elephant with a vibrant tusk and sleek fur "
+    "running swiftly towards a herd of its kind. The elephant has a calm yet determined "
+    "expression, with its ears flapping slightly as it moves at high speed. The herd consists "
+    "of several other elephants of various ages and sizes, all moving in unison. The landscape "
+    "is vast savanna with rolling hills, tall grasses, and scattered acacia trees. The sun "
+    "sets behind the horizon, casting a warm golden glow over the scene. Low-angle view, focus "
+    "on the elephant as it accelerates towards the herd."
+)
+NEG_PROMPT = "blurry, low quality, distorted"
+# The published nvidia/AnyFlow-* checkpoints are on-policy distilled with
+# fuse_guidance_scale=3.0 baked into the weights — inference uses 1.0
+# (single conditional forward, no CFG; matches AnyFlow's official demo.py).
+GUIDANCE = 1.0
+
+
+def banner(msg: str) -> None:
+    print("\n" + "=" * 80)
+    print(msg)
+    print("=" * 80)
+
+
+def free_gpu() -> None:
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.synchronize()
+    used = torch.cuda.memory_allocated() / 1e9
+    print(f"  [mem] allocated {used:.1f} GB after free")
+
+
+def init_single_rank() -> None:
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29571")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    from fastvideo.distributed import (
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+    init_distributed_environment(world_size=1, rank=0, local_rank=0, backend="nccl")
+    initialize_model_parallel(
+        tensor_model_parallel_size=1,
+        sequence_model_parallel_size=1,
+        data_parallel_size=1,
+    )
+
+
+def translate_keys(raw: dict, *, mapping: dict[str, str]) -> dict:
+    out: dict = {}
+    for k, v in raw.items():
+        new_k = k
+        for pat, repl in mapping.items():
+            new_k = re.sub(pat, repl, new_k)
+        out[new_k] = v
+    return out
+
+
+def encode_prompts():
+    banner("(1) Encode prompts via UMT5")
+    from transformers import AutoTokenizer, UMT5EncoderModel
+
+    tok = AutoTokenizer.from_pretrained(str(ANYFLOW_LOCAL), subfolder="tokenizer", use_fast=False)
+    enc = UMT5EncoderModel.from_pretrained(
+        str(ANYFLOW_LOCAL), subfolder="text_encoder", torch_dtype=DTYPE,
+    ).to(DEVICE).eval()
+
+    @torch.no_grad()
+    def encode_one(prompts):
+        out = tok(
+            prompts, padding="max_length", max_length=512, truncation=True,
+            return_attention_mask=True, return_tensors="pt")
+        ids = out.input_ids.to(DEVICE)
+        mask = out.attention_mask.to(DEVICE)
+        seq_lens = mask.gt(0).sum(dim=1).long()
+        embeds = enc(ids, mask).last_hidden_state
+        padded = []
+        for i, l in enumerate(seq_lens):
+            e = embeds[i, :l]
+            pad = torch.zeros(512 - e.size(0), e.size(1), device=DEVICE, dtype=embeds.dtype)
+            padded.append(torch.cat([e, pad], dim=0))
+        return torch.stack(padded, dim=0)
+
+    text_e = encode_one([PROMPT])
+    neg_e = encode_one([NEG_PROMPT])
+    print(f"  prompts encoded: text={tuple(text_e.shape)} neg={tuple(neg_e.shape)}")
+    del enc, tok
+    free_gpu()
+    return text_e, neg_e
+
+
+def load_transformer():
+    banner("(2) Build FastVideo Wan-14B + load AnyFlow weights")
+    from fastvideo.configs.models.dits import WanVideoConfig
+    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
+
+    cfg = WanVideoConfig()
+    arch = cfg.arch_config
+    # Wan2.1-T2V-14B arch (per AnyFlow checkpoint config.json).
+    arch.num_attention_heads = 40
+    arch.attention_head_dim = 128
+    arch.num_layers = 40
+    arch.ffn_dim = 13824
+    arch.r_embedder = True
+    arch.r_embedder_fusion = "gated"
+    arch.r_embedder_gate_value = 0.25
+    arch.r_embedder_deltatime_type = "r"
+    arch.__post_init__()
+
+    t0 = time.time()
+    model = WanTransformer3DModel(config=cfg, hf_config={}).to(DEVICE, dtype=DTYPE).eval()
+    print(f"  transformer built in {time.time() - t0:.1f}s; "
+          f"params: {sum(p.numel() for p in model.parameters())/1e9:.2f}B")
+
+    ckpt_dir = ANYFLOW_LOCAL / "transformer"
+    sd: dict = {}
+    for shard in sorted(ckpt_dir.glob("diffusion_pytorch_model-*.safetensors")):
+        sd.update(load_file(str(shard), device="cpu"))
+    print(f"  AnyFlow state dict: {len(sd)} tensors loaded")
+    sd = translate_keys(sd, mapping=arch.param_names_mapping)
+    info = model.load_state_dict(sd, strict=False)
+    print(f"  load: missing={len(info.missing_keys)} unexpected={len(info.unexpected_keys)}")
+    del sd
+    free_gpu()
+    return model
+
+
+@torch.no_grad()
+def sample(model, text_e, neg_e, nfe: int) -> torch.Tensor:
+    banner(f"(3) Sample 14B NFE={nfe}")
+    from fastvideo.forward_context import set_forward_context
+    from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import (
+        FlowMapEulerDiscreteScheduler, )
+
+    scheduler = FlowMapEulerDiscreteScheduler(num_train_timesteps=1000, shift=5.0)
+    scheduler.set_timesteps(num_inference_steps=nfe, device=DEVICE)
+    timesteps = scheduler.timesteps.to(DEVICE, dtype=DTYPE)
+
+    B, C = 1, 16
+    F = (NUM_FRAMES - 1) // 4 + 1  # temporal VAE compression = 4 (81 → 21)
+    H_l, W_l = HEIGHT // 8, WIDTH // 8
+    g = torch.Generator(device=DEVICE).manual_seed(SEED)
+    x = torch.randn(B, C, F, H_l, W_l, device=DEVICE, dtype=DTYPE, generator=g)
+
+    t0 = time.time()
+    for i, (t_cur, t_next) in enumerate(zip(timesteps[:-1], timesteps[1:])):
+        t_in = t_cur.expand(B).to(DTYPE)
+        r_in = t_next.expand(B).to(DTYPE)
+        with set_forward_context(current_timestep=t_in, attn_metadata=None):
+            flow_cond = model(
+                hidden_states=x, encoder_hidden_states=text_e,
+                timestep=t_in, r_timestep=r_in)
+            if GUIDANCE != 1.0:
+                flow_uncond = model(
+                    hidden_states=x, encoder_hidden_states=neg_e,
+                    timestep=t_in, r_timestep=r_in)
+                flow = flow_uncond + GUIDANCE * (flow_cond - flow_uncond)
+            else:
+                flow = flow_cond
+        x = scheduler.step(
+            flow, sample=x,
+            timestep=t_cur.repeat(B), r_timestep=t_next.repeat(B))
+    print(f"  NFE={nfe} sample time: {time.time() - t0:.1f}s "
+          f"({(time.time() - t0) / nfe:.1f}s/step)")
+    xf = x.float()
+    print(f"  latents mean={xf.mean().item():+.3f} std={xf.std().item():.3f} "
+          f"range=[{xf.min().item():+.2f}, {xf.max().item():+.2f}] "
+          f"finite={torch.isfinite(xf).all().item()}")
+    return x.detach()
+
+
+@torch.no_grad()
+def decode_to_mp4(latents: torch.Tensor, out_path: Path) -> tuple[Path, tuple]:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan
+    import imageio.v3 as iio
+
+    vae = AutoencoderKLWan.from_pretrained(
+        str(ANYFLOW_LOCAL), subfolder="vae", torch_dtype=DTYPE,
+    ).to(DEVICE).eval()
+    try:
+        vae.enable_tiling()
+        print(f"  VAE tiling enabled")
+    except Exception:
+        pass
+
+    mean = torch.tensor(vae.config.latents_mean, device=DEVICE, dtype=DTYPE).view(1, -1, 1, 1, 1)
+    std = torch.tensor(vae.config.latents_std, device=DEVICE, dtype=DTYPE).view(1, -1, 1, 1, 1)
+    latents_unscaled = latents * std + mean
+    t0 = time.time()
+    frames = vae.decode(latents_unscaled, return_dict=False)[0]
+    print(f"  VAE decode time: {time.time() - t0:.1f}s")
+    frames = (frames.clamp(-1, 1) + 1) / 2
+    frames = frames[0].permute(1, 2, 3, 0).float().cpu().numpy()
+    frames = (frames * 255).astype("uint8")
+    iio.imwrite(str(out_path), frames, fps=16, codec="libx264", quality=8)
+    del vae
+    free_gpu()
+    return out_path, frames.shape
+
+
+def main() -> None:
+    torch.manual_seed(SEED)
+    torch.cuda.manual_seed_all(SEED)
+    init_single_rank()
+
+    text_e, neg_e = encode_prompts()
+    model = load_transformer()
+
+    latents_list = []
+    for nfe in [4, 50]:
+        lat = sample(model, text_e, neg_e, nfe=nfe)
+        latents_list.append((nfe, lat))
+
+    del model, text_e, neg_e
+    free_gpu()
+
+    for nfe, lat in latents_list:
+        banner(f"(4) Decode 14B NFE={nfe}")
+        out_path = OUT_DIR / f"fastvideo_anyflow_14b_nfe{nfe}_seed{SEED}.mp4"
+        path, shape = decode_to_mp4(lat, out_path)
+        print(f"  decoded {shape}")
+        print(f"  saved: {path}  ({path.stat().st_size / 1e6:.2f} MB)")
+
+    banner("DONE")
+
+
+if __name__ == "__main__":
+    main()

scripts/verify_anyflow_fastvideo_parity.py

@@ -0,0 +1,404 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""FastVideo↔AnyFlow numerical parity verification.
+
+Two checks, both on a single H200:
+
+  (A) Forward parity — load nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers into
+      FastVideo's WanTransformer3DModel (with r_embedder enabled +
+      param_names_mapping handling the delta_embedder rename), forward
+      it on identical inputs against AnyFlow's reference loader, and
+      compare. Expectation: rel-mean diff < 10% in bf16 (bf16 kernel noise).
+
+  (B) Any-step end-to-end sampling — run the new
+      FlowMapEulerDiscreteScheduler through 4 Euler-flow steps on the
+      same loaded weights, confirm the final latent is finite and
+      well-scaled.
+
+Configure local checkout paths via env vars (defaults assume a sibling
+layout next to this repo):
+
+  ANYFLOW_LOCAL  — path to ``nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers``
+                   (default ``./anyflow-1.3b``)
+  ANYFLOW_REF    — path to the NVlabs/AnyFlow reference repo, used to
+                   import its loader (default ``./anyflow-ref``)
+
+Run via::
+
+    PYTHONPATH=$PWD python scripts/verify_anyflow_fastvideo_parity.py
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import sys
+import time
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+
+
+ANYFLOW_LOCAL = Path(os.environ.get("ANYFLOW_LOCAL", "./anyflow-1.3b")).expanduser()
+ANYFLOW_REF = Path(os.environ.get("ANYFLOW_REF", "./anyflow-ref")).expanduser()
+sys.path.insert(0, str(ANYFLOW_REF))
+
+SEED = 1234
+DEVICE = torch.device("cuda")
+DTYPE = torch.bfloat16
+
+
+def banner(msg: str) -> None:
+    print("\n" + "=" * 80)
+    print(msg)
+    print("=" * 80)
+
+
+# ---------------------------------------------------------------------------
+# Distributed bootstrap (single-rank). Required by WanTransformer3DModel
+# which calls get_sp_world_size() and uses ReplicatedLinear.
+# ---------------------------------------------------------------------------
+
+
+def init_single_rank() -> None:
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29551")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    from fastvideo.distributed import (
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+    init_distributed_environment(world_size=1, rank=0, local_rank=0,
+                                 backend="nccl")
+    initialize_model_parallel(
+        tensor_model_parallel_size=1,
+        sequence_model_parallel_size=1,
+        data_parallel_size=1,
+    )
+    print("  single-rank distributed environment + TP/SP/DP groups initialized")
+
+
+# ---------------------------------------------------------------------------
+# Translate AnyFlow HF safetensor keys onto FastVideo's WanTransformer3DModel
+# internal layout, applying the regex from WanVideoArchConfig.param_names_mapping.
+# ---------------------------------------------------------------------------
+
+
+def translate_keys(
+    raw: dict[str, torch.Tensor],
+    *,
+    mapping: dict[str, str],
+) -> dict[str, torch.Tensor]:
+    out: dict[str, torch.Tensor] = {}
+    for k, v in raw.items():
+        new_k = k
+        for pat, repl in mapping.items():
+            new_k = re.sub(pat, repl, new_k)
+        out[new_k] = v
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Build FastVideo WanTransformer3DModel with AnyFlow weights loaded.
+# ---------------------------------------------------------------------------
+
+
+def build_fastvideo_transformer():
+    banner("(1) Build FastVideo WanTransformer3DModel + load AnyFlow weights")
+    from fastvideo.configs.models.dits import WanVideoConfig
+    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
+
+    cfg = WanVideoConfig()
+    arch = cfg.arch_config
+    # AnyFlow Wan2.1-T2V-1.3B arch dims.
+    arch.num_attention_heads = 12
+    arch.attention_head_dim = 128
+    arch.num_layers = 30
+    arch.ffn_dim = 8960
+    # AnyFlow dual-timestep.
+    arch.r_embedder = True
+    arch.r_embedder_fusion = "gated"
+    arch.r_embedder_gate_value = 0.25
+    arch.r_embedder_deltatime_type = "r"
+    arch.__post_init__()
+
+    # WanTransformer3DModel takes (config, hf_config). hf_config is
+    # diffusers-style — we provide a minimal dict; only fields the model
+    # actually reads matter.
+    hf_config: dict = {}
+    t0 = time.time()
+    model = WanTransformer3DModel(config=cfg, hf_config=hf_config)
+    model = model.to(DEVICE, dtype=DTYPE).eval()
+    print(f"  Wan transformer built in {time.time() - t0:.1f}s; "
+          f"params: {sum(p.numel() for p in model.parameters())/1e9:.2f}B")
+
+    # Load AnyFlow checkpoint.
+    af_path = ANYFLOW_LOCAL / "transformer" / "diffusion_pytorch_model.safetensors"
+    af_raw = load_file(str(af_path), device="cpu")
+    print(f"  AnyFlow checkpoint: {len(af_raw)} tensors")
+
+    translated = translate_keys(af_raw, mapping=arch.param_names_mapping)
+    info = model.load_state_dict(translated, strict=False)
+    miss, unex = info.missing_keys, info.unexpected_keys
+    print(f"  missing_keys    : {len(miss)} (first 5: {miss[:5]})")
+    print(f"  unexpected_keys : {len(unex)} (first 5: {unex[:5]})")
+    return model, len(miss), len(unex)
+
+
+# ---------------------------------------------------------------------------
+# Build AnyFlow reference net.
+# ---------------------------------------------------------------------------
+
+
+def build_anyflow_reference():
+    banner("(2) Build AnyFlow reference loader")
+    from far.models import build_model
+
+    af_net = build_model("FAR_Wan_Transformer3DModel").from_pretrained(
+        str(ANYFLOW_LOCAL),
+        subfolder="transformer",
+        chunk_partition=None,
+        full_chunk_limit=0,
+        compressed_patch_size=[1, 4, 4],
+    ).to(DEVICE, dtype=DTYPE).eval()
+    print(f"  AnyFlow {type(af_net).__name__} ready")
+    return af_net
+
+
+# ---------------------------------------------------------------------------
+# Forward parity test.
+# ---------------------------------------------------------------------------
+
+
+def forward_compare(fv_model, af_net):
+    banner("(3) Forward output comparison")
+    B, C, F, H, W = 1, 16, 21, 60, 104
+    SEQ, DIM = 32, 4096
+    g = torch.Generator(device=DEVICE).manual_seed(SEED)
+    x = torch.randn(B, C, F, H, W, device=DEVICE, dtype=DTYPE, generator=g)
+    enc = torch.randn(B, SEQ, DIM, device=DEVICE, dtype=DTYPE, generator=g)
+
+    # AnyFlow expects per-frame (t, r) [B, F]; FastVideo T2V expects a
+    # single scalar per sample [B] (timestep.dim()==2 is reserved for
+    # Wan2.2 ti2v's per-token timestep schedule). For shared-t T2V the two
+    # are semantically equivalent.
+    t_per_frame = torch.full((B, F), 500.0, device=DEVICE, dtype=DTYPE)
+    r_per_frame = torch.full((B, F), 200.0, device=DEVICE, dtype=DTYPE)
+    t_per_sample = torch.full((B,), 500.0, device=DEVICE, dtype=DTYPE)
+    r_per_sample = torch.full((B,), 200.0, device=DEVICE, dtype=DTYPE)
+
+    with torch.no_grad():
+        # AnyFlow native: takes [B, F, C, H, W] + [B, F] (t, r).
+        x_af = x.permute(0, 2, 1, 3, 4).contiguous()
+        af_out = af_net(
+            x_af,
+            timestep=t_per_frame,
+            r_timestep=r_per_frame,
+            encoder_hidden_states=enc,
+            return_dict=False,
+            is_causal=False,
+        )[0]
+        if af_out.shape[1] != C and af_out.shape[2] == C:
+            af_out = af_out.permute(0, 2, 1, 3, 4).contiguous()
+
+        # FastVideo: [B, C, F, H, W] + [B] (t, r). Must wrap in
+        # set_forward_context so the attention layer can pick up the
+        # current timestep / attn_metadata.
+        from fastvideo.forward_context import set_forward_context
+        with set_forward_context(
+            current_timestep=t_per_sample,
+            attn_metadata=None,
+        ):
+            fv_out = fv_model(
+                hidden_states=x,
+                encoder_hidden_states=enc,
+                timestep=t_per_sample,
+                r_timestep=r_per_sample,
+            )
+
+    print(f"  AnyFlow out: shape={tuple(af_out.shape)} dtype={af_out.dtype}")
+    print(f"  FastVideo  : shape={tuple(fv_out.shape)} dtype={fv_out.dtype}")
+    if af_out.shape != fv_out.shape:
+        print("  ❌ shape mismatch")
+        return False
+    diff = (af_out.float() - fv_out.float()).abs()
+    ref = af_out.float().abs().mean().item() + 1e-12
+    print(f"  max abs diff : {diff.max().item():.3e}")
+    print(f"  mean abs diff: {diff.mean().item():.3e}")
+    print(f"  rel mean diff: {diff.mean().item() / ref:.3e}")
+    ok = diff.mean().item() / ref < 0.10
+    print(f"  >>> {'PASS' if ok else 'FAIL'} (target rel diff < 10%, bf16 noise)")
+    return ok
+
+
+# ---------------------------------------------------------------------------
+# Any-step sampling smoke via FlowMapEulerDiscreteScheduler.
+# ---------------------------------------------------------------------------
+
+
+def sample_anystep(fv_model):
+    banner("(4) Any-step 4-step Euler-flow sampling")
+    from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import (
+        FlowMapEulerDiscreteScheduler, )
+
+    scheduler = FlowMapEulerDiscreteScheduler(num_train_timesteps=1000, shift=5.0)
+    scheduler.set_timesteps(num_inference_steps=4, device=DEVICE)
+    timesteps = scheduler.timesteps.to(dtype=DTYPE)
+
+    B, C, F, H, W = 1, 16, 21, 60, 104
+    SEQ, DIM = 32, 4096
+    g = torch.Generator(device=DEVICE).manual_seed(SEED)
+    x = torch.randn(B, C, F, H, W, device=DEVICE, dtype=DTYPE, generator=g)
+    enc = torch.randn(B, SEQ, DIM, device=DEVICE, dtype=DTYPE, generator=g)
+
+    from fastvideo.forward_context import set_forward_context
+    t0 = time.time()
+    with torch.no_grad():
+        for t_cur, t_next in zip(timesteps[:-1], timesteps[1:]):
+            t_in = t_cur.expand(B).to(DTYPE)
+            r_in = t_next.expand(B).to(DTYPE)
+            with set_forward_context(current_timestep=t_in, attn_metadata=None):
+                v = fv_model(
+                    hidden_states=x,
+                    encoder_hidden_states=enc,
+                    timestep=t_in,
+                    r_timestep=r_in,
+                )
+            x = scheduler.step(
+                v, sample=x,
+                timestep=t_cur.repeat(B),
+                r_timestep=t_next.repeat(B),
+            )
+    elapsed = time.time() - t0
+    xf = x.float()
+    print(f"  elapsed: {elapsed:.1f}s")
+    print(f"  final latent: mean={xf.mean().item():+.3f} std={xf.std().item():.3f} "
+          f"range=[{xf.min().item():+.2f}, {xf.max().item():+.2f}] "
+          f"finite={torch.isfinite(xf).all().item()}")
+    ok = torch.isfinite(xf).all().item() and 0.01 < xf.std().item() < 30
+    print(f"  >>> {'PASS' if ok else 'FAIL'}")
+    return ok
+
+
+NUM_TRAIN_TIMESTEPS = 1000
+EPSILON = 5.0  # AnyFlow paper default
+
+
+@torch.no_grad()
+def training_step_compare(fv_model, af_net) -> bool:
+    """Inline replica of AnyFlow's train_bidirection central-difference loss
+    on both code paths with identical synthetic (real, noise, t, r) inputs.
+
+    Compares scalar loss + intermediate flow_pred / target tensors.
+    """
+    banner("(5) Training-step loss comparison (central-difference target)")
+    from fastvideo.forward_context import set_forward_context
+
+    B, C, F, H, W = 1, 16, 21, 60, 104
+    SEQ, DIM = 32, 4096
+    g = torch.Generator(device=DEVICE).manual_seed(SEED)
+    real = torch.randn(B, C, F, H, W, device=DEVICE, dtype=DTYPE, generator=g)
+    enc = torch.randn(B, SEQ, DIM, device=DEVICE, dtype=DTYPE, generator=g)
+    noise = torch.randn_like(real)
+    t_abs_pf = torch.full((B, F), 500.0, device=DEVICE, dtype=DTYPE)
+    r_abs_pf = torch.full((B, F), 200.0, device=DEVICE, dtype=DTYPE)
+    t_abs_ps = torch.full((B,), 500.0, device=DEVICE, dtype=DTYPE)
+    r_abs_ps = torch.full((B,), 200.0, device=DEVICE, dtype=DTYPE)
+
+    # AnyFlow inline replica: uses [B, F, C, H, W] layout.
+    real_btchw = real.permute(0, 2, 1, 3, 4).contiguous()
+    noise_btchw = noise.permute(0, 2, 1, 3, 4).contiguous()
+    t_norm_pf = (t_abs_pf / NUM_TRAIN_TIMESTEPS).view(B, F, 1, 1, 1).to(DTYPE)
+    noisy_btchw = t_norm_pf * noise_btchw + (1 - t_norm_pf) * real_btchw
+
+    def u_func_af(x_in, t_in, r_in):
+        return af_net(
+            x_in, timestep=t_in, r_timestep=r_in,
+            encoder_hidden_states=enc, return_dict=False, is_causal=False)[0]
+
+    v_pred = noise_btchw - real_btchw
+    eps = EPSILON
+    F_plus = u_func_af(noisy_btchw + v_pred * (eps / NUM_TRAIN_TIMESTEPS),
+                       t_abs_pf + eps, r_abs_pf)
+    F_minus = u_func_af(noisy_btchw - v_pred * (eps / NUM_TRAIN_TIMESTEPS),
+                        t_abs_pf - eps, r_abs_pf)
+    dF_dt_af = (F_plus - F_minus) / (2 * eps)
+    target_af = ((noise_btchw - real_btchw)
+                 - (t_abs_pf - r_abs_pf).view(B, F, 1, 1, 1) * dF_dt_af)
+    flow_af = u_func_af(noisy_btchw, t_abs_pf, r_abs_pf)
+    loss_af = (flow_af.float() - target_af.float()).pow(2).reshape(B, -1).mean(-1)
+
+    # FastVideo inline replica: uses [B, C, F, H, W] layout + [B] t/r.
+    real_bcfhw = real
+    noise_bcfhw = noise
+    t_norm_ps = (t_abs_ps / NUM_TRAIN_TIMESTEPS).view(B, 1, 1, 1, 1).to(DTYPE)
+    noisy_bcfhw = t_norm_ps * noise_bcfhw + (1 - t_norm_ps) * real_bcfhw
+
+    def u_func_fv(x_in, t_in, r_in):
+        with set_forward_context(current_timestep=t_in, attn_metadata=None):
+            return fv_model(
+                hidden_states=x_in,
+                encoder_hidden_states=enc,
+                timestep=t_in,
+                r_timestep=r_in,
+            )
+
+    v_pred_fv = noise_bcfhw - real_bcfhw
+    F_plus_fv = u_func_fv(
+        noisy_bcfhw + v_pred_fv * (eps / NUM_TRAIN_TIMESTEPS),
+        t_abs_ps + eps, r_abs_ps)
+    F_minus_fv = u_func_fv(
+        noisy_bcfhw - v_pred_fv * (eps / NUM_TRAIN_TIMESTEPS),
+        t_abs_ps - eps, r_abs_ps)
+    dF_dt_fv = (F_plus_fv - F_minus_fv) / (2 * eps)
+    target_fv = ((noise_bcfhw - real_bcfhw)
+                 - (t_abs_ps - r_abs_ps).view(B, 1, 1, 1, 1) * dF_dt_fv)
+    flow_fv = u_func_fv(noisy_bcfhw, t_abs_ps, r_abs_ps)
+    loss_fv = (flow_fv.float() - target_fv.float()).pow(2).reshape(B, -1).mean(-1)
+
+    # Compare. Align AnyFlow's [B, F, C, H, W] → [B, C, F, H, W].
+    flow_af_aligned = flow_af.permute(0, 2, 1, 3, 4)
+    target_af_aligned = target_af.permute(0, 2, 1, 3, 4)
+    flow_diff = (flow_fv.float() - flow_af_aligned.float()).abs()
+    target_diff = (target_fv.float() - target_af_aligned.float()).abs()
+
+    af_loss_v = loss_af.mean().item()
+    fv_loss_v = loss_fv.mean().item()
+    abs_diff = abs(af_loss_v - fv_loss_v)
+    rel_diff = abs_diff / abs(af_loss_v + 1e-12)
+    print(f"  AnyFlow loss : {af_loss_v:.6f}")
+    print(f"  FastVideo loss: {fv_loss_v:.6f}")
+    print(f"  abs diff     : {abs_diff:.3e}")
+    print(f"  rel diff     : {rel_diff:.3e}")
+    print(f"  flow_pred  : max abs {flow_diff.max().item():.3e}  "
+          f"mean {flow_diff.mean().item():.3e}")
+    print(f"  target     : max abs {target_diff.max().item():.3e}  "
+          f"mean {target_diff.mean().item():.3e}")
+    ok = rel_diff < 0.20
+    print(f"  >>> {'PASS' if ok else 'FAIL'} (target rel loss diff < 20%)")
+    return ok
+
+
+def main() -> None:
+    torch.manual_seed(SEED)
+    torch.cuda.manual_seed_all(SEED)
+
+    init_single_rank()
+    fv_model, n_miss, n_unex = build_fastvideo_transformer()
+    af_net = build_anyflow_reference()
+    forward_ok = forward_compare(fv_model, af_net)
+    sample_ok = sample_anystep(fv_model)
+    train_ok = training_step_compare(fv_model, af_net)
+    banner(
+        f"SUMMARY: missing_keys={n_miss} unexpected_keys={n_unex} "
+        f"forward_parity={forward_ok} sample_smoke={sample_ok} "
+        f"training_parity={train_ok}"
+    )
+    sys.exit(0 if (forward_ok and sample_ok and train_ok) else 1)
+
+
+if __name__ == "__main__":
+    main()