From a541863f433b777046bbc9d5add9d37a1a117044 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 16 Aug 2026 13:46:51 +0000 Subject: [PATCH 01/22] dit: thread a per-token projection input through the flow runner TRELLIS.2 conditions every DiT block by cross-attending over the DINOv3 token stream. Pixal3D keeps that cross-attention but wraps it: the module's weights move one level down under cross_attn.cross_attn_block, and a sibling cross_attn.proj_linear maps a per-token view-aligned feature into model space and is added to the cross-attention output, replacing it as the residual branch. Add DiTParams::proj_mode / proj_ch and a proj input to build_dit_dense, which is the whole architectural delta on the denoiser side. Conditioning now travels as a FlowCond bundle so classifier-free guidance can carry both the context and the projection; a null proj in proj mode is the negative branch and is zeroed on the device rather than materialized (it reaches 400 MB at the texture stage's token budget). The implicit FlowCond constructor keeps cross-mode call sites unchanged. (cherry picked from commit e0b76ca95e52b142a21d068aad53c2c7f75ceb11) (cherry picked from commit 8da007878feaddcf050847f6b9883533411b4b75) --- include/dit.h | 9 ++++++++- include/flow_runner.h | 24 +++++++++++++++++++----- src/dit.cpp | 18 ++++++++++++++---- src/flow_runner.cpp | 19 +++++++++++++++---- src/test_shape_flow.cpp | 2 +- src/test_slat_shape.cpp | 2 +- src/test_ss_flow.cpp | 2 +- src/test_ss_full.cpp | 2 +- src/test_ss_sample.cpp | 2 +- 9 files changed, 61 insertions(+), 19 deletions(-) diff --git a/include/dit.h b/include/dit.h index a772961..62e9594 100644 --- a/include/dit.h +++ b/include/dit.h @@ -22,6 +22,12 @@ struct DiTParams { float final_ln_eps = 1e-5f; float rms_eps = 1e-12f; bool cast_f32 = false; // cast f16 weights to f32 before matmul (precision test) + // Pixal3D "proj" image-attention mode. The block layout is identical to TRELLIS.2 except + // that the cross-attention module is wrapped: its weights sit one level deeper, under + // `cross_attn.cross_attn_block`, and a sibling `cross_attn.proj_linear` maps the per-token + // view-aligned feature into model space and is added to the cross-attention output. + bool proj_mode = false; + int proj_ch = 0; // proj_in_channels: 1024 bare, 2048 with the NAF branch }; // Build the dense SS-flow forward graph (B=1). All input tensors live in `gctx` @@ -29,11 +35,12 @@ struct DiTParams { // h0 : [in_ch, L] patchified input (channel-major) // tfreq: [256] sinusoidal timestep embedding (host-computed) // cond : [d_cond, Lc] conditioning tokens +// proj : [proj_ch, L] per-token view-aligned features (proj mode; else nullptr) // cos/sin: [1, head_dim/2, 1, L] precomputed 3D-RoPE tables // Returns the [out_ch, L] velocity; `inter` (optional) collects named intermediates. ggml_tensor* build_dit_dense(ggml_context* gctx, const Model& m, const DiTParams& p, ggml_tensor* h0, ggml_tensor* tfreq, ggml_tensor* cond, - ggml_tensor* cos, ggml_tensor* sin, + ggml_tensor* proj, ggml_tensor* cos, ggml_tensor* sin, std::map* inter = nullptr); } // namespace trellis diff --git a/include/flow_runner.h b/include/flow_runner.h index 970b412..e599ac3 100644 --- a/include/flow_runner.h +++ b/include/flow_runner.h @@ -25,6 +25,20 @@ struct SamplerParams { float sigma_min = 1e-5f; }; +// One branch of the classifier-free guidance pair. `cond` is the cross-attention context that +// every TRELLIS.2 stage uses; `proj` carries the extra per-token view-aligned features and is +// only read in Pixal3D's proj mode, where a null `proj` means the all-zero (negative) branch — +// the runner zeroes the device tensor rather than making the caller materialize the buffer. +// The implicit constructor keeps every cross-mode call site (and the reference tests) writing +// plain `cond.data()`. +struct FlowCond { + const float* cond = nullptr; // [d_cond * n_cond] + const float* proj = nullptr; // [proj_ch * N] + FlowCond() = default; + FlowCond(const float* c) : cond(c) {} + FlowCond(const float* c, const float* p) : cond(c), proj(p) {} +}; + // One DiT graph (built once for a fixed token count N), re-run per sampler step. // Token axis N = R^3 (dense) or number of active voxels (sparse); RoPE tables are // supplied by the factory (grid index math vs real voxel coords). @@ -33,13 +47,13 @@ class DitRunner { DitRunner(const Model& m, const DiTParams& p, int N, int n_cond, const std::vector& rope_cos, const std::vector& rope_sin); ~DitRunner(); - // xt: [in_ch*N] channel-major. cond: [d_cond*n_cond]. Returns velocity [out_ch*N]. - std::vector forward(const std::vector& xt, float t_scaled, const float* cond); + // xt: [in_ch*N] channel-major. Returns velocity [out_ch*N]. + std::vector forward(const std::vector& xt, float t_scaled, const FlowCond& c); int N() const { return N_; } private: const Model& m_; DiTParams p_; int N_, Lc_; ggml_context* ctx_ = nullptr; ggml_cgraph* g_ = nullptr; ggml_gallocr_t alloc_ = nullptr; - ggml_tensor *gh0_, *gtf_, *gcond_, *gcos_, *gsin_, *gout_; + ggml_tensor *gh0_, *gtf_, *gcond_, *gproj_ = nullptr, *gcos_, *gsin_, *gout_; std::vector rcos_, rsin_; // re-uploaded each forward (gallocr may reuse input buffers) std::map inter_; // [dbg] named intermediates for NaN localization bool dbg_nan_ = false, dbg_done_ = false; @@ -52,9 +66,9 @@ DitRunner* make_sparse_runner(const Model& m, const DiTParams& p, const std::vector>& coords, int n_cond); // FlowEuler guidance-interval sampler over an arbitrary forward functor. -using FlowFwd = std::function(const std::vector&, float, const float*)>; +using FlowFwd = std::function(const std::vector&, float, const FlowCond&)>; std::vector sample_flow(const FlowFwd& fwd, std::vector sample, - const float* cond, const float* neg_cond, + const FlowCond& cond, const FlowCond& neg_cond, const SamplerParams& sp, std::vector>* trace = nullptr); diff --git a/src/dit.cpp b/src/dit.cpp index bc0c413..02cef0f 100644 --- a/src/dit.cpp +++ b/src/dit.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include namespace trellis { @@ -228,7 +229,7 @@ static T* modulate(ggml_context* c, T* x, T* scale, T* shift) { return ggml_add(c, ggml_add(c, x, ggml_mul(c, x, scale)), shift); } -static T* block(ggml_context* c, const Model& m, int i, T* h, T* mod, T* cond, +static T* block(ggml_context* c, const Model& m, int i, T* h, T* mod, T* cond, T* proj, T* cos, T* sin, const DiTParams& p, std::map* inter = nullptr, T* self_mask = nullptr, T* cross_mask = nullptr) { const std::string b = "blocks." + std::to_string(i); @@ -246,7 +247,15 @@ static T* block(ggml_context* c, const Model& m, int i, T* h, T* mod, T* cond, h = ggml_add(c, h, ggml_mul(c, hh, gate_msa)); hh = layernorm(c, h, p.ln_eps, m.get(b + ".norm2.weight"), m.get(b + ".norm2.bias")); - hh = cross_attn(c, m, b + ".cross_attn", hh, cond, p, cross_mask); + if (p.proj_mode) { + // ProjectAttention: cross_attn_block(h, global_tokens) + proj_linear(view_aligned). + // The sum REPLACES the cross-attention output as the residual branch (both the dense + // and the sparse Pixal3D modules do exactly this), so the add below is unchanged. + hh = cross_attn(c, m, b + ".cross_attn.cross_attn_block", hh, cond, p, cross_mask); + hh = ggml_add(c, hh, lin(c, m, b + ".cross_attn.proj_linear", proj)); + } else { + hh = cross_attn(c, m, b + ".cross_attn", hh, cond, p, cross_mask); + } dbg("blk0_cross", hh); h = ggml_add(c, h, hh); @@ -261,9 +270,10 @@ static T* block(ggml_context* c, const Model& m, int i, T* h, T* mod, T* cond, } ggml_tensor* build_dit_dense(ggml_context* c, const Model& m, const DiTParams& p, - T* h0, T* tfreq, T* cond, T* cos, T* sin, + T* h0, T* tfreq, T* cond, T* proj, T* cos, T* sin, std::map* inter) { g_cast_f32 = p.cast_f32; + if (p.proj_mode && !proj) throw std::runtime_error("build_dit_dense: proj mode needs a proj input"); auto keep = [&](const char* n, T* t) { if (inter) (*inter)[n] = t; ggml_set_name(t, n); return t; }; T* h = lin(c, m, "input_layer", h0); // [d_model, L] @@ -281,7 +291,7 @@ ggml_tensor* build_dit_dense(ggml_context* c, const Model& m, const DiTParams& p T* self_mask = build_pad_mask(c, h0->ne[1], h0->ne[1]); T* cross_mask = build_pad_mask(c, cond->ne[1], h0->ne[1]); for (int i = 0; i < p.n_blocks; ++i) { - h = block(c, m, i, h, mod, cond, cos, sin, p, inter, self_mask, cross_mask); + h = block(c, m, i, h, mod, cond, proj, cos, sin, p, inter, self_mask, cross_mask); if (i == 0) keep("after_block0", h); if (i == 1) keep("after_block1", h); if (i == p.n_blocks - 1) keep("after_block29", h); diff --git a/src/flow_runner.cpp b/src/flow_runner.cpp index 0b5ad9e..727c0a2 100644 --- a/src/flow_runner.cpp +++ b/src/flow_runner.cpp @@ -43,10 +43,15 @@ DitRunner::DitRunner(const Model& m, const DiTParams& p, int N, int n_cond, gh0_ = ggml_new_tensor_2d(ctx_, GGML_TYPE_F32, p_.in_ch, N_); ggml_set_input(gh0_); gtf_ = ggml_new_tensor_1d(ctx_, GGML_TYPE_F32, 256); ggml_set_input(gtf_); gcond_= ggml_new_tensor_2d(ctx_, GGML_TYPE_F32, p_.d_cond, Lc_); ggml_set_input(gcond_); + if (p_.proj_mode) { + if (p_.proj_ch <= 0) throw std::runtime_error("DitRunner: proj mode needs proj_ch"); + gproj_ = ggml_new_tensor_2d(ctx_, GGML_TYPE_F32, p_.proj_ch, N_); ggml_set_input(gproj_); + } gcos_ = ggml_new_tensor_4d(ctx_, GGML_TYPE_F32, 1, half, 1, N_); ggml_set_input(gcos_); gsin_ = ggml_new_tensor_4d(ctx_, GGML_TYPE_F32, 1, half, 1, N_); ggml_set_input(gsin_); dbg_nan_ = std::getenv("TRELLIS_DBG_NAN") != nullptr; - gout_ = build_dit_dense(ctx_, m_, p_, gh0_, gtf_, gcond_, gcos_, gsin_, dbg_nan_ ? &inter_ : nullptr); + gout_ = build_dit_dense(ctx_, m_, p_, gh0_, gtf_, gcond_, gproj_, gcos_, gsin_, + dbg_nan_ ? &inter_ : nullptr); g_ = ggml_new_graph_custom(ctx_, 32768, false); ggml_build_forward_expand(g_, gout_); ggml_set_output(gout_); @@ -61,11 +66,17 @@ DitRunner::~DitRunner() { if (ctx_) ggml_free(ctx_); } -std::vector DitRunner::forward(const std::vector& xt, float t_scaled, const float* cond) { +std::vector DitRunner::forward(const std::vector& xt, float t_scaled, const FlowCond& c) { std::vector tf; timestep_embedding(t_scaled, tf); ggml_backend_tensor_set(gh0_, xt.data(), 0, xt.size() * 4); ggml_backend_tensor_set(gtf_, tf.data(), 0, tf.size() * 4); - ggml_backend_tensor_set(gcond_, cond, 0, (size_t)p_.d_cond * Lc_ * 4); + ggml_backend_tensor_set(gcond_, c.cond, 0, (size_t)p_.d_cond * Lc_ * 4); + if (gproj_) { + const size_t nb = (size_t)p_.proj_ch * N_ * 4; + // Re-written every forward: gallocr may hand the same buffer to both CFG branches. + if (c.proj) ggml_backend_tensor_set(gproj_, c.proj, 0, nb); + else ggml_backend_tensor_memset(gproj_, 0, 0, nb); + } ggml_backend_tensor_set(gcos_, rcos_.data(), 0, rcos_.size() * 4); // re-upload (buffers reused across runs) ggml_backend_tensor_set(gsin_, rsin_.data(), 0, rsin_.size() * 4); if (ggml_backend_graph_compute(m_.backend, g_) != GGML_STATUS_SUCCESS) @@ -130,7 +141,7 @@ DitRunner* make_sparse_runner(const Model& m, const DiTParams& p, } std::vector sample_flow(const FlowFwd& fwd, std::vector sample, - const float* cond, const float* neg_cond, const SamplerParams& sp, + const FlowCond& cond, const FlowCond& neg_cond, const SamplerParams& sp, std::vector>* trace) { const float sm = sp.sigma_min; const size_t Nst = sample.size(); diff --git a/src/test_shape_flow.cpp b/src/test_shape_flow.cpp index a9e0724..58787ce 100644 --- a/src/test_shape_flow.cpp +++ b/src/test_shape_flow.cpp @@ -107,7 +107,7 @@ int main(int argc, char** argv) { ggml_tensor* gsin = ggml_new_tensor_4d(c, GGML_TYPE_F32, 1, half, 1, L); ggml_set_input(gsin); std::map inter; - ggml_tensor* out = trellis::build_dit_dense(c, m, p, gh0, gtf, gcd, gcos, gsin, &inter); + ggml_tensor* out = trellis::build_dit_dense(c, m, p, gh0, gtf, gcd, nullptr, gcos, gsin, &inter); ggml_cgraph* g = ggml_new_graph_custom(c, 262144, false); ggml_build_forward_expand(g, out); diff --git a/src/test_slat_shape.cpp b/src/test_slat_shape.cpp index 6a322d9..1ff313b 100644 --- a/src/test_slat_shape.cpp +++ b/src/test_slat_shape.cpp @@ -33,7 +33,7 @@ int main(int argc, char** argv) { trellis::DiTParams p; p.in_ch = (int)Cin; p.out_ch = (int)Cin; p.d_cond = (int)Dc; if (getenv("TRELLIS_F32W")) { p.cast_f32 = true; printf("(f32 weight compute)\n"); } trellis::DitRunner* run = trellis::make_sparse_runner(m, p, coords, (int)Nimg); - trellis::FlowFwd fwd = [&](const vector& xt, float ts, const float* cd){ return run->forward(xt, ts, cd); }; + trellis::FlowFwd fwd = [&](const vector& xt, float ts, const trellis::FlowCond& cd){ return run->forward(xt, ts, cd); }; // noise [n*Cin+c] -> xt [c + Cin*n] vector xt(N * Cin); diff --git a/src/test_ss_flow.cpp b/src/test_ss_flow.cpp index d5d25c7..c92d25e 100644 --- a/src/test_ss_flow.cpp +++ b/src/test_ss_flow.cpp @@ -89,7 +89,7 @@ int main(int argc, char** argv) { ggml_tensor* gsin = ggml_new_tensor_4d(c, GGML_TYPE_F32, 1, half, 1, L); ggml_set_input(gsin); std::map inter; - ggml_tensor* out = trellis::build_dit_dense(c, m, p, gh0, gtf, gcd, gcos, gsin, &inter); + ggml_tensor* out = trellis::build_dit_dense(c, m, p, gh0, gtf, gcd, nullptr, gcos, gsin, &inter); ggml_cgraph* g = ggml_new_graph_custom(c, 32768, false); ggml_build_forward_expand(g, out); diff --git a/src/test_ss_full.cpp b/src/test_ss_full.cpp index 941fae3..5994ddd 100644 --- a/src/test_ss_full.cpp +++ b/src/test_ss_full.cpp @@ -26,7 +26,7 @@ int main(int argc, char** argv) { trellis::DiTParams p; p.in_ch = 8; p.out_ch = 8; p.d_cond = 1024; if (getenv("TRELLIS_F32W")) p.cast_f32 = true; trellis::DitRunner* run = trellis::make_dense_runner(m, p, 16, Lc); - trellis::FlowFwd fwd = [&](const vector& x, float ts, const float* c){ return run->forward(x, ts, c); }; + trellis::FlowFwd fwd = [&](const vector& x, float ts, const trellis::FlowCond& c){ return run->forward(x, ts, c); }; trellis::SamplerParams sp; sp.steps=12; sp.guidance_strength=7.5f; sp.guidance_rescale=0.7f; sp.gi0=0.6f; sp.gi1=1.0f; sp.rescale_t=5.0f; if (getenv("GS")) sp.guidance_strength = atof(getenv("GS")); if (getenv("GR")) sp.guidance_rescale = atof(getenv("GR")); diff --git a/src/test_ss_sample.cpp b/src/test_ss_sample.cpp index 7a0e548..60c3549 100644 --- a/src/test_ss_sample.cpp +++ b/src/test_ss_sample.cpp @@ -27,7 +27,7 @@ int main(int argc, char** argv) { trellis::DiTParams p; p.in_ch = (int)Cin; p.out_ch = (int)Cin; p.d_cond = (int)Dc; if (getenv("TRELLIS_F32W")) { p.cast_f32 = true; printf("(f32 weight compute)\n"); } trellis::DitRunner* run = trellis::make_dense_runner(m, p, (int)R, (int)Lc); - trellis::FlowFwd fwd = [&](const std::vector& xt, float ts, const float* cd){ return run->forward(xt, ts, cd); }; + trellis::FlowFwd fwd = [&](const std::vector& xt, float ts, const trellis::FlowCond& cd){ return run->forward(xt, ts, cd); }; // noise [c*L+sp] -> sample [c + Cin*sp] vector sample(Cin * L); From 4486c5d8d101600f18bf72d6ebe189a235c1e019 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 16 Aug 2026 13:47:05 +0000 Subject: [PATCH 02/22] pixal3d: view-aligned projection conditioning and the NAF upsampler Pixal3D cross-attends over the 5 global DINOv3 tokens only and routes the patch grid through a pixel-aligned path instead: each DiT token is a cell of a 3-D grid, projected into the image with a fixed frontal camera, and the DINOv3 feature map is sampled there. pixal3d.cpp implements the camera solve (only the FOV is free; the distance follows in closed form), the grid projection and the per-stage ProjCond assembly. test_pixal3d pins both functions to golden values evaluated from the reference implementation, because a subtly wrong projection has no runtime signal -- it just drifts the geometry off the silhouette. naf.cpp ports valeoai/NAF, which the shape and texture stages run to upsample the feature map before sampling it (hence their proj_in_channels of 2048). Only the two-branch guide encoder is learned, and it maps cleanly onto ggml; the upsampling itself is a parameter-free neighborhood cross-attention. Dilating it by exactly the upsampling factor collapses NATTEN's dilated 2-D kernel into a 9x9 window of low-resolution cells centred on the query's own cell, which is what makes the threaded CPU implementation practical -- and it only runs on the pixels the projection actually reads, not the whole upsampled map. (cherry picked from commit 9c5f948cd7304f73429900cf9316cd1428ca46d4) (cherry picked from commit e1794c3399f469cd627486cb01313f5ddbdac0ab) --- CMakeLists.txt | 6 + include/naf.h | 35 ++++++ include/pixal3d.h | 65 ++++++++++ include/preprocess.h | 5 + src/naf.cpp | 278 +++++++++++++++++++++++++++++++++++++++++++ src/pixal3d.cpp | 137 +++++++++++++++++++++ src/preprocess.cpp | 12 ++ src/test_pixal3d.cpp | 86 +++++++++++++ 8 files changed, 624 insertions(+) create mode 100644 include/naf.h create mode 100644 include/pixal3d.h create mode 100644 src/naf.cpp create mode 100644 src/pixal3d.cpp create mode 100644 src/test_pixal3d.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index a79ce9a..88dd5e6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -95,6 +95,8 @@ add_library(trellis_core STATIC src/ss_decoder.cpp src/mesh_glb.cpp src/dinov3.cpp + src/pixal3d.cpp + src/naf.cpp src/sparse.cpp src/shape_decoder.cpp src/dual_grid.cpp @@ -338,6 +340,10 @@ add_executable(trellis-test-shape-dec src/test_shape_dec.cpp) target_link_libraries(trellis-test-shape-dec PRIVATE trellis_core) set_target_properties(trellis-test-shape-dec PROPERTIES BUILD_RPATH "${GGML_RPATH}") +add_executable(trellis-test-pixal3d src/test_pixal3d.cpp) +target_link_libraries(trellis-test-pixal3d PRIVATE trellis_core) +set_target_properties(trellis-test-pixal3d PROPERTIES BUILD_RPATH "${GGML_RPATH}") + add_executable(trellis-cli src/trellis_cli_main.cpp src/trellis_cli.cpp src/trellis_args.cpp) target_link_libraries(trellis-cli PRIVATE trellis_core) set_target_properties(trellis-cli PROPERTIES BUILD_RPATH "${GGML_RPATH}") diff --git a/include/naf.h b/include/naf.h new file mode 100644 index 0000000..96bdcef --- /dev/null +++ b/include/naf.h @@ -0,0 +1,35 @@ +// NAF (Neighborhood Attention Filtering) feature upsampler — the guided upsampler Pixal3D +// runs between DINOv3 and the view-aligned projection for its shape/texture stages. +// +// Reference: valeoai/NAF, as pulled by DinoV3ProjFeatureExtractor._load_naf(). The network is +// image-guided and VFM-agnostic: the only learned part is a two-branch convolutional encoder +// over the RGB guide; the upsampling itself is a parameter-free neighborhood cross-attention +// whose queries come from the guide at high resolution, whose keys are those same guide +// features average-pooled back to the low-resolution grid, and whose values are the VFM +// (DINOv3) features. Nothing in the attention is learned, so a GGUF of the conv encoder plus +// the RoPE period buffer is the whole model. +// +// The neighborhood attention is dilated by exactly the upsampling factor, which makes each +// high-resolution query attend to a KxK window of LOW-resolution cells around its own cell — +// see naf.cpp for the derivation. That collapses NATTEN's dilated 2-D kernel into a plain +// KxK gather over the LR grid and is what makes a CPU implementation practical here. +#pragma once +#include + +namespace trellis { +struct Model; + +// Upsample `feats_lr` with `img01` as the guide, then bilinearly sample the result at +// `pts_xy`. +// img01 : RGB guide in [0,1], torch [3,S,S] memory (== ggml [S,S,3,1]). +// feats_lr : DINOv3 patch features, channel-major [C, Hf*Wf] (index c + C*(h*Wf + w)). +// out : side of the upsampled map; S must be an integer multiple of it, and `out` +// an integer multiple of Hf/Wf (both hold for every Pixal3D stage). +// pts_xy : 2*NP sample positions as (x, y) pixel coordinates in the S-sized image frame. +// Returns [C * NP] channel-major (index c + C*p), matching the projection-conditioning layout. +std::vector naf_sample(const Model& m, + const std::vector& img01, int S, + const float* feats_lr, int Hf, int Wf, int C, + int out, const std::vector& pts_xy); + +} // namespace trellis diff --git a/include/pixal3d.h b/include/pixal3d.h new file mode 100644 index 0000000..a3b5242 --- /dev/null +++ b/include/pixal3d.h @@ -0,0 +1,65 @@ +// Pixal3D view-aligned projection conditioning. +// +// Pixal3D keeps the TRELLIS.2 denoiser, sampler and decoders unchanged and swaps only how the +// image reaches the DiT. TRELLIS.2 cross-attends over all DINOv3 tokens; Pixal3D cross-attends +// over the 5 global tokens (cls + registers) only, and adds a PIXEL-ALIGNED term: every DiT +// token is a cell of a 3-D grid, that cell is projected into the image with a fixed frontal +// camera, and the DINOv3 feature map is sampled there. A per-block `proj_linear` maps the +// sampled vector into model space and adds it to the cross-attention output. +// +// Stages that set use_naf_upsample additionally sample a NAF-upsampled copy of the same feature +// map and concatenate it, which is why their proj_in_channels is 2048 rather than 1024. +#pragma once +#include +#include + +namespace trellis { +struct Model; + +struct CameraParams { + float camera_angle_x = 0.8575560450553894f; // horizontal FOV in radians + float distance = 2.0f; // camera distance along the frontal axis + float mesh_scale = 1.0f; +}; + +// The reference derives `distance` in closed form from the FOV by requiring the grid corner at +// x = -1 to project onto the image border (inference.py: distance_from_fov). Only the FOV is a +// free parameter — upstream estimates it with MoGe-2, we take it from --fov. +CameraParams pixal3d_camera(float camera_angle_x, float mesh_scale = 1.0f, + int image_resolution = 512, int extend_pixel = 0); + +// Project the centre of one grid cell of an R^3 grid to (x, y) pixel coordinates in an +// `image_resolution`-sized frame. Cells are indexed as the DiT tokenizes a dense grid: value +// along each axis is linspace(-1, 1, R)[i]. +void pixal3d_project_cell(int R, int cx, int cy, int cz, const CameraParams& cam, + int image_resolution, float& px, float& py); + +// The negative branch of classifier-free guidance is zeros_like(proj) — up to 400 MB at the +// texture stage's token budget — so it is not materialized: pass a null proj pointer and the +// runner zeroes the input tensor on the device instead. +struct ProjCond { + std::vector global; // [1024 * 5] cls + 4 register tokens, channel-major + std::vector proj; // [proj_ch * N] channel-major, one column per DiT token + int proj_ch = 1024; + int n_global = 5; +}; + +// Build the projection conditioning for one stage. +// dino : full dinov3_encode output, [1024, 5 + (S/16)^2] channel-major. +// S : the image size that produced `dino` (512 or 1024). +// grid_res : the projection grid resolution for this stage. +// proj_ch : the stage's proj_in_channels, read off its own proj_linear weight — 1024 for the +// bare feature map, 2048 when the stage concatenates a NAF-upsampled copy. +// coords : nullptr for the dense sparse-structure stage (tokens = all grid_res^3 cells in +// DiT order), otherwise the active voxel list, one token per entry. +// naf : NAF model, or nullptr. With proj_ch 2048 and no NAF model the upsampled half is +// filled with the low-resolution samples so the stage still runs (--no-naf); that +// is a degraded input, not an equivalent one. +// img01 : raw [0,1] guide image, [3,S,S] torch CHW — required when `naf` is given. +// naf_out : NAF target resolution for this stage. +ProjCond pixal3d_proj_cond(const std::vector& dino, int S, int grid_res, int proj_ch, + const CameraParams& cam, + const std::vector>* coords, + const Model* naf, const std::vector* img01, int naf_out); + +} // namespace trellis diff --git a/include/preprocess.h b/include/preprocess.h index 47744ad..df2d872 100644 --- a/include/preprocess.h +++ b/include/preprocess.h @@ -23,4 +23,9 @@ bool image_has_alpha(const std::string& path); std::vector birefnet_cutout(const std::string& path, const Model& bm, int gpu, int& sz); // Resize a square RGB/RGBA-uint8 cutout to SxS, ImageNet-normalize -> [3,S,S] torch CHW. std::vector normalize_cutout(const std::vector& rgb, int sz, int S); + +// Resize a square RGB-uint8 cutout to SxS and scale to [0,1] WITHOUT ImageNet normalization +// -> [3,S,S] torch CHW. The NAF guide branch (Pixal3D) consumes the raw [0,1] image, not the +// normalized one that feeds DINOv3. +std::vector cutout_to_chw01(const std::vector& rgb, int sz, int S); } // namespace trellis diff --git a/src/naf.cpp b/src/naf.cpp new file mode 100644 index 0000000..3a28b64 --- /dev/null +++ b/src/naf.cpp @@ -0,0 +1,278 @@ +#include "naf.h" +#include "trellis_model.h" +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-alloc.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace trellis { +using T = ggml_tensor; + +static constexpr int DIM = 256; // image-encoder output channels +static constexpr int HEADS = 4; +static constexpr int HEAD_DIM = DIM / HEADS; // 64 +static constexpr int KERNEL = 9; // neighborhood attention window (per axis) +static constexpr int GROUPS = 8; // GroupNorm groups inside EncBlock +static constexpr float GN_EPS = 1e-5f; +static constexpr int ROPE_QUARTER = HEAD_DIM / 4; // 16 periods + +// --------------------------------------------------------------------------- +// Image encoder (the only learned part of NAF), built as one ggml graph. +// --------------------------------------------------------------------------- + +// torch Conv2d(padding_mode="reflect") on a [W,H,C,1] map. ggml has no 2-D reflect pad, so the +// border rows/columns are mirrored explicitly with views + concat — p is 0 or 1 in every NAF +// conv, so this costs two small copies per axis. +static T* reflect_pad(ggml_context* c, T* x, int p) { + if (p <= 0) return x; + const int64_t W = x->ne[0], H = x->ne[1], C = x->ne[2]; + auto row = [&](int64_t r) { + return ggml_cont(c, ggml_view_4d(c, x, W, 1, C, 1, x->nb[1], x->nb[2], x->nb[3], + (size_t)r * x->nb[1])); + }; + T* y = x; + for (int i = 1; i <= p; ++i) y = ggml_concat(c, row(i), y, 1); // rows p..1 prepended + for (int i = 1; i <= p; ++i) y = ggml_concat(c, y, row(H - 1 - i), 1); // rows H-2..H-1-p appended + const int64_t H2 = y->ne[1]; + auto col = [&](int64_t k) { + return ggml_cont(c, ggml_view_4d(c, y, 1, H2, C, 1, y->nb[1], y->nb[2], y->nb[3], + (size_t)k * y->nb[0])); + }; + T* z = y; + for (int i = 1; i <= p; ++i) z = ggml_concat(c, col(i), z, 0); + for (int i = 1; i <= p; ++i) z = ggml_concat(c, z, col(W - 1 - i), 0); + return z; +} + +static T* conv2d(ggml_context* c, const Model& m, const std::string& pre, T* x, int k) { + // im2col (ggml_conv_2d) is the portable path but materializes a [k*k*C, W*H] buffer, which is + // ~2.4 GB for the 1024-guide EncBlocks. TRELLIS_NAF_CONV_DIRECT=1 switches to the fused kernel + // where the backend implements it (CUDA), trading portability for that buffer. + static const bool direct = std::getenv("TRELLIS_NAF_CONV_DIRECT") != nullptr; + x = reflect_pad(c, x, k / 2); + T* w = m.get(pre + ".weight"); + T* y = direct ? ggml_conv_2d_direct(c, w, x, 1, 1, 0, 0, 1, 1) + : ggml_conv_2d(c, w, x, 1, 1, 0, 0, 1, 1); + if (T* b = m.try_get(pre + ".bias")) y = ggml_add(c, y, ggml_reshape_4d(c, b, 1, 1, b->ne[0], 1)); + return y; +} + +static T* group_norm(ggml_context* c, const Model& m, const std::string& pre, T* x) { + x = ggml_group_norm(c, x, GROUPS, GN_EPS); + T* w = m.get(pre + ".weight"); + T* b = m.get(pre + ".bias"); + x = ggml_mul(c, x, ggml_reshape_4d(c, w, 1, 1, w->ne[0], 1)); + return ggml_add(c, x, ggml_reshape_4d(c, b, 1, 1, b->ne[0], 1)); +} + +// encoder() = Conv2d(3, dim/2, k) followed by two EncBlocks. Both blocks are built with +// residual=False and in_channels == out_channels, so neither the skip nor the 1x1 shortcut +// exists in the checkpoint — the branch is a plain chain. +static T* enc_branch(ggml_context* c, const Model& m, const std::string& pre, T* img, int k) { + T* x = conv2d(c, m, pre + ".0", img, k); + for (int i = 1; i <= 2; ++i) { + const std::string b = pre + "." + std::to_string(i); + x = group_norm(c, m, b + ".norm1", x); + x = ggml_silu(c, x); + x = conv2d(c, m, b + ".conv1", x, k); + x = group_norm(c, m, b + ".norm2", x); + x = ggml_silu(c, x); + x = conv2d(c, m, b + ".conv2", x, k); + } + return x; +} + +// Runs the guide encoder and the adaptive average pool. Returns the [out*out, DIM] feature map +// in pixel-major order (pixel p = y*out + x, channels contiguous). +static std::vector encode_guide(const Model& m, const std::vector& img01, int S, int out) { + if (S % out != 0) throw std::runtime_error("naf: guide size must be a multiple of the target size"); + const int f = S / out; + + size_t meta = ggml_tensor_overhead() * 4096 + ggml_graph_overhead_custom(8192, false) + (1 << 20); + ggml_context* c = ggml_init({ meta, nullptr, true }); + T* img = ggml_new_tensor_4d(c, GGML_TYPE_F32, S, S, 3, 1); ggml_set_input(img); + + T* a = enc_branch(c, m, "image_encoder.encoder", img, 1); + T* b = enc_branch(c, m, "image_encoder.sem_encoder", img, 3); + T* x = ggml_concat(c, a, b, 2); // [S, S, DIM, 1] + if (f > 1) x = ggml_pool_2d(c, x, GGML_OP_POOL_AVG, f, f, f, f, 0.0f, 0.0f); + ggml_set_output(x); + + ggml_cgraph* g = ggml_new_graph_custom(c, 8192, false); + ggml_build_forward_expand(g, x); + ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(m.backend)); + if (!ggml_gallocr_alloc_graph(alloc, g)) throw std::runtime_error("naf: alloc failed"); + ggml_backend_tensor_set(img, img01.data(), 0, img01.size() * 4); + if (ggml_backend_graph_compute(m.backend, g) != GGML_STATUS_SUCCESS) + throw std::runtime_error("naf: compute failed"); + std::vector planar = tensor_to_f32(x); // ggml [out, out, DIM] -> x + out*y + out*out*ch + ggml_gallocr_free(alloc); ggml_free(c); + + const size_t np = (size_t)out * out; + std::vector pix(np * DIM); + for (int ch = 0; ch < DIM; ++ch) + for (size_t p = 0; p < np; ++p) pix[p * DIM + ch] = planar[(size_t)ch * np + p]; + return pix; +} + +// --------------------------------------------------------------------------- +// Axial RoPE (the DINOv3 formulation NAF reuses), applied in place on [np, DIM]. +// Angles depend only on the pixel, so all four heads share one cos/sin table. +// --------------------------------------------------------------------------- +static void apply_rope(std::vector& pix, int H, int W, const std::vector& periods) { + const int half = HEAD_DIM / 2; // 32: rotate_half pairs j with j+32 + std::vector cs(half), sn(half); + for (int y = 0; y < H; ++y) { + const float ch = ((y + 0.5f) / H) * 2.0f - 1.0f; + for (int x = 0; x < W; ++x) { + const float cw = ((x + 0.5f) / W) * 2.0f - 1.0f; + for (int j = 0; j < half; ++j) { + const bool wj = j >= ROPE_QUARTER; + const float ang = 6.283185307179586f * (wj ? cw : ch) / periods[wj ? j - ROPE_QUARTER : j]; + cs[j] = std::cos(ang); sn[j] = std::sin(ang); + } + float* v = &pix[((size_t)y * W + x) * DIM]; + for (int h = 0; h < HEADS; ++h) { + float* q = v + (size_t)h * HEAD_DIM; + for (int j = 0; j < half; ++j) { + const float lo = q[j], hi = q[j + half]; + q[j] = lo * cs[j] - hi * sn[j]; + q[j + half] = hi * cs[j] + lo * sn[j]; + } + } + } + } +} + +// --------------------------------------------------------------------------- +// Neighborhood cross-attention. +// +// NATTEN's dilated neighborhood attention splits each axis into `dilation` interleaved +// subsequences and applies a KERNEL-wide window inside one of them. NAF sets dilation to the +// upsampling factor d, and both keys and values are nearest-upsampled from the LR grid, so tap +// t of query i reads index (i%d) + (s+t)*d whose LR cell is exactly s+t, with the window start +// s = clamp(i/d - KERNEL/2, 0, Nlr - KERNEL). The 2-D window therefore degenerates to a plain +// KERNEL x KERNEL block of LR cells, which is what this computes. +// --------------------------------------------------------------------------- +static void attend_pixel(int px, int py, int out_w, int out_h, int Hf, int Wf, int C, + const std::vector& q, const std::vector& klr, + const float* vlr, float* dst) { + const int dx = out_w / Wf, dy = out_h / Hf; + const int sx = std::min(std::max(px / dx - KERNEL / 2, 0), Wf - KERNEL); + const int sy = std::min(std::max(py / dy - KERNEL / 2, 0), Hf - KERNEL); + const float scale = 1.0f / std::sqrt((float)HEAD_DIM); + const int vhead = C / HEADS; + + const float* qv = &q[((size_t)py * out_w + px) * DIM]; + float w[KERNEL * KERNEL]; + for (int h = 0; h < HEADS; ++h) { + const float* qh = qv + (size_t)h * HEAD_DIM; + float mx = -INFINITY; + for (int a = 0; a < KERNEL; ++a) for (int b = 0; b < KERNEL; ++b) { + const float* kh = &klr[((size_t)(sy + a) * Wf + (sx + b)) * DIM + (size_t)h * HEAD_DIM]; + float s = 0; + for (int j = 0; j < HEAD_DIM; ++j) s += qh[j] * kh[j]; + s *= scale; + w[a * KERNEL + b] = s; + if (s > mx) mx = s; + } + float sum = 0; + for (int t = 0; t < KERNEL * KERNEL; ++t) { w[t] = std::exp(w[t] - mx); sum += w[t]; } + const float inv = 1.0f / sum; + float* o = dst + (size_t)h * vhead; + for (int j = 0; j < vhead; ++j) o[j] = 0.0f; + for (int a = 0; a < KERNEL; ++a) for (int b = 0; b < KERNEL; ++b) { + const float ww = w[a * KERNEL + b] * inv; + const size_t cell = (size_t)(sy + a) * Wf + (sx + b); + // values keep the DINOv3 channel-major layout: channel c of cell -> vlr[c + C*cell] + for (int j = 0; j < vhead; ++j) o[j] += ww * vlr[(size_t)h * vhead + j + (size_t)C * cell]; + } + } +} + +std::vector naf_sample(const Model& m, + const std::vector& img01, int S, + const float* feats_lr, int Hf, int Wf, int C, + int out, const std::vector& pts_xy) { + if (Hf < KERNEL || Wf < KERNEL) + throw std::runtime_error("naf: low-res feature grid smaller than the attention window"); + if (out % Hf != 0 || out % Wf != 0) + throw std::runtime_error("naf: target size must be a multiple of the feature grid"); + if (C % HEADS != 0) throw std::runtime_error("naf: channel count must be divisible by 4 heads"); + + std::vector q = encode_guide(m, img01, S, out); + std::vector periods = tensor_to_f32(m.get("image_encoder.rope.periods")); + apply_rope(q, out, out, periods); + + // keys: the same guide features average-pooled back onto the LR grid (KeyEncoder). + const int by = out / Hf, bx = out / Wf; + std::vector klr((size_t)Hf * Wf * DIM, 0.0f); + for (int y = 0; y < out; ++y) for (int x = 0; x < out; ++x) { + float* d = &klr[((size_t)(y / by) * Wf + (x / bx)) * DIM]; + const float* s = &q[((size_t)y * out + x) * DIM]; + for (int j = 0; j < DIM; ++j) d[j] += s[j]; + } + const float navg = 1.0f / (float)(by * bx); + for (float& v : klr) v *= navg; + + // Only the pixels the projection actually reads are worth attending over: collect the four + // bilinear taps of every sample point, deduplicate, and compute that set in parallel. + const size_t NP = pts_xy.size() / 2; + const float step = (float)out / (float)S; + std::vector fx(NP), fy(NP); + std::vector slot((size_t)out * out, -1); + std::vector needed; + auto want = [&](int x, int y) { + int& s = slot[(size_t)y * out + x]; + if (s < 0) { s = (int)needed.size(); needed.push_back(y * out + x); } + }; + for (size_t p = 0; p < NP; ++p) { + float gx = (pts_xy[2 * p] + 0.5f) * step - 0.5f; + float gy = (pts_xy[2 * p + 1] + 0.5f) * step - 0.5f; + gx = std::min(std::max(gx, 0.0f), (float)out - 1.0f); // grid_sample padding_mode="border" + gy = std::min(std::max(gy, 0.0f), (float)out - 1.0f); + fx[p] = gx; fy[p] = gy; + const int x0 = (int)std::floor(gx), y0 = (int)std::floor(gy); + const int x1 = std::min(x0 + 1, out - 1), y1 = std::min(y0 + 1, out - 1); + want(x0, y0); want(x1, y0); want(x0, y1); want(x1, y1); + } + + std::vector hr((size_t)needed.size() * C); + const unsigned nthreads = std::max(1u, std::thread::hardware_concurrency()); + std::vector pool; + for (unsigned t = 0; t < nthreads; ++t) { + pool.emplace_back([&, t] { + for (size_t i = t; i < needed.size(); i += nthreads) + attend_pixel(needed[i] % out, needed[i] / out, out, out, Hf, Wf, C, + q, klr, feats_lr, &hr[i * C]); + }); + } + for (auto& th : pool) th.join(); + + std::vector res((size_t)C * NP, 0.0f); + for (size_t p = 0; p < NP; ++p) { + const int x0 = (int)std::floor(fx[p]), y0 = (int)std::floor(fy[p]); + const int x1 = std::min(x0 + 1, out - 1), y1 = std::min(y0 + 1, out - 1); + const float ax = fx[p] - x0, ay = fy[p] - y0; + const float wts[4] = { (1 - ax) * (1 - ay), ax * (1 - ay), (1 - ax) * ay, ax * ay }; + const int idx[4] = { slot[(size_t)y0 * out + x0], slot[(size_t)y0 * out + x1], + slot[(size_t)y1 * out + x0], slot[(size_t)y1 * out + x1] }; + float* d = &res[(size_t)C * p]; + for (int k = 0; k < 4; ++k) { + if (wts[k] == 0.0f) continue; + const float* s = &hr[(size_t)idx[k] * C]; + for (int j = 0; j < C; ++j) d[j] += wts[k] * s[j]; + } + } + return res; +} + +} // namespace trellis diff --git a/src/pixal3d.cpp b/src/pixal3d.cpp new file mode 100644 index 0000000..f7531d0 --- /dev/null +++ b/src/pixal3d.cpp @@ -0,0 +1,137 @@ +#include "pixal3d.h" +#include "naf.h" +#include "trellis_model.h" + +#include +#include +#include + +namespace trellis { + +static constexpr int D_DINO = 1024; // DINOv3 ViT-L/16 width +static constexpr int N_GLOBAL = 5; // cls + 4 register tokens +static constexpr int PATCH = 16; + +// Blender's sensor model, as the reference reimplements it: a 32 mm sensor with the focal +// length recovered from the horizontal FOV. +static float focal_pixels(float camera_angle_x, int resolution) { + const float focal = 16.0f / std::tan(camera_angle_x * 0.5f); + return focal * (float)resolution / 32.0f; +} + +CameraParams pixal3d_camera(float camera_angle_x, float mesh_scale, + int image_resolution, int extend_pixel) { + CameraParams cam; + cam.camera_angle_x = camera_angle_x; + cam.mesh_scale = mesh_scale; + // Reference: the grid corner (-1, 0, 0) maps through the same rotation as every other grid + // point, giving world x = -0.5/mesh_scale, and is required to land on the left image border + // (target x = -extend_pixel). Solving the perspective divide for the camera distance: + // distance = f_pixels * x_world / x_ndc with x_ndc = -extend_pixel - resolution/2 + const float f_px = focal_pixels(camera_angle_x, image_resolution); + const float x_world = -0.5f / mesh_scale; + const float x_ndc = -(float)extend_pixel - (float)image_resolution * 0.5f; + cam.distance = f_px * x_world / x_ndc; + return cam; +} + +void pixal3d_project_cell(int R, int cx, int cy, int cz, const CameraParams& cam, + int image_resolution, float& px, float& py) { + const float den = R > 1 ? (float)(R - 1) : 1.0f; + const float gx = -1.0f + 2.0f * (float)cx / den; + const float gy = -1.0f + 2.0f * (float)cy / den; + const float gz = -1.0f + 2.0f * (float)cz / den; + // ProjGrid rotates the grid into Blender axes with [[1,0,0],[0,0,-1],[0,1,0]] and halves it + // by mesh_scale; the frontal view matrix then reduces the world->camera transform to + // x_cam = wx, y_cam = wz, z_cam = -wy - distance. + const float s = 1.0f / (2.0f * cam.mesh_scale); + const float wx = gx * s, wy = -gz * s, wz = gy * s; + const float depth = wy + cam.distance; + const float f_px = focal_pixels(cam.camera_angle_x, image_resolution); + const float inv = 1.0f / (depth + 1e-8f); + px = f_px * wx * inv + (float)image_resolution * 0.5f; + py = -f_px * wz * inv + (float)image_resolution * 0.5f; // image y grows downward +} + +// Bilinear sample of a channel-major [C, Hf*Wf] map at a pixel position expressed in the +// `S`-sized image frame. Mirrors F.grid_sample(align_corners=False, padding_mode="border"). +static void sample_bilinear(const float* map, int Hf, int Wf, int C, int S, + float px, float py, float* dst) { + float gx = (px + 0.5f) * (float)Wf / (float)S - 0.5f; + float gy = (py + 0.5f) * (float)Hf / (float)S - 0.5f; + gx = std::min(std::max(gx, 0.0f), (float)Wf - 1.0f); + gy = std::min(std::max(gy, 0.0f), (float)Hf - 1.0f); + const int x0 = (int)std::floor(gx), y0 = (int)std::floor(gy); + const int x1 = std::min(x0 + 1, Wf - 1), y1 = std::min(y0 + 1, Hf - 1); + const float ax = gx - x0, ay = gy - y0; + const float w00 = (1 - ax) * (1 - ay), w10 = ax * (1 - ay); + const float w01 = (1 - ax) * ay, w11 = ax * ay; + const float* p00 = map + (size_t)C * ((size_t)y0 * Wf + x0); + const float* p10 = map + (size_t)C * ((size_t)y0 * Wf + x1); + const float* p01 = map + (size_t)C * ((size_t)y1 * Wf + x0); + const float* p11 = map + (size_t)C * ((size_t)y1 * Wf + x1); + for (int j = 0; j < C; ++j) + dst[j] = w00 * p00[j] + w10 * p10[j] + w01 * p01[j] + w11 * p11[j]; +} + +ProjCond pixal3d_proj_cond(const std::vector& dino, int S, int grid_res, int proj_ch, + const CameraParams& cam, + const std::vector>* coords, + const Model* naf, const std::vector* img01, int naf_out) { + if (proj_ch != D_DINO && proj_ch != 2 * D_DINO) + throw std::runtime_error("pixal3d: unsupported proj_in_channels"); + const int Hp = S / PATCH; + const size_t ntok = dino.size() / D_DINO; + if (ntok != (size_t)N_GLOBAL + (size_t)Hp * Hp) + throw std::runtime_error("pixal3d: DINOv3 token count does not match the image size"); + + ProjCond out; + out.n_global = N_GLOBAL; + out.global.assign(dino.begin(), dino.begin() + (size_t)D_DINO * N_GLOBAL); + + // The DINOv3 output is channel-major per token and the patch tokens are row-major over the + // patch grid, which is exactly the [C, Hf*Wf] layout the samplers want — no repacking, just + // skip the global tokens. + const float* patches = dino.data() + (size_t)D_DINO * N_GLOBAL; + + const size_t N = coords ? coords->size() : (size_t)grid_res * grid_res * grid_res; + std::vector pts(2 * N); + for (size_t t = 0; t < N; ++t) { + int cx, cy, cz; + if (coords) { cx = (*coords)[t][0]; cy = (*coords)[t][1]; cz = (*coords)[t][2]; } + else { + const int R = grid_res; + cx = (int)(t / ((size_t)R * R)); cy = (int)((t / R) % R); cz = (int)(t % R); + } + pixal3d_project_cell(grid_res, cx, cy, cz, cam, S, pts[2 * t], pts[2 * t + 1]); + } + + out.proj_ch = proj_ch; + out.proj.assign((size_t)proj_ch * N, 0.0f); + + for (size_t t = 0; t < N; ++t) + sample_bilinear(patches, Hp, Hp, D_DINO, S, pts[2 * t], pts[2 * t + 1], + &out.proj[(size_t)proj_ch * t]); + + if (proj_ch == 2 * D_DINO) { + if (naf) { + if (!img01) throw std::runtime_error("pixal3d: NAF upsampling needs the raw [0,1] guide"); + std::vector hr = naf_sample(*naf, *img01, S, patches, Hp, Hp, D_DINO, naf_out, pts); + for (size_t t = 0; t < N; ++t) + std::copy(hr.begin() + (size_t)D_DINO * t, hr.begin() + (size_t)D_DINO * (t + 1), + out.proj.begin() + (size_t)proj_ch * t + D_DINO); + } else { + // --no-naf: repeat the low-resolution samples in place of the upsampled branch. The + // stage runs and stays roughly on-distribution, but every high-frequency cue the + // second branch was trained to carry is gone. + for (size_t t = 0; t < N; ++t) + std::copy(out.proj.begin() + (size_t)proj_ch * t, + out.proj.begin() + (size_t)proj_ch * t + D_DINO, + out.proj.begin() + (size_t)proj_ch * t + D_DINO); + } + } + + return out; +} + +} // namespace trellis diff --git a/src/preprocess.cpp b/src/preprocess.cpp index 6403547..bf6b6af 100644 --- a/src/preprocess.cpp +++ b/src/preprocess.cpp @@ -28,6 +28,18 @@ std::vector normalize_cutout(const std::vector& rgb, int s return out; } +std::vector cutout_to_chw01(const std::vector& rgb, int sz, int S) { + const size_t pixels = (size_t)sz * sz; + const int channels = rgb.size() == pixels * 4 ? 4 : 3; + if (sz <= 0 || S <= 0 || rgb.size() != pixels * channels) return {}; + std::vector rs((size_t)S*S*channels); + stbir_resize_uint8(rgb.data(), sz, sz, 0, rs.data(), S, S, 0, channels); + std::vector out((size_t)3*S*S); + for (int c = 0; c < 3; ++c) for (int y = 0; y < S; ++y) for (int x = 0; x < S; ++x) + out[((size_t)c*S + y)*S + x] = rs[((size_t)y*S + x)*channels + c] / 255.0f; + return out; +} + // alpha [W*H] (>0.8 = foreground) -> bbox crop (10% margin) + premultiplied square RGBA uint8. static std::vector alpha_to_cutout(const unsigned char* rgba, int W, int H, const std::vector& alpha, int& sz) { diff --git a/src/test_pixal3d.cpp b/src/test_pixal3d.cpp new file mode 100644 index 0000000..53033ce --- /dev/null +++ b/src/test_pixal3d.cpp @@ -0,0 +1,86 @@ +// Golden-value check for the Pixal3D view-aligned projection camera. +// +// The expected numbers come from the reference implementation transcribed literally out of +// Pixal3D (ProjGrid.forward + project_points_to_image_batch in +// pixal3d/trainers/flow_matching/mixins/image_conditioned_proj.py, distance_from_fov in +// inference.py) and evaluated in double precision. Everything downstream — which DINOv3 patch a +// DiT token reads, and therefore whether the generated geometry lands on the silhouette — hangs +// off these two functions, and neither has a runtime signal when it is subtly wrong. +// +// trellis-test-pixal3d +#include "pixal3d.h" + +#include +#include + +namespace { + +struct Golden { int R, res, tok; double x, y; }; + +// camera_angle_x = Pixal3D's default, mesh_scale 1, distance derived at 512 with extend_pixel 0. +constexpr double CAX = 0.8575560450553894; +constexpr double DISTANCE = 1.093750014; + +const Golden kGolden[] = { + { 16, 512, 0, 80.313726, 431.686274 }, + { 16, 512, 1, 72.643930, 439.356070 }, + { 16, 512, 16, 80.313726, 408.261438 }, + { 16, 512, 256, 103.738562, 431.686274 }, + { 16, 512, 4095, 727.578934, -215.578934 }, + { 16, 512, 2055, 272.561922, 504.428833 }, + { 32, 512, 0, 80.313726, 431.686274 }, + { 32, 512, 1, 76.684313, 435.315687 }, + { 32, 512, 32, 80.313726, 420.351676 }, + { 32, 512, 1024, 91.648324, 431.686274 }, + { 32, 512, 32767, 727.578934, -215.578934 }, + { 32, 512, 16391, 262.602800, 460.686808 }, + { 64, 1024, 0, 160.627452, 863.372548 }, + { 64, 1024, 1, 157.092739, 866.907261 }, + { 64, 1024, 64, 160.627452, 852.217864 }, + { 64, 1024, 4096, 171.782136, 863.372548 }, + { 64, 1024, 262143, 1455.157869, -431.157869 }, + { 64, 1024, 131079, 517.995316, 889.704917 }, +}; + +} // namespace + +int main() { + int bad = 0; + + trellis::CameraParams cam = trellis::pixal3d_camera((float)CAX, 1.0f, 512, 0); + const double dd = std::fabs(cam.distance - DISTANCE); + printf("distance %.9f (expected %.9f, delta %.2e)\n", cam.distance, DISTANCE, dd); + if (dd > 1e-5) { printf(" FAIL: camera distance\n"); ++bad; } + + for (const Golden& g : kGolden) { + const int cx = g.tok / (g.R * g.R), cy = (g.tok / g.R) % g.R, cz = g.tok % g.R; + float px = 0, py = 0; + trellis::pixal3d_project_cell(g.R, cx, cy, cz, cam, g.res, px, py); + const double ex = std::fabs(px - g.x), ey = std::fabs(py - g.y); + // f32 across a perspective divide at ~1e3 pixel magnitudes: 1e-3 px is the noise floor. + const bool ok = ex < 2e-3 && ey < 2e-3; + printf("R=%3d res=%4d tok=%7d x=%12.6f/%12.6f y=%12.6f/%12.6f %s\n", + g.R, g.res, g.tok, (double)px, g.x, (double)py, g.y, ok ? "ok" : "FAIL"); + if (!ok) ++bad; + } + + // The projection is defined in normalized image space, which is why one camera solved at 512 + // serves the 1024 stages too: doubling the resolution must exactly double the pixel offset + // from the image centre. + for (int tok = 0; tok < 4096; tok += 373) { + const int cx = tok / 256, cy = (tok / 16) % 16, cz = tok % 16; + float a[2], b[2]; + trellis::pixal3d_project_cell(16, cx, cy, cz, cam, 512, a[0], a[1]); + trellis::pixal3d_project_cell(16, cx, cy, cz, cam, 1024, b[0], b[1]); + for (int k = 0; k < 2; ++k) { + const double lhs = (b[k] - 512.0), rhs = 2.0 * (a[k] - 256.0); + if (std::fabs(lhs - rhs) > 2e-3) { + printf(" FAIL: scale invariance at tok=%d axis=%d (%.6f vs %.6f)\n", tok, k, lhs, rhs); + ++bad; + } + } + } + + printf(bad ? "\n%d check(s) FAILED\n" : "\nall checks passed\n", bad); + return bad ? 1 : 0; +} From 176359eb1ed4d787cab592bb77e8333b8d15fe3e Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 16 Aug 2026 13:47:20 +0000 Subject: [PATCH 03/22] cli: select the model family with --model trellis|pixal3d The pipeline is shared end to end -- same sampler schedules, same guidance, same decoders, remesh and bake -- so --model only decides how the image conditions each flow. Wire the per-stage projection parameters (grid resolution, DINOv3 image size, NAF target) from Pixal3D's stage configs, including the HR grid following the cascade's token-budget backoff. proj_in_channels is read off the checkpoint's own proj_linear rather than a config, which doubles as the guard against pointing --model at the wrong weights: both directions fail with a message instead of producing garbage. Two behaviours degrade because Pixal3D publishes only the 1024 texture flow: --res 512 has no texture model and falls back to geometry only, and the mixed-resolution texture shortcut is disabled. (cherry picked from commit c32c4f8dfc6a3f5e99a78383cffe13e9b1893d54) (cherry picked from commit a4d3d3afe33110d556d28739acb27b91577b30ea) --- include/trellis_args.h | 20 ++++++ src/trellis_args.cpp | 23 +++++++ src/trellis_cli.cpp | 146 ++++++++++++++++++++++++++++++++++++----- 3 files changed, 173 insertions(+), 16 deletions(-) diff --git a/include/trellis_args.h b/include/trellis_args.h index ffbc9c6..79a327c 100644 --- a/include/trellis_args.h +++ b/include/trellis_args.h @@ -12,6 +12,15 @@ extern bool g_no_fa; // defined in dit.cpp (TRELLIS_NOFA) extern bool g_require_gpu; // defined in trellis_model.cpp (TRELLIS_REQUIRE_GPU) extern int g_cpu_threads; // defined in trellis_model.cpp (TRELLIS_THREADS) +// Which family of flow weights the GGUF directory holds. Both share the TRELLIS.2 DiT, +// sampler and decoders; they differ only in how the image conditions the flow — see pixal3d.h. +enum class ModelFamily { + Trellis, // TRELLIS.2: cross-attention over every DINOv3 token + Pixal3D, // Pixal3D: 5 global tokens + per-token view-aligned projection +}; + +const char* model_family_name(ModelFamily f); + // Every knob for one TRELLIS.2 image->3D run. Resolved as default -> environment // (the historical TRELLIS_* / GSS / GSH names) -> CLI flag, with the CLI winning. // trellis-cli and trellis-server share the parser: the server runs it once for its @@ -26,6 +35,17 @@ struct TrellisParams { int gpu = 0; // >=0 GPU index, <0 CPU uint32_t seed = 0; + ModelFamily family = ModelFamily::Trellis; // --model trellis|pixal3d + // Pixal3D only. The projection needs the camera the image was "taken" with: upstream + // estimates the horizontal FOV with MoGe-2 and derives the distance from it in closed form. + // MoGe-2 is not ported, so the FOV is a flag; 0 keeps Pixal3D's own default (49.13 deg). + float fov_deg = 0.0f; + float mesh_scale = 1.0f; + // NAF guided upsampling of the DINOv3 feature map (the shape/texture stages' second proj + // branch). Off falls back to sampling the bare feature map twice, which halves the effective + // proj input — accepted only as a way to run without naf.gguf. + bool naf = true; + bool cascade = true; // 1024 cascade (default); --res 512 selects the light path int hr_res = 1024; // HR cascade target resolution (1024 / 1536) int max_tokens = 49152; // HR token budget (backoff floors at 1024) diff --git a/src/trellis_args.cpp b/src/trellis_args.cpp index 60fa075..508e848 100644 --- a/src/trellis_args.cpp +++ b/src/trellis_args.cpp @@ -7,6 +7,10 @@ namespace trellis { +const char* model_family_name(ModelFamily f) { + return f == ModelFamily::Pixal3D ? "pixal3d" : "trellis"; +} + void print_usage(const char* argv0, bool server) { if (server) { fprintf(stderr, @@ -24,6 +28,18 @@ void print_usage(const char* argv0, bool server) { " -o, --output PATH output .glb (default model.glb)\n" " --copyright TEXT glTF asset.copyright metadata\n" " -m, --models DIR GGUF model directory\n" + " --model FAMILY trellis (default) | pixal3d — which flow weights the model\n" + " directory holds. pixal3d swaps the DINOv3 cross-attention for\n" + " view-aligned projection conditioning; the samplers, decoders\n" + " and every postprocessing stage are shared.\n" + " --fov DEG pixal3d: horizontal field of view of the input image, which\n" + " fixes the projection camera (default 49.13, Pixal3D's own).\n" + " Upstream estimates this with MoGe-2; that model is not ported,\n" + " so a wrong FOV shows up as geometry drifting off the silhouette.\n" + " --mesh-scale F pixal3d: object scale inside the unit grid (default 1.0)\n" + " --no-naf pixal3d: skip NAF guided upsampling (needs no naf.gguf, but\n" + " the shape/texture stages then lose their high-frequency\n" + " projection branch)\n" " --gpu N GPU index, <0 = CPU (default 0)\n" " -s, --seed N RNG seed (default 42)\n" " --res 512|1024|1536 geometry resolution\n" @@ -78,6 +94,13 @@ bool parse_args(int argc, char** argv, TrellisParams& p) { else if (a == "-o" || a == "--output") { const char* v = need(a.c_str()); if (!v) return false; p.output = v; } else if (a == "--copyright") { const char* v = need(a.c_str()); if (!v) return false; p.copyright = v; } else if (a == "-m" || a == "--models") { const char* v = need(a.c_str()); if (!v) return false; p.models = v; } + else if (a == "--model") { const char* v = need(a.c_str()); if (!v) return false; + if (std::strcmp(v, "trellis") == 0) p.family = ModelFamily::Trellis; + else if (std::strcmp(v, "pixal3d") == 0) p.family = ModelFamily::Pixal3D; + else { fprintf(stderr, "[trellis] unknown model family: %s (trellis|pixal3d)\n", v); return false; } } + else if (a == "--fov") { const char* v = need(a.c_str()); if (!v) return false; p.fov_deg = (float)atof(v); } + else if (a == "--mesh-scale") { const char* v = need(a.c_str()); if (!v) return false; p.mesh_scale = (float)atof(v); } + else if (a == "--no-naf") { p.naf = false; } else if (a == "--gpu") { const char* v = need(a.c_str()); if (!v) return false; p.gpu = atoi(v); } else if (a == "-s" || a == "--seed") { const char* v = need(a.c_str()); if (!v) return false; p.seed = (uint32_t)atoi(v); } else if (a == "--res") { const char* v = need(a.c_str()); if (!v) return false; p.set_res(atoi(v)); } diff --git a/src/trellis_cli.cpp b/src/trellis_cli.cpp index 0dd031d..82bc4f7 100644 --- a/src/trellis_cli.cpp +++ b/src/trellis_cli.cpp @@ -4,6 +4,7 @@ #include "trellis_model.h" #include "preprocess.h" #include "dinov3.h" +#include "pixal3d.h" #include "flow_runner.h" #include "ss_decoder.h" #include "shape_decoder.h" @@ -14,6 +15,7 @@ #include "remesh_dc.h" #include "stb_image_write.h" #include "trellis_run.h" +#include "ggml.h" // proj_in_channels is read straight off the checkpoint's proj_linear weight #include #include @@ -23,6 +25,7 @@ #include #include #include +#include using std::vector; static double now() { return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); } @@ -69,6 +72,26 @@ int trellis_run(const trellis::TrellisParams& cfg) { const std::string& M = cfg.models; const int gpu = cfg.gpu; const bool cascade = cfg.cascade; // 1024 cascade is the TRELLIS default; --res 512 forces the light path + + // --model pixal3d. Everything below the conditioning — sampler schedules, guidance, the SS / + // shape / texture decoders, the remesh and the bake — is shared with TRELLIS.2; only how the + // image enters each DiT changes. See pixal3d.h. + const bool pix = cfg.family == trellis::ModelFamily::Pixal3D; + trellis::CameraParams cam; + if (pix) { + constexpr float PIXAL3D_DEFAULT_FOV = 0.8575560450553894f; // radians (~49.13 deg) + const float fov = cfg.fov_deg > 0.0f ? cfg.fov_deg * 3.14159265358979f / 180.0f + : PIXAL3D_DEFAULT_FOV; + // The distance is derived at 512 on purpose: the projection is resolution-independent + // once normalized, so one camera serves both the 512 and the 1024 stages. + cam = trellis::pixal3d_camera(fov, cfg.mesh_scale, 512, 0); + printf("[trellis] model family: pixal3d (fov %.2f deg, distance %.4f, mesh scale %.2f)\n", + fov * 180.0f / 3.14159265358979f, cam.distance, cam.mesh_scale); + if (cfg.fov_deg <= 0.0f) + printf(" (using Pixal3D's default FOV — MoGe-2 estimation is not ported; pass" + " --fov if the object's perspective is noticeably wider or flatter)\n"); + } + std::mt19937 rng(run_seed); std::normal_distribution randn(0.f, 1.f); auto noise = [&](size_t n){ vector v(n); for (auto& x : v) x = randn(rng); return v; }; double t0 = now(); @@ -121,12 +144,28 @@ int trellis_run(const trellis::TrellisParams& cfg) { if (cfg.bg_only) { printf("[bg-only] done (%.1fs)\n", now() - t0); return 0; } } + // Raw [0,1] guides for NAF. The DINOv3 branch wants the ImageNet-normalized tensor; NAF's + // image encoder wants the unnormalized one, so both are kept. + vector guide, guide1024; + if (pix && cfg.naf) { + guide = trellis::cutout_to_chw01(cutout, cut_sz, 512); + if (cascade) guide1024 = trellis::cutout_to_chw01(cutout, cut_sz, 1024); + } + printf("[2/6] DINOv3 conditioning\n"); - vector cond, cond1024; + vector dino, dino1024; // full token stream: 5 global + (S/16)^2 patches { trellis::Model m = trellis::Model::load(M + "/dinov3.gguf", gpu); - cond = trellis::dinov3_encode(m, chw, 512); - if (cascade) cond1024 = trellis::dinov3_encode(m, chw1024, 1024); + dino = trellis::dinov3_encode(m, chw, 512); + if (cascade) dino1024 = trellis::dinov3_encode(m, chw1024, 1024); m.free(); } + // TRELLIS.2 cross-attends over every token. Pixal3D cross-attends over the 5 global tokens + // (cls + registers) only and routes the patch grid through the projection branch instead, so + // the cross-attention context is just a prefix of the same tensor. + constexpr size_t N_GLOBAL_FLOAT = 5 * 1024; + vector cond = pix ? vector(dino.begin(), dino.begin() + N_GLOBAL_FLOAT) : dino; + vector cond1024; + if (cascade) + cond1024 = pix ? vector(dino1024.begin(), dino1024.begin() + N_GLOBAL_FLOAT) : dino1024; const int Lc = (int)(cond.size() / 1024); vector neg(cond.size(), 0.0f); const int Lc1024 = cascade ? (int)(cond1024.size() / 1024) : 0; @@ -135,15 +174,55 @@ int trellis_run(const trellis::TrellisParams& cfg) { slat_stats("cond_512 (DINOv3@512)", cond); if (cascade) slat_stats("cond_1024 (DINOv3@1024)", cond1024); + // Per-stage projection conditioning. `proj_ch` is read off the stage's own proj_linear rather + // than guessed from a config, which also doubles as the check that a --model pixal3d run is + // pointed at Pixal3D weights (and a --model trellis run is not). + auto proj_ch_of = [&](const trellis::Model& m) -> int { + ggml_tensor* w = m.try_get("blocks.0.cross_attn.proj_linear.weight"); + if (pix && !w) + throw std::runtime_error("--model pixal3d but the checkpoint has no cross_attn.proj_linear " + "(these are TRELLIS.2 weights)"); + if (!pix && w) + throw std::runtime_error("--model trellis but the checkpoint has cross_attn.proj_linear " + "(these are Pixal3D weights; pass --model pixal3d)"); + return w ? (int)w->ne[0] : 0; + }; + // grid_res: the stage's projection grid. S / naf_out: the DINOv3 image size the stage was + // trained on and its NAF target, both taken from the Pixal3D stage configs. + auto build_proj = [&](int grid_res, int S, int naf_out, int proj_ch, + const vector>* cds) { + const vector& dn = (S == 1024) ? dino1024 : dino; + const vector& gd = (S == 1024) ? guide1024 : guide; + const bool want_naf = (proj_ch == 2048) && cfg.naf; + trellis::ProjCond pc; + if (want_naf) { + trellis::Model nm = trellis::Model::load(M + "/naf.gguf", gpu); + pc = trellis::pixal3d_proj_cond(dn, S, grid_res, proj_ch, cam, cds, &nm, &gd, naf_out); + nm.free(); + } else { + pc = trellis::pixal3d_proj_cond(dn, S, grid_res, proj_ch, cam, cds, nullptr, nullptr, naf_out); + } + printf(" proj cond: grid %d^3, image %d, %d ch%s\n", grid_res, S, proj_ch, + proj_ch == 2048 ? (want_naf ? ", NAF upsampled" : ", NAF DISABLED") : ""); + return pc; + }; + printf("[3/6] sparse-structure flow + decode\n"); vector> coords; { trellis::Model m = trellis::Model::load(M + "/ss_flow.gguf", gpu); trellis::DiTParams p; p.in_ch = 8; p.out_ch = 8; p.d_cond = 1024; p.cast_f32 = F32; + p.proj_ch = proj_ch_of(m); p.proj_mode = p.proj_ch > 0; + // The sparse-structure DiT is dense at 16^3, and Pixal3D's SS stage projects a 16^3 grid, + // so its proj tokens line up one-to-one with the DiT tokens in the same x-major order. + trellis::ProjCond pc; + if (pix) pc = build_proj(16, 512, 0, p.proj_ch, nullptr); trellis::DitRunner* run = trellis::make_dense_runner(m, p, 16, Lc); - trellis::FlowFwd fwd = [&](const vector& x, float ts, const float* c){ return run->forward(x, ts, c); }; + trellis::FlowFwd fwd = [&](const vector& x, float ts, const trellis::FlowCond& c){ return run->forward(x, ts, c); }; trellis::SamplerParams sp; sp.steps=12; sp.guidance_strength=cfg.gss; sp.guidance_rescale=0.7f; sp.gi0=0.6f; sp.gi1=1.0f; sp.rescale_t=5.0f; - vector z = trellis::sample_flow(fwd, noise(8*4096), cond.data(), neg.data(), sp); // [8,4096] ne0=8 + vector z = trellis::sample_flow(fwd, noise(8*4096), + trellis::FlowCond(cond.data(), pix ? pc.proj.data() : nullptr), + trellis::FlowCond(neg.data(), nullptr), sp); // [8,4096] ne0=8 delete run; m.free(); // transpose [8,L] -> torch [8,16,16,16] memory (c*4096 + sp) vector zdec(8*4096); @@ -156,18 +235,34 @@ int trellis_run(const trellis::TrellisParams& cfg) { printf(" active voxels @res32 = %d\n", (int)coords.size()); if (coords.empty()) { fprintf(stderr, "no voxels produced\n"); return 1; } - const bool do_tex = cfg.texture; + bool do_tex = cfg.texture; + // Pixal3D publishes a single (1024) texture flow, where TRELLIS.2 publishes both. The light + // --res 512 path has no texture model to run at all, so it degrades to geometry only. + if (do_tex && pix && !cascade) { + FILE* tf = fopen((M + "/tex_flow_512.gguf").c_str(), "rb"); + if (tf) fclose(tf); + else { printf(" (pixal3d ships no res-512 texture flow -- writing geometry only)\n"); do_tex = false; } + } // one shape SLAT flow run -> normalized [32,n] (sparse, CFG 7.5, gi[0.6,1], rescale_t 3) + // `grid_res` is the resolution the sparse coords are expressed in — the same grid Pixal3D + // projects for this stage — and S / naf_out are the stage's DINOv3 image size and NAF target. + // They are ignored in TRELLIS.2 mode. auto shape_flow = [&](const std::string& path, const vector>& cds, - const float* cnd, const float* ncnd, int lc) { + const float* cnd, const float* ncnd, int lc, + int grid_res, int S, int naf_out) { const int n = (int)cds.size(); trellis::Model m = trellis::Model::load(path, gpu); trellis::DiTParams p; p.in_ch = 32; p.out_ch = 32; p.d_cond = 1024; p.cast_f32 = F32; + p.proj_ch = proj_ch_of(m); p.proj_mode = p.proj_ch > 0; + trellis::ProjCond pc; + if (pix) pc = build_proj(grid_res, S, naf_out, p.proj_ch, &cds); trellis::DitRunner* run = trellis::make_sparse_runner(m, p, cds, lc); - trellis::FlowFwd fwd = [&](const vector& x, float ts, const float* c){ return run->forward(x, ts, c); }; + trellis::FlowFwd fwd = [&](const vector& x, float ts, const trellis::FlowCond& c){ return run->forward(x, ts, c); }; trellis::SamplerParams sp; sp.steps=12; sp.guidance_strength=cfg.gsh; sp.guidance_rescale=0.5f; sp.gi0=0.6f; sp.gi1=1.0f; sp.rescale_t=3.0f; - vector sn = trellis::sample_flow(fwd, noise((size_t)32*n), cnd, ncnd, sp); // [32,n] + vector sn = trellis::sample_flow(fwd, noise((size_t)32*n), + trellis::FlowCond(cnd, pix ? pc.proj.data() : nullptr), + trellis::FlowCond(ncnd, nullptr), sp); // [32,n] delete run; m.free(); return sn; }; @@ -194,7 +289,8 @@ int trellis_run(const trellis::TrellisParams& cfg) { const int max_tok = cfg.max_tokens; printf("[4/7] shape SLAT flow (LR 512 -> upsample -> HR %d cascade, max_tok=%d)\n", hr_target, max_tok); // (1) LR shape flow @res32 with cond_512 - lr_norm = shape_flow(M + "/shape_flow_512.gguf", coords, cond.data(), neg.data(), Lc); + lr_norm = shape_flow(M + "/shape_flow_512.gguf", coords, cond.data(), neg.data(), Lc, + /*grid_res=*/32, /*S=*/512, /*naf_out=*/512); lr_dn.resize(lr_norm.size()); for (size_t n = 0; n < coords.size(); ++n) for (int c = 0; c < 32; ++c) lr_dn[(size_t)c + 32*n] = lr_norm[(size_t)c + 32*n]*SHAPE_STD[c] + SHAPE_MEAN[c]; @@ -223,13 +319,17 @@ int trellis_run(const trellis::TrellisParams& cfg) { hr_res, gi, (int)q.size(), max_tok); hr_res -= 128; } - // (4) HR shape flow @res(hr_res//16) with cond_1024 - slat_norm = shape_flow(M + "/shape_flow_1024.gguf", shc, cond1024.data(), neg1024.data(), Lc1024); + // (4) HR shape flow @res(hr_res//16) with cond_1024. The projection grid follows the same + // backoff as the token grid — Pixal3D overrides its cond model's grid_resolution to + // hr_res//16 for exactly this reason — while the NAF target stays at the stage's 512. + slat_norm = shape_flow(M + "/shape_flow_1024.gguf", shc, cond1024.data(), neg1024.data(), Lc1024, + /*grid_res=*/hr_res / 16, /*S=*/1024, /*naf_out=*/512); RES = hr_res; cond_dec = cond1024.data(); neg_dec = neg1024.data(); Lc_dec = Lc1024; } else { printf("[4/7] shape SLAT flow (512)\n"); shc = coords; - slat_norm = shape_flow(M + "/shape_flow_512.gguf", coords, cond.data(), neg.data(), Lc); + slat_norm = shape_flow(M + "/shape_flow_512.gguf", coords, cond.data(), neg.data(), Lc, + /*grid_res=*/32, /*S=*/512, /*naf_out=*/512); } const int N = (int)shc.size(); slat_dn.resize(slat_norm.size()); @@ -282,7 +382,9 @@ int trellis_run(const trellis::TrellisParams& cfg) { constexpr int DENSE_TEX = 9000000; const int tex_res = cfg.tex_res > 0 ? cfg.tex_res : (cascade && (int)so.coords.size() > DENSE_TEX ? 512 : RES); - const bool mixed = cascade && tex_res != RES; // res-1024 geometry + res-512 texture + // res-1024 geometry + res-512 texture. The shortcut needs a res-512 texture flow, which + // only TRELLIS.2 publishes, so Pixal3D always textures at the cascade resolution. + const bool mixed = cascade && tex_res != RES && !pix; printf("[6/7] texture SLAT flow + PBR decode%s\n", mixed ? " (res-512 texture on res-1024 mesh)" : ""); if (mixed) { // decode a res-512 shape (from the LR slat) to guide the res-512 tex decode @@ -301,14 +403,24 @@ int trellis_run(const trellis::TrellisParams& cfg) { const int tlc = mixed ? Lc : Lc_dec; const int tN = (int)tcoords.size(); const std::vector>& tsubs = mixed ? so_tex.subs : so.subs; + // The texture stage's projection follows whichever branch supplied its coords: the + // res-512 tex model (grid 32 @512) or the HR one (grid RES/16 @1024). Unlike the shape + // stages, Pixal3D's texture models upsample to the full image size. + const bool tex_lr = mixed || !cascade; + const int tgrid = tex_lr ? 32 : RES / 16; + const int tS = tex_lr ? 512 : 1024; + const int tnaf = tex_lr ? 256 : 1024; vector texlat; { trellis::Model m = trellis::Model::load(tflow, gpu); trellis::DiTParams p; p.in_ch = 64; p.out_ch = 32; p.d_cond = 1024; p.cast_f32 = F32; + p.proj_ch = proj_ch_of(m); p.proj_mode = p.proj_ch > 0; + trellis::ProjCond pc; + if (pix) pc = build_proj(tgrid, tS, tnaf, p.proj_ch, &tcoords); trellis::DitRunner* run = trellis::make_sparse_runner(m, p, tcoords, tlc); // state is the 32-ch noise; each forward concat [noise(32) ; shape_slat_norm(32)] -> 64ch - trellis::FlowFwd fwd = [&](const vector& st, float ts, const float* c) { + trellis::FlowFwd fwd = [&](const vector& st, float ts, const trellis::FlowCond& c) { vector x64((size_t)64 * tN); for (int n = 0; n < tN; ++n) { for (int k = 0; k < 32; ++k) x64[(size_t)k + 64*n] = st[(size_t)k + 32*n]; @@ -317,7 +429,9 @@ int trellis_run(const trellis::TrellisParams& cfg) { return run->forward(x64, ts, c); }; trellis::SamplerParams sp; sp.steps=12; sp.guidance_strength=1.0f; sp.guidance_rescale=0.0f; sp.gi0=0.6f; sp.gi1=0.9f; sp.rescale_t=3.0f; - texlat = trellis::sample_flow(fwd, noise((size_t)32*tN), tcond, tneg, sp); // [32,tN] + texlat = trellis::sample_flow(fwd, noise((size_t)32*tN), + trellis::FlowCond(tcond, pix ? pc.proj.data() : nullptr), + trellis::FlowCond(tneg, nullptr), sp); // [32,tN] delete run; m.free(); for (int n = 0; n < tN; ++n) for (int c = 0; c < 32; ++c) texlat[(size_t)c + 32*n] = texlat[(size_t)c + 32*n]*TEX_STD[c] + TEX_MEAN[c]; } From d7fdbf5b77be598e7bc05a318041ace711f1ca9d Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 16 Aug 2026 13:47:20 +0000 Subject: [PATCH 04/22] convert + docs: Pixal3D model set and integration notes TRELLIS_FAMILY=pixal3d switches the converter's manifest. Pixal3D's checkpoints carry the same filenames as TRELLIS.2's and the same tensor layout plus the two proj tensors per block, so the verbatim-name policy needs no remapping; the decoders are unchanged and can be reused. NAF ships as a torch .pth from torch.hub, so it gets its own reader. docs/pixal3d/README.md covers what actually differs, the per-stage projection table, the camera flag standing in for the unported MoGe-2 estimation, how the NAF port works and what is not verified. (cherry picked from commit d7915be1411808c80a2034e2a15c1829487b7efa) (cherry picked from commit aa467e93217ac95600fbb923a119a582d287a01e) --- README.md | 12 +++ docs/pixal3d/README.md | 228 +++++++++++++++++++++++++++++++++++++++++ tools/convert.py | 87 ++++++++++++++-- 3 files changed, 321 insertions(+), 6 deletions(-) create mode 100644 docs/pixal3d/README.md diff --git a/README.md b/README.md index 6336599..3b70fa0 100644 --- a/README.md +++ b/README.md @@ -118,8 +118,20 @@ The most useful ones: | `--atlas PX` | UV atlas size (default 2048 @1024 / 1024 @512) | | `--box-uv` | voxel-native 6-way box projection instead of the default xatlas unwrap (O(faces), faster, looser packing) | | `--seed N` | RNG seed | +| `--model trellis\|pixal3d` | which family of flow weights `--models` holds (see [Pixal3D backend](docs/pixal3d/README.md)) | | `--require-gpu` | fail instead of falling back to the (very slow, RAM-hungry) CPU path | +### Pixal3D + +`--model pixal3d` runs [TencentARC/Pixal3D](https://github.com/TencentARC/Pixal3D) on the +same engine. Pixal3D is a TRELLIS.2 fine-tune that replaces cross-attention over the DINOv3 +patch tokens with **pixel-aligned projection conditioning**: each DiT token is a grid cell, +projected into the image and sampled there. The samplers, decoders, remesh and bake are +shared, so the integration is a conditioning module plus one branch inside the DiT block. +The shape/texture stages also run the NAF guided feature upsampler, ported in +`src/naf.cpp`. See **[docs/pixal3d/README.md](docs/pixal3d/README.md)** for the model set, +the `--fov` camera flag (MoGe-2 estimation is not ported) and the known gaps. + The postprocess matches the reference pipeline op for op (see `docs/spec/27-reference-postprocess.md` / `28-divergence-matrix.md`): the raw dual-grid mesh is welded and hole-filled, **remeshed with narrow-band UDF dual diff --git a/docs/pixal3d/README.md b/docs/pixal3d/README.md new file mode 100644 index 0000000..9fcc8a9 --- /dev/null +++ b/docs/pixal3d/README.md @@ -0,0 +1,228 @@ +# Pixal3D backend for trellis.cpp + +`--model pixal3d` runs [TencentARC/Pixal3D](https://github.com/TencentARC/Pixal3D) +(*Pixal3D: Pixel-Aligned 3D Generation from Images*, SIGGRAPH 2026) on the existing +trellis.cpp engine. `--model trellis` (the default) is unchanged TRELLIS.2. + +Pixal3D is a fine-tune of TRELLIS.2, not a new architecture. Same 1.3B DiT, same 30 +blocks, same flow-Euler sampler with guidance intervals, same sparse-structure / +shape / texture cascade, and **literally the same decoder checkpoints**. One thing +changes: how the image reaches the denoiser. That is why the integration is a new +conditioning module and two lines inside the DiT block rather than a second engine. + +--- + +## What actually differs + +TRELLIS.2 cross-attends over the full DINOv3 token stream — 1 cls + 4 register + +`(S/16)²` patch tokens — and lets attention figure out which part of the image a +voxel corresponds to. + +Pixal3D splits that into two paths: + +1. **Global.** Cross-attention over the 5 global tokens only (cls + registers). The + patch tokens never enter the cross-attention. +2. **Pixel-aligned.** Every DiT token *is* a cell of a 3-D grid. That cell's centre is + projected into the image with a fixed frontal camera, the DINOv3 patch feature map + is bilinearly sampled at the projected pixel, and a per-block `proj_linear` maps + the sampled vector into model space. The result is **added to the cross-attention + output** and takes its place as the residual branch. + +In checkpoint terms (`ProjectAttention` / `SparseProjectAttention`): + +``` +TRELLIS.2 Pixal3D +blocks.N.cross_attn.to_q blocks.N.cross_attn.cross_attn_block.to_q +blocks.N.cross_attn.to_kv blocks.N.cross_attn.cross_attn_block.to_kv +blocks.N.cross_attn.to_out blocks.N.cross_attn.cross_attn_block.to_out + blocks.N.cross_attn.proj_linear.{weight,bias} +``` + +Everything else in the block — `self_attn`, `norm1..3`, `modulation`, `mlp`, +`input_layer`, `out_layer`, the RoPE tables — is bit-for-bit the same layout, which is +why `build_dit_dense` builds both with one `proj_mode` branch. + +The shape and texture stages additionally sample a **NAF-upsampled** copy of the same +feature map and concatenate it, which is why their `proj_in_channels` is 2048 instead +of 1024. `trellis.cpp` never guesses that number: it reads it off +`blocks.0.cross_attn.proj_linear.weight`, which also serves as the check that the +GGUF directory matches `--model`. + +### Per-stage projection parameters + +Taken from Pixal3D's stage configs (`configs/gen/*_proj_finetune*.json`): + +| stage | GGUF | grid | DINOv3 image | NAF target | proj ch | +|---|---|---|---|---|---| +| sparse structure | `ss_flow.gguf` | 16³ (dense, = the DiT's own token grid) | 512 | — | 1024 | +| shape LR | `shape_flow_512.gguf` | 32³ (active voxels) | 512 | 512 | 2048 | +| shape HR | `shape_flow_1024.gguf` | `hr_res/16` | 1024 | 512 | 2048 | +| texture | `tex_flow_1024.gguf` | `hr_res/16` | 1024 | 1024 | 2048 | + +The HR grid follows the cascade's token-budget backoff (`--max-tokens`): when the +`1536` target steps down, the projection grid steps with it, exactly as the reference +overrides its cond model's `grid_resolution` per call. + +### Camera + +The projection needs the camera the photo was taken with. The reference estimates the +horizontal FOV with **MoGe-2** and then derives the camera distance from it in closed +form (`distance_from_fov`), by requiring the grid corner `x = -1` to land on the image +border. + +MoGe-2 is **not ported**. The closed-form distance is, so the FOV is the only free +parameter and it is a flag: + +``` +--fov DEG horizontal field of view (default 49.13°, Pixal3D's own default) +--mesh-scale F object scale inside the unit grid (default 1.0) +``` + +The projection is resolution-independent once normalized, so one camera solved at 512 +serves the 1024 stages too. `trellis-test-pixal3d` pins both functions to golden values +taken from the reference implementation: + +``` +$ trellis-test-pixal3d +distance 1.093750000 (expected 1.093750014, delta 1.40e-08) +R= 16 res= 512 tok= 0 x= 80.313721/ 80.313726 ... ok +... +all checks passed +``` + +A wrong FOV does not crash anything — it shows up as generated geometry drifting off +the input silhouette, thicker or flatter than the object. If that is what you see, try +`--fov` a few degrees either side before blaming the weights. + +--- + +## NAF + +The shape and texture stages upsample the DINOv3 feature map before sampling it, using +[valeoai/NAF](https://github.com/valeoai/NAF) (*Neighborhood Attention Filtering*). +NAF is ported here in full (`src/naf.cpp`), which needs one extra GGUF. + +NAF is small and unusually easy to port because **the upsampling itself has no +parameters**. The only learned part is a two-branch convolutional encoder over the RGB +guide (a 1×1 branch and a 3×3 branch, each `Conv2d → 2× EncBlock(GroupNorm/SiLU)`, +reflect-padded, concatenated to 256 channels). Everything after that is a +parameter-free neighborhood cross-attention: + +- **queries** — the guide features at the target resolution, with axial RoPE applied, +- **keys** — those same features average-pooled back onto the DINOv3 grid, +- **values** — the DINOv3 features themselves, split into 4 heads of 256 channels. + +The attention is dilated by exactly the upsampling factor. That detail is what makes a +CPU implementation practical: with dilation `d`, tap `t` of query `i` reads index +`(i%d) + (s+t)·d`, whose low-resolution cell is just `s+t`. NATTEN's dilated 2-D kernel +therefore collapses to a plain **9×9 window of low-resolution cells centred on the +query's own cell**, clamped at the border. `src/naf.cpp` computes exactly that, on +CPU, threaded, and only for the pixels the projection actually reads (the four bilinear +taps of each grid point, deduplicated) rather than for the whole upsampled map. + +Two consequences worth knowing: + +- The guide encoder runs at the full image size in ggml. `ggml_conv_2d` goes through + im2col, which materializes ~2.4 GB for the 1024-guide blocks. Set + `TRELLIS_NAF_CONV_DIRECT=1` to use `ggml_conv_2d_direct` instead where the backend + implements it (CUDA), trading portability for that buffer. +- `--no-naf` runs the shape/texture stages without `naf.gguf` by repeating the + low-resolution samples in place of the upsampled branch. The stage runs and stays + roughly on-distribution, but every high-frequency cue the second branch was trained + to carry is gone. It is an escape hatch, not an equivalent. + +The one piece not verified numerically against upstream is NATTEN's border rule, which +would need `natten` installed to A/B. The window derivation above is NATTEN's +documented semantics, and the interior — where essentially every projected point lands +— is unaffected by it. + +--- + +## Models + +Pixal3D's checkpoints carry the **same filenames** as TRELLIS.2's, so keep the two +GGUF directories separate and point `--models` at the right one. + +| role | source | notes | +|------|--------|-------| +| SS flow DiT | `TencentARC/Pixal3D` `ss_flow_img_dit_1_3B_64` | proj-conditioned | +| Shape SLAT flow | `…/slat_flow_img2shape_dit_1_3B_{512,1024}` | proj, `proj_in_channels` 2048 | +| Tex SLAT flow | `…/slat_flow_imgshape2tex_dit_1_3B_1024` | proj; **no res-512 variant exists** | +| Shape / Tex / SS decoders | `…/{shape,tex}_dec_next_dc_f16c32`, `ss_dec_conv3d_16l8` | unchanged from TRELLIS.2 — reuse the GGUFs you already converted | +| Image cond | DINOv3 ViT-L/16 | unchanged | +| Feature upsampler | `valeoai/NAF` `naf_release.pth` | `naf.gguf`, ~few MB | +| BG removal | BiRefNet | unchanged | + +Convert with the same tool: + +```bash +TRELLIS_FAMILY=pixal3d python tools/convert.py # everything +TRELLIS_FAMILY=pixal3d python tools/convert.py naf # just the upsampler +``` + +Tensor names are preserved verbatim, so the extra `proj_linear` / `cross_attn_block` +tensors need no remapping. + +Because Pixal3D publishes only the 1024 texture flow: + +- `--res 512` has no texture model to run and falls back to geometry only, +- the cascade's mixed-resolution shortcut (res-1024 geometry + res-512 texture, which + suppresses the dense-decode speckle) is disabled; Pixal3D always textures at the + cascade resolution. + +Third-party GGUF conversions of Pixal3D exist on the Hub (search `Pixal3D gguf`) but +were produced for the PyTorch pipeline. They will load here only if they keep the torch +`state_dict` names; check for `blocks.0.cross_attn.proj_linear.weight` before assuming +they do. + +--- + +## Usage + +```bash +# Pixal3D, default cascade +trellis-cli in.png out.glb --model pixal3d --models /models/pixal3d-gguf + +# with a measured FOV, and BiRefNet matting +trellis-cli in.png out.glb --model pixal3d --fov 38 --bg-removal birefnet + +# no naf.gguf available +trellis-cli in.png out.glb --model pixal3d --no-naf +``` + +`trellis-server` takes the same flags at launch. + +Every other flag — `--res`, `--max-tokens`, `--gss`/`--gsh`, `--band`, `--decim`, +`--atlas`, `--box-uv`, `--tex-res`, `--seed` — behaves identically, because everything +downstream of the conditioning is shared code. + +Mismatched weights fail fast rather than producing garbage: + +``` +--model pixal3d but the checkpoint has no cross_attn.proj_linear (these are TRELLIS.2 weights) +--model trellis but the checkpoint has cross_attn.proj_linear (these are Pixal3D weights; pass --model pixal3d) +``` + +--- + +## Code map + +| file | role | +|---|---| +| `include/pixal3d.h`, `src/pixal3d.cpp` | camera solve, grid projection, feature sampling, per-stage `ProjCond` assembly | +| `include/naf.h`, `src/naf.cpp` | NAF guide encoder (ggml) + neighborhood attention (CPU, threaded) | +| `src/dit.cpp` | `proj_mode` branch in the block: deeper cross-attn prefix + `proj_linear` add | +| `include/flow_runner.h`, `src/flow_runner.cpp` | `FlowCond` carries `(cond, proj)` through CFG; `DitRunner` gains the proj input | +| `src/trellis_cli.cpp` | per-stage projection parameters, family detection, guide preparation | +| `src/test_pixal3d.cpp` | golden-value test for the camera and projection | +| `tools/convert.py` | `TRELLIS_FAMILY=pixal3d` manifest + NAF `.pth` converter | + +## Known gaps + +- **MoGe-2 FOV estimation is not ported.** Use `--fov`. +- **NAF's NATTEN border rule is unverified** against upstream (interior is unaffected). +- **No end-to-end numerical parity run.** The projection math is pinned to the + reference; the full pipeline has not been diffed against the PyTorch implementation + on real weights, because that needs the checkpoints and a CUDA box. +- **Multi-view conditioning is out of scope.** Pixal3D trains with `num_views: 2` but + the released inference pipeline is single-view, and so is this. diff --git a/tools/convert.py b/tools/convert.py index b1e364b..ba244bc 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -1,23 +1,43 @@ #!/usr/bin/env python3 -"""Convert TRELLIS.2 (and helper) safetensors checkpoints to GGUF. +"""Convert TRELLIS.2 / Pixal3D (and helper) safetensors checkpoints to GGUF. Run with the project's uv venv: /media/ilintar/D_SSD/trellis2-venv/bin/python tools/convert.py [component ...] + TRELLIS_FAMILY=pixal3d tools/convert.py # Pixal3D flows + NAF + Design: * safetensors is parsed by hand (the numpy backend can't read bf16), so we control the bf16 -> f32 -> f16 path exactly (bf16 = high 16 bits of f32). * Quantization policy: weight matrices / convs (ndim >= 2) -> f16; all 1-D params (norms, biases, gammas, the per-block `modulation` vectors) -> f32. - * Tensor names are preserved verbatim (all core models are <= 37 chars, - under GGML_MAX_NAME=64). The model config JSON is embedded as metadata. + * Tensor names are preserved verbatim (TRELLIS.2 tops out at 37 chars, + Pixal3D at 54 for blocks.N.cross_attn.cross_attn_block.q_rms_norm.gamma; + the project builds ggml with GGML_MAX_NAME=128). The model config JSON is + embedded as metadata. + +Pixal3D notes: + * The flow checkpoints carry the same filenames and the same tensor layout as + TRELLIS.2, plus two extra tensors per block: + blocks.N.cross_attn.proj_linear.{weight,bias} + and the cross-attention itself moved one level down, under + blocks.N.cross_attn.cross_attn_block.* + Both are handled by the verbatim name policy, so no remapping is needed — + trellis.cpp keys off proj_linear's presence and reads proj_in_channels + straight off its shape. + * The decoders (ss_dec, shape_dec, tex_dec) are the unchanged TRELLIS.2 ones, + so a Pixal3D model directory can reuse decoder GGUFs already converted. + * NAF is fetched by torch.hub as a .pth rather than safetensors; convert it + with the `naf` component, which reads the state dict through torch. """ import json, struct, sys, os import numpy as np import gguf +FAMILY = os.environ.get("TRELLIS_FAMILY", "trellis") MODELS = "/media/ilintar/D_SSD/models/trellis2" -OUT = f"{MODELS}/gguf" +PIXAL3D = "/media/ilintar/D_SSD/models/pixal3d" +OUT = f"{MODELS}/gguf" if FAMILY == "trellis" else f"{PIXAL3D}/gguf" # component -> (safetensors path, config json path or None, gguf arch tag) MANIFEST = { @@ -43,6 +63,36 @@ f"{MODELS}/birefnet/config.json", "birefnet-swinl"), } +# Pixal3D (TencentARC/Pixal3D). Same filenames as TRELLIS.2, different weights: the flows are +# proj-conditioned and there is no res-512 texture flow. The decoders are byte-identical to +# TRELLIS.2's, so they are converted from whichever tree is on disk. +PIXAL3D_MANIFEST = { + "ss_flow": (f"{PIXAL3D}/ckpts/ss_flow_img_dit_1_3B_64_bf16.safetensors", + f"{PIXAL3D}/ckpts/ss_flow_img_dit_1_3B_64_bf16.json", "pixal3d-ss-flow"), + "shape_flow_512": (f"{PIXAL3D}/ckpts/slat_flow_img2shape_dit_1_3B_512_bf16.safetensors", + f"{PIXAL3D}/ckpts/slat_flow_img2shape_dit_1_3B_512_bf16.json", "pixal3d-slat-flow"), + "shape_flow_1024":(f"{PIXAL3D}/ckpts/slat_flow_img2shape_dit_1_3B_1024_bf16.safetensors", + f"{PIXAL3D}/ckpts/slat_flow_img2shape_dit_1_3B_1024_bf16.json", "pixal3d-slat-flow"), + "tex_flow_1024": (f"{PIXAL3D}/ckpts/slat_flow_imgshape2tex_dit_1_3B_1024_bf16.safetensors", + f"{PIXAL3D}/ckpts/slat_flow_imgshape2tex_dit_1_3B_1024_bf16.json", "pixal3d-slat-flow"), + "shape_dec": (f"{PIXAL3D}/ckpts/shape_dec_next_dc_f16c32_fp16.safetensors", + f"{PIXAL3D}/ckpts/shape_dec_next_dc_f16c32_fp16.json", "trellis2-shape-dec"), + "tex_dec": (f"{PIXAL3D}/ckpts/tex_dec_next_dc_f16c32_fp16.safetensors", + f"{PIXAL3D}/ckpts/tex_dec_next_dc_f16c32_fp16.json", "trellis2-tex-dec"), + "ss_dec": (f"{PIXAL3D}/ckpts/ss_dec_conv3d_16l8_fp16.safetensors", + f"{PIXAL3D}/ckpts/ss_dec_conv3d_16l8_fp16.json", "trellis2-ss-dec"), + "dinov3": (f"{MODELS}/dinov3/model.safetensors", + f"{MODELS}/dinov3/config.json", "dinov3-vitl16"), + "birefnet": (f"{MODELS}/birefnet/model.safetensors", + f"{MODELS}/birefnet/config.json", "birefnet-swinl"), + "naf": (f"{PIXAL3D}/naf/naf_release.pth", None, "naf-upsampler"), +} + +if FAMILY == "pixal3d": + MANIFEST = PIXAL3D_MANIFEST +elif FAMILY != "trellis": + raise SystemExit(f"unknown TRELLIS_FAMILY {FAMILY!r} (trellis|pixal3d)") + def read_safetensors(path): """Yield (name, numpy_f32_or_f16_array) preserving natural (torch) shape.""" @@ -115,6 +165,31 @@ def convert_birefnet(w, src): return n_f16, n_f32, total +def convert_naf(w, src): + """NAF (valeoai/NAF) ships as a torch .pth from torch.hub, not safetensors: + https://github.com/valeoai/NAF/releases/download/model/naf_release.pth + Only the two-branch guide encoder and the RoPE period buffer are parameters — the + neighborhood cross-attention that performs the upsampling is parameter-free — so + everything outside image_encoder.* is dropped.""" + import torch + sd = torch.load(src, map_location="cpu") + sd = sd.get("state_dict", sd) + f16ok = os.environ.get("FORCE_F32") != "1" + n_f16 = n_f32 = total = 0 + for name, t in sd.items(): + if not name.startswith("image_encoder."): + continue + arr = t.detach().float().numpy() + if arr.ndim >= 2 and f16ok: + data = np.ascontiguousarray(arr.astype(np.float16)); n_f16 += 1 + else: + data = np.ascontiguousarray(arr.astype(np.float32)); n_f32 += 1 + w.add_tensor(name, data); total += 1 + if total == 0: + raise ValueError("naf: no image_encoder.* tensors found — is this a NAF checkpoint?") + return n_f16, n_f32, total + + def convert(component): src, cfg, arch = MANIFEST[component] os.makedirs(OUT, exist_ok=True) @@ -129,8 +204,8 @@ def convert(component): force_f32 = os.environ.get("FORCE_F32") == "1" sparse_conv = component in ("shape_dec", "tex_dec", "shape_enc") - if component == "birefnet": - n_f16, n_f32, total = convert_birefnet(w, src) + if component in ("birefnet", "naf"): + n_f16, n_f32, total = (convert_birefnet if component == "birefnet" else convert_naf)(w, src) w.write_header_to_file(); w.write_kv_data_to_file(); w.write_tensors_to_file(); w.close() sz = os.path.getsize(dst) print(f" {component:14s} -> {os.path.basename(dst):22s} {total:4d} tensors (f16={n_f16}, f32={n_f32}) {sz/1e9:.2f} GB") From 7f79f093bcb8ed1170a5a2a52406c59b8a2b7183 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 16 Aug 2026 14:06:17 +0000 Subject: [PATCH 05/22] naf: refuse the unported pre-downsample branch; record verification The reference shrinks the guide before the convolutions once it exceeds 4x the NAF target. No Pixal3D stage comes close (the largest ratio is 2), so the branch is not ported -- but fail loudly rather than silently diverge if one ever does. Also record what the port was actually checked against: an independent transcription of the reference under random weights agrees to 2.3e-6 end to end, and the collapsed 9x9 window matches NATTEN's own get_window_start exactly, so that reduction is an identity here rather than an approximation. (cherry picked from commit d57878443f5935d1ff9b284ec6495f285685f157) (cherry picked from commit b3dea88d0a418e7cad08555e8bd863cc9806fd56) --- docs/pixal3d/README.md | 27 +++++++++++++++++++-------- src/naf.cpp | 5 +++++ 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/docs/pixal3d/README.md b/docs/pixal3d/README.md index 9fcc8a9..4936932 100644 --- a/docs/pixal3d/README.md +++ b/docs/pixal3d/README.md @@ -122,6 +122,11 @@ taps of each grid point, deduplicated) rather than for the whole upsampled map. Two consequences worth knowing: +- **The texture stage is the memory peak.** It upsamples to the full 1024, so the guide + encoder runs at 1024² and the query map alone is ~1 GB host, plus up to ~800 MB of + attended pixels at the 49152-token budget, plus the ~2.4 GB im2col below — all while + the 1.3B flow model is resident. If a 16 GB device OOMs there, lower `--max-tokens` + or fall back to `--no-naf`. - The guide encoder runs at the full image size in ggml. `ggml_conv_2d` goes through im2col, which materializes ~2.4 GB for the 1024-guide blocks. Set `TRELLIS_NAF_CONV_DIRECT=1` to use `ggml_conv_2d_direct` instead where the backend @@ -131,10 +136,13 @@ Two consequences worth knowing: roughly on-distribution, but every high-frequency cue the second branch was trained to carry is gone. It is an escape hatch, not an equivalent. -The one piece not verified numerically against upstream is NATTEN's border rule, which -would need `natten` installed to A/B. The window derivation above is NATTEN's -documented semantics, and the interior — where essentially every projected point lands -— is unaffected by it. +The port was checked end to end against an independent transcription of the reference +(`naf.py` + `layers/{convolutions,rope,attentions}.py`) driven by random weights, on +both a pooled and an unpooled configuration, with sample points outside the frame to +exercise the border clamp: **max absolute error 2.3e-6**. The collapsed window and +NATTEN's own `get_window_start` on the upsampled grid agreed exactly, so the reduction +above is an identity, not an approximation, for every stage here (it needs +`length % dilation == 0`, which all four satisfy). --- @@ -220,9 +228,12 @@ Mismatched weights fail fast rather than producing garbage: ## Known gaps - **MoGe-2 FOV estimation is not ported.** Use `--fov`. -- **NAF's NATTEN border rule is unverified** against upstream (interior is unaffected). -- **No end-to-end numerical parity run.** The projection math is pinned to the - reference; the full pipeline has not been diffed against the PyTorch implementation - on real weights, because that needs the checkpoints and a CUDA box. +- **No end-to-end numerical parity run on real weights.** The projection is pinned to + golden values from the reference and NAF was diffed against a transcription of it + under random weights, but the full pipeline has not been run against the PyTorch + implementation on the released checkpoints — that needs the weights and a CUDA box. +- **NAF's pre-downsample branch is not implemented.** It only triggers when the guide + exceeds 4× the NAF target, which no Pixal3D stage does; the code refuses rather than + diverging if a stage ever would. - **Multi-view conditioning is out of scope.** Pixal3D trains with `num_views: 2` but the released inference pipeline is single-view, and so is this. diff --git a/src/naf.cpp b/src/naf.cpp index 3a28b64..26ec610 100644 --- a/src/naf.cpp +++ b/src/naf.cpp @@ -95,6 +95,11 @@ static T* enc_branch(ggml_context* c, const Model& m, const std::string& pre, T* static std::vector encode_guide(const Model& m, const std::vector& img01, int S, int out) { if (S % out != 0) throw std::runtime_error("naf: guide size must be a multiple of the target size"); const int f = S / out; + // ImageEncoder.forward bilinearly shrinks the guide before the convolutions once it exceeds + // 4x the target size. No Pixal3D stage gets anywhere near that (the largest ratio is 2), so + // the branch is not implemented — but refuse rather than silently diverge if one ever does. + if (f > 4) throw std::runtime_error("naf: guide more than 4x the target size needs the " + "reference's pre-downsample branch, which is not ported"); size_t meta = ggml_tensor_overhead() * 4096 + ggml_graph_overhead_custom(8192, false) + (1 << 20); ggml_context* c = ggml_init({ meta, nullptr, true }); From 061673164cb72cac3f4f42e366136317039e59b4 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 16 Aug 2026 14:21:23 +0000 Subject: [PATCH 06/22] tools: probe a GGUF's tensor table without downloading the weights Third-party Pixal3D conversions exist on the Hub, but whether they load here comes down to two things buried in the header: whether the tensor names are the verbatim torch state_dict names, and what dtype the 1-D parameters carry. The second is the sharp edge -- dit.cpp adds `modulation` to an f32 timestep embedding and multiplies `norm2.weight` into an f32 activation, so f16 there is a silent mismatch rather than a slow path. GGUF puts its metadata and tensor index at the head of the file, so a few MB of HTTP Range answers both without pulling 2.7 GB per flow. Stdlib only. (cherry picked from commit 70772d6b3a109bc3b9264e7c91e0cf6f28c33e5c) (cherry picked from commit 0dbb0478f239e8522c373ec95d01182272e56ab7) --- tools/gguf_probe.py | 183 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 183 insertions(+) create mode 100755 tools/gguf_probe.py diff --git a/tools/gguf_probe.py b/tools/gguf_probe.py new file mode 100755 index 0000000..75d40d1 --- /dev/null +++ b/tools/gguf_probe.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +"""Print a GGUF's tensor table without downloading the weights. + +GGUF puts its metadata and full tensor index at the head of the file, so a few MB of +HTTP Range is enough to answer the only two questions that decide whether a third-party +conversion will load in trellis.cpp: + + 1. are the tensor names the verbatim torch state_dict names, and + 2. what dtype are the 1-D parameters (the per-block `modulation` vector, the norm + weights) stored as — trellis.cpp adds `modulation` to an f32 timestep embedding and + multiplies `norm2.weight` into an f32 activation, so f16 there is a silent + mismatch, not a slow path. + +Usage: + tools/gguf_probe.py [name-substring ...] + + tools/gguf_probe.py https://huggingface.co/USER/REPO/resolve/main/model.gguf + +With no substrings it prints a summary plus the tensors trellis.cpp is picky about. +Only the standard library is used, so it runs anywhere python does. +""" +import struct +import sys +import urllib.request + +GGUF_MAGIC = 0x46554747 + +# ggml_type -> name, for the types a checkpoint conversion realistically emits. +TYPES = { + 0: "F32", 1: "F16", 2: "Q4_0", 3: "Q4_1", 6: "Q5_0", 7: "Q5_1", 8: "Q8_0", + 9: "Q8_1", 10: "Q2_K", 11: "Q3_K", 12: "Q4_K", 13: "Q5_K", 14: "Q6_K", + 15: "Q8_K", 30: "BF16", +} + +# What trellis.cpp requires beyond "the tensor exists". See src/dit.cpp. +PROBES = [ + ("blocks.0.cross_attn.proj_linear.weight", None, "Pixal3D marker; ne[0] is proj_in_channels"), + ("blocks.0.cross_attn.cross_attn_block.to_q.weight", None, "proj-mode cross-attn nesting"), + ("blocks.0.modulation", "F32", "added to the f32 timestep embedding"), + ("blocks.0.norm2.weight", "F32", "multiplied into an f32 activation"), + ("blocks.0.norm2.bias", "F32", "added to an f32 activation"), + ("blocks.0.self_attn.q_rms_norm.gamma", None, "f16 tolerated (cast at graph build)"), + ("input_layer.weight", None, "matmul weight, f16/quant fine"), +] + + +class Head: + """Reads a local file or an HTTP URL, fetching in chunks and caching what it read.""" + + def __init__(self, src, chunk=1 << 22): + self.src, self.chunk, self.buf = src, chunk, b"" + self.remote = src.startswith("http://") or src.startswith("https://") + if not self.remote: + self.fh = open(src, "rb") + + def _grow(self, need): + while len(self.buf) < need: + if self.remote: + lo, hi = len(self.buf), len(self.buf) + max(self.chunk, need - len(self.buf)) - 1 + req = urllib.request.Request(self.src, headers={"Range": f"bytes={lo}-{hi}"}) + with urllib.request.urlopen(req) as r: + if r.status != 206: + raise SystemExit("server ignored the Range request; download the file instead") + part = r.read() + else: + self.fh.seek(len(self.buf)) + part = self.fh.read(max(self.chunk, need - len(self.buf))) + if not part: + raise SystemExit("unexpected end of file while reading the GGUF header") + self.buf += part + + def at(self, off, n): + self._grow(off + n) + return self.buf[off:off + n] + + +def main(argv): + if len(argv) < 2: + raise SystemExit(__doc__) + src, wanted = argv[1], argv[2:] + h = Head(src) + pos = 0 + + def take(n): + nonlocal pos + b = h.at(pos, n) + pos += n + return b + + def u32(): + return struct.unpack("5} ne={ne}") + return 0 + + print("\n trellis.cpp compatibility probes:") + bad = 0 + for name, want, why in PROBES: + got = by_name.get(name) + if got is None: + print(f" MISSING {name:<52} ({why})") + bad += 1 + continue + ne, t = got + note = "" + if want and ty(t) != want: + note = f" <-- expected {want}: {why}" + bad += 1 + print(f" {ty(t):>5} {name:<52} ne={ne}{note}") + + proj = by_name.get("blocks.0.cross_attn.proj_linear.weight") + print() + if proj: + print(f" => Pixal3D weights, proj_in_channels = {proj[0][0]} (--model pixal3d)") + else: + print(" => no proj_linear: TRELLIS.2 weights (--model trellis)") + print(f" => {bad} compatibility problem(s)" if bad else " => no compatibility problems found") + return 1 if bad else 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) From 47ae389f4fe099fad587f6fe196b3d409b14de5f Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 16 Aug 2026 15:17:23 +0000 Subject: [PATCH 07/22] tools: gguf_probe checks block structure and matmul dtypes Two things the first version missed on a real third-party conversion. Structure: a DiT is N identical blocks plus a fixed preamble, so every block must carry the same suffix set. A converter that renamed, merged or dropped tensors shows up as a ragged block here instead of as a failure at inference time. For a Pixal3D SLat flow the arithmetic is 30 x 23 + 10 = 700. Dtypes: report the matmul weights separately. BF16 loads and runs -- ggml has a dedicated cuBLAS BF16 path and CPU type traits -- but --f32 only casts F16, so that escape hatch does not apply to BF16 weights and it is worth saying so before someone reaches for it to debug a bad sm_120-family matmul. (cherry picked from commit 1c545a12185c3a8578021d92c3583a8870e3bc12) (cherry picked from commit 5b82d1f664fbd9c7b3d01077eaf5b691643ab2c9) --- tools/gguf_probe.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tools/gguf_probe.py b/tools/gguf_probe.py index 75d40d1..837e52d 100755 --- a/tools/gguf_probe.py +++ b/tools/gguf_probe.py @@ -169,6 +169,45 @@ def ty(t): bad += 1 print(f" {ty(t):>5} {name:<52} ne={ne}{note}") + # Structural completeness. A DiT is 30 identical blocks plus a fixed preamble, so every + # block must carry the same set of suffixes. A converter that renamed, merged or dropped + # tensors shows up here as a ragged block rather than as a crash at inference time. + blocks = {} + top = [] + for n, ne, t in tensors: + if n.startswith("blocks."): + i, _, suffix = n[len("blocks."):].partition(".") + if i.isdigit(): + blocks.setdefault(int(i), set()).add(suffix) + continue + top.append(n) + if blocks: + nb = max(blocks) + 1 + ref = blocks.get(0, set()) + ragged = [i for i in range(nb) if blocks.get(i) != ref] + print(f"\n structure: {nb} blocks x {len(ref)} tensors + {len(top)} top-level " + f"= {nb * len(ref) + len(top)} (file has {n_tensors})") + if len(blocks) != nb or ragged: + print(f" RAGGED: block indices missing or inconsistent: {ragged[:8]}") + bad += 1 + if nb * len(ref) + len(top) != n_tensors: + print(" COUNT MISMATCH: some blocks share suffixes unevenly") + bad += 1 + + # BF16 matmul weights load and run (ggml has a dedicated cuBLAS BF16 path and CPU type + # traits), but --f32 only casts F16, so that escape hatch does not apply to them. + mm = [t for n, ne, t in tensors if len(ne) >= 2 and "rms_norm" not in n] + hist = {} + for t in mm: + hist[ty(t)] = hist.get(ty(t), 0) + 1 + if hist: + print(f" matmul weights: " + ", ".join(f"{k}={v}" for k, v in sorted(hist.items()))) + if "BF16" in hist: + print(" note: BF16 runs, but --f32 only casts F16 — it will not affect these") + quant = [k for k in hist if k.startswith("Q")] + if quant: + print(f" note: k-quants/{'/'.join(quant)} are untested in trellis.cpp's graphs") + proj = by_name.get("blocks.0.cross_attn.proj_linear.weight") print() if proj: From f371be69a2cabcc2e5d1d94036febbb6d1a8e1c9 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 16 Aug 2026 18:06:34 +0000 Subject: [PATCH 08/22] models: prefix the Pixal3D-specific GGUFs and share one model directory Pixal3D's checkpoints carry the same upstream filenames as TRELLIS.2's, which forced either two model directories or a tree of symlinks -- friction for anyone testing both families. Name the family-specific files pixal3d_*.gguf instead: the four flows and NAF. The decoders, DINOv3 and BiRefNet are byte-identical between the families and keep their plain names, so both --model values read one directory and adding Pixal3D to a working TRELLIS.2 set is 5 new files rather than a second copy of everything. The converter applies the prefix automatically. (cherry picked from commit e51a731f7478de4dd5a421ea05dd113e2da21c9f) (cherry picked from commit a8b2f63f594038e899be4412ef0f9f7f92e80771) --- README.md | 9 +++++-- docs/pixal3d/README.md | 60 ++++++++++++++++++++++++++---------------- src/trellis_cli.cpp | 19 ++++++++----- tools/convert.py | 19 ++++++++++--- 4 files changed, 72 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index 3b70fa0..94571de 100644 --- a/README.md +++ b/README.md @@ -129,8 +129,13 @@ patch tokens with **pixel-aligned projection conditioning**: each DiT token is a projected into the image and sampled there. The samplers, decoders, remesh and bake are shared, so the integration is a conditioning module plus one branch inside the DiT block. The shape/texture stages also run the NAF guided feature upsampler, ported in -`src/naf.cpp`. See **[docs/pixal3d/README.md](docs/pixal3d/README.md)** for the model set, -the `--fov` camera flag (MoGe-2 estimation is not ported) and the known gaps. +`src/naf.cpp`. + +Both families use the **same model directory**: the Pixal3D flows and NAF are named +`pixal3d_*.gguf`, while the decoders, DINOv3 and BiRefNet are byte-identical and shared, +so adding Pixal3D to a working TRELLIS.2 set is 5 new files. See +**[docs/pixal3d/README.md](docs/pixal3d/README.md)** for the model list, the `--fov` +camera flag (MoGe-2 estimation is not ported) and the known gaps. The postprocess matches the reference pipeline op for op (see `docs/spec/27-reference-postprocess.md` / `28-divergence-matrix.md`): the raw diff --git a/docs/pixal3d/README.md b/docs/pixal3d/README.md index 4936932..2b359ad 100644 --- a/docs/pixal3d/README.md +++ b/docs/pixal3d/README.md @@ -54,10 +54,10 @@ Taken from Pixal3D's stage configs (`configs/gen/*_proj_finetune*.json`): | stage | GGUF | grid | DINOv3 image | NAF target | proj ch | |---|---|---|---|---|---| -| sparse structure | `ss_flow.gguf` | 16³ (dense, = the DiT's own token grid) | 512 | — | 1024 | -| shape LR | `shape_flow_512.gguf` | 32³ (active voxels) | 512 | 512 | 2048 | -| shape HR | `shape_flow_1024.gguf` | `hr_res/16` | 1024 | 512 | 2048 | -| texture | `tex_flow_1024.gguf` | `hr_res/16` | 1024 | 1024 | 2048 | +| sparse structure | `pixal3d_ss_flow.gguf` | 16³ (dense, = the DiT's own token grid) | 512 | — | 1024 | +| shape LR | `pixal3d_shape_flow_512.gguf` | 32³ (active voxels) | 512 | 512 | 2048 | +| shape HR | `pixal3d_shape_flow_1024.gguf` | `hr_res/16` | 1024 | 512 | 2048 | +| texture | `pixal3d_tex_flow_1024.gguf` | `hr_res/16` | 1024 | 1024 | 2048 | The HR grid follows the cascade's token-budget backoff (`--max-tokens`): when the `1536` target steps down, the projection grid steps with it, exactly as the reference @@ -148,20 +148,24 @@ above is an identity, not an approximation, for every stage here (it needs ## Models -Pixal3D's checkpoints carry the **same filenames** as TRELLIS.2's, so keep the two -GGUF directories separate and point `--models` at the right one. - -| role | source | notes | -|------|--------|-------| -| SS flow DiT | `TencentARC/Pixal3D` `ss_flow_img_dit_1_3B_64` | proj-conditioned | -| Shape SLAT flow | `…/slat_flow_img2shape_dit_1_3B_{512,1024}` | proj, `proj_in_channels` 2048 | -| Tex SLAT flow | `…/slat_flow_imgshape2tex_dit_1_3B_1024` | proj; **no res-512 variant exists** | -| Shape / Tex / SS decoders | `…/{shape,tex}_dec_next_dc_f16c32`, `ss_dec_conv3d_16l8` | unchanged from TRELLIS.2 — reuse the GGUFs you already converted | -| Image cond | DINOv3 ViT-L/16 | unchanged | -| Feature upsampler | `valeoai/NAF` `naf_release.pth` | `naf.gguf`, ~few MB | -| BG removal | BiRefNet | unchanged | - -Convert with the same tool: +**One model directory serves both families.** Pixal3D's checkpoints carry the same +upstream filenames as TRELLIS.2's, so the family-specific ones are written with a +`pixal3d_` prefix. Everything that is byte-identical between the two keeps its plain +name and is shared — adding Pixal3D to a working TRELLIS.2 set is **5 new files**, not +a second copy of everything. + +| file | source | shared? | +|------|--------|---------| +| `pixal3d_ss_flow.gguf` | `TencentARC/Pixal3D` `ss_flow_img_dit_1_3B_64` | no — proj-conditioned | +| `pixal3d_shape_flow_512.gguf` | `…/slat_flow_img2shape_dit_1_3B_512` | no — `proj_in_channels` 2048 | +| `pixal3d_shape_flow_1024.gguf` | `…/slat_flow_img2shape_dit_1_3B_1024` | no | +| `pixal3d_tex_flow_1024.gguf` | `…/slat_flow_imgshape2tex_dit_1_3B_1024` | no; **no res-512 variant exists** | +| `pixal3d_naf.gguf` | `valeoai/NAF` `naf_release.pth` | no — Pixal3D only, a few MB | +| `shape_dec.gguf`, `tex_dec.gguf`, `ss_dec.gguf` | `…/{shape,tex}_dec_next_dc_f16c32`, `ss_dec_conv3d_16l8` | **yes** — unchanged from TRELLIS.2 | +| `dinov3.gguf` | DINOv3 ViT-L/16 | **yes** | +| `birefnet.gguf` | BiRefNet | **yes** | + +Convert with the same tool; the prefix is applied automatically: ```bash TRELLIS_FAMILY=pixal3d python tools/convert.py # everything @@ -179,17 +183,27 @@ Because Pixal3D publishes only the 1024 texture flow: cascade resolution. Third-party GGUF conversions of Pixal3D exist on the Hub (search `Pixal3D gguf`) but -were produced for the PyTorch pipeline. They will load here only if they keep the torch -`state_dict` names; check for `blocks.0.cross_attn.proj_linear.weight` before assuming -they do. +were produced for the PyTorch pipeline. They load here only if they keep the torch +`state_dict` names — `tools/gguf_probe.py` answers that from ~4 MB of HTTP Range, +without downloading the weights: + +```bash +tools/gguf_probe.py https://huggingface.co/USER/REPO/resolve/main/some_flow.gguf +``` + +It reports the tensor names, the per-block structure (a SLat flow is 30 × 23 + 10 = +700 tensors), `proj_in_channels`, and the dtypes trellis.cpp is picky about — the 1-D +parameters must be F32, because `dit.cpp` adds `modulation` to an f32 timestep +embedding and multiplies `norm2.weight` into an f32 activation. BF16 matmul weights +are fine, but note `--f32` only casts F16 and so will not affect them. --- ## Usage ```bash -# Pixal3D, default cascade -trellis-cli in.png out.glb --model pixal3d --models /models/pixal3d-gguf +# Pixal3D, default cascade — same model directory as TRELLIS.2 +trellis-cli in.png out.glb --model pixal3d --models /models # with a measured FOV, and BiRefNet matting trellis-cli in.png out.glb --model pixal3d --fov 38 --bg-removal birefnet diff --git a/src/trellis_cli.cpp b/src/trellis_cli.cpp index 82bc4f7..16bb7e9 100644 --- a/src/trellis_cli.cpp +++ b/src/trellis_cli.cpp @@ -77,6 +77,11 @@ int trellis_run(const trellis::TrellisParams& cfg) { // shape / texture decoders, the remesh and the bake — is shared with TRELLIS.2; only how the // image enters each DiT changes. See pixal3d.h. const bool pix = cfg.family == trellis::ModelFamily::Pixal3D; + // Both families live in ONE model directory. The flows and NAF are family-specific and carry + // a `pixal3d_` prefix; the decoders, DINOv3 and BiRefNet are byte-identical between the two + // and keep their plain names, so adding Pixal3D to an existing TRELLIS.2 set is 5 new files + // rather than a second copy of everything. + const std::string FP = pix ? "/pixal3d_" : "/"; trellis::CameraParams cam; if (pix) { constexpr float PIXAL3D_DEFAULT_FOV = 0.8575560450553894f; // radians (~49.13 deg) @@ -196,7 +201,7 @@ int trellis_run(const trellis::TrellisParams& cfg) { const bool want_naf = (proj_ch == 2048) && cfg.naf; trellis::ProjCond pc; if (want_naf) { - trellis::Model nm = trellis::Model::load(M + "/naf.gguf", gpu); + trellis::Model nm = trellis::Model::load(M + FP + "naf.gguf", gpu); pc = trellis::pixal3d_proj_cond(dn, S, grid_res, proj_ch, cam, cds, &nm, &gd, naf_out); nm.free(); } else { @@ -210,7 +215,7 @@ int trellis_run(const trellis::TrellisParams& cfg) { printf("[3/6] sparse-structure flow + decode\n"); vector> coords; { - trellis::Model m = trellis::Model::load(M + "/ss_flow.gguf", gpu); + trellis::Model m = trellis::Model::load(M + FP + "ss_flow.gguf", gpu); trellis::DiTParams p; p.in_ch = 8; p.out_ch = 8; p.d_cond = 1024; p.cast_f32 = F32; p.proj_ch = proj_ch_of(m); p.proj_mode = p.proj_ch > 0; // The sparse-structure DiT is dense at 16^3, and Pixal3D's SS stage projects a 16^3 grid, @@ -239,7 +244,7 @@ int trellis_run(const trellis::TrellisParams& cfg) { // Pixal3D publishes a single (1024) texture flow, where TRELLIS.2 publishes both. The light // --res 512 path has no texture model to run at all, so it degrades to geometry only. if (do_tex && pix && !cascade) { - FILE* tf = fopen((M + "/tex_flow_512.gguf").c_str(), "rb"); + FILE* tf = fopen((M + FP + "tex_flow_512.gguf").c_str(), "rb"); if (tf) fclose(tf); else { printf(" (pixal3d ships no res-512 texture flow -- writing geometry only)\n"); do_tex = false; } } @@ -289,7 +294,7 @@ int trellis_run(const trellis::TrellisParams& cfg) { const int max_tok = cfg.max_tokens; printf("[4/7] shape SLAT flow (LR 512 -> upsample -> HR %d cascade, max_tok=%d)\n", hr_target, max_tok); // (1) LR shape flow @res32 with cond_512 - lr_norm = shape_flow(M + "/shape_flow_512.gguf", coords, cond.data(), neg.data(), Lc, + lr_norm = shape_flow(M + FP + "shape_flow_512.gguf", coords, cond.data(), neg.data(), Lc, /*grid_res=*/32, /*S=*/512, /*naf_out=*/512); lr_dn.resize(lr_norm.size()); for (size_t n = 0; n < coords.size(); ++n) for (int c = 0; c < 32; ++c) @@ -322,13 +327,13 @@ int trellis_run(const trellis::TrellisParams& cfg) { // (4) HR shape flow @res(hr_res//16) with cond_1024. The projection grid follows the same // backoff as the token grid — Pixal3D overrides its cond model's grid_resolution to // hr_res//16 for exactly this reason — while the NAF target stays at the stage's 512. - slat_norm = shape_flow(M + "/shape_flow_1024.gguf", shc, cond1024.data(), neg1024.data(), Lc1024, + slat_norm = shape_flow(M + FP + "shape_flow_1024.gguf", shc, cond1024.data(), neg1024.data(), Lc1024, /*grid_res=*/hr_res / 16, /*S=*/1024, /*naf_out=*/512); RES = hr_res; cond_dec = cond1024.data(); neg_dec = neg1024.data(); Lc_dec = Lc1024; } else { printf("[4/7] shape SLAT flow (512)\n"); shc = coords; - slat_norm = shape_flow(M + "/shape_flow_512.gguf", coords, cond.data(), neg.data(), Lc, + slat_norm = shape_flow(M + FP + "shape_flow_512.gguf", coords, cond.data(), neg.data(), Lc, /*grid_res=*/32, /*S=*/512, /*naf_out=*/512); } const int N = (int)shc.size(); @@ -395,7 +400,7 @@ int trellis_run(const trellis::TrellisParams& cfg) { } // tex flow + decode inputs: HR path (shc/slat_norm/cond_dec/so.subs) vs res-512 mixed path // (coords/lr_norm/cond_512/so_tex.subs). The tex decoder upsamples via the guide subdivision. - const std::string tflow = M + (mixed ? "/tex_flow_512.gguf" : (cascade ? "/tex_flow_1024.gguf" : "/tex_flow_512.gguf")); + const std::string tflow = M + FP + (mixed ? "tex_flow_512.gguf" : (cascade ? "tex_flow_1024.gguf" : "tex_flow_512.gguf")); const vector>& tcoords = mixed ? coords : shc; const vector& tslat = mixed ? lr_norm : slat_norm; const float* tcond = mixed ? cond.data() : cond_dec; diff --git a/tools/convert.py b/tools/convert.py index ba244bc..66971e5 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -25,8 +25,10 @@ Both are handled by the verbatim name policy, so no remapping is needed — trellis.cpp keys off proj_linear's presence and reads proj_in_channels straight off its shape. - * The decoders (ss_dec, shape_dec, tex_dec) are the unchanged TRELLIS.2 ones, - so a Pixal3D model directory can reuse decoder GGUFs already converted. + * The decoders (ss_dec, shape_dec, tex_dec), DINOv3 and BiRefNet are the + unchanged TRELLIS.2 ones and keep their plain names; the flows and NAF are + written as `pixal3d_*.gguf`. Both families therefore share ONE model + directory, and adding Pixal3D to an existing set is 5 new files. * NAF is fetched by torch.hub as a .pth rather than safetensors; convert it with the `naf` component, which reads the state dict through torch. """ @@ -190,10 +192,21 @@ def convert_naf(w, src): return n_f16, n_f32, total +# Components that are byte-identical between the two families and therefore keep their plain +# name, so one model directory can serve both. Everything else gets a `pixal3d_` prefix. +SHARED = {"shape_dec", "tex_dec", "ss_dec", "dinov3", "birefnet"} + + +def out_name(component): + if FAMILY == "pixal3d" and component not in SHARED: + return f"pixal3d_{component}.gguf" + return f"{component}.gguf" + + def convert(component): src, cfg, arch = MANIFEST[component] os.makedirs(OUT, exist_ok=True) - dst = f"{OUT}/{component}.gguf" + dst = f"{OUT}/{out_name(component)}" w = gguf.GGUFWriter(dst, arch) if cfg and os.path.exists(cfg): with open(cfg) as f: From eb24b94ef959e27d004da8ba4eb55882752bc4ae Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 16 Aug 2026 20:39:51 +0000 Subject: [PATCH 09/22] convert: take model paths from the environment The checkpoint trees live outside the repo, so the hardcoded defaults are only ever right on one machine -- everyone else had to edit the script before running it, which is both friction and a stray diff waiting to be committed by accident. TRELLIS_MODELS, PIXAL3D_MODELS and TRELLIS_GGUF_OUT override them. The old values remain the defaults, so existing invocations are unaffected. (cherry picked from commit 9bb6e7360e55f1336da67a96cb4416fb378d139e) (cherry picked from commit 8e73f224d46195b02a87f0d41afa0675ec169de9) --- tools/convert.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tools/convert.py b/tools/convert.py index 66971e5..7127542 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -6,6 +6,11 @@ TRELLIS_FAMILY=pixal3d tools/convert.py # Pixal3D flows + NAF +Paths come from the environment, so nothing here needs editing per machine: + TRELLIS_MODELS TRELLIS.2 checkpoint tree + PIXAL3D_MODELS Pixal3D checkpoint tree (and naf/naf_release.pth) + TRELLIS_GGUF_OUT where the .gguf files are written + Design: * safetensors is parsed by hand (the numpy backend can't read bf16), so we control the bf16 -> f32 -> f16 path exactly (bf16 = high 16 bits of f32). @@ -37,9 +42,13 @@ import gguf FAMILY = os.environ.get("TRELLIS_FAMILY", "trellis") -MODELS = "/media/ilintar/D_SSD/models/trellis2" -PIXAL3D = "/media/ilintar/D_SSD/models/pixal3d" -OUT = f"{MODELS}/gguf" if FAMILY == "trellis" else f"{PIXAL3D}/gguf" +# Source trees and output directory. The defaults are the author's layout; override them from +# the environment rather than editing this file (the checkpoints live outside the repo, so +# every machine — and every Windows contributor — needs different paths). +MODELS = os.environ.get("TRELLIS_MODELS", "/media/ilintar/D_SSD/models/trellis2") +PIXAL3D = os.environ.get("PIXAL3D_MODELS", "/media/ilintar/D_SSD/models/pixal3d") +OUT = os.environ.get("TRELLIS_GGUF_OUT", + f"{MODELS}/gguf" if FAMILY == "trellis" else f"{PIXAL3D}/gguf") # component -> (safetensors path, config json path or None, gguf arch tag) MANIFEST = { From 0c6161cd9298ad354fa04b482712e45f31a1c87c Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 16 Aug 2026 20:55:07 +0000 Subject: [PATCH 10/22] convert: let NAF_CKPT point straight at naf_release.pth NAF is a separate project's GitHub release, not part of the Pixal3D checkpoint tree, so requiring it under PIXAL3D_MODELS/naf/ was an arbitrary constraint on where it may be downloaded. The default is unchanged. (cherry picked from commit d2fd02f504ec939ee795c5310369678ed44e0ac2) (cherry picked from commit b6941f44130a4dadc0a99332b796507ddf83d766) --- tools/convert.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tools/convert.py b/tools/convert.py index 7127542..5d986a1 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -10,6 +10,7 @@ TRELLIS_MODELS TRELLIS.2 checkpoint tree PIXAL3D_MODELS Pixal3D checkpoint tree (and naf/naf_release.pth) TRELLIS_GGUF_OUT where the .gguf files are written + NAF_CKPT full path to naf_release.pth (default PIXAL3D_MODELS/naf/) Design: * safetensors is parsed by hand (the numpy backend can't read bf16), so we @@ -96,7 +97,9 @@ f"{MODELS}/dinov3/config.json", "dinov3-vitl16"), "birefnet": (f"{MODELS}/birefnet/model.safetensors", f"{MODELS}/birefnet/config.json", "birefnet-swinl"), - "naf": (f"{PIXAL3D}/naf/naf_release.pth", None, "naf-upsampler"), + # NAF is a separate project's release, not part of the Pixal3D checkpoint tree, so point + # NAF_CKPT straight at the .pth wherever it was downloaded. + "naf": (os.environ.get("NAF_CKPT", f"{PIXAL3D}/naf/naf_release.pth"), None, "naf-upsampler"), } if FAMILY == "pixal3d": From 05f41cd58445f35c0c997c4e16b32e9aa40fd4c1 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 16 Aug 2026 21:45:23 +0000 Subject: [PATCH 11/22] docs: point at a pre-built Pixal3D GGUF set The five Pixal3D-specific files are published at vegax87/Pixal3D under the exact names the loader expects, so testing the backend no longer means converting from the TencentARC safetensors first. The conversion route stays documented, now with the environment variables rather than the edit-the-script instructions. (cherry picked from commit e94e88169a4bf5a6a456bfc66db13ada85d358f7) (cherry picked from commit 464e56e3c99294eb5b2ce663d3e8adda15fe6eaa) --- README.md | 3 ++- docs/pixal3d/README.md | 27 ++++++++++++++++++++++++--- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 94571de..d1c511e 100644 --- a/README.md +++ b/README.md @@ -133,7 +133,8 @@ The shape/texture stages also run the NAF guided feature upsampler, ported in Both families use the **same model directory**: the Pixal3D flows and NAF are named `pixal3d_*.gguf`, while the decoders, DINOv3 and BiRefNet are byte-identical and shared, -so adding Pixal3D to a working TRELLIS.2 set is 5 new files. See +so adding Pixal3D to a working TRELLIS.2 set is 5 new files — pre-built at +[`vegax87/Pixal3D`](https://huggingface.co/vegax87/Pixal3D). See **[docs/pixal3d/README.md](docs/pixal3d/README.md)** for the model list, the `--fov` camera flag (MoGe-2 estimation is not ported) and the known gaps. diff --git a/docs/pixal3d/README.md b/docs/pixal3d/README.md index 2b359ad..e685357 100644 --- a/docs/pixal3d/README.md +++ b/docs/pixal3d/README.md @@ -148,6 +148,17 @@ above is an identity, not an approximation, for every stage here (it needs ## Models +**Pre-built GGUFs:** [`vegax87/Pixal3D`](https://huggingface.co/vegax87/Pixal3D) — the five +Pixal3D-specific files, already named as the loader expects. Drop them into the model +directory you already use for TRELLIS.2: + +```bash +for f in ss_flow shape_flow_512 shape_flow_1024 tex_flow_1024 naf; do + curl -fL -o "$MODELS/pixal3d_$f.gguf" \ + "https://huggingface.co/vegax87/Pixal3D/resolve/main/pixal3d_$f.gguf" +done +``` + **One model directory serves both families.** Pixal3D's checkpoints carry the same upstream filenames as TRELLIS.2's, so the family-specific ones are written with a `pixal3d_` prefix. Everything that is byte-identical between the two keeps its plain @@ -165,13 +176,23 @@ a second copy of everything. | `dinov3.gguf` | DINOv3 ViT-L/16 | **yes** | | `birefnet.gguf` | BiRefNet | **yes** | -Convert with the same tool; the prefix is applied automatically: +To convert from the source checkpoints instead, the same tool handles both families and +applies the prefix automatically. Paths come from the environment, so nothing in the +script needs editing: ```bash -TRELLIS_FAMILY=pixal3d python tools/convert.py # everything -TRELLIS_FAMILY=pixal3d python tools/convert.py naf # just the upsampler +export TRELLIS_FAMILY=pixal3d +export PIXAL3D_MODELS=/path/to/TencentARC-Pixal3D # the ckpts/ tree +export TRELLIS_GGUF_OUT=/path/to/models +export NAF_CKPT=/path/to/naf_release.pth # optional; defaults to $PIXAL3D_MODELS/naf/ + +python tools/convert.py # everything +python tools/convert.py naf # just the upsampler ``` +NAF is the one component read through torch rather than safetensors, because it ships as +a `.pth` GitHub release. It converts to ~1.3 MB: only `image_encoder.*` carries weights. + Tensor names are preserved verbatim, so the extra `proj_linear` / `cross_attn_block` tensors need no remapping. From 9b8dc63118ff73a55c2fb589bf41d581bd8eb490 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 17 Aug 2026 20:54:55 +0000 Subject: [PATCH 12/22] dit: name the tensor when a checkpoint's shapes do not fit the graph A GGUF whose layout differs from what build_dit_dense assumes surfaces as a bare GGML_ASSERT(ggml_can_mul_mat) inside ggml, followed by a core dump that names neither the tensor nor the shapes -- the backtrace only reaches trellis::lin, which is called a dozen times per block. Third-party conversions are a normal way to obtain these weights, so check the width in lin() and report the offending tensor with both shapes instead. (cherry picked from commit 84825d2f0cf0469b4304c2bce298e9fb42aa8133) (cherry picked from commit a8b2fe9727fe50a70a47fd4f27ec346e85fbe26a) --- src/dit.cpp | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/dit.cpp b/src/dit.cpp index 02cef0f..9f9a261 100644 --- a/src/dit.cpp +++ b/src/dit.cpp @@ -19,8 +19,22 @@ static bool g_cast_f32 = false; // set per build_dit_dense call static constexpr int64_t kAttnChunkBytes = 1024ll * 1024 * 1024; bool g_no_fa = false; // --no-fa; set by trellis_run +static std::string ne_str(const T* t) { + std::string s = "["; + for (int i = 0; i < 4 && t->ne[i] > 1; ++i) s += (i ? ", " : "") + std::to_string(t->ne[i]); + return s + "]"; +} + static T* lin(ggml_context* c, const Model& m, const std::string& p, T* x) { T* w = m.get(p + ".weight"); + // A GGUF whose layout does not match what this graph assumes reaches ggml as a bare + // GGML_ASSERT(ggml_can_mul_mat) and a core dump, naming neither the tensor nor the shapes. + // Since third-party conversions are a normal way to obtain these weights, say what broke. + if (w->ne[0] != x->ne[0]) + throw std::runtime_error("dit: " + p + ".weight expects an input width of " + + std::to_string(w->ne[0]) + " but the activation is " + + std::to_string(x->ne[0]) + " wide (weight ne=" + ne_str(w) + + ", input ne=" + ne_str(x) + ")"); if (g_cast_f32 && w->type == GGML_TYPE_F16) w = ggml_cast(c, w, GGML_TYPE_F32); T* y = ggml_mul_mat(c, w, x); if (T* b = m.try_get(p + ".bias")) y = ggml_add(c, y, b); From 6f3c3954eab63725277cc53eb1015cbdd7a30e81 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 17 Aug 2026 21:01:50 +0000 Subject: [PATCH 13/22] tools: gguf_probe surfaces foreign-runtime metadata Converters aimed at other runtimes reshape tensors into a quantization-friendly 2-D form and stash the real shape in metadata -- a 1536-element RMS-norm gamma becomes [256, 6] rather than [128, 12], and input_layer becomes [256, 48] rather than [8, 1536]. Element counts match, so nothing looks wrong until ggml asserts on the first matmul. Report the non-standard metadata keys, which name the convention and make the reshaping obvious. Also decode metadata scalars properly so array values print as numbers instead of raw bytes. (cherry picked from commit 0ead0a0cb5f1273c779d2283fe95c33bce5b03cc) (cherry picked from commit da9ea57e0fcfb1cb488fe1c97fd692a40ff045de) --- tools/gguf_probe.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/tools/gguf_probe.py b/tools/gguf_probe.py index 837e52d..5f4566f 100755 --- a/tools/gguf_probe.py +++ b/tools/gguf_probe.py @@ -104,11 +104,13 @@ def string(): # Value readers by gguf_metadata_value_type. Only the scalar payloads need real # decoding; everything else just has to be skipped by the right number of bytes. - fixed = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8} + fixed = {0: "'}") + if len(extra) > 4: + print(f" ... and {len(extra) - 4} more") + by_name = {n: (ne, t) for n, ne, t in tensors} if wanted: From 1a05d8117909c292a40c5cfb7bf8d15cfc739b51 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 17 Aug 2026 21:04:44 +0000 Subject: [PATCH 14/22] model: restore tensor shapes recorded by foreign-runtime converters ComfyUI-style GGUF tooling stores tensors in a quantization-friendly 2-D layout instead of their natural shape -- a 1536-element RMS-norm gamma becomes [256, 6] rather than [128, 12], an [8, 1536] input projection becomes [256, 48] -- and records the real shape under a `.orig_shape.` metadata key. Element order is untouched, so restoring the shape reinterprets the same bytes. Without it nothing looks wrong: names match, element counts match, the buffer loads, and the mismatch only surfaces when a matmul deep in the graph finally compares widths. Since these conversions are a normal way to obtain the weights, read the metadata rather than requiring everyone to reconvert. Quantized tensors are skipped (their row length must stay a multiple of the block size), as is any tensor whose recorded shape does not match its element count -- that is padding, not reshaping, and silently reinterpreting it would be wrong. (cherry picked from commit 8e18e689bc7a11e78632a195de7fa5baec12fca9) (cherry picked from commit 2a246fbfd7b50822830124fee985e51dd2bf1c75) --- docs/pixal3d/README.md | 24 +++++++++++------ src/trellis_model.cpp | 58 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 8 deletions(-) diff --git a/docs/pixal3d/README.md b/docs/pixal3d/README.md index e685357..a0551c0 100644 --- a/docs/pixal3d/README.md +++ b/docs/pixal3d/README.md @@ -203,20 +203,28 @@ Because Pixal3D publishes only the 1024 texture flow: suppresses the dense-decode speckle) is disabled; Pixal3D always textures at the cascade resolution. -Third-party GGUF conversions of Pixal3D exist on the Hub (search `Pixal3D gguf`) but -were produced for the PyTorch pipeline. They load here only if they keep the torch -`state_dict` names — `tools/gguf_probe.py` answers that from ~4 MB of HTTP Range, -without downloading the weights: +Third-party GGUF conversions of Pixal3D exist on the Hub (search `Pixal3D gguf`). Most +were produced for the PyTorch pipeline with ComfyUI-style tooling, which stores tensors +in a quantization-friendly 2-D layout — a 1536-element RMS-norm gamma becomes `[256, 6]` +rather than `[128, 12]` — and records the real shape in a `*.orig_shape.` +metadata key. The loader restores those shapes on load, so such files work; it logs +`restored N reshaped tensor(s)` when it does. Element counts must match exactly, so a +padded (rather than merely reshaped) tensor is left alone and will fail loudly. + +What still has to hold is the naming: the tensors must carry the verbatim torch +`state_dict` names. `tools/gguf_probe.py` answers that from ~4 MB of HTTP Range, without +downloading the weights: ```bash tools/gguf_probe.py https://huggingface.co/USER/REPO/resolve/main/some_flow.gguf ``` It reports the tensor names, the per-block structure (a SLat flow is 30 × 23 + 10 = -700 tensors), `proj_in_channels`, and the dtypes trellis.cpp is picky about — the 1-D -parameters must be F32, because `dit.cpp` adds `modulation` to an f32 timestep -embedding and multiplies `norm2.weight` into an f32 activation. BF16 matmul weights -are fine, but note `--f32` only casts F16 and so will not affect them. +700 tensors), `proj_in_channels`, any foreign-runtime metadata, and the dtypes +trellis.cpp is picky about — the 1-D parameters must be F32, because `dit.cpp` adds +`modulation` to an f32 timestep embedding and multiplies `norm2.weight` into an f32 +activation. BF16 matmul weights are fine, but note `--f32` only casts F16 and so will +not affect them. --- diff --git a/src/trellis_model.cpp b/src/trellis_model.cpp index 4240b46..751e0d6 100644 --- a/src/trellis_model.cpp +++ b/src/trellis_model.cpp @@ -120,6 +120,60 @@ static ggml_backend* make_backend(int gpu) { return cpu_backend(); } +// Conversions aimed at other runtimes store tensors in a quantization-friendly 2-D layout +// rather than their natural shape — a 1536-element RMS-norm gamma becomes [256, 6] instead of +// [128, 12], an [8, 1536] input projection becomes [256, 48] — and record the real shape in a +// metadata key ending `.orig_shape.` (ComfyUI-GGUF and its derivatives do this). +// The element order is untouched, so restoring the shape is a pure reinterpretation of the same +// bytes. Without it the layout looks plausible everywhere and only fails when a matmul finally +// compares widths, deep inside the graph. +// +// Only non-quantized types are restored: a quantized tensor's row length must stay a multiple +// of its block size, which reshaping would break. +static int restore_orig_shapes(gguf_context* gguf, ggml_context* meta) { + static const char* MARK = ".orig_shape."; + int fixed = 0; + for (int64_t k = 0, nk = gguf_get_n_kv(gguf); k < nk; ++k) { + const std::string key = gguf_get_key(gguf, k); + const size_t p = key.find(MARK); + if (p == std::string::npos) continue; + if (gguf_get_kv_type(gguf, k) != GGUF_TYPE_ARRAY) continue; + + ggml_tensor* t = ggml_get_tensor(meta, key.substr(p + strlen(MARK)).c_str()); + if (!t || ggml_blck_size(t->type) != 1) continue; + + const size_t nd = gguf_get_arr_n(gguf, k); + if (nd < 1 || nd > GGML_MAX_DIMS) continue; + const void* raw = gguf_get_arr_data(gguf, k); + const gguf_type at = gguf_get_arr_type(gguf, k); + + // The recorded shape is in torch order (outermost dimension first); ggml's ne is the + // reverse, so read it back to front. + int64_t ne[GGML_MAX_DIMS] = { 1, 1, 1, 1 }; + int64_t elems = 1; + for (size_t d = 0; d < nd; ++d) { + int64_t v; + switch (at) { + case GGUF_TYPE_INT32: v = ((const int32_t*) raw)[d]; break; + case GGUF_TYPE_UINT32: v = ((const uint32_t*) raw)[d]; break; + case GGUF_TYPE_INT64: v = ((const int64_t*) raw)[d]; break; + case GGUF_TYPE_UINT64: v = (int64_t)((const uint64_t*) raw)[d]; break; + default: v = -1; + } + if (v <= 0) { elems = -1; break; } + ne[nd - 1 - d] = v; + elems *= v; + } + if (elems != ggml_nelements(t)) continue; // padded, not merely reshaped — leave it + + for (int d = 0; d < GGML_MAX_DIMS; ++d) t->ne[d] = ne[d]; + t->nb[0] = ggml_type_size(t->type); + for (int d = 1; d < GGML_MAX_DIMS; ++d) t->nb[d] = t->nb[d - 1] * t->ne[d - 1]; + ++fixed; + } + return fixed; +} + Model Model::load(const std::string& path, int gpu) { Model m; @@ -137,6 +191,10 @@ Model Model::load(const std::string& path, int gpu) { if (int64_t k = gguf_find_key(m.gguf, "trellis.config_json"); k >= 0) m.config_json = gguf_get_val_str(m.gguf, k); + if (int fixed = restore_orig_shapes(m.gguf, meta)) + fprintf(stderr, "[trellis] %s: restored %d reshaped tensor(s) from orig_shape metadata\n", + path.c_str(), fixed); + m.backend = make_backend(gpu); m.on_gpu = gpu >= 0; From 6b2ebf2011ad802d730f4bb3fa1f223261f890da Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 17 Aug 2026 22:01:15 +0000 Subject: [PATCH 15/22] pixal3d: report projection coverage and the two proj halves The projection conditioning has no runtime failure mode. A wrong camera, a wrong grid or a broken upsampler all yield finite numbers and a mesh whose silhouette still looks roughly right; the damage appears only as high-frequency noise in the generated surface, which is indistinguishable from a bad seed by eye. Print three figures per stage instead. Coverage catches a collapsed camera (the cube's corners fall outside the frame legitimately, so this is a ratio, not a pass/fail). The per-half mean/std and their cosine similarity catch a broken NAF branch: both halves describe the same points at different detail, so they must share a scale and correlate strongly -- and nothing else in the pipeline would reveal it. (cherry picked from commit 9d4098749a8c45a3069b4bc2e5bcb65738ddb4a3) (cherry picked from commit 4d8de19090bcb9be29675ad72a46570a646542c8) --- src/pixal3d.cpp | 51 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/src/pixal3d.cpp b/src/pixal3d.cpp index f7531d0..a9e93e2 100644 --- a/src/pixal3d.cpp +++ b/src/pixal3d.cpp @@ -4,6 +4,7 @@ #include #include +#include #include namespace trellis { @@ -74,6 +75,55 @@ static void sample_bilinear(const float* map, int Hf, int Wf, int C, int S, dst[j] = w00 * p00[j] + w10 * p10[j] + w01 * p01[j] + w11 * p11[j]; } +// The projection conditioning has no runtime failure mode: a wrong camera, a wrong grid or a +// broken upsampler all produce finite numbers and a plausible-looking mesh silhouette, with the +// damage showing up only as high-frequency noise in the generated surface. These three figures +// separate those causes. +// - coverage: how many grid cells land inside the frame. The cube's corners legitimately fall +// outside (the object is inscribed in it), but a badly wrong camera collapses this. +// - the two halves of a 2048-channel vector: the low-resolution DINOv3 samples and the +// NAF-upsampled ones describe the SAME points at different detail, so they must share a +// scale and correlate strongly. A broken upsampler shows up here and nowhere else. +static void proj_stats(const ProjCond& pc, size_t N, const std::vector& pts, int S) { + if (N == 0) return; + size_t inside = 0; + for (size_t t = 0; t < N; ++t) + if (pts[2*t] >= 0 && pts[2*t] < S && pts[2*t+1] >= 0 && pts[2*t+1] < S) ++inside; + + auto half = [&](int off, double& mean, double& sd) { + double s = 0, s2 = 0; + for (size_t t = 0; t < N; ++t) + for (int c = 0; c < D_DINO; ++c) { + const double v = pc.proj[(size_t)pc.proj_ch * t + off + c]; + s += v; s2 += v * v; + } + const double n = (double)N * D_DINO; + mean = s / n; sd = std::sqrt(std::max(0.0, s2 / n - mean * mean)); + }; + double lm, ls; + half(0, lm, ls); + printf(" [stats] proj: %.0f%% of cells inside frame; lr mean=%.3f std=%.3f", + 100.0 * (double)inside / (double)N, lm, ls); + if (pc.proj_ch == 2 * D_DINO) { + double hm, hs; + half(D_DINO, hm, hs); + double dot = 0, nl = 0, nh = 0; + for (size_t t = 0; t < N; ++t) { + const float* p = &pc.proj[(size_t)pc.proj_ch * t]; + double d = 0, a = 0, b = 0; + for (int c = 0; c < D_DINO; ++c) { + d += (double)p[c] * p[D_DINO + c]; + a += (double)p[c] * p[c]; + b += (double)p[D_DINO + c] * p[D_DINO + c]; + } + dot += d; nl += a; nh += b; + } + const double cs = (nl > 0 && nh > 0) ? dot / std::sqrt(nl * nh) : 0.0; + printf("; hr mean=%.3f std=%.3f; cos(lr,hr)=%.3f", hm, hs, cs); + } + printf("\n"); +} + ProjCond pixal3d_proj_cond(const std::vector& dino, int S, int grid_res, int proj_ch, const CameraParams& cam, const std::vector>* coords, @@ -131,6 +181,7 @@ ProjCond pixal3d_proj_cond(const std::vector& dino, int S, int grid_res, } } + proj_stats(out, N, pts, S); return out; } From a13146c9af442fd31f5eecd34aeddee1a351d2e6 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 17 Aug 2026 22:39:02 +0000 Subject: [PATCH 16/22] pixal3d: use Pixal3D's HR coord quantizer, not TRELLIS.2's The cascade quantizes the upsampled res-512 coords down to the HR token grid. TRELLIS.2 floors u * grid; Pixal3D rounds u * (grid - 1). This port followed the former, because Pixal3D's own sample_shape_slat_cascade still carries TRELLIS.2's version -- but run() never calls that helper and inlines the other formula. The difference is not cosmetic for Pixal3D. There the token index also selects which node of the projection grid the token samples the image at, and that grid is the endpoint-inclusive linspace(-1, 1, grid): rounding to the nearest node is what keeps the pixel-aligned sample registered. Flooring u * grid shifts by half a cell and dilates by grid/(grid-1), and because the index is an integer the error is a step function, not a smooth warp -- at grid 64 it moves a quarter of the coordinates a full cell along each axis, in stripes of period 8, so 58% of tokens end up misregistered on at least one axis. A cell is 16 res-1024 voxels, one ViT-L/16 patch at the 1024 conditioning size. That is exactly the observed failure: the silhouette survives, since it rides on the five global cross-attention tokens, while the high-frequency detail does not and the decode comes out speckled -- face/vertex ratio 1.24 against 2.00 for the TRELLIS.2 baseline on the same image. The sparse-structure and LR shape stages were unaffected because neither requantizes: SS is a dense self-consistent grid, and the LR coords come straight from the SS decode at the grid resolution the stage already expects. Kept family-conditional -- flooring remains correct for TRELLIS.2, whose coords only feed RoPE, where a half-cell offset is a smooth reparameterization. (cherry picked from commit 67b56acf8d33c78a8e810cbcc541a45649f0040f) (cherry picked from commit 9aef5d4e6cfa73a6684cd032dafa65b4317e1e17) --- src/trellis_cli.cpp | 29 +++++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/src/trellis_cli.cpp b/src/trellis_cli.cpp index 16bb7e9..7c55f41 100644 --- a/src/trellis_cli.cpp +++ b/src/trellis_cli.cpp @@ -304,16 +304,33 @@ int trellis_run(const trellis::TrellisParams& cfg) { vector> hr_coords; { trellis::Model m = trellis::Model::load(M + "/shape_dec.gguf", gpu); hr_coords = trellis::shape_upsample(m, lr_dn, coords); m.free(); } - // (3) quantize res512 -> res(hr_res//16) with the reference's adaptive token-budget backoff - // (sample_shape_slat_cascade): start at hr_target, step -128 toward the 1024 floor while - // the unique token count would exceed max_num_tokens. grid = hr_res//16 is integral since - // 128/16 = 8 (1536->96, 1408->88, ..., 1024->64). + // (3) quantize res512 -> res(hr_res//16) with the reference's adaptive token-budget backoff: + // start at hr_target, step -128 toward the 1024 floor while the unique token count + // would exceed max_num_tokens. grid = hr_res//16 is integral since 128/16 = 8 + // (1536->96, 1408->88, ..., 1024->64). + // + // The two families quantize DIFFERENTLY, and it matters far more than it looks. + // TRELLIS.2 floors `u * grid` — a cell quantizer, and the coords only feed RoPE, where + // a half-cell offset is a smooth reparameterization. Pixal3D rounds `u * (grid - 1)`, + // because for it the token index ALSO selects which projection-grid node the token + // samples the image at, and that grid is the endpoint-inclusive linspace(-1, 1, grid): + // rounding to the nearest node is what keeps the pixel-aligned sample registered. + // (Pixal3D's own sample_shape_slat_cascade still carries the TRELLIS.2 form, but run() + // never calls it — that dead helper is what this port originally followed.) + // + // Using the wrong one is not a smooth warp but a step function: at grid 64 it moves + // 25% of coordinates a full cell along each axis, in stripes of period 8, so 58% of + // tokens are misregistered on at least one axis. A cell is 16 res-1024 voxels — one + // ViT-L/16 patch at S=1024. The silhouette survives on the global tokens; the fine + // detail does not, and the decode comes out speckled. int hr_res = hr_target; for (;;) { const int gi = hr_res / 16; // integral grid (ref's hr_resolution//16) - const float g = (float)gi; + const float g = pix ? (float)(gi - 1) : (float)gi; + auto qz = [&](float c) { return pix ? (int)lrintf((c + 0.5f) / 512.f * g) + : (int)((c + 0.5f) / 512.f * g); }; std::set> q; - for (auto& c : hr_coords) q.insert({ (int)((c[0]+0.5f)/512.f*g), (int)((c[1]+0.5f)/512.f*g), (int)((c[2]+0.5f)/512.f*g) }); + for (auto& c : hr_coords) q.insert({ qz((float)c[0]), qz((float)c[1]), qz((float)c[2]) }); if ((int)q.size() < max_tok || hr_res <= 1024) { shc.assign(q.begin(), q.end()); printf(" upsampled coords @res512=%d -> quantized @res%d (grid %d) = %d tokens\n", From 4d6676c6a2ff6f421f0290f0d2e6532e554a5d36 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 17 Aug 2026 23:07:25 +0000 Subject: [PATCH 17/22] pixal3d: check the 1024 conditioning against the 512 one The HR shape stage is the only one conditioned on DINOv3 at 1024, and it is the only stage producing a speckled decode -- the sparse-structure and LR shape stages, which read the 512 map through the same code, are clean (face/vertex 2.04 against the TRELLIS.2 baseline's 2.01). Both maps encode the same image, so the features sampled at a given 3-D point must agree across them. Report that agreement. Nothing else in the pipeline can distinguish "the 1024 map is being read with the wrong spatial layout" from "the latent is bad": the projection statistics stay well-scaled, the two proj halves stay correlated, and every value stays finite either way. (cherry picked from commit fc3b8a3e14531fd38855c4d162bf8fa9e7995227) (cherry picked from commit 8d52f0349319d7f5f12991ebfd7c857b3fc1df4b) --- include/pixal3d.h | 10 ++++++++++ src/pixal3d.cpp | 26 ++++++++++++++++++++++++++ src/trellis_cli.cpp | 8 ++++++++ 3 files changed, 44 insertions(+) diff --git a/include/pixal3d.h b/include/pixal3d.h index a3b5242..f574fc5 100644 --- a/include/pixal3d.h +++ b/include/pixal3d.h @@ -57,6 +57,16 @@ struct ProjCond { // is a degraded input, not an equivalent one. // img01 : raw [0,1] guide image, [3,S,S] torch CHW — required when `naf` is given. // naf_out : NAF target resolution for this stage. +// Sample the same 3-D points from two DINOv3 feature maps produced at different input sizes and +// report the mean cosine similarity. Both maps encode the same image, so a point's features must +// agree strongly across them; the value only collapses if one map is being read with the wrong +// spatial layout. This isolates "the conditioning is wrong at 1024 but right at 512", which no +// other statistic in the pipeline can distinguish from a bad latent. +double pixal3d_cross_res_agreement(const std::vector& dino_a, int Sa, + const std::vector& dino_b, int Sb, + int grid_res, const CameraParams& cam, + const std::vector>& coords); + ProjCond pixal3d_proj_cond(const std::vector& dino, int S, int grid_res, int proj_ch, const CameraParams& cam, const std::vector>* coords, diff --git a/src/pixal3d.cpp b/src/pixal3d.cpp index a9e93e2..fde79d4 100644 --- a/src/pixal3d.cpp +++ b/src/pixal3d.cpp @@ -75,6 +75,32 @@ static void sample_bilinear(const float* map, int Hf, int Wf, int C, int S, dst[j] = w00 * p00[j] + w10 * p10[j] + w01 * p01[j] + w11 * p11[j]; } +double pixal3d_cross_res_agreement(const std::vector& dino_a, int Sa, + const std::vector& dino_b, int Sb, + int grid_res, const CameraParams& cam, + const std::vector>& coords) { + if (coords.empty()) return 0.0; + const int Ha = Sa / PATCH, Hb = Sb / PATCH; + const float* pa = dino_a.data() + (size_t)D_DINO * N_GLOBAL; + const float* pb = dino_b.data() + (size_t)D_DINO * N_GLOBAL; + std::vector va(D_DINO), vb(D_DINO); + double acc = 0; + // Every 37th point: the figure is a mean over thousands of samples either way, and the full + // sweep would cost more than the flow step it is diagnosing. + size_t n = 0; + for (size_t t = 0; t < coords.size(); t += 37, ++n) { + float xa, ya, xb, yb; + pixal3d_project_cell(grid_res, coords[t][0], coords[t][1], coords[t][2], cam, Sa, xa, ya); + pixal3d_project_cell(grid_res, coords[t][0], coords[t][1], coords[t][2], cam, Sb, xb, yb); + sample_bilinear(pa, Ha, Ha, D_DINO, Sa, xa, ya, va.data()); + sample_bilinear(pb, Hb, Hb, D_DINO, Sb, xb, yb, vb.data()); + double d = 0, na = 0, nb = 0; + for (int c = 0; c < D_DINO; ++c) { d += (double)va[c]*vb[c]; na += (double)va[c]*va[c]; nb += (double)vb[c]*vb[c]; } + if (na > 0 && nb > 0) acc += d / std::sqrt(na * nb); + } + return n ? acc / (double)n : 0.0; +} + // The projection conditioning has no runtime failure mode: a wrong camera, a wrong grid or a // broken upsampler all produce finite numbers and a plausible-looking mesh silhouette, with the // damage showing up only as high-frequency noise in the generated surface. These three figures diff --git a/src/trellis_cli.cpp b/src/trellis_cli.cpp index 7c55f41..041245d 100644 --- a/src/trellis_cli.cpp +++ b/src/trellis_cli.cpp @@ -344,6 +344,14 @@ int trellis_run(const trellis::TrellisParams& cfg) { // (4) HR shape flow @res(hr_res//16) with cond_1024. The projection grid follows the same // backoff as the token grid — Pixal3D overrides its cond model's grid_resolution to // hr_res//16 for exactly this reason — while the NAF target stays at the stage's 512. + // The HR stage is the only one conditioned on DINOv3 at 1024. If its feature map were read + // with a different spatial layout than the 512 one — which the LR stage uses successfully — + // nothing else in the pipeline would show it: the features would still be well-scaled, + // well-correlated across the proj halves, and finite. + if (pix) + printf(" [stats] proj: DINOv3@1024 vs @512 agreement on the same points = %.3f\n", + trellis::pixal3d_cross_res_agreement(dino1024, 1024, dino, 512, + hr_res / 16, cam, shc)); slat_norm = shape_flow(M + FP + "shape_flow_1024.gguf", shc, cond1024.data(), neg1024.data(), Lc1024, /*grid_res=*/hr_res / 16, /*S=*/1024, /*naf_out=*/512); RES = hr_res; cond_dec = cond1024.data(); neg_dec = neg1024.data(); Lc_dec = Lc1024; From b2ae9faa4146aec3d900cbdfa1987af33405ee6e Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 17 Aug 2026 23:23:29 +0000 Subject: [PATCH 18/22] dit: skip FlashAttention when the key sequence is shorter than one tile FlashAttention is here to avoid materialising the [Lk, Lq, nh] score matrix, which reaches terabytes at the HR flow. When Lk is shorter than a single FA key tile that matrix is a few megabytes and FA buys nothing, while still paying the whole cost of its padding machinery. Pixal3D's proj mode cross-attends over 5 global tokens, so the key dimension is zero-padded from 5 to 256 -- 98% padding. TRELLIS.2 never gets near that regime (1029 keys at 512, 4101 at 1024), and the comments in this function record that ggml's CUDA FlashAttention has repeatedly mishandled padded key tiles in ways that depend on the token count. Route short-KV attention through the exact chunked path, where at Lk = 5 it is cheaper than FA regardless. (cherry picked from commit b35d91d92908fdffe0af72f8c9f1660b8c6a8c19) (cherry picked from commit 3ac026496914b43bc6940339487a22619c1325b3) --- src/dit.cpp | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/dit.cpp b/src/dit.cpp index 9f9a261..e6264e2 100644 --- a/src/dit.cpp +++ b/src/dit.cpp @@ -131,7 +131,14 @@ static T* sdpa(ggml_context* c, T* q, T* k, T* v, int d_model, T* mask = nullptr // reads only -0.0019 vs -0.00014 (its oracle is -0.00010). MMA already accumulated KQ in // FP32, so only the VKQ sum stagnated there -- same bug, ~9x milder. // --no-fa falls back to the exact chunked path (correct on any backend, ~2.7x slower). - const bool no_fa = g_no_fa; + // FlashAttention exists here to avoid materialising the [Lk, Lq, nh] score matrix, which is + // terabytes at the HR flow. When Lk is shorter than one FA key tile that matrix is trivially + // small and FA buys nothing — while costing a great deal of risk. Pixal3D's proj mode + // cross-attends over 5 global tokens, so the key dim gets zero-padded 5 -> 256: 98% padding, + // a regime TRELLIS.2 (1029 or 4101 keys) never reaches, and precisely the shape whose mask + // handling the comments below record as fragile and token-count dependent. Take the exact + // path instead; at Lk = 5 it is cheaper than FA anyway. + const bool no_fa = g_no_fa || k->ne[2] < 256; if (!no_fa) { // TRELLIS_FA_FAST=1: F16 K/V + default (F16) accumulation — the shapes // the Vulkan coopmat FA shaders are specialized for. A/B only: F16 K/V From a222b7419403a143cd18dc793c3b5c26d2f115d5 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 18 Aug 2026 20:35:31 +0000 Subject: [PATCH 19/22] pixal3d: expose --extend-pixel for subjects cropped by the frame The camera solve registers the projection grid so its corner lands on the image border, which silently assumes the subject is entirely inside the frame. When the input is cropped -- an image generator asked for too close a shot, a photo framed tight -- that assumption is wrong and the object gets squeezed into the grid: the visible part comes out compressed and the occluded part is invented at the wrong scale. The reference carries an extend_pixel parameter for exactly this and never exposes it, since MoGe-2 estimates from the whole frame. pixal3d_camera already took it; this only plumbs it through to a flag. At 49.13 degrees and a 512 frame: 0 px gives distance 1.094, 64 px gives 0.875, 128 px gives 0.729. (cherry picked from commit 539c7047a88a0d81191d69ac91bc99437baa9731) (cherry picked from commit 5dc4f6e33ba413cfa7fcb1472a78cee4ab08ecee) --- docs/pixal3d/README.md | 13 +++++++++++-- include/trellis_args.h | 6 ++++++ src/trellis_args.cpp | 5 +++++ src/trellis_cli.cpp | 6 +++--- 4 files changed, 25 insertions(+), 5 deletions(-) diff --git a/docs/pixal3d/README.md b/docs/pixal3d/README.md index a0551c0..aafa59b 100644 --- a/docs/pixal3d/README.md +++ b/docs/pixal3d/README.md @@ -74,10 +74,19 @@ MoGe-2 is **not ported**. The closed-form distance is, so the FOV is the only fr parameter and it is a flag: ``` ---fov DEG horizontal field of view (default 49.13°, Pixal3D's own default) ---mesh-scale F object scale inside the unit grid (default 1.0) +--fov DEG horizontal field of view (default 49.13°, Pixal3D's own default) +--mesh-scale F object scale inside the unit grid (default 1.0) +--extend-pixel N how far the subject continues past the image border (default 0) ``` +`--extend-pixel` matters more than it sounds. The solve registers the grid corner onto +the image border, which silently assumes the subject is *fully inside the frame*. Feed it +an image cropped at the edges — a generator asked for a close-up, say — and the object is +squeezed into the grid: the visible part comes out compressed and the unseen part is +invented at the wrong scale. Raising it moves the virtual border outward, and the camera +steps back to match (49.13° at 512: 0 px → distance 1.094, 64 px → 0.875, 128 px → 0.729). +Upstream has the same parameter but never exposes it, since MoGe-2 sees the whole frame. + The projection is resolution-independent once normalized, so one camera solved at 512 serves the 1024 stages too. `trellis-test-pixal3d` pins both functions to golden values taken from the reference implementation: diff --git a/include/trellis_args.h b/include/trellis_args.h index 79a327c..cd720c3 100644 --- a/include/trellis_args.h +++ b/include/trellis_args.h @@ -41,6 +41,12 @@ struct TrellisParams { // MoGe-2 is not ported, so the FOV is a flag; 0 keeps Pixal3D's own default (49.13 deg). float fov_deg = 0.0f; float mesh_scale = 1.0f; + // The camera solve assumes the object is fully inside the frame: it registers the projection + // grid so that the grid corner lands on the image border. When the subject is cropped by the + // frame that assumption fails, and the object gets compressed into the grid. --extend-pixel + // moves the virtual border outward by N pixels, telling the solve the object continues past + // the edge. It is the reference's own `extend_pixel`, which upstream never exposes. + int extend_pixel = 0; // NAF guided upsampling of the DINOv3 feature map (the shape/texture stages' second proj // branch). Off falls back to sampling the bare feature map twice, which halves the effective // proj input — accepted only as a way to run without naf.gguf. diff --git a/src/trellis_args.cpp b/src/trellis_args.cpp index 508e848..06d9fe4 100644 --- a/src/trellis_args.cpp +++ b/src/trellis_args.cpp @@ -37,6 +37,10 @@ void print_usage(const char* argv0, bool server) { " Upstream estimates this with MoGe-2; that model is not ported,\n" " so a wrong FOV shows up as geometry drifting off the silhouette.\n" " --mesh-scale F pixal3d: object scale inside the unit grid (default 1.0)\n" + " --extend-pixel N pixal3d: how far the subject continues past the image border,\n" + " in pixels of a 512 frame (default 0 = fully inside). Raise it\n" + " when the input is cropped: the camera solve otherwise assumes\n" + " the whole object fits, and squeezes it into the grid.\n" " --no-naf pixal3d: skip NAF guided upsampling (needs no naf.gguf, but\n" " the shape/texture stages then lose their high-frequency\n" " projection branch)\n" @@ -100,6 +104,7 @@ bool parse_args(int argc, char** argv, TrellisParams& p) { else { fprintf(stderr, "[trellis] unknown model family: %s (trellis|pixal3d)\n", v); return false; } } else if (a == "--fov") { const char* v = need(a.c_str()); if (!v) return false; p.fov_deg = (float)atof(v); } else if (a == "--mesh-scale") { const char* v = need(a.c_str()); if (!v) return false; p.mesh_scale = (float)atof(v); } + else if (a == "--extend-pixel") { const char* v = need(a.c_str()); if (!v) return false; p.extend_pixel = atoi(v); } else if (a == "--no-naf") { p.naf = false; } else if (a == "--gpu") { const char* v = need(a.c_str()); if (!v) return false; p.gpu = atoi(v); } else if (a == "-s" || a == "--seed") { const char* v = need(a.c_str()); if (!v) return false; p.seed = (uint32_t)atoi(v); } diff --git a/src/trellis_cli.cpp b/src/trellis_cli.cpp index 041245d..483daeb 100644 --- a/src/trellis_cli.cpp +++ b/src/trellis_cli.cpp @@ -89,9 +89,9 @@ int trellis_run(const trellis::TrellisParams& cfg) { : PIXAL3D_DEFAULT_FOV; // The distance is derived at 512 on purpose: the projection is resolution-independent // once normalized, so one camera serves both the 512 and the 1024 stages. - cam = trellis::pixal3d_camera(fov, cfg.mesh_scale, 512, 0); - printf("[trellis] model family: pixal3d (fov %.2f deg, distance %.4f, mesh scale %.2f)\n", - fov * 180.0f / 3.14159265358979f, cam.distance, cam.mesh_scale); + cam = trellis::pixal3d_camera(fov, cfg.mesh_scale, 512, cfg.extend_pixel); + printf("[trellis] model family: pixal3d (fov %.2f deg, distance %.4f, mesh scale %.2f, extend %d px)\n", + fov * 180.0f / 3.14159265358979f, cam.distance, cam.mesh_scale, cfg.extend_pixel); if (cfg.fov_deg <= 0.0f) printf(" (using Pixal3D's default FOV — MoGe-2 estimation is not ported; pass" " --fov if the object's perspective is noticeably wider or flatter)\n"); From 78a2a360c976132a36351d0532035709a5255b26 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 18 Aug 2026 20:56:51 +0000 Subject: [PATCH 20/22] docs: correct what --extend-pixel is for The previous text claimed it rescues an input whose subject is cropped by the frame. It does not, and recommending it that way produces a worse result than leaving it alone. Background removal already reframes: it crops a square around the visible alpha bbox with a 10% margin, so the cutout reaching the camera always has the subject inscribed regardless of the original framing. Pushing the border past that stretches the grid over territory with no pixels behind it -- the projection lands on the border clamp, those cells go unconditioned, and the model fills them with whatever it likes. Observed on a cropped photograph as a tail growing out of the object. No parameter recovers information the file does not contain. Say so, and point at re-framing the source instead. (cherry picked from commit 846cb07b76761e7c87ac042b81a91d3794936908) (cherry picked from commit 7246f6302975afd0ff4cc0cd19c244fc6306f834) --- docs/pixal3d/README.md | 22 +++++++++++++++------- include/trellis_args.h | 10 +++++----- src/trellis_args.cpp | 8 ++++---- 3 files changed, 24 insertions(+), 16 deletions(-) diff --git a/docs/pixal3d/README.md b/docs/pixal3d/README.md index aafa59b..9a0bdeb 100644 --- a/docs/pixal3d/README.md +++ b/docs/pixal3d/README.md @@ -79,13 +79,21 @@ parameter and it is a flag: --extend-pixel N how far the subject continues past the image border (default 0) ``` -`--extend-pixel` matters more than it sounds. The solve registers the grid corner onto -the image border, which silently assumes the subject is *fully inside the frame*. Feed it -an image cropped at the edges — a generator asked for a close-up, say — and the object is -squeezed into the grid: the visible part comes out compressed and the unseen part is -invented at the wrong scale. Raising it moves the virtual border outward, and the camera -steps back to match (49.13° at 512: 0 px → distance 1.094, 64 px → 0.875, 128 px → 0.729). -Upstream has the same parameter but never exposes it, since MoGe-2 sees the whole frame. +`--extend-pixel` moves the virtual border outward, so the camera steps back to match +(49.13° at 512: 0 px → distance 1.094, 64 px → 0.875, 128 px → 0.729). Upstream carries the +same parameter and never exposes it. + +**Leave it at 0 unless you know why you are raising it.** In particular it does *not* +rescue an input whose subject is cropped by the frame: background removal already reframes, +cropping a square around the visible alpha bbox with a 10% margin, so the cutout that +reaches the camera always has the subject inscribed. Extending past that stretches the grid +over territory with no pixels behind it — the projection falls onto the border clamp, those +cells go unconditioned, and the model fills them with whatever it likes. On a cropped +photograph the observed result is a tail growing out of the object. + +There is no parameter for a cropped input, because the information is not in the file. +Re-frame the source image with air around the subject instead; this model is pixel-aligned +by construction, so framing buys more than any camera flag. The projection is resolution-independent once normalized, so one camera solved at 512 serves the 1024 stages too. `trellis-test-pixal3d` pins both functions to golden values diff --git a/include/trellis_args.h b/include/trellis_args.h index cd720c3..0d45331 100644 --- a/include/trellis_args.h +++ b/include/trellis_args.h @@ -41,11 +41,11 @@ struct TrellisParams { // MoGe-2 is not ported, so the FOV is a flag; 0 keeps Pixal3D's own default (49.13 deg). float fov_deg = 0.0f; float mesh_scale = 1.0f; - // The camera solve assumes the object is fully inside the frame: it registers the projection - // grid so that the grid corner lands on the image border. When the subject is cropped by the - // frame that assumption fails, and the object gets compressed into the grid. --extend-pixel - // moves the virtual border outward by N pixels, telling the solve the object continues past - // the edge. It is the reference's own `extend_pixel`, which upstream never exposes. + // Pushes the camera's virtual image border outward, so the projection grid spans more than + // the frame. The reference's own `extend_pixel`, which upstream never exposes. Rarely useful + // here: background removal already reframes around the subject's alpha bbox, and grid cells + // pushed past the cutout sample the border clamp — unconditioned, and the model fills them + // arbitrarily. int extend_pixel = 0; // NAF guided upsampling of the DINOv3 feature map (the shape/texture stages' second proj // branch). Off falls back to sampling the bare feature map twice, which halves the effective diff --git a/src/trellis_args.cpp b/src/trellis_args.cpp index 06d9fe4..d2f5530 100644 --- a/src/trellis_args.cpp +++ b/src/trellis_args.cpp @@ -37,10 +37,10 @@ void print_usage(const char* argv0, bool server) { " Upstream estimates this with MoGe-2; that model is not ported,\n" " so a wrong FOV shows up as geometry drifting off the silhouette.\n" " --mesh-scale F pixal3d: object scale inside the unit grid (default 1.0)\n" - " --extend-pixel N pixal3d: how far the subject continues past the image border,\n" - " in pixels of a 512 frame (default 0 = fully inside). Raise it\n" - " when the input is cropped: the camera solve otherwise assumes\n" - " the whole object fits, and squeezes it into the grid.\n" + " --extend-pixel N pixal3d: push the camera's virtual image border outward by N\n" + " pixels of a 512 frame (default 0). Rarely useful: background\n" + " removal already reframes around the subject, and extending\n" + " past the cutout projects onto pixels that do not exist.\n" " --no-naf pixal3d: skip NAF guided upsampling (needs no naf.gguf, but\n" " the shape/texture stages then lose their high-frequency\n" " projection branch)\n" From 5dfb0749a3ddf0d4a75e71a1333dc4d7ee3b1ab3 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 18 Aug 2026 21:21:50 +0000 Subject: [PATCH 21/22] server: accept the model family and the camera per request --model already worked as a launch flag, since the server runs the same parser for its defaults and every request copies them. But one model directory now holds both families -- that is what the pixal3d_ prefix bought -- so a single resident server can serve either, and only the request field was missing. The camera fields matter more. fov, mesh_scale and extend_pixel describe the image rather than the run, so pinning them at launch is close to useless on a server taking arbitrary uploads: every image wants its own FOV. (cherry picked from commit 41054689e039806de3d87260bd3e1f08ce4af2f6) (cherry picked from commit 1725a23751d3a4b941ea31f2b40fa689756c90b0) --- docs/pixal3d/README.md | 12 +++++++++++- src/trellis-server.cpp | 19 +++++++++++++++++-- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/docs/pixal3d/README.md b/docs/pixal3d/README.md index 9a0bdeb..52aa7f5 100644 --- a/docs/pixal3d/README.md +++ b/docs/pixal3d/README.md @@ -258,7 +258,17 @@ trellis-cli in.png out.glb --model pixal3d --fov 38 --bg-removal birefnet trellis-cli in.png out.glb --model pixal3d --no-naf ``` -`trellis-server` takes the same flags at launch. +`trellis-server` takes the same flags at launch, and since one model directory holds both +families a single resident server can serve either — `POST /generate` accepts a `model` +field (`trellis` | `pixal3d`) alongside the existing ones. The camera is per-request too, +which matters more than it sounds: `fov`, `mesh_scale` and `extend_pixel` describe the +*image*, not the run, so a launch-time default is close to useless on a server taking +arbitrary uploads. + +```bash +curl -F image=@in.png -F model=pixal3d -F fov=38 -F resolution=1024 \ + http://127.0.0.1:8080/generate -o out.glb +``` Every other flag — `--res`, `--max-tokens`, `--gss`/`--gsh`, `--band`, `--decim`, `--atlas`, `--box-uv`, `--tex-res`, `--seed` — behaves identically, because everything diff --git a/src/trellis-server.cpp b/src/trellis-server.cpp index d03afb8..71d2058 100644 --- a/src/trellis-server.cpp +++ b/src/trellis-server.cpp @@ -5,8 +5,10 @@ // fields "seed", "resolution" (512/1024/1536), "bg_removal" // (threshold|birefnet), "uv" (xatlas = default, unique // chart space; box = faster projection), "band" (narrow-band -// DC remesh band width, default 1 — see --band). Returns -// model/gltf-binary. +// DC remesh band width, default 1 — see --band), "model" +// (trellis|pixal3d — one directory holds both), and for +// pixal3d the camera: "fov" (degrees), "mesh_scale", +// "extend_pixel". Returns model/gltf-binary. // // Launch-time defaults come from CLI flags (see trellis::parse_args); // each request copies those defaults and applies its own overrides. The model @@ -99,6 +101,19 @@ int main(int argc, char** argv) { if (req.has_file("bg_removal")) p.birefnet = (req.get_file_value("bg_removal").content == "birefnet") ? 1 : 0; if (req.has_file("uv")) p.xatlas = (req.get_file_value("uv").content == "xatlas"); if (req.has_file("band")) p.band = atoi(req.get_file_value("band").content.c_str()); + // Both families share one model directory (the Pixal3D flows carry a pixal3d_ prefix), + // so a single resident server can serve either -- the client just has to say which. + if (req.has_file("model")) { + const std::string& m = req.get_file_value("model").content; + if (m == "pixal3d") p.family = trellis::ModelFamily::Pixal3D; + else if (m == "trellis") p.family = trellis::ModelFamily::Trellis; + else { res.status = 400; res.set_content("model must be trellis or pixal3d\n", "text/plain"); return; } + } + // Camera, for pixal3d. Unlike every other knob these describe the IMAGE rather than the + // run, so a launch-time default is close to useless on a server taking arbitrary uploads. + if (req.has_file("fov")) p.fov_deg = (float) atof(req.get_file_value("fov").content.c_str()); + if (req.has_file("mesh_scale")) p.mesh_scale = (float) atof(req.get_file_value("mesh_scale").content.c_str()); + if (req.has_file("extend_pixel")) p.extend_pixel = atoi(req.get_file_value("extend_pixel").content.c_str()); if (req.has_file("webp")) { const std::string& w = req.get_file_value("webp").content; p.webp = (w == "off" || w == "0" || w == "false") ? 0 From bb4a367341fd27d4b630e3dc6a08673bc3fefe9f Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 20 Aug 2026 12:10:14 +0000 Subject: [PATCH 22/22] docs: record that proj mode depends on the FlashAttention V-range fix Proj mode adds proj_linear(proj) into every block's residual, so its activations run larger than TRELLIS.2's. CUDA's tensor-core FA kernels convert BF16 K/V to F16 internally, and those larger values overflow: before the V pre-scaling in sdpa(), the res-1024 cascade decoded 1.19M voxels at a face/vertex ratio of 1.25 while TRELLIS.2 was unaffected at the same token count. With it, and an identical seed, FA gives 3.47M at 2.31 against 3.22M at 2.11 on the exact path, with the sparse-structure stage agreeing to 0.4%. Worth writing down because the failure is silent, resolution dependent, and looks like a bad seed rather than a numerical fault. --- docs/pixal3d/README.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/pixal3d/README.md b/docs/pixal3d/README.md index 52aa7f5..5bbe9f4 100644 --- a/docs/pixal3d/README.md +++ b/docs/pixal3d/README.md @@ -297,6 +297,14 @@ Mismatched weights fail fast rather than producing garbage: ## Known gaps +- **Needs the FlashAttention V-range fix.** Proj mode adds `proj_linear(proj)` into every + block's residual, so its activations run larger than TRELLIS.2's. CUDA's tensor-core FA + kernels convert BF16 K/V to F16 internally, and those larger values overflow: before the + V pre-scaling in `sdpa()`, the res-1024 cascade decoded 1.19M voxels at a face/vertex + ratio of 1.25, a surface full of holes, while TRELLIS.2 was unaffected at the same token + count. With it, the same run gives 3.47M at 2.31 against 3.22M at 2.11 on the exact path + (identical seed), and the sparse-structure stage agrees to 0.4%. Build against an older + ggml and the holes come back; `--no-fa` is the escape hatch if they ever do. - **MoGe-2 FOV estimation is not ported.** Use `--fov`. - **No end-to-end numerical parity run on real weights.** The projection is pinned to golden values from the reference and NAF was diffed against a transcription of it