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

Update attention.py

Browse Source
This commit is contained in:
kijai
2025-07-20 16:22:27 +03:00
parent e5b125cd89
commit 59add151eb
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wanvideo/modules/attention.py
+4 -4
View File
@@ -16,11 +16,11 @@ except ModuleNotFoundError:
try:
from sageattention import sageattn
@torch.compiler.disable()
def sageattn_func(q, k, v, attn_mask=None, dropout_p=0, is_causal=False):
def sageattn_func(q, k, v, attn_mask=None, dropout_p=0, is_causal=False, tensor_layout="HND"):
if q.dtype == torch.float32:
return sageattn(q.to(torch.float16), k.to(torch.float16), v.to(torch.float16), attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, tensor_layout="NHD").to(torch.float32)
return sageattn(q.to(torch.float16), k.to(torch.float16), v.to(torch.float16), attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, tensor_layout=tensor_layout).to(torch.float32)
else:
return sageattn(q, k, v, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, tensor_layout="NHD")
return sageattn(q, k, v, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, tensor_layout=tensor_layout)
except Exception as e:
print(f"Warning: Could not load sageattention: {str(e)}")
if isinstance(e, ModuleNotFoundError):
@@ -183,4 +183,4 @@ def attention(
elif attention_mode == 'sdpa':
return torch.nn.functional.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2).contiguous()
else:
return sageattn_func(q, k, v).contiguous()
return sageattn_func(q, k, v, tensor_layout="NHD").contiguous()