[Admin maintenance] Support new ZeroGPU hardware

README.md

@@ -4,7 +4,7 @@ emoji: 🦀
 colorFrom: red
 colorTo: indigo
 sdk: gradio
-sdk_version: 4.32.1
+sdk_version: 5.49.1
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -1,18 +1,115 @@
 import os
-import shlex
+import sys
 import subprocess
+import tempfile
 
-import gradio as gr
-import numpy as np
 import spaces
 import torch
+import ctypes
+
+# Monkey-patch xformers.ops.memory_efficient_attention to fall back to
+# torch SDPA on Blackwell (sm_120). On the new ZeroGPU GPUs, neither
+# the FA3 nor the cutlass xformers kernel supports compute capability
+# 12.0, so any call into xformers' MEA raises NotImplementedError.
+# imagedream's mv_unet imports xformers.ops directly, so we have to
+# patch it before that import happens.
+import xformers
+import xformers.ops as _xops
+
+
+def _xformers_mea_sdpa(query, key, value, attn_bias=None, p=0.0, scale=None,
+                      op=None, **kwargs):
+    # xformers MEA accepts (B, M, K) or (B, M, H, K). Torch SDPA wants
+    # (B, H, M, K). Reshape appropriately and reverse on output.
+    if query.dim() == 3:
+        # Single-head: add an H=1 axis.
+        q = query.unsqueeze(1)
+        k = key.unsqueeze(1)
+        v = value.unsqueeze(1)
+        squeeze_out = True
+    else:
+        # (B, M, H, K) -> (B, H, M, K)
+        q = query.transpose(1, 2)
+        k = key.transpose(1, 2)
+        v = value.transpose(1, 2)
+        squeeze_out = False
+    attn_mask = attn_bias
+    if hasattr(attn_mask, "materialize"):
+        try:
+            attn_mask = attn_mask.materialize(
+                shape=(q.shape[0], q.shape[1], q.shape[2], k.shape[2]),
+                dtype=q.dtype,
+                device=q.device,
+            )
+        except Exception:
+            attn_mask = None
+    out = torch.nn.functional.scaled_dot_product_attention(
+        q, k, v, attn_mask=attn_mask, dropout_p=p, scale=scale,
+    )
+    if squeeze_out:
+        return out.squeeze(1)
+    return out.transpose(1, 2)
+
+
+_xops.memory_efficient_attention = _xformers_mea_sdpa
+xformers.ops.memory_efficient_attention = _xformers_mea_sdpa
+
+import gradio as gr
+import numpy as np
 from diffusers import DiffusionPipeline
 
-subprocess.run(
-    shlex.split(
-        "pip install wheel/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl"
+
+CUDA_HOME = "/cuda-image/usr/local/cuda-13.0"
+CUDA_LIBDIR = os.path.join(CUDA_HOME, "lib64")
+
+
+@spaces.GPU(duration=600)
+def _first_gpu_setup():
+    try:
+        import diff_gaussian_rasterization  # noqa
+        return
+    except ImportError:
+        pass
+
+    patch_dir = tempfile.mkdtemp(prefix="torch_cuda_patch_")
+    with open(os.path.join(patch_dir, "sitecustomize.py"), "w") as f:
+        f.write(
+            "try:\n"
+            "    import torch.utils.cpp_extension as _c\n"
+            "    _c._check_cuda_version = lambda *a, **k: None\n"
+            "except Exception:\n"
+            "    pass\n"
+        )
+
+    env = os.environ.copy()
+    env["CUDA_HOME"] = CUDA_HOME
+    env["CUDA_PATH"] = CUDA_HOME
+    env["PATH"] = os.path.join(CUDA_HOME, "bin") + os.pathsep + env.get("PATH", "")
+    env["PYTHONPATH"] = patch_dir + os.pathsep + env.get("PYTHONPATH", "")
+    env["TORCH_CUDA_ARCH_LIST"] = "12.0"
+
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install", "--no-deps",
+         "setuptools", "wheel", "ninja", "packaging"],
     )
-)
+
+    # The PyPI release of `diff-gaussian-rasterization` is the original
+    # graphdeco-inria version; the LGM model uses that one.
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install",
+         "--no-build-isolation", "--no-deps",
+         "git+https://github.com/graphdeco-inria/diff-gaussian-rasterization.git"],
+        env=env,
+    )
+
+
+_first_gpu_setup()
+try:
+    ctypes.CDLL(os.path.join(CUDA_LIBDIR, "libcudart.so.13"), mode=ctypes.RTLD_GLOBAL)
+    os.environ["LD_LIBRARY_PATH"] = CUDA_LIBDIR + os.pathsep + os.environ.get("LD_LIBRARY_PATH", "")
+except OSError:
+    pass
+
 
 TMP_DIR = "/tmp"
 os.makedirs(TMP_DIR, exist_ok=True)

requirements.txt

@@ -1,4 +1,3 @@
-wheel
 numpy
 tyro
 diffusers
@@ -27,8 +26,5 @@ kiui >= 0.2.3
 xatlas
 roma
 plyfile
-torch == 2.2.0
-torchvision == 0.17.0
-torchaudio == 2.2.0
 xformers
 ushlex

wheel/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:957eac9a321f881c46e8d8a9206bd5f253b64685f5cc35f07d943b908ef2263f
-size 3134028