From 59add151eb02efb64819381214dfd8fd4b2f4593 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sun, 20 Jul 2025 16:22:27 +0300 Subject: [PATCH] Update attention.py --- wanvideo/modules/attention.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/wanvideo/modules/attention.py b/wanvideo/modules/attention.py index 39a3a80..006b4da 100644 --- a/wanvideo/modules/attention.py +++ b/wanvideo/modules/attention.py @@ -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()