wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-09 07:57:30 +08:00
Code Issues Actions 19 Packages Projects Releases Wiki Activity

Fix.

Browse Source
This commit is contained in:
comfyanonymous 2026-07-04 20:53:17 -04:00 committed by GitHub
parent 501b808481
commit 2131e630e0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
comfy/ldm/modules/attention.py
View File

@ -643,9 +643,6 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
skip_output_reshape=skip_output_reshape,
**kwargs
)
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, H, -3)
q_s, k_s, v_s = q, k, v
N = q.shape[2]
dim_head = D
else:
@ -671,7 +668,13 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
**kwargs
)
if not skip_reshape:
if skip_reshape:
q_s = q
if kwargs.get("enable_gqa", False):
k_s, v_s = _repeat_kv_for_gqa(k, v, H, -3)
else:
k_s, v_s = k, v
else:
q_s, k_s, v_s = _reshape_qkv_to_heads(q, k, v, B, heads, dim_head, kwargs.get("enable_gqa", False))
q_s, k_s, v_s = map(lambda t: t.permute(0, 2, 1, 3).contiguous(), (q_s, k_s, v_s))
B, H, L, D = q_s.shape