[Admin maintenance] Support new ZeroGPU hardware

README.md

@@ -4,7 +4,7 @@ emoji: 🖥️
 colorFrom: indigo
 colorTo: blue
 sdk: gradio
-sdk_version: 5.34.2
+sdk_version: 5.49.1
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -1,12 +1,158 @@
-import gradio as gr
+import os
+os.environ.setdefault("SPCONV_ALGO", "native")
+os.environ.setdefault("ATTN_BACKEND", "xformers")
+os.environ.setdefault("SPARSE_ATTN_BACKEND", "xformers")
+
+import subprocess
+import sys
+
+# Install gradio_litmodel3d ignoring its over-restrictive gradio<5 cap.
+try:
+    import gradio_litmodel3d  # noqa: F401
+except ImportError:
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install", "--no-deps", "gradio_litmodel3d==0.0.1"],
+    )
+
 import spaces
+import torch
+import ctypes
+import tempfile
+
+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"),
+    ]:
+        try:
+            __import__(modname)
+        except ImportError:
+            need[name] = True
+    if not need:
+        return
+
+    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"],
+    )
+
+    if "nvdiffrast" in need:
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "install",
+             "--no-build-isolation", "--no-deps",
+             "git+https://github.com/NVlabs/nvdiffrast/"],
+            env=env,
+        )
+    if "diff_gaussian_rasterization" in need:
+        mip = tempfile.mkdtemp(prefix="mip_")
+        subprocess.check_call(
+            ["git", "clone", "--recursive", "--depth=1",
+             "https://github.com/autonomousvision/mip-splatting.git", mip],
+        )
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "install",
+             "--no-build-isolation", "--no-deps",
+             os.path.join(mip, "submodules", "diff-gaussian-rasterization")],
+            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
+
+# xformers on the Blackwell (sm_120) ZeroGPU container is built without CUDA
+# extensions for any FwOp: cutlassF-pt rejects compute capability >= (9, 0)
+# ("too new") and FlashAttn3 is Hopper-only. Reroute xformers.ops.memory_efficient_attention
+# (used by DINOv2, VGGT, trellis dense+sparse paths) to torch.nn.functional.scaled_dot_product_attention,
+# which is CUDA-native on torch 2.10/2.11 and supports sm_120. Must be patched BEFORE
+# anything that calls memory_efficient_attention is imported.
+try:
+    import xformers.ops as _xops
+    import torch.nn.functional as _F
+    from xformers.ops.fmha.attn_bias import BlockDiagonalMask as _BlockDiagonalMask
+
+    def _bdm_starts(seqinfo):
+        # xformers' BlockDiagonalMask sub-attribute. Try the public python-list view first;
+        # otherwise pull from the tensor and tolist().
+        for attr in ("seqstart_py", "_seqstart_py"):
+            v = getattr(seqinfo, attr, None)
+            if v is not None:
+                return list(v)
+        t = getattr(seqinfo, "seqstart", None)
+        if t is not None:
+            return t.detach().cpu().tolist()
+        raise AttributeError("BlockDiagonalMask seqinfo has no seqstart_py / seqstart")
+
+    def _mea_sdpa(q, k, v, attn_bias=None, p=0.0, scale=None, op=None):
+        # q, k, v shapes: [B, M, H, K] (xformers convention). SDPA wants [B, H, M, K].
+        if isinstance(attn_bias, _BlockDiagonalMask):
+            # Block-diagonal mask used by trellis sparse attention to batch
+            # variable-length sequences in a single dense tensor. Materialize each
+            # block separately and concat — SDPA has no block-diagonal kernel.
+            q_starts = _bdm_starts(attn_bias.q_seqinfo)
+            k_starts = _bdm_starts(attn_bias.k_seqinfo)
+            outs = []
+            # q,k,v come in as [1, total_tokens, H, K]
+            for i in range(len(q_starts) - 1):
+                qs, qe = q_starts[i], q_starts[i + 1]
+                ks, ke = k_starts[i], k_starts[i + 1]
+                qi = q[:, qs:qe].transpose(1, 2)  # [1, H, Lq, K]
+                ki = k[:, ks:ke].transpose(1, 2)  # [1, H, Lk, K]
+                vi = v[:, ks:ke].transpose(1, 2)
+                oi = _F.scaled_dot_product_attention(qi, ki, vi, dropout_p=p, scale=scale)
+                outs.append(oi.transpose(1, 2))   # back to [1, Li, H, K]
+            return torch.cat(outs, dim=1)
+
+        attn_mask = None
+        if attn_bias is not None and hasattr(attn_bias, "materialize"):
+            attn_mask = attn_bias.materialize((q.shape[0], q.shape[2], q.shape[1], k.shape[1]),
+                                              dtype=q.dtype, device=q.device)
+        elif attn_bias is not None:
+            attn_mask = attn_bias
+
+        qh = q.transpose(1, 2)  # [B, H, M, K]
+        kh = k.transpose(1, 2)
+        vh = v.transpose(1, 2)
+        out = _F.scaled_dot_product_attention(qh, kh, vh, attn_mask=attn_mask, dropout_p=p, scale=scale)
+        return out.transpose(1, 2)  # [B, M, H, K]
+
+    _xops.memory_efficient_attention = _mea_sdpa
+    print("[blackwell] xformers.memory_efficient_attention rerouted to torch SDPA")
+except Exception as _e:
+    print(f"[blackwell] xformers SDPA shim skipped: {_e}")
+
+import gradio as gr
 from gradio_litmodel3d import LitModel3D
 
-import os
 import shutil
-os.environ['SPCONV_ALGO'] = 'native'
 from typing import *
-import torch
 import numpy as np
 import imageio
 from easydict import EasyDict as edict

requirements.txt

@@ -1,9 +1,5 @@
---extra-index-url https://download.pytorch.org/whl/cu121
-
 kornia==0.8.0
 
-torch==2.4.0
-torchvision==0.19.0
 pillow==10.4.0
 imageio==2.36.1
 imageio-ffmpeg==0.5.1
@@ -19,17 +15,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.46.3
-gradio_litmodel3d==0.0.1
-pydantic==2.10.6
 einops==0.8.1
-# huggingface_hub==0.25.0
-huggingface_hub==0.33.4
 lpips==0.1.4
-spaces==0.37.1
 timm==1.0.23
-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