From ce23c5bfb1491f3970e9867d722b9abebd652efa Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 19 May 2026 14:26:39 +0800 Subject: [PATCH] [feat] AnyFlowPretrainMethod skeleton + (t, r) per-batch sampling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the single-student AnyFlow flow-map pretrain method as a TrainingMethod subclass. The __init__ parses and validates all method config knobs (diffusion_ratio, consistency_ratio, epsilon, weight_type, fuse_guidance_scale, shift), builds a FlowMapEulerDiscreteScheduler, and wires a single optimizer+scheduler over the student. The (t, r) sampling helper _sample_pair_timesteps lives at module scope so it can be unit-tested without instantiating the full method. It matches the AnyFlow paper formulation: - two uniform draws u1, u2; t = max, r = min - first diffusion_ratio * B: r := t (plain flow matching) - next consistency_ratio * B: r := 0 (consistency to clean data) - remainder: free reconstruction range single_train_step raises NotImplementedError for now — filled in by the next commit which adds the central-difference target and the loss assembly (Task 7). --- .../methods/distribution_matching/__init__.py | 3 + .../distribution_matching/anyflow_pretrain.py | 226 ++++++++++++++++++ 2 files changed, 229 insertions(+) create mode 100644 fastvideo/train/methods/distribution_matching/anyflow_pretrain.py diff --git a/fastvideo/train/methods/distribution_matching/__init__.py b/fastvideo/train/methods/distribution_matching/__init__.py index 015f8e124..97118c3f0 100644 --- a/fastvideo/train/methods/distribution_matching/__init__.py +++ b/fastvideo/train/methods/distribution_matching/__init__.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 +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 +9,7 @@ from fastvideo.train.methods.distribution_matching.streaming_long_tuning import StreamingLongTuningMethod, ) __all__ = [ + "AnyFlowPretrainMethod", "DMD2Method", "SelfForcingMethod", "StreamingLongTuningMethod", 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 000000000..810bb9e4a --- /dev/null +++ b/fastvideo/train/methods/distribution_matching/anyflow_pretrain.py @@ -0,0 +1,226 @@ +# 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, 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 + + +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. + raise NotImplementedError("filled in Task 7") + + 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), + )