[Admin maintenance] Support new ZeroGPU hardware (#20)

- [Admin maintenance] Support new ZeroGPU hardware (93d601f7ad6176bd32925886e0257f8c889db3dc)

README.md

@@ -5,7 +5,7 @@ colorFrom: red
 colorTo: purple
 sdk: gradio
 python_version: 3.10.8
-sdk_version: 4.12.0
+sdk_version: 5.49.1
 app_file: app.py
 pinned: true
 short_description: Create a 1M faces 3D colored model from an image!

app.py

@@ -1,24 +1,210 @@
-import shlex
+import os
+import sys
 import subprocess
-subprocess.run(
-    shlex.split(
-        "pip install package/onnxruntime_gpu-1.17.0-cp310-cp310-manylinux_2_28_x86_64.whl --force-reinstall --no-deps"
+import tempfile
+import spaces
+
+# Build pytorch3d and nvdiffrast against the installed torch the first time a GPU is allocated.
+# Both packages publish no prebuilt wheels for torch >= 2.8 (Blackwell / CUDA 13 stack).
+CUDA_HOME = "/cuda-image/usr/local/cuda-13.0"
+
+
+def _build_env():
+    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["TORCH_CUDA_ARCH_LIST"] = "12.0"
+    return env
+
+
+def _ensure_pkg(name, install_args):
+    try:
+        __import__(name)
+        return True
+    except ImportError:
+        pass
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install",
+         "--no-build-isolation",
+         "--no-deps",
+         *install_args],
+        env=_build_env(),
     )
-)
-subprocess.run(
-    shlex.split(
-        "pip install package/nvdiffrast-0.3.1.torch-cp310-cp310-linux_x86_64.whl --force-reinstall --no-deps"
+
+
+def _strip_pulsar_from_pytorch3d_src(src):
+    """Drop pulsar from a cloned pytorch3d source tree.
+
+    Pulsar fails to link on modern GCC because the explicit template
+    instantiations of pulsar::Renderer::fill_bg<true>/forward<true>/...
+    end up as hidden symbols and ld errors with:
+
+        hidden symbol `..._fill_bg<true>...' isn't defined
+        final link failed: bad value
+
+    Unique3D never imports pulsar (only pytorch3d.renderer.{cameras,mesh,...}
+    and pytorch3d.structures). Strip pulsar entirely before building.
+    """
+    import shutil
+
+    # 1. Remove the C++/CUDA sources so setup.py's glob skips them.
+    pulsar_csrc = os.path.join(src, "pytorch3d", "csrc", "pulsar")
+    if os.path.isdir(pulsar_csrc):
+        shutil.rmtree(pulsar_csrc)
+
+    # 2. Remove the python-side pulsar package.
+    pulsar_py = os.path.join(src, "pytorch3d", "renderer", "points", "pulsar")
+    if os.path.isdir(pulsar_py):
+        shutil.rmtree(pulsar_py)
+
+    # 3. Rewrite ext.cpp to drop every pulsar-related include and the entire
+    # pulsar pybind11 registration block. Use the well-defined "Pulsar." comment
+    # marker that starts the registration block and run to the matching #endif.
+    ext_cpp = os.path.join(src, "pytorch3d", "csrc", "ext.cpp")
+    if os.path.isfile(ext_cpp):
+        with open(ext_cpp, "r") as f:
+            lines = f.readlines()
+
+        # Phase A: drop the three pulsar include lines (they sit inside
+        # `#if !defined(USE_ROCM)` blocks together with non-pulsar includes? No —
+        # check the file: the top of ext.cpp has two #if !defined(USE_ROCM)
+        # blocks that only contain pulsar includes. Drop the whole #if/#endif
+        # pairs that contain only pulsar includes.
+        out = []
+        i = 0
+        while i < len(lines):
+            ln = lines[i]
+            if ln.strip().startswith("#if !defined(USE_ROCM)"):
+                # Look ahead for the matching #endif. Collect block body.
+                j = i + 1
+                block = []
+                while j < len(lines) and not lines[j].strip().startswith("#endif"):
+                    block.append(lines[j])
+                    j += 1
+                # Decide: if every non-blank/non-comment line in block mentions
+                # "pulsar", drop the whole #if..#endif. Otherwise keep as-is.
+                meaningful = [
+                    b for b in block
+                    if b.strip() and not b.strip().startswith("//")
+                ]
+                if meaningful and all("pulsar" in b.lower() for b in meaningful):
+                    # Drop block (and matching #endif).
+                    i = j + 1
+                    continue
+            out.append(ln)
+            i += 1
+
+        # Phase B: drop the pulsar pybind registration block.
+        # Find the "// Pulsar." comment, walk back to its preceding blank line,
+        # then forward to the matching #endif (the last #endif before `}`).
+        text = "".join(out)
+        marker = "  // Pulsar."
+        idx = text.find(marker)
+        if idx >= 0:
+            # Find the start of the line containing the marker.
+            line_start = text.rfind("\n", 0, idx) + 1
+            # Find the trailing `#endif` of the pulsar block. Walk forward; the
+            # block ends with a `#endif` on a line by itself, right before `}`.
+            end_marker = "\n#endif\n}"
+            end_idx = text.find(end_marker, line_start)
+            if end_idx < 0:
+                # Fall back: just drop to the next `}` on its own line.
+                end_idx = text.find("\n}\n", line_start)
+                if end_idx < 0:
+                    raise RuntimeError("could not locate pulsar block end in ext.cpp")
+                # Keep the `}`.
+                text = text[:line_start] + text[end_idx + 1:]
+            else:
+                # Drop the `#endif` line but keep the `}`.
+                text = text[:line_start] + text[end_idx + len("\n#endif"):]
+
+        with open(ext_cpp, "w") as f:
+            f.write(text)
+
+    # 4. Patch renderer/points/__init__.py AND renderer/__init__.py: both have a
+    # `if not torch.version.hip:` block that imports PulsarPointsRenderer; replace
+    # the import with `pass` so the symbol is just absent.
+    for init_path in [
+        os.path.join(src, "pytorch3d", "renderer", "points", "__init__.py"),
+        os.path.join(src, "pytorch3d", "renderer", "__init__.py"),
+    ]:
+        if not os.path.isfile(init_path):
+            continue
+        with open(init_path, "r") as f:
+            init_lines = f.readlines()
+        new_lines = []
+        skip_block = False
+        for ln in init_lines:
+            if "torch.version.hip" in ln and ln.lstrip().startswith("if "):
+                indent = " " * (len(ln) - len(ln.lstrip()))
+                new_lines.append(ln)
+                new_lines.append(indent + "    pass  # pulsar disabled (Blackwell linker)\n")
+                skip_block = True
+                continue
+            if skip_block:
+                # Drop any indented continuation lines of the if-block, including
+                # multi-line `from .points import (` parenthesized imports.
+                stripped = ln.strip()
+                if stripped == "" or ln.startswith(" ") or ln.startswith("\t"):
+                    if "pulsar" in ln.lower() or stripped == "" or stripped.startswith(")") or stripped.startswith("from ") or stripped.startswith("import "):
+                        # part of the pulsar import block — drop
+                        continue
+                    skip_block = False
+                else:
+                    skip_block = False
+            new_lines.append(ln)
+        with open(init_path, "w") as f:
+            f.writelines(new_lines)
+
+
+@spaces.GPU(duration=1500)
+def _build_pytorch3d():
+    # Probe the actual import path Unique3D uses; if the previously installed
+    # pytorch3d still has a broken `renderer/__init__.py` (e.g. left a dangling
+    # `from .points import PulsarPointsRenderer`), uninstall and rebuild.
+    try:
+        from pytorch3d.renderer import TexturesVertex  # noqa: F401
+        return
+    except Exception:
+        pass
+    subprocess.call(
+        [sys.executable, "-m", "pip", "uninstall", "-y", "pytorch3d"],
+    )
+
+    src = tempfile.mkdtemp(prefix="pytorch3d_src_")
+    subprocess.check_call(
+        ["git", "clone", "--depth=1", "--branch=stable",
+         "https://github.com/facebookresearch/pytorch3d.git", src],
+    )
+    _strip_pulsar_from_pytorch3d_src(src)
+
+    # Force a fresh wheel build (don't reuse the cached broken wheel from a
+    # previous run, which still has the pulsar import in renderer/__init__.py).
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install",
+         "--no-build-isolation", "--no-deps",
+         "--force-reinstall", "--no-cache-dir", src],
+        env=_build_env(),
     )
-)
+
+
+@spaces.GPU(duration=600)
+def _build_nvdiffrast():
+    _ensure_pkg("nvdiffrast", ["git+https://github.com/NVlabs/nvdiffrast.git"])
+
+
+# Run setup before importing anything that needs these packages.
+_build_pytorch3d()
+_build_nvdiffrast()
+
 
 if __name__ == "__main__":
-    import os
     from huggingface_hub import login
     hf_token = os.environ.get("HF_TOKEN")
-    login(token=hf_token)
-    
-    import os
-    import sys
+    if hf_token:
+        login(token=hf_token)
+
     sys.path.append(os.curdir)
     import torch
     torch.set_float32_matmul_precision('medium')
@@ -51,7 +237,7 @@ _DESCRIPTION = '''
 
 def launch():
     model_zoo.init_models()
-        
+
     with gr.Blocks(
         title=_TITLE,
         # theme=gr.themes.Monochrome(),
@@ -63,6 +249,6 @@ def launch():
         create_3d_ui("wkl")
 
     demo.queue().launch(share=True)
-    
+
 if __name__ == '__main__':
     fire.Fire(launch)

gradio_app/all_models.py

@@ -6,7 +6,7 @@ from diffusers import StableDiffusionControlNetImg2ImgPipeline, StableDiffusionP
 class MyModelZoo:
     _pipe_disney_controlnet_lineart_ipadapter_i2i: StableDiffusionControlNetImg2ImgPipeline = None
     
-    base_model = "benjamin-paine/stable-diffusion-v1-5"
+    base_model = "stable-diffusion-v1-5/stable-diffusion-v1-5"
 
     def __init__(self, base_model=None) -> None:
         if base_model is not None:

gradio_app/gradio_3dgen.py

@@ -11,14 +11,14 @@ from scripts.utils import save_glb_and_video
 # from scripts.multiview_inference import geo_reconstruct
 from scripts.multiview_inference import geo_reconstruct_part1, geo_reconstruct_part2, geo_reconstruct_part3
 
-@spaces.GPU(duration=100)
+@spaces.GPU(duration=600)
 def run_mv(preview_img, input_processing, seed):
     if preview_img.size[0] <= 512:
         preview_img = run_sr_fast([preview_img])[0]
     rgb_pils, front_pil = run_mvprediction(preview_img, remove_bg=input_processing, seed=int(seed)) # 6s
     return rgb_pils, front_pil
 
-@spaces.GPU(duration=100) # seems split into multiple part will leads to `RuntimeError`, before fix it, still initialize here
+@spaces.GPU(duration=600) # seems split into multiple part will leads to `RuntimeError`, before fix it, still initialize here
 def generate3dv2(preview_img, input_processing, seed, render_video=True, do_refine=True, expansion_weight=0.1, init_type="std"):
     if preview_img is None:
         raise gr.Error("The input image is none!")

packages.txt

@@ -0,0 +1,4 @@
+libopengl0
+libegl1
+libxkbcommon-x11-0
+libdbus-1-3

requirements.txt

@@ -1,7 +1,5 @@
-pytorch3d @ https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu121_pyt221/pytorch3d-0.7.6-cp310-cp310-linux_x86_64.whl
-ort_nightly_gpu @ https://aiinfra.pkgs.visualstudio.com/2692857e-05ef-43b4-ba9c-ccf1c22c437c/_packaging/d3daa2b0-aa56-45ac-8145-2c3dc0661c87/pypi/download/ort-nightly-gpu/1.17.dev20240118002/ort_nightly_gpu-1.17.0.dev20240118002-cp310-cp310-manylinux_2_28_x86_64.whl
-onnxruntime_gpu @ https://pkgs.dev.azure.com/onnxruntime/2a773b67-e88b-4c7f-9fc0-87d31fea8ef2/_packaging/7fa31e42-5da1-4e84-a664-f2b4129c7d45/pypi/download/onnxruntime-gpu/1.17/onnxruntime_gpu-1.17.0-cp310-cp310-manylinux_2_28_x86_64.whl
-torch==2.2.0
+onnxruntime-gpu
+torch_scatter @ https://data.pyg.org/whl/torch-2.11.0+cu130/torch_scatter-2.1.2+pt211cu130-cp310-cp310-linux_x86_64.whl
 accelerate
 datasets
 diffusers>=0.26.3
@@ -9,7 +7,7 @@ fire
 gradio
 jaxtyping
 numba
-numpy
+numpy<2
 omegaconf>=2.3.0
 opencv_python
 opencv_python_headless
@@ -18,12 +16,10 @@ Pillow
 pygltflib
 pymeshlab>=2023.12
 rembg[gpu]
-torch>=2.0.1
-torch_scatter @ https://data.pyg.org/whl/torch-2.2.0%2Bcu121/torch_scatter-2.1.2%2Bpt22cu121-cp310-cp310-linux_x86_64.whl
 tqdm
 transformers
 trimesh
 typeguard
 wandb
 xformers
-ninja
+ninja

scripts/load_onnx.py

@@ -1,20 +1,9 @@
 import onnxruntime
 import torch
 
-providers = [
-    # ('TensorrtExecutionProvider', {
-    #     'device_id': 0,
-    #     'trt_max_workspace_size': 8 * 1024 * 1024 * 1024,
-    #     'trt_fp16_enable': True,
-    #     'trt_engine_cache_enable': True,
-    # }),
-    ('CUDAExecutionProvider', {
-        'device_id': 0,
-        'arena_extend_strategy': 'kSameAsRequested',
-        'gpu_mem_limit': 8 * 1024 * 1024 * 1024,
-        'cudnn_conv_algo_search': 'HEURISTIC',
-    })
-]
+# ZeroGPU containers ship CUDA 13; the onnxruntime-gpu wheel links against
+# libcublasLt.so.12 (CUDA 12) → CUDAExecutionProvider fails to load. Use CPU.
+providers = ['CPUExecutionProvider']
 
 def load_onnx(file_path: str):
     assert file_path.endswith(".onnx")

scripts/sd_model_zoo.py

@@ -4,7 +4,7 @@ import torch
 from copy import deepcopy
 
 ENABLE_CPU_CACHE = False
-DEFAULT_BASE_MODEL = "benjamin-paine/stable-diffusion-v1-5"
+DEFAULT_BASE_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5"
 
 cached_models = {}  # cache for models to avoid repeated loading, key is model name
 def cache_model(func):

scripts/utils.py

@@ -13,14 +13,9 @@ from typing import List, Tuple
 from PIL import Image
 import trimesh
 
-providers = [
-    ('CUDAExecutionProvider', {
-        'device_id': 0,
-        'arena_extend_strategy': 'kSameAsRequested',
-        'gpu_mem_limit': 8 * 1024 * 1024 * 1024,
-        'cudnn_conv_algo_search': 'HEURISTIC',
-    })
-]
+# onnxruntime-gpu wheel links against CUDA 12 libs; ZeroGPU containers ship
+# CUDA 13 → CUDAExecutionProvider can't load. Use CPU for rembg.
+providers = ['CPUExecutionProvider']
 
 session = new_session(providers=providers)