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.
Classes
Bases: 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.
Source code in fastvideo/train/methods/distribution_matching/anyflow_pretrain.py
| 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()
|
Methods:
backward(loss_map: dict[str, 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.
Source code in fastvideo/train/methods/distribution_matching/anyflow_pretrain.py
| 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,
)
|
Functions: