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

Update multitalk.py

Browse Source
This commit is contained in:
kijai
2025-08-17 19:33:20 +03:00
parent aa62d59d9e
commit e4f422cf99
1 changed files with 5 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
multitalk/multitalk.py
+5 -7
View File
@@ -256,13 +256,11 @@ class SingleStreamAttention(nn.Module):
N_t, N_h, N_w = shape
x_extra = None
try:
x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t)
except:
if x.shape[0] * N_t != encoder_hidden_states.shape[0]:
x_extra = x[:, -N_h * N_w:, :]
x = x[:, :-N_h * N_w, :]
N_t = N_t - 1
x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t)
x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t)
# get q for hidden_state
B, N, C = x.shape
@@ -369,13 +367,13 @@ class SingleStreamMultiAttention(SingleStreamAttention):
return super().forward(x, encoder_hidden_states, shape)
N_t, N_h, N_w = shape
x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t)
x_extra = None
if x.shape[0] != encoder_hidden_states.shape[0]:
if x.shape[0] * N_t != encoder_hidden_states.shape[0]:
x_extra = x[:, -N_h * N_w:, :]
x = x[:, :-N_h * N_w, :]
N_t = N_t - 1
x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t)
# Query projection
B, N, C = x.shape