diff --git a/.gitignore b/.gitignore index b600824..ed59cc9 100644 --- a/.gitignore +++ b/.gitignore @@ -36,3 +36,24 @@ tests/ss_dec_ref.bin # clangd index .cache/ +corpus_*/ +crash-* + +# fuzzing +/corpus_*/ +crash-* + +# demo server build artifacts +/server/trellis2-server +/server/trellis2-server-linux +/server/vendor/ + +# python bytecode +__pycache__/ +*.pyc + +# model / reference artifacts +/models/ +/ggufs/ +/dumps/ +slow-unit-* diff --git a/CMakeLists.txt b/CMakeLists.txt index f1c48c5..d3aa9d7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -20,6 +20,17 @@ endif() # ───────────────────────────────────────────────────────────────────────── option(TRELLIS2_USE_EXTERNAL_GGML "Consume external ggml via find_package() instead of the bundled submodule" OFF) +# libFuzzer harnesses (clang only). The sanitizer flags apply globally — set +# BEFORE any target so the whole tree (ggml, stb decode, preprocessing, +# loaders) is instrumented. Harness targets live in fuzz/ (added at the end). +option(TRELLIS2_FUZZ "Build libFuzzer harnesses (requires clang)" OFF) +set(TRELLIS2_FUZZ_SANITIZERS "address,undefined" CACHE STRING + "Sanitizers for the fuzz build") +if(TRELLIS2_FUZZ) + add_compile_options(-fsanitize=${TRELLIS2_FUZZ_SANITIZERS} -fsanitize=fuzzer-no-link -fno-omit-frame-pointer -g -O1) + add_link_options(-fsanitize=${TRELLIS2_FUZZ_SANITIZERS}) +endif() + # ggml options — enable Metal on Apple (bundled-ggml path only) option(TRELLIS2_METAL "Enable Metal backend" ON) if(APPLE AND TRELLIS2_METAL AND NOT TRELLIS2_USE_EXTERNAL_GGML) @@ -35,17 +46,33 @@ else() set(TRELLIS2_GGML_TARGET ggml) endif() -add_library(trellis2 trellis2.cpp trellis2.h) +add_library(trellis2 + trellis2.cpp trellis2.h + trellis2_capi.cpp trellis2_capi.h + mesh_export.cpp mesh_export.h + third_party/xatlas/xatlas.cpp + third_party/meshoptimizer/simplifier.cpp + third_party/meshoptimizer/indexgenerator.cpp + third_party/meshoptimizer/vfetchoptimizer.cpp + third_party/meshoptimizer/allocator.cpp) target_include_directories(trellis2 PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/stb ) +# trellis2_capi.cpp reuses the example isosurface extractor for the demo path; +# mesh_export.cpp vendors xatlas (UV unwrap) and Simplify.h (QEM decimation). +target_include_directories(trellis2 PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/examples + ${CMAKE_CURRENT_SOURCE_DIR}/third_party/xatlas + ${CMAKE_CURRENT_SOURCE_DIR}/third_party/meshoptimizer) + +find_package(Threads REQUIRED) # xatlas uses std::thread if(TRELLIS2_USE_EXTERNAL_GGML AND DEFINED GGML_INCLUDE_DIR) target_include_directories(trellis2 PUBLIC ${GGML_INCLUDE_DIR}) endif() -target_link_libraries(trellis2 PUBLIC ${TRELLIS2_GGML_TARGET}) +target_link_libraries(trellis2 PUBLIC ${TRELLIS2_GGML_TARGET} PRIVATE Threads::Threads) target_compile_features(trellis2 PUBLIC cxx_std_14) set_property(TARGET trellis2 PROPERTY POSITION_INDEPENDENT_CODE ON) @@ -62,5 +89,10 @@ endif() option(TRELLIS2_BUILD_TESTS "Build test executables" OFF) if(TRELLIS2_BUILD_TESTS) + enable_testing() add_subdirectory(tests) endif() + +if(TRELLIS2_FUZZ) + add_subdirectory(fuzz) +endif() diff --git a/README.md b/README.md index cbc7f64..0e3bb2d 100644 --- a/README.md +++ b/README.md @@ -1,23 +1,88 @@ # trellis2.cpp A C++/[ggml](https://github.com/ggml-org/ggml) implementation of the -**TRELLIS.2** image-to-3D pipeline (stage 1 first: sparse-structure flow). +**TRELLIS.2** image-to-3D pipeline: an image goes in, a 3D mesh with per-vertex +PBR textures comes out, with all inference in C++/ggml (no PyTorch at runtime). +The demo can also bake the result into a portable, image-textured **GLB** — +CUDA-free UV atlas + PBR bake, no reference container required. Modeled structurally after [sam3.cpp](https://github.com/rms80/sam3.cpp): single-file library (`trellis2.h` / `trellis2.cpp`), bundled ggml as a submodule (Metal on by default on Apple), DLL-export decoration, and a -CMake build with example executables. +CMake build with example executables. A flat C ABI (`trellis2_capi.h`) drives +a Go demo server with a browser mesh viewer. -## Status +## Quick start (demo) -Early scaffolding. Implemented so far: +```sh +git submodule update --init --depth 1 # ggml +scripts/download_models.sh # HF checkpoints -> models/ (~7 GB) +docker build -f docker/Dockerfile.ref -t trellis2-ref docker # convert weights / gen refs +docker build -f docker/Dockerfile.demo -t trellis2-demo docker # CUDA runtime + Go +# convert every checkpoint to GGUF (f16 for the demo, f32 for validation) +docker run --rm -v "$PWD":/work -w /work trellis2-ref bash scripts/convert_all.sh +# build the CUDA shared lib + Go server, then run +docker run --rm -v "$PWD":/work -w /work trellis2-demo bash -c ' + cmake -B build-cuda-shared -G Ninja -DCMAKE_BUILD_TYPE=Release -DGGML_CUDA=ON \ + -DCMAKE_CUDA_ARCHITECTURES=120 -DBUILD_SHARED_LIBS=ON && cmake --build build-cuda-shared -j + cd server && go build -o trellis2-server-linux .' +docker run --rm --device nvidia.com/gpu=all -v "$PWD":/work -w /work/server -p 8742:8742 \ + trellis2-demo ./trellis2-server-linux -lib /work/build-cuda-shared/libtrellis2.so -ggufs /work/ggufs +# open http://localhost:8742 , drop an image (transparent background works best) +``` + +The browser UI has a **quality** selector: coarse preview (64³ marching cubes), +512³ fine, or **1024³ cascade** (the TRELLIS.2 default — sharpest). Coarse falls +back automatically if the shape-SLAT models are absent (`-coarse`); the 1024 +cascade needs the extra 1024 model (`-no-1024` disables it). + +On a 16 GB RTX 50-series: the 512 fine path runs image→mesh in ~110 s (~1M-vertex +512³ mesh); the 1024 cascade adds a second 1.3B-model pass and the 1024³ decoder +for a ~5M-vertex mesh (~5 min, ~10 GB VRAM, and a ~14 GB host-RAM spike for the +1024³ sparse-conv decode). + +## Pipeline + +``` +image (RGBA) + → preprocess alpha bbox crop, premultiply, PIL-exact Lanczos-512 [C++, byte-exact] + → DINOv3 ViT-L/16 [1, 1029, 1024] conditioning tokens [C++/ggml] + → SS-flow DiT 1.3B dense DiT, 12-step CFG flow-Euler → z_s [C++/ggml] + → SS decoder dense 3D-conv → 64³ occupancy → 32³ voxel scaffold [C++/ggml] + → shape-SLAT DiT 1.3B sparse DiT over active voxels, 12-step CFG [C++/ggml] + → shape VAE decoder sparse ConvNeXt U-Net, 16× up → 512³ dual-grid fields [C++/ggml] + → flexible dual grid → triangle mesh [C++] +``` -- **`.dinodata` loader** — reads the DINOv3 conditioning tensor that the - sparse-structure flow DiT consumes as cross-attention K/V. This is the full - DINOv3 ViT-L/16 token sequence (`[1, 1029, 1024]` = 1 CLS + 4 register + - 1024 patch) of the **last transformer layer** with an **affine-free - LayerNorm** (NOT HF's `last_hidden_state`). `neg_cond = zeros_like(cond)`, - so it is never stored. Files are produced by `trellis2-shiv/dump_dinodata.py`. +The **1024 cascade** (default in TRELLIS.2) adds a second pass on top: the 512 +result's decoder `.upsample(×4)` predicts a denser coordinate scaffold, which is +quantized to 64³ and fed to a second 1.3B shape-SLAT flow (the 1024 model, +conditioned on a 1024-res DINOv3 encode) and the same decoder at 1024³ — a +~5M-vertex mesh. The ~49k-token HR attention only fits in VRAM via flash +attention (`sdpa_auto`); see [docs/VERIFICATION.md](docs/VERIFICATION.md). + +Every stage is validated tap-by-tap against the PyTorch reference — see +[docs/VERIFICATION.md](docs/VERIFICATION.md). Highlights: preprocessing is +byte-exact, the DINOv3 encoder matches to rel-L2 ≤ 7e-7 across 40 taps, and the +sparse U-Net decoder is numerically exact through all four conv levels. + +## Components + +- **Image preprocessing + DINOv3 encoder** — `trellis2_preprocess_rgba()` + reproduces `pipeline.preprocess_image` (the has-alpha path) with a + PIL-compatible fixed-point Lanczos-3 resampler (byte-exact vs Pillow). + `trellis2_dino_encode()` runs the full DINOv3 ViT-L/16 (axial-2D RoPE, + LayerScale, exact-GELU MLP) and applies the affine-free final LayerNorm the + flow models expect — the `[1, 1029, 1024]` conditioning that used to come + from an external `dump_dinodata.py`. `dino_encode` chains them: + + ```sh + ./build/examples/dino_encode ggufs/dino_f16.gguf image.png cond.dinodata + ``` + +- **`.dinodata` loader** — `trellis2_load_dinodata()` still reads/writes the + precomputed conditioning tensor (1 CLS + 4 register + 1024 patch, last layer, + affine-free LN; `neg_cond = zeros_like(cond)`), for testing and CLI chaining. - **SS-flow DiT weights** — `convert_ss_flow_to_gguf.py` converts the stage-1 `ss_flow_img_dit_1_3B_64_bf16` checkpoint to GGUF; `trellis2_ss_flow_load()` @@ -58,10 +123,10 @@ Early scaffolding. Implemented so far: # -> logits [1,64,64,64], occupied(>0) grid (the coarse voxel scaffold) ``` -- **Occupancy → mesh** — `ss_mesh` decodes a z_s latent and exports the +- **Occupancy → coarse mesh** — `ss_mesh` decodes a z_s latent and exports the `{logit = 0}` isosurface as a watertight OBJ via a self-contained marching cubes (`examples/marching_cubes.h`, the tetrahedral / Freudenthal variant — no - 256-row table, provably manifold). Chain it after sampling to *see* the result: + 256-row table, provably manifold). This is the fast preview path: ```sh ./build/examples/ss_sample ss_flow_dit_f16.gguf /path/img.dinodata z_s.latent @@ -69,6 +134,18 @@ Early scaffolding. Implemented so far: # -> watertight shape.obj in the centered unit cube; open in any 3D viewer ``` +- **Shape-SLAT flow + decoder (fine geometry)** — `trellis2_slat_flow_sample()` + runs the sparse 1.3B DiT over the active voxels of the 32³ scaffold (same + block structure as the SS-flow DiT, 3D RoPE over each voxel's coords), + denormalized with `shape_slat_normalization` baked into the GGUF. + `trellis2_shape_dec_decode()` runs `FlexiDualGridVaeDecoder` — a sparse + ConvNeXt U-Net whose 3×3×3 submanifold convolutions are expressed as 27 + gather+GEMM steps, with each level's learned subdivision growing the active + set (32³ → 512³, 16×). `examples/flexible_dual_grid.h` turns the 7-channel + per-voxel output (dual-vertex offset, per-axis intersection flags, quad split + weight) into the triangle mesh. This is the real TRELLIS.2 geometry, driven + end-to-end by the demo server. + ## Validate the forward pass ```sh @@ -154,11 +231,16 @@ reduction. | `trellis2.cpp` | implementation | | `convert_ss_flow_to_gguf.py` | stage-1 DiT checkpoint → GGUF converter | | `convert_ss_dec_to_gguf.py` | stage-1 decoder checkpoint → GGUF converter | -| `examples/` | CLI tools (`dino_info`, `ss_flow_info`, `ss_sample`, `ss_decode`, `ss_mesh`) | +| `mesh_export.{h,cpp}` | CUDA-free textured-GLB bake: decimate → UV atlas → PBR bake → glTF | +| `examples/` | CLI tools (`dino_info`, `ss_flow_info`, `ss_sample`, `ss_decode`, `ss_mesh`, `mesh2glb`) | | `examples/marching_cubes.h` | single-file isosurface → OBJ extractor | +| `third_party/` | vendored `meshoptimizer` (QEM decimate) + `xatlas` (opt-in charts) | | `ggml/` | submodule, pinned to the same commit as sam3.cpp | | `stb/` | `stb_image.h` / `stb_image_write.h` for image I/O | ## License -MIT. See [LICENSE](LICENSE). +MIT. See [LICENSE](LICENSE). Vendored third-party code is also MIT: +[meshoptimizer](https://github.com/zeux/meshoptimizer) (Arseny Kapoulkine) and +[xatlas](https://github.com/jpcy/xatlas) (Jonathan Young) under `third_party/`, +and `stb` (public domain / MIT). diff --git a/convert_dino_to_gguf.py b/convert_dino_to_gguf.py new file mode 100644 index 0000000..d5a102f --- /dev/null +++ b/convert_dino_to_gguf.py @@ -0,0 +1,201 @@ +#!/usr/bin/env python3 +""" +Convert the DINOv3 ViT-L/16 image-conditioning encoder (HF format, +facebook/dinov3-vitl16-pretrain-lvd1689m or a mirror) to a GGUF file for +trellis2.cpp. + +TRELLIS.2 uses this model (transformers DINOv3ViTModel) to turn the +preprocessed 512x512 input image into the [1, 1029, 1024] conditioning tensor +consumed by the flow models: 1 CLS + 4 register + 32x32 patch tokens, last +transformer layer output passed through an affine-free layer norm (the model's +own final `norm` layer is NOT applied). + +Architecture (config.json): hidden 1024, 24 layers, 16 heads, MLP 4096 (exact +GELU, not gated), patch 16, axial 2D RoPE over patch-center coords with +theta=100, LayerScale on both residual branches, q/v/o biases but no k bias. + +Self-contained like the other converters (safetensors + numpy). Tensors keep +their HF state-dict names; hyperparameters travel as `trellis2.dino.*` KV. + +Usage: + python convert_dino_to_gguf.py --model models/dinov3-vitl16/model.safetensors \ + --output ggufs/dino_f32.gguf --ftype 0 + +ftype: 0 = f32 (validation), 1 = f16 (matrices f16; norms/biases/1-D f32). +""" + +import argparse +import json +import os +import struct +import sys + +import numpy as np + +# ── GGUF / GGML constants (must match the bundled ggml) ────────────────────── +GGUF_MAGIC = b"GGUF" +GGUF_VERSION = 3 +GGUF_ALIGNMENT = 32 + +GGML_TYPE_F32 = 0 +GGML_TYPE_F16 = 1 + +GGUF_VT_UINT32 = 4 +GGUF_VT_FLOAT32 = 6 +GGUF_VT_BOOL = 7 +GGUF_VT_STRING = 8 + +ARCH = "trellis2-dino" +KV_PREFIX = "trellis2.dino." + + +def _gguf_str(s: str) -> bytes: + b = s.encode("utf-8") + return struct.pack(" bytes: + return _gguf_str(key) + struct.pack(" int: + return (n + a - 1) // a * a + + +def choose_type(name, shape, ftype: int) -> int: + if ftype == 0: + return GGML_TYPE_F32 + # matrices f16; 1-D (biases, norms, layerscale) and token embeddings f32 + if len(shape) < 2: + return GGML_TYPE_F32 + if "cls_token" in name or "register_tokens" in name or "mask_token" in name: + return GGML_TYPE_F32 + return GGML_TYPE_F16 + + +def to_bytes(arr_f32: np.ndarray, ggml_type: int) -> bytes: + if ggml_type == GGML_TYPE_F32: + return arr_f32.astype("