Merge 5f7e7ed289 into 827906c30f

Adding support for cards that aren't Ampere architecture

README.md

@@ -36,6 +36,7 @@ In this repository, we present **Wan2.1**, a comprehensive and open suite of vid
 
 ## Community Works
 If your work has improved **Wan2.1** and you would like more people to see it, please inform us.
+- [ATI](https://github.com/bytedance/ATI), built on **Wan2.1-I2V-14B**, is a trajectory-based motion-control framework that unifies object, local, and camera movements in video generation. Refer to [their website](https://anytraj.github.io/) for more examples.
 - [Phantom](https://github.com/Phantom-video/Phantom) has developed a unified video generation framework for single and multi-subject references based on both **Wan2.1-T2V-1.3B** and **Wan2.1-T2V-14B**. Please refer to [their examples](https://github.com/Phantom-video/Phantom).
 - [UniAnimate-DiT](https://github.com/ali-vilab/UniAnimate-DiT), based on **Wan2.1-14B-I2V**, has trained a Human image animation model and has open-sourced the inference and training code. Feel free to enjoy it!
 - [CFG-Zero](https://github.com/WeichenFan/CFG-Zero-star) enhances **Wan2.1** (covering both T2V and I2V models) from the perspective of CFG.

wan/modules/attention.py

@@ -79,6 +79,7 @@ def flash_attention(
         k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
         v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
 
+    try:
         q = q.to(v.dtype)
         k = k.to(v.dtype)
     
@@ -126,6 +127,53 @@ def flash_attention(
                 window_size=window_size,
                 deterministic=deterministic).unflatten(0, (b, lq))
     
+    except RuntimeError as e:
+        if "FlashAttention only supports Ampere GPUs or newer" in str(e):
+            #for cards like 2080ti that aren't Ampere structure
+            from torch import nn
+            import torch.nn.functional as F
+        
+            q = q.to(half(k).dtype)
+
+            # 转置维度，保证形状为 [B, N, L, C]
+            q = q.view(b, lq, q.size(1), q.size(2)).transpose(1, 2)
+            k = k.view(b, lk, k.size(1), k.size(2)).transpose(1, 2)
+            v = v.view(b, lk, v.size(1), v.size(2)).transpose(1, 2)
+
+            # 计算注意力
+            # 注意：确保 Q、K、V 的形状为 [B, N, L, C]
+            # 设置默认缩放因子
+            if softmax_scale is None:
+                softmax_scale = 1.0 / q.size(-1) ** 0.5
+
+            # 如果 q_scale 存在，则应用缩放
+            if q_scale is not None:
+                q = q * q_scale
+
+            # 创建掩码
+            if causal:
+                attn_mask = torch.triu(torch.full((q.size(2), k.size(2)), -torch.inf), diagonal=1).to(q.device)
+            else:
+                attn_mask = None
+
+            # 计算注意力
+            # 使用 scaled_dot_product_attention
+            x = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=attn_mask,
+                dropout_p=dropout_p,
+                is_causal=causal,
+            )
+
+            # 转换回原形状 [B, L, N, C]
+            x = x.transpose(1, 2).contiguous()
+
+            # 对输出应用 Dropout
+            dropout = nn.Dropout(dropout_p)
+            x = dropout(x)            
+        else:
+            raise
+            
     # output
     return x.type(out_dtype)