Initial support for Pixal3D

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}")

README.md

@@ -118,8 +118,26 @@ 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`.
+
+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 — 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.
+
 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

include/dit.h

@@ -22,18 +22,25 @@ 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`
 // and must be flagged ggml_set_input by the caller; weights come from `m`.
 //   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<std::string, ggml_tensor*>* inter = nullptr);
 
 } // namespace trellis

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<float>& rope_cos, const std::vector<float>& rope_sin);
     ~DitRunner();
-    // xt: [in_ch*N] channel-major. cond: [d_cond*n_cond]. Returns velocity [out_ch*N].
-    std::vector<float> forward(const std::vector<float>& xt, float t_scaled, const float* cond);
+    // xt: [in_ch*N] channel-major. Returns velocity [out_ch*N].
+    std::vector<float> forward(const std::vector<float>& 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<float> rcos_, rsin_;   // re-uploaded each forward (gallocr may reuse input buffers)
     std::map<std::string, ggml_tensor*> 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<std::array<int,3>>& coords, int n_cond);
 
 // FlowEuler guidance-interval sampler over an arbitrary forward functor.
-using FlowFwd = std::function<std::vector<float>(const std::vector<float>&, float, const float*)>;
+using FlowFwd = std::function<std::vector<float>(const std::vector<float>&, float, const FlowCond&)>;
 std::vector<float> sample_flow(const FlowFwd& fwd, std::vector<float> sample,
-                               const float* cond, const float* neg_cond,
+                               const FlowCond& cond, const FlowCond& neg_cond,
                                const SamplerParams& sp,
                                std::vector<std::vector<float>>* trace = nullptr);
 

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 <vector>
+
+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<float> naf_sample(const Model& m,
+                              const std::vector<float>& img01, int S,
+                              const float* feats_lr, int Hf, int Wf, int C,
+                              int out, const std::vector<float>& pts_xy);
+
+} // namespace trellis

include/pixal3d.h

@@ -0,0 +1,75 @@
+// 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 <array>
+#include <vector>
+
+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<float> global;      // [1024 * 5] cls + 4 register tokens, channel-major
+    std::vector<float> 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.
+// 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<float>& dino_a, int Sa,
+                                   const std::vector<float>& dino_b, int Sb,
+                                   int grid_res, const CameraParams& cam,
+                                   const std::vector<std::array<int, 3>>& coords);
+
+ProjCond pixal3d_proj_cond(const std::vector<float>& dino, int S, int grid_res, int proj_ch,
+                           const CameraParams& cam,
+                           const std::vector<std::array<int, 3>>* coords,
+                           const Model* naf, const std::vector<float>* img01, int naf_out);
+
+} // namespace trellis

include/preprocess.h

@@ -23,4 +23,9 @@ bool image_has_alpha(const std::string& path);
 std::vector<unsigned char> 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<float> normalize_cutout(const std::vector<unsigned char>& 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<float> cutout_to_chw01(const std::vector<unsigned char>& rgb, int sz, int S);
 } // namespace trellis

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,23 @@ 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;
+    // 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
+    // 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)