xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-10-10 11:47:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix signature

Browse Source
This commit is contained in:
contentis 2024-10-22 02:58:07 -07:00 committed by GitHub
parent 92bd8fbc45
commit 713430d863
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
comfy/ldm/modules/diffusionmodules/mmdit.py
View File

@ -326,7 +326,7 @@ class SelfAttention(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
q, k, v = self.pre_attention(x) q, k, v = self.pre_attention(x)
x = optimized_attention( x = optimized_attention(
q, k, v, num_heads=self.num_heads q, k, v, heads=self.num_heads
) )
x = self.post_attention(x) x = self.post_attention(x)
return x return x
@ -529,8 +529,8 @@ class DismantledBlock(nn.Module):
assert not self.pre_only assert not self.pre_only
qkv, intermediates = self.pre_attention(x, c) qkv, intermediates = self.pre_attention(x, c)
attn = optimized_attention( attn = optimized_attention(
qkv, q[0], k[1], v[2],
num_heads=self.attn.num_heads, heads=self.attn.num_heads,
) )
return self.post_attention(attn, *intermediates) return self.post_attention(attn, *intermediates)
@ -555,8 +555,8 @@ def _block_mixing(context, x, context_block, x_block, c):
qkv = tuple(o) qkv = tuple(o)
attn = optimized_attention( attn = optimized_attention(
qkv, q[0], k[1], v[2],
num_heads=x_block.attn.num_heads, heads=x_block.attn.num_heads,
) )
context_attn, x_attn = ( context_attn, x_attn = (
attn[:, : context_qkv[0].shape[1]], attn[:, : context_qkv[0].shape[1]],
@ -640,7 +640,7 @@ class SelfAttentionContext(nn.Module):
def forward(self, x): def forward(self, x):
qkv = self.qkv(x) qkv = self.qkv(x)
q, k, v = split_qkv(qkv, self.dim_head) q, k, v = split_qkv(qkv, self.dim_head)
x = optimized_attention((q.reshape(q.shape[0], q.shape[1], -1), k, v), self.heads) x = optimized_attention(q.reshape(q.shape[0], q.shape[1], -1), k, v, heads=self.heads)
return self.proj(x) return self.proj(x)
class ContextProcessorBlock(nn.Module): class ContextProcessorBlock(nn.Module):