nodes/kijai-ComfyUI-WanVideoWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix prompt splitting

Browse Source
This commit is contained in:
kijai
2025-12-07 18:34:21 +02:00
parent 123c9ca312
commit 2369cdbbe9
1 changed files with 21 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wanvideo/modules/model.py
+4 -12
View File
@@ -522,14 +522,6 @@ class WanSelfAttention(nn.Module):
def forward_split(self, q, k, v, seq_lens, grid_sizes, seq_chunks): 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 # Split by frames if multiple prompts are provided
frames, height, width = grid_sizes[0] frames, height, width = grid_sizes[0]
tokens_per_frame = height * width tokens_per_frame = height * width
@@ -1284,7 +1276,7 @@ class WanAttentionBlock(nn.Module):
x_ip = x_ip.addcmul(y_ip, gate_mlp_ip) x_ip = x_ip.addcmul(y_ip, gate_mlp_ip)
return x, x_ip, lynx_ref_feature, x_ovi 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): def split_cross_attn_ffn(self, x, context, shift_mlp, scale_mlp, gate_mlp, clip_embed=None, grid_sizes=None):
# Get number of prompts # Get number of prompts
num_prompts = context.shape[0] num_prompts = context.shape[0]
@@ -1333,9 +1325,9 @@ class WanAttentionBlock(nn.Module):
# Continue with FFN # Continue with FFN
x = x + x_combined x = x + x_combined
y = self.ffn_chunked(x, shift_mlp, scale_mlp) mod_x = torch.addcmul(shift_mlp, self.norm2(x.to(shift_mlp.dtype)), 1 + scale_mlp)
x = x.addcmul(y, gate_mlp) y = self.ffn_chunked(mod_x, num_chunks=1)
return x return x.addcmul(y, gate_mlp)
class VaceWanAttentionBlock(WanAttentionBlock): class VaceWanAttentionBlock(WanAttentionBlock):
def __init__( def __init__(