Initial release
This commit is contained in:
+197
@@ -0,0 +1,197 @@
|
||||
# trellis2.cpp completion plan
|
||||
|
||||
Goal: image in → 3D mesh out, pure C++/ggml inference behind a Go demo server with a
|
||||
browser mesh viewer. No PBR textures for now. Process mirrors depth-anything.cpp
|
||||
(layer-by-layer parity vs the PyTorch reference) and privacy-filter.cpp (libFuzzer
|
||||
harnesses on untrusted inputs, sanitizer builds).
|
||||
|
||||
## Pipeline target
|
||||
|
||||
TRELLIS.2 "512" pipeline type (non-cascade), geometry only:
|
||||
|
||||
```
|
||||
image (RGBA preferred)
|
||||
→ preprocess (alpha crop, premultiply, resize 512, ImageNet norm) [C++ port]
|
||||
→ DINOv3 ViT-L/16 cond tokens [1, 1+4+1024, 1024] [C++ port — was external .dinodata]
|
||||
→ SS-flow DiT (1.3B dense, 12-step CFG flow Euler) [already ported + validated]
|
||||
→ SS decoder → 64³ occupancy → max_pool 32³ → coords [already ported; pool/coords new]
|
||||
→ shape-SLAT flow (1.3B sparse DiT, res 32, 32ch) [to port]
|
||||
→ FlexiDualGrid VAE decoder (sparse ConvNeXt U-Net, 16× up) [to port]
|
||||
→ flexible dual grid → triangle mesh [to port, CPU]
|
||||
→ (postprocess: fill holes — best effort, CuMesh not portable)
|
||||
```
|
||||
|
||||
Fallback shipped at every point in time: marching-cubes preview mesh from the 64³
|
||||
occupancy (already works), so the demo is demoable before the SLAT stages land.
|
||||
|
||||
The upstream neural background-removal net (BiRefNet/RMBG-2.0, a separate ~1 GB
|
||||
model) remains out of scope. The port does include deterministic removal of
|
||||
near-black/near-white backgrounds connected to the image border, with soft alpha
|
||||
edges and browser controls to force or disable it. Texture and cascade status is
|
||||
tracked below.
|
||||
|
||||
## Weights
|
||||
|
||||
`scripts/download_models.sh` → `models/`. DINOv3 comes from the ungated
|
||||
`camenduru/dinov3-vitl16-pretrain-lvd1689m` mirror because the official
|
||||
`facebook/dinov3-vitl16-pretrain-lvd1689m` is license-gated and this account has no
|
||||
access yet (403). Request access and re-download officially when possible.
|
||||
|
||||
## Validation (depth-anything.cpp process)
|
||||
|
||||
- Reference activations dumped from Python via forward hooks into
|
||||
`dumps/reference_<stage>.gguf` + manifest (GGUF as the dump format so C++ reads it
|
||||
with the ggml API it already links).
|
||||
- Reference env: docker `pytorch/pytorch:2.7.1-cuda12.8-cudnn9-devel` + stock pip
|
||||
deps only. Custom CUDA backends (FlexGEMM/flash-attn/o-voxel) are NOT installed;
|
||||
sparse attention (batch=1) is monkeypatched to dense SDPA and sparse conv to a
|
||||
pure-torch gather-GEMM — slow but runs on CPU too, and doubles as an executable
|
||||
spec for the C++ port. 16 GB VRAM is only a constraint for the reference runs;
|
||||
stage-at-a-time + low_vram-style model swapping keeps us under it.
|
||||
- C++ side: per-layer taps compared with `tests/parity.hpp`-style
|
||||
`atol+rtol*|ref|` gate (2e-3 default), ctest with `SKIP_RETURN_CODE 77` when
|
||||
fixtures are absent.
|
||||
- Existing whole-stage tests (ss_flow/ss_sample/ss_dec) keep working; their
|
||||
ref generators get path fixes (the `trellis2-shiv` layout doesn't exist here).
|
||||
|
||||
## Fuzzing (privacy-filter.cpp process)
|
||||
|
||||
Non-GGUF untrusted inputs, libFuzzer + ASan/UBSan (clang):
|
||||
- image bytes → stb_image decode → preprocess (the demo upload path)
|
||||
- `.dinodata` loader
|
||||
- `.latent`/occupancy readers used by the example CLIs
|
||||
GGUF model files are trusted assets (same threat model as privacy-filter).
|
||||
|
||||
## Demo
|
||||
|
||||
- `server/` Go (stdlib http + purego dlopen of `libtrellis2.so`), depth-anything
|
||||
server pattern: single inference mutex, `POST /api/generate` (multipart image,
|
||||
params) → job id, progress polling, `GET /api/mesh/<id>` binary mesh,
|
||||
self-contained embedded `web/index.html` WebGL viewer (no CDN, no build step).
|
||||
- C API additions in `trellis2.h` for: dino encode from RGBA buffer, full pipeline
|
||||
run with progress callback, mesh buffer accessors.
|
||||
|
||||
## Order of work
|
||||
|
||||
1. infra: ref container, weights, fix existing ref scripts, regenerate refs,
|
||||
existing 3 tests green on this machine (CPU + asan)
|
||||
2. DINOv3 encoder port + per-layer validation → image→coarse-mesh e2e in C++
|
||||
3. Go server + web viewer on coarse pipeline (user-visible milestone)
|
||||
4. fuzz harnesses + short campaigns
|
||||
5. shape-SLAT flow port (sparse DiT) + validation
|
||||
6. FDG VAE decoder + flexible-dual-grid mesher + validation → real mesh in demo
|
||||
7. docs (VERIFICATION.md), CUDA build check, quantized variants, benchmarks
|
||||
|
||||
## Status (done)
|
||||
|
||||
All of 1–7 complete. Image→mesh works end-to-end, C++/ggml only, coarse
|
||||
(marching-cubes preview) and fine (512³ dual-grid) paths, in the Go demo. Every
|
||||
stage validated tap-by-tap (docs/VERIFICATION.md): preprocessing byte-exact,
|
||||
DINOv3 rel-L2 ≤7e-7, sparse U-Net decoder exact through 4 levels. CUDA build
|
||||
(sm_120) works; 3D-conv decoders pinned to CPU. One fuzz bug found+fixed.
|
||||
|
||||
Runtime on the 16 GB RTX 5070 Ti (GPU shape decode auto-enabled): fine (512³)
|
||||
**~39 s**, 1024 cascade **~133 s**. (Was 85 s for 512 with the decoder on CPU.)
|
||||
|
||||
### Where the fine-path time goes (measured, `TRELLIS2_TIMING=1`)
|
||||
|
||||
Per-stage wall-clock for one 512³ generation (f16, ~1.09 M-vertex mesh), CPU
|
||||
decoder vs the auto GPU decoder:
|
||||
|
||||
| stage | CPU decode | GPU decode | device (GPU mode) |
|
||||
|------------|-----------:|-----------:|-------------------|
|
||||
| dino | 0.1 s | 0.1 s | GPU |
|
||||
| ss_flow | 21 s | 21 s | GPU |
|
||||
| ss_dec | 3 s | 3 s | CPU (dense CONV_3D)|
|
||||
| slat_flow | 11 s | 10 s | GPU |
|
||||
| shape_dec | **49 s** | **2.5 s** | **GPU** (was 59 % of the run) |
|
||||
| mesh | 0.2 s | 0.2 s | CPU |
|
||||
| **total** | **85 s** | **39 s** | |
|
||||
|
||||
The original "GPU ~30 %" reading was the *average*: the biggest stage (the
|
||||
FlexiDualGrid VAE decoder) ran on the CPU with the GPU idle. It now runs on the
|
||||
GPU — **~2.5 s vs ~49 s, ~20×** — placed there automatically when VRAM allows
|
||||
(see the shape-decoder note under "Not done"). The flow DiTs (`ss_flow` +
|
||||
`slat_flow`) are now the bulk of the run. Two facts about them, from
|
||||
`TRELLIS2_TIMING`:
|
||||
|
||||
- Attention is now `ggml_flash_attn_ext` for **all** flow forwards, not just the
|
||||
cascade's huge token counts (`sdpa_auto`, `TRELLIS2_SDPA_EXACT` restores the
|
||||
old materialized softmax). Flash is bit-identical to full softmax on the CPU
|
||||
(both flow forwards still validate at ~2.4e-4 / 2.9e-4 rel-L2) and on the GPU
|
||||
it is ~30 % faster *and* O(L) memory — the exact path's 805 MB score matrix
|
||||
`alloc_graph`-fails once the resident pipeline leaves <~2 GB free, so flash is
|
||||
what keeps the flow fitting under host-VRAM pressure, not only a speedup.
|
||||
- Each forward still runs at only ~17 % of the card's f16 tensor-core peak; the
|
||||
matmuls already hit the f16 cores (ggml converts the f32 activations), so the
|
||||
remaining slack is the f32 elementwise/permute traffic between matmuls. CUDA
|
||||
graphs do not help (they `cudaGraphInstantiate`-OOM with the pipeline
|
||||
resident) and are left off.
|
||||
|
||||
The shape decoder is a *sparse* net that keeps millions of voxels in the tensor
|
||||
**row** dimension (`ne1`). ggml's CUDA CONCAT/PAD kernels launch a grid
|
||||
dimension equal to that count and abort (`invalid configuration argument`) above
|
||||
the 65535 grid-Y/Z limit — which is why the original port used the CPU. But
|
||||
CONCAT/PAD were only used to append one missing-neighbor zero row; replacing
|
||||
that with a clamp-index + 0/1-mask formulation (get_rows a valid row, multiply
|
||||
the missing ones by zero — byte-identical) keeps every op inside ggml kernels
|
||||
that tile the voxel dimension (`get_rows`, broadcast `mul`, `mul_mat`, `add`,
|
||||
`norm`, `silu` all verified OK at 3 M voxels). The decoder then runs on the GPU
|
||||
in ~2.5 s (512³) / ~12.5 s (1024³).
|
||||
|
||||
### VRAM-aware auto-placement of the shape decoder
|
||||
|
||||
The decoder's placement is decided from **measured free VRAM**
|
||||
(`trellis2_gpu_free_vram()` → `cudaMemGetInfo`), not a hardcoded flag:
|
||||
|
||||
- At load, `t2_pipeline_load` records the free VRAM *before* the flow DiTs are
|
||||
loaded (= what freeing them again reclaims) and places the decoder on the GPU
|
||||
when that covers the target tier's decode peak (~3 GB at 512³, ~7.5 GB at
|
||||
1024³) plus the decoder weights and a margin; otherwise CPU. `TRELLIS2_SHAPE_DEC_GPU`
|
||||
/ `_CPU` force it; a CPU-only build (no GPU device) always picks CPU.
|
||||
- At decode time, `ensure_decode_vram` checks free VRAM again and, if the decode
|
||||
would not fit, frees the flow DiTs (all finished by then) to reclaim their
|
||||
~5–7 GB; `reload_flows` brings them back on the next `t2_generate` (a few
|
||||
seconds from page cache). So 512³ decodes with the flows resident (no reload
|
||||
cost) while the big 1024³ decode transparently frees-and-reloads. This is the
|
||||
`T2_LOAD_LOW_VRAM` idea, applied automatically only when the numbers require it.
|
||||
|
||||
Validated end-to-end on the 16 GB card with **no env vars**: 512 → 39 s, 1024
|
||||
cascade → 133 s (the 1024³ GPU decode is 12.5 s), and three back-to-back cascade
|
||||
generations exercise the free/reload path with no OOM.
|
||||
|
||||
## Not done (out of original scope / future)
|
||||
|
||||
- **PBR textures — DONE.** The validated standalone texturing path re-encodes
|
||||
the decoded dual grid to condition texture flow, then trilinearly samples the
|
||||
decoded six-channel PBR volume at the surface. The browser and GLB path
|
||||
preserve alpha and use the correct base-color space. The first integrated
|
||||
subdivision-guide path was removed after it produced collapsed materials.
|
||||
- **`1024_cascade` — DONE.** Full high-resolution path: LR flow (512 model) →
|
||||
`trellis2_shape_dec_upsample(×4)` → quantize+dedup to the 64³ HR scaffold →
|
||||
HR flow (1024 model, cond_1024) → 1024³ decode → dual-grid mesh. Enabled by
|
||||
flash attention (`sdpa_auto`) for the HR token counts. Validated by
|
||||
`test_cascade`: upsample coords + HR scaffold exact (to 1 voxel of ~995k),
|
||||
HR flow forward rel-L2 3.1e-4 (CPU). The 1024³ decode is the same decoder
|
||||
validated exactly at the 512 tier; its explicit parity gate is behind
|
||||
`TRELLIS2_CASCADE_DECODE` because it transiently needs ~14 GB host RAM.
|
||||
- **Still to do:** `1536_cascade` (add the token-reduction loop + res 1536). The
|
||||
`<8 GB` case is now partly covered: the shape decoder auto-frees the flow DiTs
|
||||
around a GPU decode (see below), so the pieces of the `T2_LOAD_LOW_VRAM`
|
||||
lifecycle exist; a full per-stage load/free mode would extend that to the flow
|
||||
DiTs themselves for cards that can't hold even one 1.3B DiT + a decode.
|
||||
- **Background removal** (BiRefNet/RMBG-2.0) — the demo instructs a transparent
|
||||
PNG and uses the image as-is otherwise. A separate ~1 GB seg model.
|
||||
- **GPU shape decoder — DONE.** The mask-conv change plus the VRAM-aware
|
||||
placement + free/reload lifecycle (both described above) make the decoder run
|
||||
on the GPU automatically when it fits: 512³ ~2.5 s and 1024³ ~12.5 s vs ~49 s /
|
||||
minutes on the CPU. No env var required; `TRELLIS2_SHAPE_DEC_{GPU,CPU}` still
|
||||
force it, and it stays on the CPU on cards too small or in a CPU-only build.
|
||||
- **CUDA CONV_3D** — the SS *occupancy* decoder (`ss_dec`, 3 s) uses a genuine
|
||||
dense `ggml_conv_3d_direct` with no CUDA kernel, so it stays on CPU regardless.
|
||||
- **Quantized shipping variants / benchmarks / CI** — f16 is the demo default;
|
||||
q8/q4 conversion + a benchmark table + a two-tier CI workflow are the natural
|
||||
next hardening steps (privacy-filter.cpp has the template).
|
||||
- **CPU-only SLAT sampler tightness in ctest** — the `slat` ctest is labeled
|
||||
slow and validated in-container; the sampler validates tightly on CPU
|
||||
(TRELLIS2_SLAT_STRICT=1) but that CPU run is minutes-long.
|
||||
@@ -0,0 +1,86 @@
|
||||
# Verification
|
||||
|
||||
Every ported stage is validated numerically against the PyTorch reference, per
|
||||
component, not just end-to-end — the depth-anything.cpp approach. Reference
|
||||
activations are dumped from `microsoft/TRELLIS.2` (via `scripts/refgen.sh`
|
||||
inside the CUDA container) into GGUF files under `dumps/`, and the C++ side is
|
||||
compared tap-by-tap with `tests/parity.hpp` (gate: `|got-ref| <= atol +
|
||||
rtol*|ref|`, reported as max-abs / rel-L2 / per-row).
|
||||
|
||||
## One-time setup
|
||||
|
||||
```sh
|
||||
docker build -f docker/Dockerfile.ref -t trellis2-ref docker # PyTorch reference env
|
||||
scripts/download_models.sh # HF checkpoints -> models/
|
||||
docker run --rm -v "$PWD":/work -w /work trellis2-ref bash -c '
|
||||
python convert_dino_to_gguf.py --output ggufs/dino_f32.gguf --ftype 0
|
||||
python convert_ss_flow_to_gguf.py --model models/TRELLIS.2-4B/ckpts/ss_flow_img_dit_1_3B_64_bf16.safetensors --output ggufs/ss_flow_f32.gguf --ftype 0
|
||||
python convert_ss_dec_to_gguf.py --model models/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16.safetensors --output ggufs/ss_dec_f32.gguf --ftype 0
|
||||
python convert_slat_flow_to_gguf.py --model models/TRELLIS.2-4B/ckpts/slat_flow_img2shape_dit_1_3B_512_bf16.safetensors --pipeline-json models/TRELLIS.2-4B/pipeline.json --output ggufs/slat_flow_f32.gguf --ftype 0
|
||||
python convert_shape_dec_to_gguf.py --output ggufs/shape_dec_f32.gguf --ftype 0'
|
||||
scripts/refgen.sh # dump reference activations
|
||||
```
|
||||
|
||||
## Run
|
||||
|
||||
```sh
|
||||
cmake -B build -DTRELLIS2_BUILD_TESTS=ON && cmake --build build -j
|
||||
ctest --test-dir build -LE model # fast, no assets (marching cubes, preprocess)
|
||||
ctest --test-dir build # full parity (needs ggufs/ + dumps/)
|
||||
```
|
||||
|
||||
## Parity table (f32 GGUF vs true-fp32 reference)
|
||||
|
||||
| stage | test | tap coverage | result |
|
||||
|---|---|---|---|
|
||||
| image preprocess (alpha crop, premultiply, Lanczos-512) | `test_preprocess` | full 512×512 RGB | **byte-exact** (0/786432 differ) |
|
||||
| DINOv3 ViT-L/16 encoder | `test_dino` | 40 taps: embeddings, RoPE, per-layer output + first/last-layer detail (norm/attn/layerscale/mlp), final affine-free LN | **PASS**, rel-L2 ≤ 7e-7 all taps |
|
||||
| SS-flow DiT forward | `test_ss_flow_forward` | full output | **PASS**, rel-L2 2.4e-4 |
|
||||
| SS-flow Euler sampler (12-step CFG) | `test_ss_sample` | z_s latent | **PASS**, rel-L2 5.7e-3, sign 99.85% (CPU) |
|
||||
| SS decoder (dense 3D-conv → 64³ occupancy) | `test_ss_dec` | occupancy logits | **PASS**, rel-L2 2e-5 |
|
||||
| shape-SLAT flow forward | `test_slat` | full output | **PASS**, rel-L2 2.9e-4 (CPU) / 8e-4 (GPU) |
|
||||
| shape-SLAT VAE decoder (sparse ConvNeXt U-Net, 4 levels, 16× up) | `test_slat` | per-level features + subdivision logits + final 7-ch output, all 5 levels | **PASS**, rel-L2 ≤ 6e-7 (levels 0–3 exact; final set within 0.0001%) |
|
||||
| integrated subdivision guide | `test_slat` | all decoder levels; final guide coordinates equal decoded shape coordinates | **PASS** |
|
||||
| standalone shape encoder → texture flow → texture decoder | `test_texture` | shape latent, flow forward/sampler, guided 6-channel PBR decode | parity-gated; sampler backend drift uses the documented loose gate |
|
||||
| sparse PBR surface sampling | `test_pbr_sampling` | dense trilinear interpolation + sparse-boundary normalization | **PASS** |
|
||||
| GLB PBR/alpha export | `test_mesh_export` | direct vertex RGBA, retained metallic/roughness, glTF alpha mode | **PASS** |
|
||||
| dual-grid mesh extraction | `test_marching_cubes` (invariants) + visual | watertight-manifold, Euler characteristic, winding | **PASS** |
|
||||
| **1024 cascade** — decoder upsample(×4) → 512³ coords | `test_cascade` | full coord set + quantized 64³ HR scaffold | **PASS**, set match to 0.0001% (1 voxel of 995k) |
|
||||
| **1024 cascade** — HR (1024-model) flow forward | `test_cascade` | full output | **PASS**, rel-L2 ~3e-4 (CPU); ~1e-2 on GPU flash |
|
||||
| **1024 cascade** — final 1024³ decode (3.97M voxels) | `test_cascade` | per-level features + subdivision + 7-ch output | **PASS**, rel-L2 ≤ 2e-2, set within 0.0001% |
|
||||
|
||||
Notes:
|
||||
- **Flash attention (default for every flow forward).** `sdpa_auto()` uses
|
||||
`ggml_flash_attn_ext` for both flow DiTs at all token counts;
|
||||
`TRELLIS2_SDPA_EXACT` restores the old materialized `[L_k, L_q, heads]`
|
||||
softmax. Flash is bit-faithful to full softmax on CPU (SS-flow 2.4e-4, SLAT
|
||||
2.9e-4 — identical to exact, so the tap parity above is unaffected) but incurs
|
||||
~3e-3 rel-L2 on the CUDA F16-MMA kernel. It was already required for the HR
|
||||
cascade (49,152-token attention fits on 16 GB vs a 108 GiB exact matrix); it
|
||||
is now the default because on the GPU it is also ~30 % faster per forward and
|
||||
O(L) memory — the exact path's 805 MB SS-flow score matrix `alloc_graph`-fails
|
||||
when the resident pipeline leaves little free VRAM. Was previously gated to
|
||||
score matrices >1 GiB (i.e. only the HR stage). See docs/PLAN.md for the
|
||||
per-stage runtime profile.
|
||||
- **TF32 matters.** PyTorch's default CUDA matmul/attention uses TF32 (≈10-bit
|
||||
mantissa) and reduced-precision flash SDPA, which shows up as ~1e-3 relative
|
||||
error versus true fp32. `scripts/ref_common.py` disables it so the golden
|
||||
dumps are real fp32; otherwise a correct port looks like it has a 0.08-rel-L2
|
||||
bug (this exact trap cost a debugging session — see the flow-forward gate).
|
||||
- **Sampler drift.** The 12-step Euler + CFG-rescale loop chaotically amplifies
|
||||
per-step fp differences between backends; it validates tightly on CPU and
|
||||
drifts to ~0.1 rel-L2 on GPU. The decoder gate therefore decodes the
|
||||
*reference* SLAT so decoder parity is independent of sampler trajectory.
|
||||
- **Subdivision boundary.** A handful of level-3 subdivision logits sit within
|
||||
fp-noise of zero; the `>0` threshold can flip them, so the final active-voxel
|
||||
set differs by ~4 voxels out of ~4 million (0.0001%) run-to-run and
|
||||
hardware-to-hardware. This is inherent to a hard threshold, not a port bug.
|
||||
|
||||
## GPU
|
||||
|
||||
`-DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=120` (Blackwell / RTX 50-series). The
|
||||
flow DiTs and DINOv3 run on CUDA; the 3D-conv decoders (SS decoder CONV_3D and
|
||||
the sparse-conv gather-GEMM) run on CPU because the bundled ggml has no CUDA
|
||||
CONV_3D kernel — they are a small fraction of total inference time. GPU f32
|
||||
matmul is fp16-class, so tap parity on CUDA is ~1e-3 (deterministic, not device
|
||||
noise); CPU is the tight-tolerance reference backend.
|
||||
Reference in New Issue
Block a user