Update hymm_sp/modules/models_audio.py

hymm_sp/modules/models_audio.py

@@ -7,7 +7,7 @@ import torch.nn.functional as F
 from diffusers.models import ModelMixin
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 
-from flash_attn.flash_attn_interface import flash_attn_varlen_func
+#from flash_attn.flash_attn_interface import flash_attn_varlen_func
 
 
 
@@ -173,29 +173,30 @@ class DoubleStreamBlock(nn.Module):
                 x.view(x.shape[0] * x.shape[1], *x.shape[2:])
                 for x in [q, k, v]
             ]
+            attn = None
             
-            attn = flash_attn_varlen_func(
-                q,
-                k,
-                v,
-                cu_seqlens_q,
-                cu_seqlens_kv,
-                max_seqlen_q,
-                max_seqlen_kv,
-            )
+            # attn = flash_attn_varlen_func(
+            #     q,
+            #     k,
+            #     v,
+            #     cu_seqlens_q,
+            #     cu_seqlens_kv,
+            #     max_seqlen_q,
+            #     max_seqlen_kv,
+            # )
             attn = attn.view(img_k.shape[0], max_seqlen_q, -1).contiguous()
         else:
-                attn, _ = parallel_attention(
-                (img_q, txt_q),
-                (img_k, txt_k),
-                (img_v, txt_v),
-                img_q_len=img_q.shape[1],
-                img_kv_len=img_k.shape[1],
-                cu_seqlens_q=cu_seqlens_q,
-                cu_seqlens_kv=cu_seqlens_kv,
-                max_seqlen_q=max_seqlen_q,
-                max_seqlen_kv=max_seqlen_kv,
-            )
+            #     attn, _ = parallel_attention(
+            #     (img_q, txt_q),
+            #     (img_k, txt_k),
+            #     (img_v, txt_v),
+            #     img_q_len=img_q.shape[1],
+            #     img_kv_len=img_k.shape[1],
+            #     cu_seqlens_q=cu_seqlens_q,
+            #     cu_seqlens_kv=cu_seqlens_kv,
+            #     max_seqlen_q=max_seqlen_q,
+            #     max_seqlen_kv=max_seqlen_kv,
+            # )
         img_attn, txt_attn = attn[:, :img.shape[1]], attn[:, img.shape[1]:]
 
         if CPU_OFFLOAD: torch.cuda.empty_cache()