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

Allow Phantom to work with MultiTalk

Browse Source
This commit is contained in:
kijai
2025-08-17 16:46:50 +03:00
parent 689797bdb7
commit 1d7e0848e0
1 changed files with 19 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
multitalk/multitalk.py
+19 -4
View File
@@ -253,8 +253,14 @@ class SingleStreamAttention(nn.Module):
self.attention_mode = attention_mode
def forward(self, x: torch.Tensor, encoder_hidden_states: torch.Tensor, shape=None, enable_sp=False, kv_seq=None) -> torch.Tensor:
N_t, N_h, N_w = shape
x_extra = None
if x.shape[0] != 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)
# get q for hidden_state
@@ -282,17 +288,18 @@ class SingleStreamAttention(nn.Module):
encoder_v.transpose(1, 2),
attention_mode=self.attention_mode
)
#x = torch.nn.functional.scaled_dot_product_attention(
# q, encoder_k, encoder_v, attn_mask=None, is_causal=False, dropout_p=0.0)
# linear transform
x_output_shape = (B, N, C)
#x = x.transpose(1, 2)
x = x.reshape(x_output_shape)
x = self.proj(x)
x = self.proj_drop(x)
x = self.proj_drop(x)
x = rearrange(x, "(B N_t) S C -> B (N_t S) C", N_t=N_t)
if x_extra is not None:
x = torch.cat([x, torch.zeros_like(x_extra)], dim=1)
return x
@@ -363,6 +370,12 @@ class SingleStreamMultiAttention(SingleStreamAttention):
N_t, _, _ = 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]:
x_extra = x[:, -N_h * N_w:, :]
x = x[:, :-N_h * N_w, :]
N_t = N_t - 1
# Query projection
B, N, C = x.shape
q = self.q_linear(x)
@@ -473,5 +486,7 @@ class SingleStreamMultiAttention(SingleStreamAttention):
# Restore original layout
x = rearrange(x, "(B N_t) S C -> B (N_t S) C", N_t=N_t)
if x_extra is not None:
x = torch.cat([x, torch.zeros_like(x_extra)], dim=1)
return x