[feat] Thread optional r_timestep through WanTransformer3DModel.forward

The forward signature now accepts r_timestep as an explicit kwarg (not
silently swallowed by **kwargs), gets flattened in parallel with timestep
when the input is 2-D (Wan 2.2 ti2v case), and is forwarded to the
condition_embedder.

The constructor now propagates the four arch_config flags
(r_embedder/r_embedder_fusion/r_embedder_gate_value/r_embedder_deltatime_type)
into the WanTimeTextImageEmbedding so a YAML can opt into AnyFlow's
dual-timestep conditioning through training.dit_config overrides.

Defaults preserve byte-identity with the legacy single-timestep path —
when r_timestep is omitted the delta_embedder branch is skipped entirely
(test_embedder_enabled_without_r_timestep_is_bit_identical_to_legacy).

fastvideo/models/dits/wanvideo.py

@@ -644,6 +644,10 @@ class WanTransformer3DModel(BaseDiT):
             time_freq_dim=config.freq_dim,
             text_embed_dim=config.text_dim,
             image_embed_dim=config.image_dim,
+            r_embedder=config.r_embedder,
+            r_embedder_fusion=config.r_embedder_fusion,
+            r_embedder_gate_value=config.r_embedder_gate_value,
+            r_embedder_deltatime_type=config.r_embedder_deltatime_type,
         )
 
         # 3. Transformer blocks
@@ -685,6 +689,7 @@ class WanTransformer3DModel(BaseDiT):
                 encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
                 | None = None,
                 guidance=None,
+                r_timestep: torch.Tensor | None = None,
                 **kwargs) -> torch.Tensor:
         orig_dtype = hidden_states.dtype
         if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
@@ -733,8 +738,17 @@ class WanTransformer3DModel(BaseDiT):
         else:
             ts_seq_len = None
 
+        # AnyFlow dual-timestep — match timestep's flattening so embedder
+        # sees aligned shapes.
+        if r_timestep is not None and r_timestep.dim() == 2:
+            r_timestep = r_timestep.flatten()
+
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len)
+            timestep,
+            encoder_hidden_states,
+            encoder_hidden_states_image,
+            timestep_seq_len=ts_seq_len,
+            r_timestep=r_timestep)
         if ts_seq_len is not None:
             # batch_size, seq_len, 6, inner_dim
             timestep_proj = timestep_proj.unflatten(2, (6, -1))

fastvideo/tests/training/distill/test_anyflow_pretrain.py

@@ -168,3 +168,54 @@ def test_embedder_gate_not_in_state_dict() -> None:
     emb = _make_embedder(r_embedder=True, fusion="gated", gate=0.25)
     keys = list(emb.state_dict().keys())
     assert not any("_r_embedder_gate" in k for k in keys)
+
+
+# ---------------------------------------------------------------------------
+# Task 3: WanTransformer3DModel threads r_timestep through.
+# ---------------------------------------------------------------------------
+
+
+def test_wan_transformer_forward_signature_has_r_timestep() -> None:
+    """The forward signature must declare r_timestep explicitly (not
+    swallowed by **kwargs) so callers and type checkers can see it."""
+    import inspect
+
+    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
+
+    sig = inspect.signature(WanTransformer3DModel.forward)
+    assert "r_timestep" in sig.parameters
+    param = sig.parameters["r_timestep"]
+    assert param.default is None
+
+
+def test_wan_transformer_init_propagates_r_embedder_config() -> None:
+    """When the arch config sets r_embedder=True the WanTransformer3DModel
+    constructor must instantiate the embedder with the delta path active.
+
+    We avoid full WanTransformer3DModel instantiation (which requires
+    distributed init) by reading the source's __init__ to confirm it
+    forwards the four arch flags to WanTimeTextImageEmbedding.
+    """
+    import inspect
+
+    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
+
+    src = inspect.getsource(WanTransformer3DModel.__init__)
+    # All four arch config fields must be passed to WanTimeTextImageEmbedding.
+    assert "r_embedder=config.r_embedder" in src
+    assert "r_embedder_fusion=config.r_embedder_fusion" in src
+    assert "r_embedder_gate_value=config.r_embedder_gate_value" in src
+    assert "r_embedder_deltatime_type=config.r_embedder_deltatime_type" in src
+
+
+def test_wan_transformer_forward_threads_r_timestep_to_embedder() -> None:
+    """The forward must pass r_timestep into the embedder call (verified via
+    source inspection to avoid heavyweight distributed bring-up)."""
+    import inspect
+
+    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
+
+    src = inspect.getsource(WanTransformer3DModel.forward)
+    assert "r_timestep=r_timestep" in src, (
+        "WanTransformer3DModel.forward must forward r_timestep into "
+        "self.condition_embedder")