xinyun/ComfyUI-KJNodes
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git synced 2026-01-29 09:37:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix sage3 for Wan

Browse Source
This commit is contained in:
kijai 2025-10-06 21:23:47 +03:00
parent b1fffd33ba
commit 08164edab3
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
nodes/model_optimization_nodes.py
View File

@ -122,10 +122,16 @@ class BaseLoaderKJ:
from sageattn3 import sageattn3_blackwell
if sage_attention == "sageattn3_per_block_mean":
def func(q, k, v, is_causal=False, attn_mask=None, **kwargs):
return sageattn3_blackwell(q, k, v, is_causal=is_causal, attn_mask=attn_mask, per_block_mean=True)
if q.shape == k.shape and q.shape == v.shape:
return sageattn3_blackwell(q, k, v, is_causal=is_causal, attn_mask=attn_mask, per_block_mean=True)
else:
return sageattn(q, k, v, is_causal=is_causal, attn_mask=attn_mask, tensor_layout="NHD")
else:
def func(q, k, v, is_causal=False, attn_mask=None, **kwargs):
return sageattn3_blackwell(q, k, v, is_causal=is_causal, attn_mask=attn_mask, per_block_mean=False)
if q.shape == k.shape and q.shape == v.shape:
return sageattn3_blackwell(q, k, v, is_causal=is_causal, attn_mask=attn_mask, per_block_mean=False)
else:
return sageattn(q, k, v, is_causal=is_causal, attn_mask=attn_mask, tensor_layout="NHD")
return func
sage_func = set_sage_func(sage_attention)