[Admin maintenance] Support new ZeroGPU hardware

README.md

@@ -4,7 +4,7 @@ emoji: 🦀
 colorFrom: red
 colorTo: indigo
 sdk: gradio
-sdk_version: 4.20.1
+sdk_version: 5.49.1
 python_version: 3.10.13
 app_file: app.py
 pinned: false

app.py

@@ -1,11 +1,135 @@
 import os
+import sys
 import shlex
 import subprocess
+import tempfile
+import ctypes
+
+import spaces
+import torch
+
+
+# ---------------------------------------------------------------------------
+# Blackwell (sm_120) shim: xformers' FA3/FA2/CutlassF kernels in the prebuilt
+# torch-2.10/2.11 wheel all reject compute capability 12.0, so any call into
+# `xformers.ops.memory_efficient_attention(...)` raises NotImplementedError.
+# The LGM stack calls MEA directly in two places:
+#   - core/attention.py: 4D shape (B, M, H, K) (the dino-style Attention class)
+#   - mvdream/mv_unet.py: 3D shape (B*H, M, K) (the cross-attention block)
+# Route both through torch SDPA. Must be installed BEFORE the imports that
+# pull in core/attention.py and mvdream/mv_unet.py.
+# ---------------------------------------------------------------------------
+import xformers
+import xformers.ops as _xops
+
+
+def _xformers_mea_sdpa(query, key, value, attn_bias=None, p=0.0, scale=None,
+                       op=None, **kwargs):
+    if query.dim() == 3:
+        # (B, M, K) -> single-head; add an H=1 axis.
+        q = query.unsqueeze(1)
+        k = key.unsqueeze(1)
+        v = value.unsqueeze(1)
+        squeeze_out = True
+    else:
+        # (B, M, H, K) -> (B, H, M, K)
+        q = query.transpose(1, 2)
+        k = key.transpose(1, 2)
+        v = value.transpose(1, 2)
+        squeeze_out = False
+    attn_mask = attn_bias
+    if hasattr(attn_mask, "materialize"):
+        try:
+            attn_mask = attn_mask.materialize(
+                shape=(q.shape[0], q.shape[1], q.shape[2], k.shape[2]),
+                dtype=q.dtype,
+                device=q.device,
+            )
+        except Exception:
+            attn_mask = None
+    out = torch.nn.functional.scaled_dot_product_attention(
+        q, k, v, attn_mask=attn_mask, dropout_p=p, scale=scale,
+    )
+    if squeeze_out:
+        return out.squeeze(1)
+    return out.transpose(1, 2)
+
+
+_xops.memory_efficient_attention = _xformers_mea_sdpa
+xformers.ops.memory_efficient_attention = _xformers_mea_sdpa
+
+
+# ---------------------------------------------------------------------------
+# Build `diff_gaussian_rasterization` from source against torch 2.10/2.11/cu128
+# on the first GPU call. The vendored wheel/diff_gaussian_rasterization-...whl
+# in this Space was built against torch 2.4 (`libcudart.so.11.0`) and won't
+# load on Blackwell. We build the original graphdeco-inria fork (what
+# core/gs.py imports).
+# ---------------------------------------------------------------------------
+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():
+    try:
+        import diff_gaussian_rasterization  # noqa: F401
+        return
+    except ImportError:
+        pass
+
+    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"],
+    )
+
+    # Build the vendored `diff-gaussian-rasterization/` source in this repo.
+    # This Space ships a custom fork that returns 4 outputs (image, radii,
+    # depth, alpha); the upstream graphdeco-inria release returns only 2 and
+    # would crash `core/gs.py` at unpack time. The vendored tree is what
+    # LGM's training/inference code was written against.
+    vendored_dgr = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                                "diff-gaussian-rasterization")
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install",
+         "--no-build-isolation", "--no-deps",
+         vendored_dgr],
+        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
+
+
+# ---------------------------------------------------------------------------
+# Now the usual app imports.
+# ---------------------------------------------------------------------------
 import tyro
 import imageio
 import numpy as np
 import tqdm
-import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torchvision.transforms.functional as TF
@@ -13,13 +137,9 @@ from safetensors.torch import load_file
 import rembg
 import gradio as gr
 
-import spaces
-
-# download checkpoints
 from huggingface_hub import hf_hub_download
-ckpt_path = hf_hub_download(repo_id="ashawkey/LGM", filename="model_fp16_fixrot.safetensors")
 
-subprocess.run(shlex.split("pip install wheel/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl"))
+ckpt_path = hf_hub_download(repo_id="ashawkey/LGM", filename="model_fp16_fixrot.safetensors")
 
 import kiui
 from kiui.op import recenter
@@ -34,7 +154,6 @@ IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
 GRADIO_VIDEO_PATH = 'gradio_output.mp4'
 GRADIO_PLY_PATH = 'gradio_output.ply'
 
-# opt = tyro.cli(AllConfigs)
 opt = Options(
     input_size=256,
     up_channels=(1024, 1024, 512, 256, 128), # one more decoder
@@ -79,7 +198,6 @@ pipe_text = MVDreamPipeline.from_pretrained(
     'ashawkey/mvdream-sd2.1-diffusers', # remote weights
     torch_dtype=torch.float16,
     trust_remote_code=True,
-    # local_files_only=True,
 )
 pipe_text = pipe_text.to(device)
 
@@ -87,7 +205,6 @@ pipe_image = MVDreamPipeline.from_pretrained(
     "ashawkey/imagedream-ipmv-diffusers", # remote weights
     torch_dtype=torch.float16,
     trust_remote_code=True,
-    # local_files_only=True,
 )
 pipe_image = pipe_image.to(device)
 
@@ -256,7 +373,7 @@ with block:
         inputs=[input_image],
         outputs=[output_image, output_video, output_file],
         fn=lambda x: process(input_image=x, prompt=''),
-        cache_examples=True,
+        cache_examples=False,
         label='Image-to-3D Examples'
     )
 
@@ -273,7 +390,7 @@ with block:
         inputs=[input_text],
         outputs=[output_image, output_video, output_file],
         fn=lambda x: process(input_image=None, prompt=x),
-        cache_examples=True,
+        cache_examples=False,
         label='Text-to-3D Examples'
     )
 

requirements.txt

@@ -1,4 +1,3 @@
-torch==2.4.0
 xformers
 
 numpy

wheel/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:42bf718442ba764469170abc09d99a70b7c1d891dc290f2e1247db09c95a0e88
-size 3021758