Update hymm_sp/modules/models_audio.py

hymm_sp/modules/models_audio.py

@@ -166,39 +166,7 @@ class DoubleStreamBlock(nn.Module):
         v = torch.cat((img_v, txt_v), dim=1)
 
         # Compute attention.
-        if CPU_OFFLOAD or DISABLE_SP:
-            assert cu_seqlens_q.shape[0] == 2 * img.shape[0] + 1
-
-            q, k, v = [
-                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 = 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,
-            # )
-        img_attn, txt_attn = attn[:, :img.shape[1]], attn[:, img.shape[1]:]
-
+       
         if CPU_OFFLOAD: torch.cuda.empty_cache()
 
         # Calculate the img bloks.