[Admin maintenance] Support new ZeroGPU hardware

Thank you so much for having shared this Space with the community on this demo. We have upgraded the ZeroGPU infra-structure to run on modern blackwell architecture.
For that, we need to upgrade your demo to support that. This PR fixes your demo to work with the new architecture. As this is something we broke on our end, we may merge this PR autonomously. If this breaks unexpectedly or brings unintended consequences, feel free to revert, modify or otherwise. Any issues you can email apolinario@huggingface.co

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