Fix prompt splitting

wanvideo/modules/model.py

@@ -522,14 +522,6 @@ class WanSelfAttention(nn.Module):
 
 
     def forward_split(self, q, k, v, seq_lens, grid_sizes, seq_chunks):
-        r"""
-        Args:
-            x(Tensor): Shape [B, L, num_heads, C / num_heads]
-            seq_lens(Tensor): Shape [B]
-            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
-            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
-        """
-
         # Split by frames if multiple prompts are provided
         frames, height, width = grid_sizes[0]
         tokens_per_frame = height * width
@@ -1284,7 +1276,7 @@ class WanAttentionBlock(nn.Module):
             x_ip = x_ip.addcmul(y_ip, gate_mlp_ip)
         return x, x_ip, lynx_ref_feature, x_ovi
 
-    @torch.compiler.disable()
+
     def split_cross_attn_ffn(self, x, context, shift_mlp, scale_mlp, gate_mlp, clip_embed=None, grid_sizes=None):
         # Get number of prompts
         num_prompts = context.shape[0]
@@ -1333,9 +1325,9 @@ class WanAttentionBlock(nn.Module):
 
         # Continue with FFN
         x = x + x_combined
-        y = self.ffn_chunked(x, shift_mlp, scale_mlp)
-        x = x.addcmul(y, gate_mlp)
-        return x
+        mod_x = torch.addcmul(shift_mlp, self.norm2(x.to(shift_mlp.dtype)), 1 + scale_mlp)
+        y = self.ffn_chunked(mod_x, num_chunks=1)
+        return x.addcmul(y, gate_mlp)
 
 class VaceWanAttentionBlock(WanAttentionBlock):
     def __init__(