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/README.md b/README.md index 6336599..d1c511e 100644 --- a/README.md +++ b/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 diff --git a/docs/pixal3d/README.md b/docs/pixal3d/README.md new file mode 100644 index 0000000..5bbe9f4 --- /dev/null +++ b/docs/pixal3d/README.md @@ -0,0 +1,317 @@ +# 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 | `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 +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) +--extend-pixel N how far the subject continues past the image border (default 0) +``` + +`--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 +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 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 + 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 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). + +--- + +## 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 +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** | + +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 +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. + +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`). 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`, 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. + +--- + +## Usage + +```bash +# 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 + +# no naf.gguf available +trellis-cli in.png out.glb --model pixal3d --no-naf +``` + +`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 +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 + +- **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 + 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/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/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..f574fc5 --- /dev/null +++ b/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 +#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. +// 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, + 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/include/trellis_args.h b/include/trellis_args.h index ffbc9c6..0d45331 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,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) diff --git a/src/dit.cpp b/src/dit.cpp index bc0c413..e6264e2 100644 --- a/src/dit.cpp +++ b/src/dit.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include namespace trellis { @@ -18,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); @@ -116,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 @@ -228,7 +250,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 +268,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 +291,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 +312,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/naf.cpp b/src/naf.cpp new file mode 100644 index 0000000..26ec610 --- /dev/null +++ b/src/naf.cpp @@ -0,0 +1,283 @@ +#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; + // 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 }); + 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..fde79d4 --- /dev/null +++ b/src/pixal3d.cpp @@ -0,0 +1,214 @@ +#include "pixal3d.h" +#include "naf.h" +#include "trellis_model.h" + +#include +#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]; +} + +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 +// 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, + 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); + } + } + + proj_stats(out, N, pts, S); + 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; +} 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); 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 diff --git a/src/trellis_args.cpp b/src/trellis_args.cpp index 60fa075..d2f5530 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,22 @@ 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" + " --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" " --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 +98,14 @@ 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 == "--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); } 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..483daeb 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,31 @@ 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; + // 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) + 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, 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"); + } + 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 +149,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 +179,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 + FP + "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::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, + // 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 +240,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 + 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; } + } // 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 +294,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 + 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) lr_dn[(size_t)c + 32*n] = lr_norm[(size_t)c + 32*n]*SHAPE_STD[c] + SHAPE_MEAN[c]; @@ -203,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", @@ -223,13 +341,25 @@ 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. + // 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; } 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(); slat_dn.resize(slat_norm.size()); @@ -282,7 +412,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 @@ -293,7 +425,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; @@ -301,14 +433,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 +459,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]; } 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; diff --git a/tools/convert.py b/tools/convert.py index b1e364b..5d986a1 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -1,23 +1,55 @@ #!/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 + +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 + 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 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), 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. """ import json, struct, sys, os import numpy as np import gguf -MODELS = "/media/ilintar/D_SSD/models/trellis2" -OUT = f"{MODELS}/gguf" +FAMILY = os.environ.get("TRELLIS_FAMILY", "trellis") +# 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 = { @@ -43,6 +75,38 @@ 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 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": + 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,10 +179,46 @@ 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 + + +# 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: @@ -129,8 +229,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") diff --git a/tools/gguf_probe.py b/tools/gguf_probe.py new file mode 100755 index 0000000..5f4566f --- /dev/null +++ b/tools/gguf_probe.py @@ -0,0 +1,238 @@ +#!/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("'}") + if len(extra) > 4: + print(f" ... and {len(extra) - 4} more") + + by_name = {n: (ne, t) for n, ne, t in tensors} + + if wanted: + for w in wanted: + hits = [(n, ne, t) for n, ne, t in tensors if w in n] + print(f"\n '{w}': {len(hits)} match(es)") + for n, ne, t in hits[:20]: + print(f" {n:<62} {ty(t):>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}") + + # 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: + 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))