From fba7170ed01ec3d2ca42a6852864d2a078abfe13 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 19 May 2026 14:23:24 +0800 Subject: [PATCH] [feat] Thread optional r_timestep through WanTransformer3DModel.forward MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 16 +++++- .../training/distill/test_anyflow_pretrain.py | 51 +++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index a6bb45361..6b925c5f0 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/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)) diff --git a/fastvideo/tests/training/distill/test_anyflow_pretrain.py b/fastvideo/tests/training/distill/test_anyflow_pretrain.py index 285312096..f5544ffed 100644 --- a/fastvideo/tests/training/distill/test_anyflow_pretrain.py +++ b/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")