Update mvdiffusion/models_unclip/attn_processors.py

mvdiffusion/models_unclip/attn_processors.py

@@ -19,31 +19,35 @@ else:
 
 
 def _memory_efficient_attention(query, key, value, attn_bias=None):
-    """Ruft xformers.ops.memory_efficient_attention auf – mit per-call Fallback.
+    """Speichereffiziente Attention mit xformers → SDPA → bmm Fallback-Kette.
 
-    Warum per-call statt einmaligem hasattr-Check:
-    Auf ZeroGPU wird die GPU erst beim @spaces.GPU-Decorator allokiert.
-    xformers._C.so initialisiert sich erst wenn wirklich CUDA verfügbar ist.
-    Ein hasattr-Check beim Import (ohne GPU) gibt False zurück, obwohl die
-    Funktion später im GPU-Kontext problemlos verfügbar wäre.
+    Tensoren nach head_to_batch_dim: (batch*heads, seq, dim)
+
+    1. xformers.ops.memory_efficient_attention – Flash Attention, O(N) Speicher
+       Wird per-call versucht: auf ZeroGPU lädt _C.so erst im GPU-Kontext.
+    2. torch.nn.functional.scaled_dot_product_attention (SDPA) – PyTorch 2.x
+       Flash Attention 2 intern, O(N) Speicher, kein quadratischer Bedarf.
+    3. torch.bmm – nur als letzter Fallback, quadratischer Speicher (44 GB bei PSHuman!)
     """
+    # 1. xformers (per-call: funktioniert im GPU-Kontext auch wenn beim Start False)
     if xformers is not None:
         try:
             return xformers.ops.memory_efficient_attention(
                 query, key, value, attn_bias=attn_bias
             )
         except AttributeError:
-            # xformers geladen, aber _C.so noch nicht initialisiert (kein GPU-Kontext)
-            # → torch.bmm-Fallback
-            pass
-
-    # Fallback: skaliertes Dot-Product
-    scale = query.shape[-1] ** -0.5
-    attn_scores = torch.bmm(query * scale, key.transpose(-1, -2))
-    if attn_bias is not None:
-        attn_scores = attn_scores + attn_bias
-    attn_probs = torch.nn.functional.softmax(attn_scores, dim=-1)
-    return torch.bmm(attn_probs, value)
+            pass  # _C.so noch nicht initialisiert → SDPA
+
+    # 2. PyTorch SDPA mit Flash Attention 2 (memory-effizient, seit torch 2.0)
+    # Reshape: (batch*heads, seq, dim) → (batch*heads, 1, seq, dim)
+    q = query.unsqueeze(1)
+    k = key.unsqueeze(1)
+    v = value.unsqueeze(1)
+    mask = attn_bias.unsqueeze(1) if attn_bias is not None else None
+    out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+    return out.squeeze(1)
+
+    # 3. torch.bmm – NICHT verwendet (44 GB für PSHuman Multi-View → OOM)
 
 class RowwiseMVAttention(Attention):
     def set_use_memory_efficient_attention_xformers(