diff --git a/docs/distillation/anyflow.md b/docs/distillation/anyflow.md new file mode 100644 index 0000000000..85904d5abd --- /dev/null +++ b/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). diff --git a/examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml b/examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml new file mode 100644 index 0000000000..2fe7b928b5 --- /dev/null +++ b/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 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: + 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 diff --git a/examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml b/examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml new file mode 100644 index 0000000000..0e870a4a1b --- /dev/null +++ b/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 diff --git a/fastvideo/configs/models/dits/wanvideo.py b/fastvideo/configs/models/dits/wanvideo.py index 8ab9c7b31e..0b76d8c5fa 100644 --- a/fastvideo/configs/models/dits/wanvideo.py +++ b/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 diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 2c32a5053f..6b925c5f05 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__(self, config: WanVideoConfig, hf_config: dict[str, 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 @@ def forward(self, 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 @@ def forward(self, 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)) diff --git a/fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py new file mode 100644 index 0000000000..95f6e808ec --- /dev/null +++ b/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 diff --git a/fastvideo/tests/training/distill/test_anyflow_onpolicy.py b/fastvideo/tests/training/distill/test_anyflow_onpolicy.py new file mode 100644 index 0000000000..7040a25488 --- /dev/null +++ b/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")) diff --git a/fastvideo/tests/training/distill/test_anyflow_pretrain.py b/fastvideo/tests/training/distill/test_anyflow_pretrain.py new file mode 100644 index 0000000000..fee6f5878f --- /dev/null +++ b/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}") diff --git a/fastvideo/tests/training/distill/test_anyflow_smoke.py b/fastvideo/tests/training/distill/test_anyflow_smoke.py new file mode 100644 index 0000000000..8769901429 --- /dev/null +++ b/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) diff --git a/fastvideo/train/methods/distribution_matching/__init__.py b/fastvideo/train/methods/distribution_matching/__init__.py index 015f8e124b..2016423903 100644 --- a/fastvideo/train/methods/distribution_matching/__init__.py +++ b/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 @@ StreamingLongTuningMethod, ) __all__ = [ + "AnyFlowMethod", + "AnyFlowPretrainMethod", "DMD2Method", "SelfForcingMethod", "StreamingLongTuningMethod", diff --git a/fastvideo/train/methods/distribution_matching/anyflow.py b/fastvideo/train/methods/distribution_matching/anyflow.py new file mode 100644 index 0000000000..ebdc110b96 --- /dev/null +++ b/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 diff --git a/fastvideo/train/methods/distribution_matching/anyflow_pretrain.py b/fastvideo/train/methods/distribution_matching/anyflow_pretrain.py new file mode 100644 index 0000000000..821e021f7f --- /dev/null +++ b/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), + ) diff --git a/fastvideo/train/models/wan/wan.py b/fastvideo/train/models/wan/wan.py index 7e11c88d65..316dafd1e9 100644 --- a/fastvideo/train/models/wan/wan.py +++ b/fastvideo/train/models/wan/wan.py @@ -358,6 +358,53 @@ def predict_noise( 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, diff --git a/mkdocs.yml b/mkdocs.yml index df1d9b2990..e7515381e9 100644 --- a/mkdocs.yml +++ b/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 diff --git a/scripts/demo_anyflow_14b.py b/scripts/demo_anyflow_14b.py new file mode 100644 index 0000000000..53afad76b1 --- /dev/null +++ b/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() diff --git a/scripts/verify_anyflow_fastvideo_parity.py b/scripts/verify_anyflow_fastvideo_parity.py new file mode 100644 index 0000000000..2f7b717ad7 --- /dev/null +++ b/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()