Use prebuilt diff-gaussian-rasterization wheel (splatting_cuda decoder); drop gsplat JIT

README.md

@@ -30,13 +30,17 @@ How the Space works:
    runs the **complete bitstream path** from `ModelWrapper.test_step`: encode context →
    compress to a real bitstream at the chosen λ → decompress → predict Gaussians, using bf16
    autocast like `trainer.precision=bf16-mixed`.
-4. The authors' gsplat decoder (antialiased, eps2d=0.1) renders held-out views (PSNR is reported
+4. The authors' `splatting_cuda` decoder (Inria `diff-gaussian-rasterization`, prebuilt sm_120
+   wheel) renders held-out views (PSNR is reported
    against the real photos) and a fly-through video. The decoded Gaussians are shown in an
    interactive `gr.Model3D` splat viewer, and the `.bin` bitstream is available to download.
 
 Notes:
-* gsplat's CUDA kernels are JIT-compiled at startup against the CUDA toolkit in the ZeroGPU
-  image, following `hugging-apps/resplat-recurrent-gs`.
+* The authors' evaluation scripts render with gsplat (antialiased). This Space uses the
+  repo's other supported decoder, `model/decoder=splatting_cuda`, because a prebuilt Blackwell
+  wheel exists for it. Renders can differ very slightly from the paper's numbers.
+* COLMAP runs in a separate process with 16 threads. On the ~192-CPU ZeroGPU hosts COLMAP's
+  default thread count makes the mapper fail to register images.
 * The viewer gets a compact copy of the Gaussians: DC colour only, near-transparent splats
   pruned, re-centred so the scene is framed.
 

app.py

@@ -3,15 +3,13 @@
 Upload overlapping photos of a static scene. The Space runs COLMAP (pycolmap, CPU) to get
 camera poses, then CodecSplat (authors' DL3DV 8-view 256x448 checkpoint):
   context views -> 2D latent features -> learned entropy coder -> real bitstream
-  -> decoder -> 3D Gaussians -> gsplat rendering.
+  -> decoder -> 3D Gaussians -> Gaussian-splatting rasterizer.
 You get an interactive splat viewer, a fly-through video, held-out novel views, the
 size of the compressed scene bitstream (and the bitstream itself).
 """
 
 from __future__ import annotations
 
-import ctypes
-import glob
 import os
 import sys
 import tempfile
@@ -25,8 +23,6 @@ sys.path.insert(0, str(APP_DIR))
 # threads each spin up a multi-threaded BLAS; keep BLAS single-threaded.
 os.environ["OPENBLAS_NUM_THREADS"] = "1"
 os.environ.setdefault("OMP_NUM_THREADS", "4")
-os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
-os.environ.setdefault("MAX_JOBS", str(min(32, (os.cpu_count() or 8))))
 
 import spaces  # noqa: E402  (must precede any CUDA-touching import)
 import torch  # noqa: E402
@@ -34,67 +30,6 @@ import torch  # noqa: E402
 torch.set_float32_matmul_precision("high")
 
 
-# ---------------------------------------------------------------------------
-# 1. gsplat CUDA backend: JIT-compile once at startup against the CUDA toolkit
-#    shipped in the ZeroGPU image (same approach as hugging-apps/resplat-recurrent-gs).
-# ---------------------------------------------------------------------------
-def _pick_cuda_home() -> str | None:
-    torch_cuda = (torch.version.cuda or "").split(".")[0]
-    candidates: list[str] = []
-    for pattern in ("/cuda-image/usr/local/cuda-*", "/usr/local/cuda-*"):
-        candidates += sorted(glob.glob(pattern), reverse=True)
-    candidates += ["/cuda-image/usr/local/cuda", "/usr/local/cuda"]
-    usable = [c for c in candidates if os.path.isfile(os.path.join(c, "bin", "nvcc"))]
-    if not usable:
-        return None
-    for c in usable:
-        if torch_cuda and f"cuda-{torch_cuda}" in c:
-            return c
-    return usable[0]
-
-
-os.environ.setdefault("_ORIG_LD_LIBRARY_PATH", os.environ.get("LD_LIBRARY_PATH", ""))
-CUDA_HOME = _pick_cuda_home()
-print(f"[setup] torch {torch.__version__} (cuda {torch.version.cuda}) | CUDA_HOME={CUDA_HOME}", flush=True)
-if CUDA_HOME:
-    os.environ["CUDA_HOME"] = CUDA_HOME
-    os.environ["CUDA_PATH"] = CUDA_HOME
-    os.environ["PATH"] = os.path.join(CUDA_HOME, "bin") + os.pathsep + os.environ.get("PATH", "")
-    _libdir = os.path.join(CUDA_HOME, "lib64")
-    os.environ["LD_LIBRARY_PATH"] = _libdir + os.pathsep + os.environ.get("LD_LIBRARY_PATH", "")
-    for _so in sorted(glob.glob(os.path.join(_libdir, "libcudart.so*")), reverse=True):
-        try:
-            ctypes.CDLL(_so, mode=ctypes.RTLD_GLOBAL)
-            break
-        except OSError:
-            continue
-
-def _cuda_diag(stage: str) -> None:
-    """Log whether the CUDA driver got loaded into the main (GPU-less) process."""
-    try:
-        maps = open("/proc/self/maps").read()
-        libcuda = "libcuda.so" in maps
-    except Exception:
-        libcuda = "?"
-    print(f"[diag] {stage}: libcuda mapped={libcuda} "
-          f"torch.cuda.is_initialized={torch.cuda.is_initialized()}", flush=True)
-
-
-_cuda_diag("after torch import")
-import torch.utils.cpp_extension as _cppext  # noqa: E402
-
-_cppext._check_cuda_version = lambda *a, **k: None
-
-print("[setup] compiling gsplat CUDA kernels (one-off, several minutes)...", flush=True)
-_t0 = time.time()
-from gsplat.cuda._backend import _C as _GSPLAT_C  # noqa: E402
-
-if _GSPLAT_C is None:
-    raise RuntimeError(f"gsplat CUDA backend unavailable (CUDA_HOME={CUDA_HOME}).")
-print(f"[setup] gsplat backend ready in {time.time() - _t0:.1f}s", flush=True)
-_cuda_diag("after gsplat build")
-
-
 # ---------------------------------------------------------------------------
 # 2. Model
 # ---------------------------------------------------------------------------
@@ -106,7 +41,6 @@ from PIL import Image  # noqa: E402
 
 import codecsplat_pipeline as cp  # noqa: E402
 
-_cuda_diag("after pipeline import")
 from colmap_sfm import IMAGE_EXT, run_sfm_subprocess  # noqa: E402
 
 MODEL_REPO = "pengpeng-yu/CodecSplat"
@@ -121,7 +55,6 @@ print("[setup] building CodecSplat...", flush=True)
 ENCODER = cp.build_encoder(ckpt_path).to("cuda")
 DECODER = cp.build_decoder().to("cuda")
 print("[setup] model ready", flush=True)
-_cuda_diag("after model build")
 
 
 def _u8(t: torch.Tensor) -> np.ndarray:
@@ -237,7 +170,6 @@ def reconstruct(images, lmb: float = 1024.0, matcher: str = "exhaustive",
         raise gr.Error(str(exc)) from exc
     t_sfm = time.time() - t_start
     img_paths, c2w, K = _scene_arrays(sfm["frames"])
-    _cuda_diag("before GPU call")
 
     progress(0.55, desc="CodecSplat: encode -> bitstream -> decode -> render")
     r = codecsplat_gpu(img_paths, c2w, K, float(lmb))

codecsplat_pipeline.py

@@ -76,19 +76,21 @@ def build_encoder(checkpoint_path: str):
 
 
 def build_decoder():
-    from src.model.decoder.decoder_splatting_gsplat import (
-        DecoderSplattingGsplat,
-        DecoderSplattingGsplatCfg,
+    """The authors' `model/decoder=splatting_cuda` decoder (Inria diff-gaussian-rasterization).
+
+    The evaluation scripts use the gsplat decoder; on ZeroGPU we use the repo's other
+    supported decoder because a prebuilt sm_120 wheel exists for it (gsplat would have to
+    be JIT-compiled against a CUDA toolkit at startup).
+    """
+    from src.model.decoder.decoder_splatting_cuda import (
+        DecoderSplattingCUDA,
+        DecoderSplattingCUDACfg,
     )
 
-    dec_cfg = OmegaConf.load(APP_DIR / "config/model/decoder/splatting_gsplat.yaml")
-    dec_cfg = OmegaConf.merge(dec_cfg, DECODER_OVERRIDES)
-    dec_cfg = from_dict(
-        DecoderSplattingGsplatCfg,
-        OmegaConf.to_container(dec_cfg, resolve=True),
-        config=DaciteConfig(strict=False),
+    decoder = DecoderSplattingCUDA(
+        DecoderSplattingCUDACfg(name="splatting_cuda"),
+        SimpleNamespace(background_color=[0.0, 0.0, 0.0]),
     )
-    decoder = DecoderSplattingGsplat(dec_cfg, SimpleNamespace(background_color=[0.0, 0.0, 0.0]))
     return decoder.eval()
 
 

requirements.txt

@@ -1,7 +1,6 @@
 torch==2.11.0
 torchvision==0.26.0
-gsplat==1.5.3
-ninja
+https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt211-cu130-cp312/diff_gaussian_rasterization-0.0.0-cp312-cp312-linux_x86_64.whl
 compressai==1.2.8
 numpy<2
 pycolmap==4.1.1

src/model/decoder/__init__.py

@@ -16,15 +16,22 @@ except ImportError as exc:
 else:
     CUDA_SPLATTING_IMPORT_ERROR = None
 
-from .decoder_splatting_gsplat import (
-    DecoderSplattingGsplat,
-    DecoderSplattingGsplatCfg,
-)
-
-DECODERS = {
-    "splatting_gsplat": DecoderSplattingGsplat,
-    "gsplat": DecoderSplattingGsplat,
-}
+try:  # gsplat is optional in this Space (we render with diff-gaussian-rasterization)
+    from .decoder_splatting_gsplat import (
+        DecoderSplattingGsplat,
+        DecoderSplattingGsplatCfg,
+    )
+except ImportError:
+    DecoderSplattingGsplat = None
+
+    @dataclass
+    class DecoderSplattingGsplatCfg:
+        name: Literal["splatting_gsplat"]
+
+DECODERS = {}
+if DecoderSplattingGsplat is not None:
+    DECODERS["splatting_gsplat"] = DecoderSplattingGsplat
+    DECODERS["gsplat"] = DecoderSplattingGsplat
 
 DecoderCfg = DecoderSplattingGsplatCfg | DecoderSplattingCUDACfg