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

Update attention.py

Browse Source 
Fix https://github.com/comfyanonymous/ComfyUI/issues/8634
This commit is contained in:
Sakura-Luna 2025-06-29 16:21:57 +08:00 committed by GitHub
parent 2a0b138feb
commit 9a33651bb5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -446,6 +446,7 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
# add a heads dimension if there isn't already one # add a heads dimension if there isn't already one
if mask.ndim == 3: if mask.ndim == 3:
mask = mask.unsqueeze(1) mask = mask.unsqueeze(1)
mask = mask.contiguous()
if SDP_BATCH_LIMIT >= b: if SDP_BATCH_LIMIT >= b:
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)