diff --git a/.gitignore b/.gitignore index b600824..0c505ed 100644 --- a/.gitignore +++ b/.gitignore @@ -36,3 +36,25 @@ 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/ +/generations/ +slow-unit-* diff --git a/CMakeLists.txt b/CMakeLists.txt index f1c48c5..f6ed085 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,28 @@ 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) 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 for the optional chart-based UV unwrap. +target_include_directories(trellis2 PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/examples + ${CMAKE_CURRENT_SOURCE_DIR}/third_party/xatlas) + +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 +84,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..35aee87 100644 --- a/README.md +++ b/README.md @@ -1,23 +1,150 @@ # 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 export the result into a portable, full-density **GLB** with +standard interpolated vertex colour and retained PBR attributes—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_ggufs.sh # prebuilt f16 GGUFs -> ggufs/ (~14 GB) +docker build -f docker/Dockerfile.demo -t trellis2-demo docker # CUDA runtime + Go +# 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 -store /work/generations -unload-idle +# open http://localhost:8742 and drop an image +``` + +Or just run `scripts/demo.sh`, which builds the lib + server, auto-downloads any +missing GGUFs, and launches the container. + +### Prebuilt GGUFs + +`scripts/download_ggufs.sh` pulls the ready-made f16 GGUFs from three public repos +under the [LocalAI-io](https://huggingface.co/LocalAI-io) org, so you can skip the +safetensors download and the conversion step entirely: + +- [`TRELLIS.2-4B-GGUF`](https://huggingface.co/LocalAI-io/TRELLIS.2-4B-GGUF) — MIT +- [`TRELLIS-image-large-GGUF`](https://huggingface.co/LocalAI-io/TRELLIS-image-large-GGUF) — MIT +- [`dinov3-vitl16-pretrain-lvd1689m-GGUF`](https://huggingface.co/LocalAI-io/dinov3-vitl16-pretrain-lvd1689m-GGUF) — DINOv3 License (Built with DINOv3) + +**Developers** who need the f32 validation variants, or who want to regenerate the +GGUFs from source, use the original flow instead: `scripts/download_models.sh` (HF +safetensors → `models/`, ~7 GB) then `docker run … trellis2-ref bash +scripts/convert_all.sh` (safetensors → GGUF, f16 + f32). The card + license sources +for the published repos live in `scripts/hf/`, and `scripts/upload_ggufs.sh` +(re)publishes them. + +Completed generations are committed atomically under `generations/` (final +mesh, replay frames, and manifest) and restored with the same job IDs after a +server restart. Pass `-store ''` to disable persistence or `-store PATH` to use +a different durable location. Incomplete writes are ignored on startup. With +`-unload-idle`, the HTTP server starts without allocating model VRAM, loads the +pipeline on the first generation, and releases it again when the queue is idle. + +The browser UI has a **quality** selector: coarse preview (64³ marching cubes), +512³ fine, or **1024³ cascade** (the TRELLIS.2 default and highest resolution +currently supported here). Upstream's optional 1536 cascade is not yet ported. Coarse falls +back automatically if the shape-SLAT models are absent (`-coarse`); the 1024 +cascade needs the extra 1024 model (`-no-1024` disables it). +Enable **free VRAM when idle** to unload the resident model pipeline between +generations; the next queued generation reloads it automatically. +**Live steps** is off by default because each sparse-structure frame requires an +extra CPU occupancy decode between GPU inference steps. Its button always says +`on` or `off`; enabling it records the frames used by replay and showcase mode. +Completed jobs expose `durationMs`, `livePreview`, and per-stage `stageTimings` +through `/api/job/{id}` and persist those diagnostics in their manifest. +The always-visible **asset export** panel preserves the generated polygon count, +can preview component cleanup, optionally keep only the largest connected piece, +restore the original preview, and download a Three.js-ready GLB. Dense generated +materials are stored as standard interpolated vertex colours rather than a +sub-texel per-triangle atlas; original metallic/roughness values are also retained +in the custom `_METALLIC_ROUGHNESS` attribute. All components are preserved by +default; destructive cleanup must be selected explicitly. Showcase mode likewise loads the original +full-density mesh. Open `/showcase` for the separate full-screen storyboard: it +starts each generation with its saved source image centered, moves that image to +the upper-right, replays the recorded stages, and then lingers on a slowly +rotating final model. New generations retain the original upload byte-for-byte +for display, plus the full-resolution processed PNG actually used by TRELLIS for +repeatable server-side regeneration without another upload. Select a saved mesh +and use **regenerate from saved image** to run it again with the current settings. +Older manifests fall back to their thumbnail and can only regenerate when their +older processed source file is available. + +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 (RGB/RGBA) + → background cleanup border-connected black/white → feathered alpha [C++/browser] + → 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 → decoded dual grid [C++/ggml] + ├→ flexible dual grid → triangle mesh [C++] + └→ shape VAE encoder validated dual-grid → shape SLat + subdivision guide [C++/ggml] + → texture-SLAT DiT shape-SLat concat conditioning [C++/ggml] + → texture decoder replay subdivision → sparse 6-channel PBR volume [C++/ggml] + → material sampling trilinear PBR at surface vertices [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). + +The neural components are validated tap-by-tap against the PyTorch reference, +with separate integration regressions for subdivision guidance, sparse material +sampling, and GLB alpha preservation — 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_remove_solid_background_rgba()` first converts a detected near-black + or near-white background connected to the image border into softly feathered + alpha, while preserving enclosed black/white subject details and existing + alpha masks. The demo exposes automatic, forced-black, forced-white, and keep + original modes. + `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 +185,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 +196,27 @@ 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. + +- **PBR texture generation** — the decoded dual grid is encoded to the shape + SLat used to condition texture flow, using the numerically validated + standalone texturing path. The decoded six-channel volume (base color, + metallic, roughness, alpha) is sampled trilinearly at the actual dual-grid + surface positions. Collapsed all-saturated outputs are rejected instead of + being persisted as apparently successful textures. The browser linearizes base + color before PBR lighting and preserves opacity. Material sampling steps are + controlled separately from geometry steps, matching upstream's defaults. + ## Validate the forward pass ```sh @@ -154,11 +302,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 full-density GLB export with direct vertex colour/PBR attributes | +| `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 `xatlas` (opt-in chart-based UV unwrap) | | `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("