[Admin maintenance] Support new ZeroGPU hardware

app.py

@@ -1,8 +1,130 @@
 import gradio as gr
-import spaces
 import os
-import shutil
 os.environ['SPCONV_ALGO'] = 'native'
+os.environ['ATTN_BACKEND'] = 'xformers'
+os.environ['SPARSE_ATTN_BACKEND'] = 'xformers'
+
+import spaces
+
+# Build the CUDA extensions that used to ship as torch-2.4-ABI prebuilt wheels
+# (nvdiffrast, diff_gaussian_rasterization) and the torch-2.4-cu121-only
+# torch-scatter, all from source against the actually-installed torch.
+import subprocess, sys, tempfile, ctypes
+
+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():
+    need = {}
+    for name, modname in [
+        ("nvdiffrast", "nvdiffrast"),
+        ("diff_gaussian_rasterization", "diff_gaussian_rasterization"),
+        ("torch_scatter", "torch_scatter"),
+        ("detectron2", "detectron2"),
+    ]:
+        try:
+            __import__(modname)
+        except ImportError:
+            need[name] = True
+    if not need:
+        print("CUDA extensions already present.")
+        return
+
+    if not os.path.exists(os.path.join(CUDA_HOME, "bin", "nvcc")):
+        raise RuntimeError(f"nvcc not at {CUDA_HOME}/bin/nvcc; update CUDA_HOME.")
+
+    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"],
+    )
+
+    if "nvdiffrast" in need:
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "install", "--no-build-isolation",
+             "git+https://github.com/NVlabs/nvdiffrast/"],
+            env=env,
+        )
+    if "diff_gaussian_rasterization" in need:
+        mip_dir = tempfile.mkdtemp(prefix="mip_splatting_")
+        subprocess.check_call(
+            ["git", "clone", "--recursive", "--depth=1",
+             "https://github.com/autonomousvision/mip-splatting.git", mip_dir],
+        )
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "install", "--no-build-isolation",
+             os.path.join(mip_dir, "submodules", "diff-gaussian-rasterization")],
+            env=env,
+        )
+    if "torch_scatter" in need:
+        # Prebuilt wheel for torch 2.10.0 + cu128 (matches the torch wheel's
+        # CUDA), avoids the CUDA-version mismatch you get when building against
+        # the container's CUDA 13 toolkit.
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "install", "--no-deps",
+             "torch-scatter",
+             "-f", "https://data.pyg.org/whl/torch-2.10.0+cu128.html"],
+        )
+    if "detectron2" in need:
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "install", "--no-build-isolation",
+             "git+https://github.com/facebookresearch/detectron2.git"],
+            env=env,
+        )
+
+
+_first_gpu_setup()
+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", "")
+
+# torch 2.6 flipped torch.load weights_only default to True; the bbox_gen
+# partfield checkpoint pickles yacs.config.CfgNode and other non-tensor
+# objects. Trusted upstream checkpoint — fall back to weights_only=False.
+import torch as _torch
+_orig_load = _torch.load
+_torch.load = lambda *a, **k: _orig_load(*a, **{**k, "weights_only": k.get("weights_only", False)})
+
+# xformers on Blackwell (sm_120) has no operator for fp32 memory_efficient_attention:
+# FA3 / Cutlass cap at compute capability 9.0, FA2 only supports fp16/bf16. DINOv2
+# (loaded via torch.hub) calls xops.memory_efficient_attention with fp32 → no
+# dispatch. Route through torch SDPA, which handles sm_120 fp32 natively.
+import xformers.ops as _xops
+import torch.nn.functional as _F
+
+def _mea_via_sdpa(q, k, v, attn_bias=None, p=0.0, scale=None, op=None):
+    q = q.transpose(1, 2); k = k.transpose(1, 2); v = v.transpose(1, 2)
+    attn_mask = None
+    if attn_bias is not None:
+        if hasattr(attn_bias, "materialize"):
+            try:
+                attn_mask = attn_bias.materialize((q.shape[-2], k.shape[-2]), device=q.device, dtype=q.dtype)
+            except Exception:
+                attn_mask = None
+        elif isinstance(attn_bias, _torch.Tensor):
+            attn_mask = attn_bias
+    out = _F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=p, scale=scale)
+    return out.transpose(1, 2).contiguous()
+
+_xops.memory_efficient_attention = _mea_via_sdpa
+
+import shutil
 from huggingface_hub import hf_hub_download
 
 from app_utils import (

modules/bbox_gen/models/autogressive_bbox_gen.py

@@ -105,7 +105,7 @@ class BboxGen(nn.Module):
             self.decoder: BBoxOPT = AutoModelForCausalLM.from_config(
                 self.decoder_config, 
                 torch_dtype=torch.bfloat16,
-                attn_implementation="flash_attention_2"
+                attn_implementation="sdpa"
             )
         else:
             self.decoder: BBoxOPT = AutoModelForCausalLM.from_config(

requirements.txt

@@ -1,7 +1,7 @@
---extra-index-url https://download.pytorch.org/whl/cu121
-
-torch==2.4.0
-torchvision==0.19.0
+torch==2.10.0
+torchvision==0.25.0
+einops
+psutil
 pillow==10.4.0
 imageio==2.36.1
 imageio-ffmpeg==0.5.1
@@ -17,10 +17,9 @@ pyvista==0.44.2
 pymeshfix==0.17.0
 igraph==0.11.8
 git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8
-xformers==0.0.27.post2
+xformers
 spconv-cu120==2.3.6
 transformers==4.50.3
-pydantic==2.10.6
 diffusers==0.32.0
 lightning==2.2
 mesh2sdf
@@ -35,10 +34,3 @@ timm
 h5py
 boto3
 git+https://github.com/facebookresearch/segment-anything.git
-git+https://github.com/facebookresearch/detectron2.git
---find-links https://data.pyg.org/whl/torch-2.4.0+cu121.html
-torch-scatter
-
-https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.0.post2/flash_attn-2.7.0.post2+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
-https://huggingface.co/spaces/JeffreyXiang/TRELLIS/resolve/main/wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl?download=true
-https://huggingface.co/spaces/JeffreyXiang/TRELLIS/resolve/main/wheels/nvdiffrast-0.3.3-cp310-cp310-linux_x86_64.whl?download=true