From df95c85283d7625fbdf664d0133a2e1c114ba14a Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sun, 27 Apr 2025 17:31:06 +0300 Subject: [PATCH] Fix sageattn when using fp32 --- wanvideo/modules/attention.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/wanvideo/modules/attention.py b/wanvideo/modules/attention.py index 923a9e5..f8266f8 100644 --- a/wanvideo/modules/attention.py +++ b/wanvideo/modules/attention.py @@ -17,7 +17,10 @@ try: from sageattention import sageattn @torch.compiler.disable() def sageattn_func(q, k, v, attn_mask=None, dropout_p=0, is_causal=False): - return sageattn(q, k, v, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal) + 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).to(torch.float32) + else: + return sageattn(q, k, v, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal) except Exception as e: print(f"Warning: Could not load sageattention: {str(e)}") if isinstance(e, ModuleNotFoundError):