Update skyreels_v2_infer/modules/transformer.py

skyreels_v2_infer/modules/transformer.py

@@ -13,7 +13,7 @@ from torch.nn.attention.flex_attention import BlockMask
 from torch.nn.attention.flex_attention import create_block_mask
 from torch.nn.attention.flex_attention import flex_attention
 
-from .attention import flash_attention
+from .attention import flash_attention, attention
 
 
 flex_attention = torch.compile(flex_attention, dynamic=False, mode="max-autotune")
@@ -160,7 +160,7 @@ class WanSelfAttention(nn.Module):
         if not self._flag_ar_attention:
             q = rope_apply(q, grid_sizes, freqs)
             k = rope_apply(k, grid_sizes, freqs)
-            x = flash_attention(q=q, k=k, v=v, window_size=self.window_size)
+            x = attention(q=q, k=k, v=v, window_size=self.window_size)
         else:
             q = rope_apply(q, grid_sizes, freqs)
             k = rope_apply(k, grid_sizes, freqs)
@@ -199,7 +199,7 @@ class WanT2VCrossAttention(WanSelfAttention):
         v = self.v(context).view(b, -1, n, d)
 
         # compute attention
-        x = flash_attention(q, k, v)
+        x = attention(q, k, v)
 
         # output
         x = x.flatten(2)