jinx/ComfyUI
1
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 23:14:49 +08:00
Code Issues Activity

Fix mask issue in some attention functions.

Browse Source
This commit is contained in:
comfyanonymous
2024-11-22 02:10:09 -05:00
parent 8f0009aad0
commit 2fd9c1308a
2 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
comfy/ldm/modules/sub_quadratic_attention.py
View File

@@ -234,6 +234,8 @@ def efficient_dot_product_attention(
def get_mask_chunk(chunk_idx: int) -> Tensor:
if mask is None:
return None
if mask.shape[1] == 1:
return mask
chunk = min(query_chunk_size, q_tokens)
return mask[:,chunk_idx:chunk_idx + chunk]