Skip to content

anyflow

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.

Classes

fastvideo.train.methods.distribution_matching.anyflow.AnyFlowMethod 

AnyFlowMethod(*, cfg: Any, role_models: dict[str, Any])

Bases: DMD2Method

AnyFlow on-policy distillation (multi-step rollout).

Source code in fastvideo/train/methods/distribution_matching/anyflow.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
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