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

Fix sageattn_varlen

Browse Source
This commit is contained in:
kijai
2025-12-01 01:11:06 +02:00
parent e5be3e5263
commit 196d39695f
1 changed files with 13 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wanvideo/modules/attention.py
+3 -3
View File
@@ -51,10 +51,10 @@ except Exception as e:
try:
from sageattention import sageattn_varlen
from typing import List
@torch.library.custom_op("wanvideo::sageattn_varlen", mutates_args=())
def sageattn_varlen_func(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_lens: list, k_lens: list, max_seqlen_q: int, max_seqlen_k: int, dropout_p: float = 0.0, is_causal: bool = False
) -> torch.Tensor:
def sageattn_varlen_func(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_lens: List[int], k_lens: List[int], max_seqlen_q: int, max_seqlen_k: int, dropout_p: float = 0.0, is_causal: bool = False) -> torch.Tensor:
cu_seqlens_q = torch.tensor([0] + list(torch.cumsum(torch.tensor(q_lens), dim=0)), device=q.device, dtype=torch.int32)
cu_seqlens_k = torch.tensor([0] + list(torch.cumsum(torch.tensor(k_lens), dim=0)), device=q.device, dtype=torch.int32)
if not (q.dtype == k.dtype == v.dtype):
@@ -68,7 +68,7 @@ try:
def _(q, k, v, q_lens, k_lens, max_seqlen_q, max_seqlen_k, dropout_p=0.0, is_causal=False):
# Return tensor with same shape as q
return q.clone()
except:
except Exception as e:
sageattn_varlen_func = None
try: