[Admin maintenance] Fix for new ZeroGPU hardware (Blackwell sm_120)

README.md

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

app.py

@@ -1,16 +1,114 @@
-import gradio as gr
+import os
+os.environ['SPCONV_ALGO'] = 'native'
+os.environ['ATTN_BACKEND'] = 'xformers'
+os.environ['SPARSE_ATTN'] = 'xformers'
+
+import subprocess, sys, tempfile, ctypes
+
+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
+
+CUDA_HOME = "/cuda-image/usr/local/cuda-13.0"
+CUDA_LIBDIR = os.path.join(CUDA_HOME, "lib64")
+REPO_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+@spaces.GPU(duration=600)
+def first_gpu_setup():
+    need_nvdiffrast = need_dgr = False
+    try:
+        import nvdiffrast  # noqa: F401
+    except ImportError:
+        need_nvdiffrast = True
+    try:
+        import diff_gaussian_rasterization  # noqa: F401
+    except ImportError:
+        need_dgr = True
+    if not (need_nvdiffrast or need_dgr):
+        print("CUDA extensions already installed; skipping build.")
+        return
+
+    if not os.path.exists(os.path.join(CUDA_HOME, "bin", "nvcc")):
+        raise RuntimeError(
+            f"nvcc not found at {CUDA_HOME}/bin/nvcc on the ZeroGPU worker. "
+            "The new-hardware CUDA path may have moved; please 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 need_nvdiffrast:
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "install", "--no-build-isolation",
+             os.path.join(REPO_DIR, "extensions", "nvdiffrast")],
+            env=env,
+        )
+
+    if need_dgr:
+        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,
+        )
+
+
+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", "")
+
+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
 from PIL import Image
+
+# xformers on Blackwell (sm_120) picks Flash-Attn-3 (Hopper-only) and crashes
+# with "invalid argument". Force it to use Cutlass kernels instead.
+import xformers.ops as _xops
+_orig_mea = _xops.memory_efficient_attention
+_cutlass_op = (_xops.fmha.cutlass.FwOp, _xops.fmha.cutlass.BwOp)
+def _mea_cutlass(*args, **kwargs):
+    kwargs.setdefault("op", _cutlass_op)
+    return _orig_mea(*args, **kwargs)
+_xops.memory_efficient_attention = _mea_cutlass
+
 from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils

requirements.txt

@@ -18,8 +18,3 @@ git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057
 xformers
 spconv-cu120==2.3.6
 transformers==4.46.3
-gradio_litmodel3d==0.0.1
-pydantic==2.10.6
-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