[Admin maintenance] Support new ZeroGPU hardware

app.py

@@ -1,15 +1,164 @@
-import gradio as gr
+import os
+# Force attention backends compatible with the new ZeroGPU (Blackwell) stack.
+# Must be set BEFORE any trellis / dinov2 import.
+# Trellis's dense attention has a native SDPA path; use it.
+os.environ.setdefault('ATTN_BACKEND', 'sdpa')
+# Sparse attention only knows 'xformers' or 'flash_attn'; keep 'xformers' but
+# monkey-patch xformers.ops.memory_efficient_attention to SDPA below (none of
+# the prebuilt xformers ops support sm_120 / Blackwell).
+os.environ.setdefault('SPARSE_ATTN_BACKEND', 'xformers')
+os.environ.setdefault('SPCONV_ALGO', 'native')
+# Force dinov2 (loaded via torch.hub for image conditioning) to take its pure
+# torch.nn.functional.scaled_dot_product_attention path instead of importing
+# xformers.ops.memory_efficient_attention (which raises on sm_120).
+os.environ.setdefault('XFORMERS_DISABLED', '1')
+
+import sys
+import subprocess
+import tempfile
+import ctypes
+
 import spaces
-from gradio_litmodel3d import LitModel3D
+import torch
+import gradio as gr
+
+# ---------------------------------------------------------------------------
+# xformers -> SDPA shim for Blackwell (sm_120).
+# The prebuilt xformers wheel ships FA3, FA2 and CutlassF ops that all assert
+# device capability <= (9, 0); none load on sm_120, so any call to
+# memory_efficient_attention raises NotImplementedError. dinov2 (image
+# conditioning model in trellis) and trellis's own sparse paths both call it.
+# Replace memory_efficient_attention with an SDPA-backed implementation that
+# also handles xformers.fmha.BlockDiagonalMask (used by sparse attention).
+# ---------------------------------------------------------------------------
+try:
+    import xformers  # noqa: F401
+    import xformers.ops as _xops
+    from torch.nn.functional import scaled_dot_product_attention as _sdpa
+
+    try:
+        _BlockDiagonalMask = _xops.fmha.BlockDiagonalMask
+    except Exception:
+        _BlockDiagonalMask = None
+
+    def _mea_sdpa(q, k, v, attn_bias=None, p=0.0, scale=None, *args, **kwargs):
+        # q, k, v: [B, N, H, C] (xformers layout). SDPA expects [B, H, N, C].
+        if attn_bias is None:
+            qh = q.transpose(1, 2)
+            kh = k.transpose(1, 2)
+            vh = v.transpose(1, 2)
+            out = _sdpa(qh, kh, vh, dropout_p=p, scale=scale)
+            return out.transpose(1, 2).contiguous()
+
+        if _BlockDiagonalMask is not None and isinstance(attn_bias, _BlockDiagonalMask):
+            # BlockDiagonal: q, k, v come as [1, T, H, C] where T is the
+            # concatenation of variable-length blocks. Split, apply SDPA per
+            # block, concatenate. q and kv can have different seqlens.
+            q_info = attn_bias.q_seqinfo
+            kv_info = attn_bias.k_seqinfo
+            q_starts = q_info.seqstart_py
+            kv_starts = kv_info.seqstart_py
+            outs = []
+            for i in range(len(q_starts) - 1):
+                qs, qe = q_starts[i], q_starts[i + 1]
+                ks, ke = kv_starts[i], kv_starts[i + 1]
+                qi = q[:, qs:qe].transpose(1, 2)
+                ki = k[:, ks:ke].transpose(1, 2)
+                vi = v[:, ks:ke].transpose(1, 2)
+                oi = _sdpa(qi, ki, vi, dropout_p=p, scale=scale)
+                outs.append(oi.transpose(1, 2))
+            return torch.cat(outs, dim=1).contiguous()
+
+        # Fallback: dense additive bias.
+        qh = q.transpose(1, 2)
+        kh = k.transpose(1, 2)
+        vh = v.transpose(1, 2)
+        out = _sdpa(qh, kh, vh, attn_mask=attn_bias, dropout_p=p, scale=scale)
+        return out.transpose(1, 2).contiguous()
+
+    _xops.memory_efficient_attention = _mea_sdpa
+    print("[xformers-shim] Replaced memory_efficient_attention with SDPA backend (Blackwell sm_120 fallback).")
+except Exception as _e:
+    print(f"[xformers-shim] Skipped: {_e}")
 
-import os
 import shutil
-os.environ['SPCONV_ALGO'] = 'native'
 from typing import *
-import torch
 import numpy as np
 import imageio
 from PIL import Image
+
+# Build nvdiffrast and diff_gaussian_rasterization from source on first GPU call.
+CUDA_HOME = "/cuda-image/usr/local/cuda-13.0"
+CUDA_LIBDIR = os.path.join(CUDA_HOME, "lib64")
+_NVDIFFRAST_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "extensions", "nvdiffrast")
+
+
+@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"  # Blackwell sm_120
+
+    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",
+             _NVDIFFRAST_DIR],
+            env=env,
+        )
+    if "diff_gaussian_rasterization" in need:
+        # Hi3DGen actually uses the mip-splatting submodule fork; not the
+        # original graphdeco-inria release on PyPI.
+        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
+
 from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.utils import render_utils
 import trimesh

requirements.txt

@@ -1,11 +1,10 @@
---extra-index-url https://download.pytorch.org/whl/cu121
-huggingface-hub==0.36.0
+huggingface-hub
 diffusers==0.35.0
 accelerate==1.2.1
 kornia==0.8.0
-timm==0.6.7
-torch==2.4.0
-torchvision==0.19.0
+timm
+torch==2.10.0
+torchvision==0.25.0
 pillow==10.4.0
 imageio==2.36.1
 imageio-ffmpeg==0.5.1
@@ -21,15 +20,8 @@ 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
-spconv-cu120==2.3.6
+xformers
+spconv-cu126==2.3.8
 transformers==4.46.3
-gradio_litmodel3d==0.0.1
-triton==3.0.0
-nvidia-cudnn-cu12==9.1.0.70
-nvidia-nccl-cu12==2.20.5
-tokenizers==0.20.3
-spaces==0.42.1
-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
+einops
+spaces