Finalize integration and configs

Signed-off-by: Julius Berner <mail@jberner.info>

Author: juliusberner

fastgen/configs/experiments/WanT2V/config_anyflow.py

@@ -67,7 +67,7 @@ def create_config():
     # learns the guided flow directly. guidance_scale stays None — MeanFlow's
     # target-side eq. 19 fusion is a different mechanism.
     config.model.guidance_scale = None
-    config.model.loss_config.guidance_fuse_scale = 3.0
+    config.model.guidance_fuse_scale = 3.0
     config.model.cond_dropout_prob = 0.1
 
     # ------ (t, r) sampling: shifted uniform pairs + AnyFlow buckets ------

fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py

@@ -112,7 +112,7 @@ def create_config():
     config.model.loss_config.use_jvp_finite_diff = True
     config.model.loss_config.jvp_finite_diff_eps = 5e-3
     config.model.loss_config.rebalance_to_flow_matching = True
-    config.model.loss_config.guidance_fuse_scale = 3.0
+    config.model.guidance_fuse_scale = 3.0
     config.model.cond_dropout_prob = 0.1
     config.model.precision_amp_jvp = "float32"
     config.model.sample_t_cfg.flow_matching_ratio = 0.5

fastgen/configs/methods/config_anyflow.py

@@ -59,13 +59,24 @@ class ModelConfig(DMD2ModelConfig):
     # the fixed student_sample_steps.
     student_sample_steps_list: Optional[List[int]] = None
 
+    # Prediction-side guidance fusion for the co-trained flow-map loss; see
+    # `config_mean_flow.ModelConfig`.
+    guidance_fuse_scale: Optional[float] = None
+
     # Text dropout for the co-trained flow-map loss (reference drop_text_ratio).
     cond_dropout_prob: Optional[float] = None
     cond_keys_no_dropout: List[str] = attrs.field(factory=list)
 
     # Precision for autocast in the co-trained loss JVP (None = training precision).
     precision_amp_jvp: str | None = None
 
+    # MeanFlow's target-side guidance knobs. AnyFlow typixcally guides on the 
+    # PREDICTION side (`guidance_fuse_scale`), so these stay at their
+    # no-op defaults
+    guidance_mixture_ratio: Optional[float] = None
+    guidance_t_start: float = 0.0
+    guidance_t_end: float = 1.0
+
 
 @attrs.define(slots=False)
 class Config(BaseConfig):

fastgen/configs/methods/config_mean_flow.py

@@ -86,11 +86,6 @@ class LossConfig:
     # ("beta08", "gaussian", "uniform"). Multiplies the adaptive norm_method 
     # weight above; None disables it.
     weight_type: Optional[str] = None
-    # prediction-side guidance fusion scale (AnyFlow guidance distillation):
-    # the conditional output is trained to be the guided flow directly via
-    # (u_cond + (g-1) * u_uncond) / g against the raw data velocity. None 
-    # keeps MeanFlow's target-side guidance (guidance_scale).
-    guidance_fuse_scale: Optional[float] = None
     # rebalance the flow-map / consistency (r < t) sample losses to the global
     # flow-matching (r = t) loss mean via a detached per-sample factor
     rebalance_to_flow_matching: bool = False
@@ -116,6 +111,12 @@ class ModelConfig(BaseModelConfig):
     # optimizer
     net_optimizer: DictConfig = attrs.field(factory=lambda: copy.deepcopy(RAdamOptimizerConfig))
 
+    # prediction-side guidance fusion scale (AnyFlow guidance distillation):
+    # the conditional output is trained to be the guided flow directly via
+    # (u_cond + (g-1) * u_uncond) / g against the raw data velocity. None
+    # keeps MeanFlow's target-side guidance (guidance_scale).
+    guidance_fuse_scale: Optional[float] = None
+
     # condition dropout probability
     cond_dropout_prob: Optional[float] = None
 

fastgen/methods/consistency_model/README.md

@@ -73,7 +73,7 @@ Two-stage training: frozen Stage-1 CM for `t < transition_t`, trainable student
 
 ## MeanFlow
 
-**File:** [`mean_flow.py`](mean_flow.py) | **Reference:** [Geng et al., 2025](https://arxiv.org/abs/2505.13447)
+**File:** [`mean_flow.py`](mean_flow.py) | **References:** [Geng et al., 2025](https://arxiv.org/abs/2505.13447), [Sabour et al., 2025](https://arxiv.org/abs/2506.14603)
 
 Learns average velocity between trajectory points: `x_r = x_t - (t-r) · u(x_t, t, r)`.
 
@@ -82,6 +82,7 @@ Learns average velocity between trajectory points: `x_r = x_t - (t-r) · u(x_t,
 - `loss_config.use_jvp_finite_diff`: Use finite difference for JVP (e.g., for compatibility with Flash Attention and FSDP)
 - `sample_t_cfg`, `sample_r_cfg`: Configs of the distributions for sampling `t` and `r`
 - `sample_t_cfg.flow_matching_ratio`: Fraction of the batch with `r = t` (flow matching loss)
+- `guidance_scale`, `guidance_fuse_scale`: If `guidance_fuse_scale` is `None`, use target-side guidance (`guidance_scale`, 3-way mixture via `guidance_mixture_ratio`, t-window via `guidance_t_start` / `guidance_t_end`), otherwise use prediction-side guidance (conditional output learns the guided flow directly)
 
 
 **Configs:** [`EDM/config_mf_cifar10.py`](../../configs/experiments/EDM/config_mf_cifar10.py), [`DiT/config_mf_b.py`](../../configs/experiments/DiT/config_mf_b.py), [`DiT/config_mf_xl.py`](../../configs/experiments/DiT/config_mf_xl.py), [`WanT2V/config_mf.py`](../../configs/experiments/WanT2V/config_mf.py), [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) (Stage 1 of [AnyFlow](../distribution_matching/README.md#anyflow))

fastgen/methods/consistency_model/mean_flow.py

@@ -125,6 +125,90 @@ def _drop_condition(self, condition: Any, neg_condition: Any) -> Tuple[Any, Opti
             return condition, keep
         raise TypeError(f"Unsupported condition type: {type(condition)}")
 
+    @torch.no_grad()
+    def _get_velocity(
+        self,
+        x: torch.Tensor,
+        z: torch.Tensor,
+        t: torch.Tensor,
+        condition: Optional[Any] = None,
+        neg_condition: Optional[Any] = None,
+    ) -> Tuple[Any, torch.Tensor]:
+        """Regression target for the flow-map loss, plus the condition it was built from.
+
+        Two independent choices:
+
+        * ``loss_config.use_cd`` picks the target SOURCE -- the teacher, or the
+          conditional data velocity.
+        * ``guidance_fuse_scale`` guides the *prediction* in
+          ``_compute_mf_loss``, so the target needs no guidance of its own and we only
+          drop the condition. Otherwise we guide the target too: through the teacher, or
+          -- without one -- by the net's own cond/uncond pass.
+        """
+        fuse_scale = self.config.guidance_fuse_scale
+        if fuse_scale is not None:
+            assert fuse_scale > 0, f"guidance_fuse_scale must be > 0, got {fuse_scale} (None disables fusion)"
+            condition, _ = self._drop_condition(condition, neg_condition)
+
+        x_t = self.net.noise_scheduler.forward_process(x, z, t)
+
+        if self.loss_config.use_cd:
+            dxt_dt = self.teacher(x_t, t, condition=condition, fwd_pred_type="flow")
+            # Under fusion the target stays unguided, or it would be guided twice.
+            if self.config.guidance_scale is not None and fuse_scale is None:
+                guidance_scale = torch.where(
+                    ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)),
+                    self.config.guidance_scale,
+                    1.0,
+                )
+                guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype)
+                neg_dxt_dt = self.teacher(x_t, t, condition=neg_condition, fwd_pred_type="flow")
+                dxt_dt = dxt_dt + (guidance_scale - 1.0) * (dxt_dt - neg_dxt_dt)
+        else:
+            dxt_dt = self.net.noise_scheduler.cond_velocity(x=x, eps=z, t=t)
+
+            # unconditional score estimation from meanflow eq (19). Skipped under
+            # prediction-side fusion: that is this same guidance, moved onto the
+            # prediction, and the dropout it needs already ran above.
+            if fuse_scale is None and (
+                self.config.guidance_scale is not None or self.config.guidance_mixture_ratio is not None
+            ):
+                # Turn off dropout
+                self.net.eval()
+                neg_dxt_dt = self.net(x_t, t, r=t, condition=neg_condition, fwd_pred_type="flow")
+                guidance_scale = self.config.guidance_scale or 1.0
+                guidance_scale = torch.where(
+                    ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)),
+                    guidance_scale,
+                    1.0,
+                )
+                guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype)
+
+                if self.config.guidance_mixture_ratio is None:
+                    guided_dxt_dt = neg_dxt_dt + guidance_scale * (dxt_dt - neg_dxt_dt)
+                else:
+                    guidance_mixture_ratio = torch.where(
+                        ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)),
+                        self.config.guidance_mixture_ratio,
+                        0.0,
+                    )
+                    guidance_mixture_ratio = expand_like(guidance_mixture_ratio, x_t).to(dtype=x_t.dtype)
+                    cond_dxt_dt = self.net(x_t, t, r=t, condition=condition, fwd_pred_type="flow")
+                    guided_dxt_dt = (
+                        guidance_scale * dxt_dt
+                        + (1.0 - guidance_scale - guidance_mixture_ratio) * neg_dxt_dt
+                        + guidance_mixture_ratio * cond_dxt_dt
+                    )
+
+                self.net.train()
+                condition, keep = self._drop_condition(condition, neg_condition)
+                if keep is not None:
+                    # Same subset: a kept sample is conditional + guided, a
+                    # dropped one unconditional + unguided.
+                    dxt_dt = torch.where(expand_like(keep, dxt_dt), guided_dxt_dt, dxt_dt)
+
+        return condition, dxt_dt
+
     def _estimate_jvp_finite_difference(
         self,
         net_wrapper: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor],
@@ -489,50 +573,34 @@ def _compute_mf_loss(
         z = torch.randn_like(real_data)
         x_t = self.net.noise_scheduler.forward_process(real_data, z, t)
 
-        guidance_fuse_scale = getattr(self.loss_config, "guidance_fuse_scale", None)
+        condition, dxt_dt = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition)
+        # prevent JVP to use cached conversions (which can break the computational graph) that were created in the no_grad context of _get_velocity
+        torch.clear_autocast_cache()
+        u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition)
+        assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients"
+
+        # additional forward pass to get u_theta with gradient; see also https://github.com/Gsunshine/py-meanflow?tab=readme-ov-file#note-on-jvp
+        assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}"
+        u_theta = self.net(
+            x_t,
+            t,
+            r=r,
+            condition=condition,
+            fwd_pred_type="flow",
+        )
+
+        guidance_fuse_scale = self.config.guidance_fuse_scale
         if guidance_fuse_scale is not None:
-            # Guidance distillation fused on the PREDICTION side: the
-            # conditional output is trained to be the guided flow directly.
-            # The regression target stays the raw data velocity, the
-            # unconditional branch is queried at the SAME (t, r) flow-map
-            # slice, and the effective prediction is
-            # (u_cond + (g - 1) * u_uncond) / g. Text dropout replaces the
-            # condition with neg_condition for a random subset beforehand.
-            g = float(guidance_fuse_scale)
-            assert g > 0, f"guidance_fuse_scale must be > 0, got {g} (set it to None to disable guidance fusion)"
-            assert (
-                neg_condition is not None
-            ), "guidance_fuse_scale requires neg_condition: the unconditional branch is queried at the same (t, r)"
-            condition, _ = self._drop_condition(condition, neg_condition)
-            dxt_dt = self.net.noise_scheduler.cond_velocity(x=real_data, eps=z, t=t)
-            torch.clear_autocast_cache()
-            # dF/dt of the fused prediction: finite difference of the
-            # conditional output divided by g (the unconditional derivative is
-            # dropped).
-            u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) / g
-            assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients"
-
-            assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}"
-            u_theta = self.net(x_t, t, r=r, condition=condition, fwd_pred_type="flow")
+            # Guidance distillation on the PREDICTION side (see `_get_velocity`): the
+            # conditional output learns the guided flow directly, so only the prediction
+            # changes. The uncond branch is queried at the SAME (t, r) flow-map slice,
+            # giving (u_cond + (g - 1) * u_uncond) / g; dF/dt is then the conditional
+            # finite difference over g, with the unconditional derivative dropped.
+            u_theta_jvp = u_theta_jvp / guidance_fuse_scale
             with torch.no_grad():
                 u_uncond = self.net(x_t, t, r=r, condition=neg_condition, fwd_pred_type="flow")
-            u_theta = (u_theta + (g - 1.0) * u_uncond) / g
-        else:
-            condition, dxt_dt = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition)
-            # prevent JVP to use cached conversions (which can break the computational graph) that were created in the no_grad context of _get_velocity
-            torch.clear_autocast_cache()
-            u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition)
-            assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients"
-
-            # additional forward pass to get u_theta with gradient; see also https://github.com/Gsunshine/py-meanflow?tab=readme-ov-file#note-on-jvp
-            assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}"
-            u_theta = self.net(
-                x_t,
-                t,
-                r=r,
-                condition=condition,
-                fwd_pred_type="flow",
-            )
+            u_theta = (u_theta + (guidance_fuse_scale - 1.0) * u_uncond) / guidance_fuse_scale
+
         mf_loss, tangent, loss_weight, warmup_weight = self._mf_pred_to_loss(
             u_theta=u_theta, u_theta_jvp=u_theta_jvp, x_t=x_t, dxt_dt=dxt_dt, t=t, r=r, iteration=iteration
         )
@@ -564,69 +632,6 @@ def __init__(self, config: ModelConfig):
         self.config = config
         self._init_flow_map_loss()
 
-    @torch.no_grad()
-    def _get_velocity(
-        self,
-        x: torch.Tensor,
-        z: torch.Tensor,
-        t: torch.Tensor,
-        condition: Optional[torch.Tensor] = None,
-        neg_condition: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        x_t = self.net.noise_scheduler.forward_process(x, z, t)
-
-        if self.loss_config.use_cd:
-            dxt_dt = self.teacher(x_t, t, condition=condition, fwd_pred_type="flow")
-            if self.config.guidance_scale is not None:
-                guidance_scale = torch.where(
-                    ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)),
-                    self.config.guidance_scale,
-                    1.0,
-                )
-                guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype)
-                neg_dxt_dt = self.teacher(x_t, t, condition=neg_condition, fwd_pred_type="flow")
-                dxt_dt = dxt_dt + (guidance_scale - 1.0) * (dxt_dt - neg_dxt_dt)
-        else:
-            dxt_dt = self.net.noise_scheduler.cond_velocity(x=x, eps=z, t=t)
-
-            # unconditional score estimation from meanflow eq (19)
-            if self.config.guidance_scale is not None or self.config.guidance_mixture_ratio is not None:
-                # Turn off dropout
-                self.net.eval()
-                neg_dxt_dt = self.net(x_t, t, r=t, condition=neg_condition, fwd_pred_type="flow")
-                guidance_scale = self.config.guidance_scale or 1.0
-                guidance_scale = torch.where(
-                    ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)),
-                    guidance_scale,
-                    1.0,
-                )
-                guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype)
-
-                if self.config.guidance_mixture_ratio is None:
-                    guided_dxt_dt = neg_dxt_dt + guidance_scale * (dxt_dt - neg_dxt_dt)
-                else:
-                    guidance_mixture_ratio = torch.where(
-                        ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)),
-                        self.config.guidance_mixture_ratio,
-                        0.0,
-                    )
-                    guidance_mixture_ratio = expand_like(guidance_mixture_ratio, x_t).to(dtype=x_t.dtype)
-                    cond_dxt_dt = self.net(x_t, t, r=t, condition=condition, fwd_pred_type="flow")
-                    guided_dxt_dt = (
-                        guidance_scale * dxt_dt
-                        + (1.0 - guidance_scale - guidance_mixture_ratio) * neg_dxt_dt
-                        + guidance_mixture_ratio * cond_dxt_dt
-                    )
-
-                self.net.train()
-                condition, keep = self._drop_condition(condition, neg_condition)
-                if keep is not None:
-                    # Same subset: a kept sample is conditional + guided, a
-                    # dropped one unconditional + unguided.
-                    dxt_dt = torch.where(expand_like(keep, dxt_dt), guided_dxt_dt, dxt_dt)
-
-        return condition, dxt_dt
-
     def single_train_step(
         self, data: Dict[str, Any], iteration: int
     ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: