Update app.py

app.py

@@ -1,8 +1,10 @@
+import os
+# PYTORCH_CUDA_ALLOC_CONF MUSS vor torch-Import stehen (GPT-Fix)
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
 import spaces
 import sys
 import torch
-import os
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 import shutil
 import tempfile
 import gradio as gr
@@ -18,20 +20,17 @@ from pathlib import Path
 
 
 def _get_subprocess_env() -> dict:
-    """Baut die Subprocess-Umgebung mit korrekten CUDA- und Library-Pfaden.
+    """Subprocess-Umgebung mit CUDA-Library-Pfaden und NVML-freiem Allocator.
 
-    Wichtige Fixes:
-    - LD_LIBRARY_PATH: libcudart.so via nvidia pip-Pakete findbar machen
-    - PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True nutzt NVML, das im
-      Subprocess auf ZeroGPU Blackwell + torch 2.11.0 nicht stabil ist.
-      Im Subprocess max_split_size_mb:512 verwenden (NVML-frei, stabil).
+    backend:cudaMallocAsync ersetzt CUDACachingAllocator komplett –
+    kein NVML nötig, stabil auf ZeroGPU Blackwell + torch 2.11.0 + CUDA 13.
     """
     import site
     import glob as _glob
 
     env = os.environ.copy()
 
-    # 1. CUDA-Library-Pfade für libcudart
+    # CUDA-Library-Pfade für libcudart
     nvidia_paths: list[str] = []
     all_site = site.getsitepackages()
     try:
@@ -53,16 +52,47 @@ def _get_subprocess_env() -> dict:
         nvidia_paths + [torch_lib] + extra + ([existing_ld] if existing_ld else [])
     )
 
-    # 2. CUDA-Allocator: expandable_segments:True braucht NVML, das im
-    #    ZeroGPU-Subprocess mit torch 2.11.0 zu CUDACachingAllocator-Assertions führt.
-    #    Im Subprocess stabilen Fallback verwenden.
-    env["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
+    # NVML-freier Allocator für Subprocess (CUDACachingAllocator-Assertions vermeiden)
+    env["PYTORCH_CUDA_ALLOC_CONF"] = "backend:cudaMallocAsync"
 
     return env
 
 
+def _validate_xformers() -> bool:
+    """Prüft ob xformers.ops.memory_efficient_attention wirklich funktioniert.
+    Gibt True zurück wenn ein Mini-Test mit echten CUDA-Tensoren besteht.
+    """
+    try:
+        import xformers
+        import xformers.ops
+        print(f"[DIAG] xformers version   : {xformers.__version__}", flush=True)
+
+        has_mea = hasattr(xformers.ops, 'memory_efficient_attention')
+        print(f"[DIAG] xformers MEA attr  : {has_mea}", flush=True)
+
+        if not has_mea:
+            print("[DIAG] xformers MEA test  : SKIP (attr fehlt)", flush=True)
+            return False
+
+        if not torch.cuda.is_available():
+            print("[DIAG] xformers MEA test  : SKIP (kein CUDA)", flush=True)
+            return False
+
+        # Mini-Test mit kleinen CUDA-Tensoren
+        q = torch.randn(2, 16, 64, device="cuda", dtype=torch.float16)
+        k = torch.randn(2, 16, 64, device="cuda", dtype=torch.float16)
+        v = torch.randn(2, 16, 64, device="cuda", dtype=torch.float16)
+        _ = xformers.ops.memory_efficient_attention(q, k, v)
+        torch.cuda.synchronize()
+        print("[DIAG] xformers MEA test  : PASS ✅", flush=True)
+        return True
+
+    except Exception as e:
+        print(f"[DIAG] xformers MEA test  : FAIL – {e}", flush=True)
+        return False
+
+
 def _log_cuda_diagnostics():
-    """Loggt torch-Version, CUDA, libcudart-Pfad und pytorch3d beim Start."""
     import glob as _glob
     import diffusers
     print(f"[DIAG] diffusers version   : {diffusers.__version__}", flush=True)
@@ -71,17 +101,13 @@ def _log_cuda_diagnostics():
     print(f"[DIAG] CUDA available     : {torch.cuda.is_available()}", flush=True)
     print(f"[DIAG] sys.executable     : {sys.executable}", flush=True)
     env = _get_subprocess_env()
-    ld_path = env.get("LD_LIBRARY_PATH", "")
-    print(f"[DIAG] LD_LIBRARY_PATH    : {ld_path}", flush=True)
-    found = []
-    for d in ld_path.split(":"):
-        if d:
-            found.extend(_glob.glob(os.path.join(d, "libcudart.so*")))
+    ld = env.get("LD_LIBRARY_PATH", "")
+    found = [h for d in ld.split(":") if d for h in _glob.glob(os.path.join(d, "libcudart.so*"))]
     print(f"[DIAG] libcudart found    : {found}", flush=True)
     print(f"[DIAG] subprocess ALLOC   : {env.get('PYTORCH_CUDA_ALLOC_CONF')}", flush=True)
+    _validate_xformers()
     import importlib.util
-    p3d = importlib.util.find_spec("pytorch3d")
-    print(f"[DIAG] pytorch3d spec     : {p3d}", flush=True)
+    print(f"[DIAG] pytorch3d spec     : {importlib.util.find_spec('pytorch3d')}", flush=True)
 
 _log_cuda_diagnostics()