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")