Address comments

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

Author: juliusberner

fastgen/configs/experiments/WanT2V/config_anyflow.py

@@ -4,8 +4,8 @@
 """AnyFlow flow-map pretrain config on Wan-1.3B T2V (paper Stage 1).
 
 AnyFlow's pretrain objective is MeanFlow's with a fixed ``beta08`` per-timestep
-weighting, a finite-difference JVP, shifted timestep sampling, and a 
-``consistency_ratio`` fraction of the batch pinned to ``r = 0`` — so this config 
+weighting, a finite-difference JVP, shifted timestep sampling, and a
+``consistency_ratio`` fraction of the batch pinned to ``r = 0`` — so this config
 runs``MeanFlowModel`` directly. The values below mirror the reference recipe
 ``train_wan1b_student_shift5_81f_480p_lr5e-5_6k_b32.yml``.
 Known deviations from the reference: full-rank fine-tuning instead of the paper's
@@ -103,7 +103,7 @@ def create_config():
     config.model.sample_t_cfg.t_list = [1.0, 0.9375, 0.8333333333333334, 0.625, 0.0]
 
     # ------ data / trainer ------
-    config.dataloader_train = VideoLoaderConfig
+    config.dataloader_train = copy.deepcopy(VideoLoaderConfig)
     config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8)
     config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1
     config.dataloader_train.batch_size = 1

fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py

@@ -137,7 +137,7 @@ def create_config():
     config.trainer.callbacks.ema.start_iter = 6399
 
     # ------ data / trainer ------
-    config.dataloader_train = VideoLoaderConfig
+    config.dataloader_train = copy.deepcopy(VideoLoaderConfig)
     config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8)
     config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1
     config.dataloader_train.batch_size = 1

fastgen/configs/methods/config_anyflow.py

@@ -11,8 +11,6 @@
 
 from typing import List, Optional
 
-from typing import Optional
-
 import attrs
 from omegaconf import DictConfig
 
@@ -61,6 +59,11 @@ class ModelConfig(DMD2ModelConfig):
 
     # Prediction-side guidance fusion for the co-trained flow-map loss; see
     # `config_mean_flow.ModelConfig`.
+    #
+    # Deviates from the AnyFlow reference, which divides dF/dt by g for the whole
+    # batch. We divide only on samples that kept their condition: a dropped sample's
+    # fused prediction is plain `u_uncond`, so an ungated 1/g would regress it onto
+    # a different fixed point.
     guidance_fuse_scale: Optional[float] = None
 
     # Text dropout for the co-trained flow-map loss (reference drop_text_ratio).
@@ -70,7 +73,7 @@ class ModelConfig(DMD2ModelConfig):
     # 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 
+    # MeanFlow's target-side guidance knobs. AnyFlow typically guides on the
     # PREDICTION side (`guidance_fuse_scale`), so these stay at their
     # no-op defaults
     guidance_mixture_ratio: Optional[float] = None

fastgen/configs/methods/config_mean_flow.py

@@ -83,7 +83,7 @@ class LossConfig:
     # loss type (choice between l2 and opt_grad)
     loss_type: str = "opt_grad"
     # optional fixed per-timestep loss weighting evaluated as a function of t
-    # ("beta08", "gaussian", "uniform"). Multiplies the adaptive norm_method 
+    # ("beta08", "gaussian", "uniform"). Multiplies the adaptive norm_method
     # weight above; None disables it.
     weight_type: Optional[str] = None
     # rebalance the flow-map / consistency (r < t) sample losses to the global

fastgen/methods/consistency_model/mean_flow.py

@@ -59,9 +59,9 @@ class FlowMapLossMixin:
     training objective; distribution-matching methods can co-train it alongside
     their own objective.
 
-    The host class must provide ``net``, ``device``, ``config``,
-    ``precision_amp``, a ``_get_velocity`` implementation, and call
-    ``_init_flow_map_loss`` from its ``__init__``.
+    The host class must provide ``net``, ``device``, ``config`` and ``precision_amp``
+    (plus ``teacher`` when ``loss_config.use_cd``), and call ``_init_flow_map_loss``
+    from its ``__init__``.
     """
 
     def _init_flow_map_loss(self) -> None:
@@ -70,7 +70,8 @@ def _init_flow_map_loss(self) -> None:
         self.sample_r_cfg = self.config.sample_r_cfg
         self.loss_config = self.config.loss_config
 
-        # Precision for JVP
+        # Drop the JVP autocast when it matches the outer one, since a nested region in the
+        # same dtype is a no-op.
         if self.config.precision_amp_jvp is None or self.config.precision_amp_jvp == self.precision_amp:
             self.precision_amp_jvp = None
         else:
@@ -89,28 +90,29 @@ def _init_flow_map_loss(self) -> None:
             grid = shift * grid / (1 + (shift - 1) * grid)
             self._timestep_weight_scale = float(num_steps / self._timestep_weight_raw(grid).sum())
 
-    def _drop_condition(self, condition: Any, neg_condition: Any) -> Tuple[Any, Optional[torch.Tensor]]:
+    def _drop_condition(
+        self, condition: Any, neg_condition: Any, batch_size: int, device: torch.device
+    ) -> Tuple[Any, torch.Tensor]:
         """Replace the condition with neg_condition for a per-sample subset.
 
-        Returns ``(condition, keep)``; ``keep`` is the ``[B]`` bool mask (None if
-        no dropout), so callers can reuse the same subset. Keys in
-        ``cond_keys_no_dropout`` are never replaced.
+        Returns ``(condition, keep)``; ``keep`` is the ``[B]`` bool mask of the
+        samples that stayed conditional, so callers can reuse the same subset.
 
         ``deterministic_buckets`` decides whether an index carries bucket
         information: if so the buckets are cut on the GLOBAL index and an
         index-based rule would only hit flow matching on rank 0, so draw per
         sample; otherwise drop the first ``num_to_drop``.
         """
+        # Dropout disabled, or no negative condition to swap in: every sample
+        # stays conditional.
         if self.config.cond_dropout_prob is None or neg_condition is None:
-            return condition, None
+            return condition, torch.ones(batch_size, dtype=torch.bool, device=device)
 
-        ref = neg_condition if isinstance(neg_condition, torch.Tensor) else next(iter(neg_condition.values()))
-        batch_size = ref.shape[0]
         if self.sample_t_cfg.deterministic_buckets:
-            keep = torch.rand(batch_size, device=ref.device) >= self.config.cond_dropout_prob
+            keep = torch.rand(batch_size, device=device) >= self.config.cond_dropout_prob
         else:
-            num_to_drop = (torch.rand(batch_size, device=ref.device) < self.config.cond_dropout_prob).sum()
-            keep = torch.arange(batch_size, device=ref.device) >= num_to_drop
+            num_to_drop = (torch.rand(batch_size, device=device) < self.config.cond_dropout_prob).sum()
+            keep = torch.arange(batch_size, device=device) >= num_to_drop
 
         if isinstance(condition, torch.Tensor):
             return torch.where(expand_like(keep, condition), condition, neg_condition), keep
@@ -133,8 +135,9 @@ def _get_velocity(
         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.
+    ) -> Tuple[Any, torch.Tensor, torch.Tensor]:
+        """Regression target for the flow-map loss, the condition it was built from,
+        and the ``[B]`` mask of samples that stayed conditional.
 
         Two independent choices:
 
@@ -148,7 +151,9 @@ def _get_velocity(
         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)
+            condition, keep = self._drop_condition(condition, neg_condition, x.shape[0], x.device)
+        else:
+            keep = torch.ones(x.shape[0], dtype=torch.bool, device=x.device)
 
         x_t = self.net.noise_scheduler.forward_process(x, z, t)
 
@@ -201,13 +206,12 @@ def _get_velocity(
                     )
 
                 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)
+                condition, keep = self._drop_condition(condition, neg_condition, x_t.shape[0], x_t.device)
+                # 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
+        return condition, dxt_dt, keep
 
     def _estimate_jvp_finite_difference(
         self,
@@ -573,7 +577,7 @@ def _compute_mf_loss(
         z = torch.randn_like(real_data)
         x_t = self.net.noise_scheduler.forward_process(real_data, z, t)
 
-        condition, dxt_dt = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition)
+        condition, dxt_dt, keep = 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)
@@ -594,9 +598,10 @@ def _compute_mf_loss(
             # 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
+            # giving (u_cond + (g - 1) * u_uncond) / g; dF/dt is then the finite
+            # difference over g on conditional samples, with the unconditional
+            # derivative dropped.
+            u_theta_jvp = torch.where(expand_like(keep, u_theta_jvp), u_theta_jvp / guidance_fuse_scale, u_theta_jvp)
             with torch.no_grad():
                 u_uncond = self.net(x_t, t, r=r, condition=neg_condition, fwd_pred_type="flow")
             u_theta = (u_theta + (guidance_fuse_scale - 1.0) * u_uncond) / guidance_fuse_scale

fastgen/methods/distribution_matching/anyflow.py

@@ -38,8 +38,7 @@ class AnyFlowModel(FlowMapLossMixin, DMD2Model):
     """AnyFlow on-policy stage: DMD2 with a flow-map rollout student.
 
     ``FlowMapLossMixin`` supplies the co-trained Stage-1 objective and the
-    flow-map validation sample loop, which integrates with ``r = t_next`` —
-    FastGenModel's default x0-prediction loop never passes ``r``.
+    flow-map validation sample loop, which integrates with ``r = t_next``.
     """
 
     def __init__(self, config: ModelConfig):
@@ -133,7 +132,7 @@ def gen_data_from_net(
         grad_step = self._broadcast_choice(num_steps)
         t_list = self._rollout_t_list(num_steps)
 
-        # The leading jump exists only for grad_step > 0 and the trailing one 
+        # The leading jump exists only for grad_step > 0 and the trailing one
         # only for grad_step + 1 < num_steps.
         seg_t = [t_list[0]] if grad_step > 0 else []
         seg_t += [t_list[grad_step], t_list[grad_step + 1]]

fastgen/networks/Wan/utils.py

@@ -129,6 +129,13 @@ def remap_anyflow_keys(state_dict: Mapping[str, Any]) -> Mapping[str, Any]:
         # [transformer.]condition_embedder.delta_embedder.linear_1.weight
         #   -> [transformer.]r_embedder.time_embedder.linear_1.weight
         prefix, _, suffix = k.partition(delta_marker)
-        new_sd[f"{prefix}r_embedder.time_embedder.{suffix}"] = new_sd.pop(k)
+        target = f"{prefix}r_embedder.time_embedder.{suffix}"
+        if target in new_sd:
+            raise ValueError(
+                f"remap_anyflow_keys: rewriting {k!r} would overwrite the existing {target!r}. "
+                "This checkpoint carries both the AnyFlow and the FastGen r-pathway layouts; "
+                "drop one of them before loading."
+            )
+        new_sd[target] = new_sd.pop(k)
     logger.info(f"remap_anyflow_keys: rewrote {len(delta_keys)} delta_embedder tensors into r_embedder.")
     return new_sd

tests/test_anyflowmodel.py

@@ -343,8 +343,10 @@ def counting_forward(*args, **kwargs):
         return orig_forward(*args, **kwargs)
 
     model.net.forward = counting_forward
-    gen = model.gen_data_from_net(input_student, t_student, condition=cond)
-    model.net.forward = orig_forward
+    try:
+        gen = model.gen_data_from_net(input_student, t_student, condition=cond)
+    finally:
+        model.net.forward = orig_forward
 
     assert len(calls) <= 3, f"rollout must compress to <= 3 forwards, got {len(calls)}"
     assert gen.requires_grad

tests/test_meanflowmodel.py

@@ -120,3 +120,137 @@ def test_single_train_step_update_fp32_jvp():
     assert "mf_loss" in loss_map
     assert "gen_rand" in outputs
     assert isinstance(outputs["gen_rand"], Callable)
+
+
+@pytest.mark.parametrize("cond_dropout_prob, expect_guided", [(None, True), (0.0, True), (1.0, False)])
+def test_target_side_guidance_applies_when_no_cond_dropout(cond_dropout_prob, expect_guided):
+    """`guidance_scale` must reach the target for every conditional sample.
+
+    No dropout means every sample is conditional, so it must guide exactly like
+    `p=0.0`. It used not to: `_drop_condition` returned `keep=None` when
+    `cond_dropout_prob is None` and the caller skipped the update entirely --
+    computing `guided_dxt_dt` at the cost of an extra forward pass and then
+    discarding it, so `guidance_scale` was silently inert on that path.
+    (Pre-existing: `_mix_condition` returned early on `cond_dropout_prob is None`
+    before the AnyFlow work.) `_drop_condition` now always returns a mask -- all-True
+    here -- so the caller has no special case left to forget.
+
+    The net's forward is stubbed to depend only on the condition, so the guided
+    velocity is exact rather than initialization-dependent -- EDM's `SongUNet`
+    zero-inits `out_conv`, which would otherwise make the cond and uncond passes
+    bitwise identical and guidance a provable no-op.
+    """
+    config = create_config()
+    instance = config.model
+    opts = ["-", "img_resolution=2", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True"]
+    instance.net = override_config_with_opts(instance.net, opts)
+    instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    instance.precision = "float32"
+    instance.pretrained_model_path = ""
+    instance.input_shape = [3, 2, 2]
+    instance.cond_dropout_prob = cond_dropout_prob
+    instance.guidance_scale = 2.0
+    instance.guidance_fuse_scale = None
+    model = MeanFlowModel(instance)
+    model.on_train_begin()
+
+    batch_size = 4
+    real = torch.randn(batch_size, 3, 2, 2, device=model.device, dtype=torch.float32)
+    z = torch.randn_like(real)
+    t = torch.full((batch_size,), 0.5, device=model.device, dtype=model.net.noise_scheduler.t_precision)
+    condition = torch.nn.functional.one_hot(torch.arange(batch_size) % 10, num_classes=10)
+    condition = condition.to(model.device, torch.float32)
+    neg_condition = torch.zeros(batch_size, 10, device=model.device, dtype=torch.float32)
+
+    # Depends ONLY on the condition: the all-zero neg pass returns exactly 0.
+    def condition_only_forward(x_t, t, **kwargs):
+        val = kwargs["condition"].sum(dim=1).reshape(-1, 1, 1, 1).to(x_t.dtype)
+        return torch.ones_like(x_t) * val
+
+    orig_forward = model.net.forward
+    model.net.forward = condition_only_forward
+    try:
+        cond_out, dxt_dt, _ = model._get_velocity(real, z, t, condition=condition, neg_condition=neg_condition)
+    finally:
+        model.net.forward = orig_forward
+
+    # neg_dxt_dt == 0, so guided == neg + scale * (plain - neg) == 2 * plain.
+    plain = model.net.noise_scheduler.cond_velocity(x=real, eps=z, t=t)
+    dropped = (cond_out == neg_condition).all(dim=1)
+    assert bool(dropped.all()) is not expect_guided
+
+    expected = 2.0 * plain if expect_guided else plain
+    assert torch.allclose(dxt_dt, expected), (dxt_dt - expected).abs().max()
+
+
+def test_fused_jvp_scaling_is_gated_on_kept_samples():
+    """Under `guidance_fuse_scale`, dF/dt must be divided by g only for samples that
+    stayed conditional.
+
+    A dropped sample's fused prediction collapses to plain `u_uncond` (the fusion
+    self-cancels), so scaling its derivative too would regress it onto
+    `v - (t - r) * d(u_uncond)/dt / g` instead of the unconditional MeanFlow identity.
+    The AnyFlow reference scales the whole batch
+    (`compute_central_difference(..., guidance)` in
+    `far/trainers/trainer_wan_anyflow_pretrain.py`); we gate on `keep`.
+
+    The net is stubbed to ignore `condition`, so dropping it changes nothing about the
+    raw derivative -- making the 1/g gate the ONLY difference between the two runs.
+    """
+    g = 3.0
+    config = create_config()
+    instance = config.model
+    opts = ["-", "img_resolution=2", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True"]
+    instance.net = override_config_with_opts(instance.net, opts)
+    instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    instance.precision = "float32"
+    instance.pretrained_model_path = ""
+    instance.input_shape = [3, 2, 2]
+    instance.guidance_fuse_scale = g
+    instance.loss_config.use_jvp_finite_diff = True
+    instance.sample_t_cfg.deterministic_buckets = False
+    model = MeanFlowModel(instance)
+    model.on_train_begin()
+
+    batch_size = 4
+    real = torch.randn(batch_size, 3, 2, 2, device=model.device, dtype=torch.float32)
+    t_prec = model.net.noise_scheduler.t_precision
+    t = torch.full((batch_size,), 0.6, device=model.device, dtype=t_prec)
+    r = torch.full((batch_size,), 0.2, device=model.device, dtype=t_prec)
+    condition = torch.nn.functional.one_hot(torch.arange(batch_size) % 10, num_classes=10)
+    condition = condition.to(model.device, torch.float32)
+    neg_condition = torch.zeros(batch_size, 10, device=model.device, dtype=torch.float32)
+
+    # `+ 0.0 * param.sum()` leaves the value untouched but ties the output to the
+    # autograd graph: `_mf_pred_to_loss` asserts
+    # `u_theta.requires_grad is torch.is_grad_enabled()`. The JVP runs under `_jvp`'s
+    # `@torch.no_grad()`, so `u_theta_jvp` stays grad-free as that code also asserts.
+    param = next(model.net.parameters())
+
+    def condition_independent_forward(x_t, t, **kwargs):
+        tt = t.reshape(-1, *([1] * (x_t.ndim - 1))).to(x_t.dtype)
+        return x_t * (1.0 + tt) + 0.0 * param.sum().to(x_t.dtype)
+
+    orig_forward = model.net.forward
+
+    def run(cond_dropout_prob):
+        model.config.cond_dropout_prob = cond_dropout_prob
+        model.net.forward = condition_independent_forward
+        try:
+            torch.manual_seed(11)
+            return model._compute_mf_loss(
+                real_data=real,
+                t=t,
+                r=r,
+                iteration=0,
+                condition=condition,
+                neg_condition=neg_condition,
+            )[2]
+        finally:
+            model.net.forward = orig_forward
+
+    jvp_kept = run(0.0)  # every sample conditional -> every derivative divided by g
+    jvp_dropped = run(1.0)  # every sample unconditional -> none divided
+
+    assert torch.allclose(jvp_dropped, g * jvp_kept, atol=1e-5, rtol=1e-4), (jvp_dropped - g * jvp_kept).abs().max()
+    assert not torch.allclose(jvp_dropped, jvp_kept, atol=1e-6)