Initial support for Pixal3D - #38
Open
vegax87 wants to merge 22 commits into
Open
Conversation
Owner
|
The flash attention precision is something that's a common bug, I have to look into it more carefully. |
vegax87
marked this pull request as draft
August 20, 2026 10:17
TRELLIS.2 conditions every DiT block by cross-attending over the DINOv3 token stream. Pixal3D keeps that cross-attention but wraps it: the module's weights move one level down under cross_attn.cross_attn_block, and a sibling cross_attn.proj_linear maps a per-token view-aligned feature into model space and is added to the cross-attention output, replacing it as the residual branch. Add DiTParams::proj_mode / proj_ch and a proj input to build_dit_dense, which is the whole architectural delta on the denoiser side. Conditioning now travels as a FlowCond bundle so classifier-free guidance can carry both the context and the projection; a null proj in proj mode is the negative branch and is zeroed on the device rather than materialized (it reaches 400 MB at the texture stage's token budget). The implicit FlowCond constructor keeps cross-mode call sites unchanged. (cherry picked from commit e0b76ca95e52b142a21d068aad53c2c7f75ceb11) (cherry picked from commit 8da007878feaddcf050847f6b9883533411b4b75)
Pixal3D cross-attends over the 5 global DINOv3 tokens only and routes the patch grid through a pixel-aligned path instead: each DiT token is a cell of a 3-D grid, projected into the image with a fixed frontal camera, and the DINOv3 feature map is sampled there. pixal3d.cpp implements the camera solve (only the FOV is free; the distance follows in closed form), the grid projection and the per-stage ProjCond assembly. test_pixal3d pins both functions to golden values evaluated from the reference implementation, because a subtly wrong projection has no runtime signal -- it just drifts the geometry off the silhouette. naf.cpp ports valeoai/NAF, which the shape and texture stages run to upsample the feature map before sampling it (hence their proj_in_channels of 2048). Only the two-branch guide encoder is learned, and it maps cleanly onto ggml; the upsampling itself is a parameter-free neighborhood cross-attention. Dilating it by exactly the upsampling factor collapses NATTEN's dilated 2-D kernel into a 9x9 window of low-resolution cells centred on the query's own cell, which is what makes the threaded CPU implementation practical -- and it only runs on the pixels the projection actually reads, not the whole upsampled map. (cherry picked from commit 9c5f948cd7304f73429900cf9316cd1428ca46d4) (cherry picked from commit e1794c3399f469cd627486cb01313f5ddbdac0ab)
The pipeline is shared end to end -- same sampler schedules, same guidance, same decoders, remesh and bake -- so --model only decides how the image conditions each flow. Wire the per-stage projection parameters (grid resolution, DINOv3 image size, NAF target) from Pixal3D's stage configs, including the HR grid following the cascade's token-budget backoff. proj_in_channels is read off the checkpoint's own proj_linear rather than a config, which doubles as the guard against pointing --model at the wrong weights: both directions fail with a message instead of producing garbage. Two behaviours degrade because Pixal3D publishes only the 1024 texture flow: --res 512 has no texture model and falls back to geometry only, and the mixed-resolution texture shortcut is disabled. (cherry picked from commit c32c4f8dfc6a3f5e99a78383cffe13e9b1893d54) (cherry picked from commit a4d3d3afe33110d556d28739acb27b91577b30ea)
TRELLIS_FAMILY=pixal3d switches the converter's manifest. Pixal3D's checkpoints carry the same filenames as TRELLIS.2's and the same tensor layout plus the two proj tensors per block, so the verbatim-name policy needs no remapping; the decoders are unchanged and can be reused. NAF ships as a torch .pth from torch.hub, so it gets its own reader. docs/pixal3d/README.md covers what actually differs, the per-stage projection table, the camera flag standing in for the unported MoGe-2 estimation, how the NAF port works and what is not verified. (cherry picked from commit d7915be1411808c80a2034e2a15c1829487b7efa) (cherry picked from commit aa467e93217ac95600fbb923a119a582d287a01e)
The reference shrinks the guide before the convolutions once it exceeds 4x the NAF target. No Pixal3D stage comes close (the largest ratio is 2), so the branch is not ported -- but fail loudly rather than silently diverge if one ever does. Also record what the port was actually checked against: an independent transcription of the reference under random weights agrees to 2.3e-6 end to end, and the collapsed 9x9 window matches NATTEN's own get_window_start exactly, so that reduction is an identity here rather than an approximation. (cherry picked from commit d57878443f5935d1ff9b284ec6495f285685f157) (cherry picked from commit b3dea88d0a418e7cad08555e8bd863cc9806fd56)
Third-party Pixal3D conversions exist on the Hub, but whether they load here comes down to two things buried in the header: whether the tensor names are the verbatim torch state_dict names, and what dtype the 1-D parameters carry. The second is the sharp edge -- dit.cpp adds `modulation` to an f32 timestep embedding and multiplies `norm2.weight` into an f32 activation, so f16 there is a silent mismatch rather than a slow path. GGUF puts its metadata and tensor index at the head of the file, so a few MB of HTTP Range answers both without pulling 2.7 GB per flow. Stdlib only. (cherry picked from commit 70772d6b3a109bc3b9264e7c91e0cf6f28c33e5c) (cherry picked from commit 0dbb0478f239e8522c373ec95d01182272e56ab7)
Two things the first version missed on a real third-party conversion. Structure: a DiT is N identical blocks plus a fixed preamble, so every block must carry the same suffix set. A converter that renamed, merged or dropped tensors shows up as a ragged block here instead of as a failure at inference time. For a Pixal3D SLat flow the arithmetic is 30 x 23 + 10 = 700. Dtypes: report the matmul weights separately. BF16 loads and runs -- ggml has a dedicated cuBLAS BF16 path and CPU type traits -- but --f32 only casts F16, so that escape hatch does not apply to BF16 weights and it is worth saying so before someone reaches for it to debug a bad sm_120-family matmul. (cherry picked from commit 1c545a12185c3a8578021d92c3583a8870e3bc12) (cherry picked from commit 5b82d1f664fbd9c7b3d01077eaf5b691643ab2c9)
Pixal3D's checkpoints carry the same upstream filenames as TRELLIS.2's, which forced either two model directories or a tree of symlinks -- friction for anyone testing both families. Name the family-specific files pixal3d_*.gguf instead: the four flows and NAF. The decoders, DINOv3 and BiRefNet are byte-identical between the families and keep their plain names, so both --model values read one directory and adding Pixal3D to a working TRELLIS.2 set is 5 new files rather than a second copy of everything. The converter applies the prefix automatically. (cherry picked from commit e51a731f7478de4dd5a421ea05dd113e2da21c9f) (cherry picked from commit a8b2f63f594038e899be4412ef0f9f7f92e80771)
The checkpoint trees live outside the repo, so the hardcoded defaults are only ever right on one machine -- everyone else had to edit the script before running it, which is both friction and a stray diff waiting to be committed by accident. TRELLIS_MODELS, PIXAL3D_MODELS and TRELLIS_GGUF_OUT override them. The old values remain the defaults, so existing invocations are unaffected. (cherry picked from commit 9bb6e7360e55f1336da67a96cb4416fb378d139e) (cherry picked from commit 8e73f224d46195b02a87f0d41afa0675ec169de9)
NAF is a separate project's GitHub release, not part of the Pixal3D checkpoint tree, so requiring it under PIXAL3D_MODELS/naf/ was an arbitrary constraint on where it may be downloaded. The default is unchanged. (cherry picked from commit d2fd02f504ec939ee795c5310369678ed44e0ac2) (cherry picked from commit b6941f44130a4dadc0a99332b796507ddf83d766)
The five Pixal3D-specific files are published at vegax87/Pixal3D under the exact names the loader expects, so testing the backend no longer means converting from the TencentARC safetensors first. The conversion route stays documented, now with the environment variables rather than the edit-the-script instructions. (cherry picked from commit e94e88169a4bf5a6a456bfc66db13ada85d358f7) (cherry picked from commit 464e56e3c99294eb5b2ce663d3e8adda15fe6eaa)
A GGUF whose layout differs from what build_dit_dense assumes surfaces as a bare GGML_ASSERT(ggml_can_mul_mat) inside ggml, followed by a core dump that names neither the tensor nor the shapes -- the backtrace only reaches trellis::lin, which is called a dozen times per block. Third-party conversions are a normal way to obtain these weights, so check the width in lin() and report the offending tensor with both shapes instead. (cherry picked from commit 84825d2f0cf0469b4304c2bce298e9fb42aa8133) (cherry picked from commit a8b2fe9727fe50a70a47fd4f27ec346e85fbe26a)
Converters aimed at other runtimes reshape tensors into a quantization-friendly 2-D form and stash the real shape in metadata -- a 1536-element RMS-norm gamma becomes [256, 6] rather than [128, 12], and input_layer becomes [256, 48] rather than [8, 1536]. Element counts match, so nothing looks wrong until ggml asserts on the first matmul. Report the non-standard metadata keys, which name the convention and make the reshaping obvious. Also decode metadata scalars properly so array values print as numbers instead of raw bytes. (cherry picked from commit 0ead0a0cb5f1273c779d2283fe95c33bce5b03cc) (cherry picked from commit da9ea57e0fcfb1cb488fe1c97fd692a40ff045de)
ComfyUI-style GGUF tooling stores tensors in a quantization-friendly 2-D layout instead of their natural shape -- a 1536-element RMS-norm gamma becomes [256, 6] rather than [128, 12], an [8, 1536] input projection becomes [256, 48] -- and records the real shape under a `.orig_shape.<tensor>` metadata key. Element order is untouched, so restoring the shape reinterprets the same bytes. Without it nothing looks wrong: names match, element counts match, the buffer loads, and the mismatch only surfaces when a matmul deep in the graph finally compares widths. Since these conversions are a normal way to obtain the weights, read the metadata rather than requiring everyone to reconvert. Quantized tensors are skipped (their row length must stay a multiple of the block size), as is any tensor whose recorded shape does not match its element count -- that is padding, not reshaping, and silently reinterpreting it would be wrong. (cherry picked from commit 8e18e689bc7a11e78632a195de7fa5baec12fca9) (cherry picked from commit 2a246fbfd7b50822830124fee985e51dd2bf1c75)
The projection conditioning has no runtime failure mode. A wrong camera, a wrong grid or a broken upsampler all yield finite numbers and a mesh whose silhouette still looks roughly right; the damage appears only as high-frequency noise in the generated surface, which is indistinguishable from a bad seed by eye. Print three figures per stage instead. Coverage catches a collapsed camera (the cube's corners fall outside the frame legitimately, so this is a ratio, not a pass/fail). The per-half mean/std and their cosine similarity catch a broken NAF branch: both halves describe the same points at different detail, so they must share a scale and correlate strongly -- and nothing else in the pipeline would reveal it. (cherry picked from commit 9d4098749a8c45a3069b4bc2e5bcb65738ddb4a3) (cherry picked from commit 4d8de19090bcb9be29675ad72a46570a646542c8)
The cascade quantizes the upsampled res-512 coords down to the HR token grid. TRELLIS.2 floors u * grid; Pixal3D rounds u * (grid - 1). This port followed the former, because Pixal3D's own sample_shape_slat_cascade still carries TRELLIS.2's version -- but run() never calls that helper and inlines the other formula. The difference is not cosmetic for Pixal3D. There the token index also selects which node of the projection grid the token samples the image at, and that grid is the endpoint-inclusive linspace(-1, 1, grid): rounding to the nearest node is what keeps the pixel-aligned sample registered. Flooring u * grid shifts by half a cell and dilates by grid/(grid-1), and because the index is an integer the error is a step function, not a smooth warp -- at grid 64 it moves a quarter of the coordinates a full cell along each axis, in stripes of period 8, so 58% of tokens end up misregistered on at least one axis. A cell is 16 res-1024 voxels, one ViT-L/16 patch at the 1024 conditioning size. That is exactly the observed failure: the silhouette survives, since it rides on the five global cross-attention tokens, while the high-frequency detail does not and the decode comes out speckled -- face/vertex ratio 1.24 against 2.00 for the TRELLIS.2 baseline on the same image. The sparse-structure and LR shape stages were unaffected because neither requantizes: SS is a dense self-consistent grid, and the LR coords come straight from the SS decode at the grid resolution the stage already expects. Kept family-conditional -- flooring remains correct for TRELLIS.2, whose coords only feed RoPE, where a half-cell offset is a smooth reparameterization. (cherry picked from commit 67b56acf8d33c78a8e810cbcc541a45649f0040f) (cherry picked from commit 9aef5d4e6cfa73a6684cd032dafa65b4317e1e17)
The HR shape stage is the only one conditioned on DINOv3 at 1024, and it is the only stage producing a speckled decode -- the sparse-structure and LR shape stages, which read the 512 map through the same code, are clean (face/vertex 2.04 against the TRELLIS.2 baseline's 2.01). Both maps encode the same image, so the features sampled at a given 3-D point must agree across them. Report that agreement. Nothing else in the pipeline can distinguish "the 1024 map is being read with the wrong spatial layout" from "the latent is bad": the projection statistics stay well-scaled, the two proj halves stay correlated, and every value stays finite either way. (cherry picked from commit fc3b8a3e14531fd38855c4d162bf8fa9e7995227) (cherry picked from commit 8d52f0349319d7f5f12991ebfd7c857b3fc1df4b)
FlashAttention is here to avoid materialising the [Lk, Lq, nh] score matrix, which reaches terabytes at the HR flow. When Lk is shorter than a single FA key tile that matrix is a few megabytes and FA buys nothing, while still paying the whole cost of its padding machinery. Pixal3D's proj mode cross-attends over 5 global tokens, so the key dimension is zero-padded from 5 to 256 -- 98% padding. TRELLIS.2 never gets near that regime (1029 keys at 512, 4101 at 1024), and the comments in this function record that ggml's CUDA FlashAttention has repeatedly mishandled padded key tiles in ways that depend on the token count. Route short-KV attention through the exact chunked path, where at Lk = 5 it is cheaper than FA regardless. (cherry picked from commit b35d91d92908fdffe0af72f8c9f1660b8c6a8c19) (cherry picked from commit 3ac026496914b43bc6940339487a22619c1325b3)
The camera solve registers the projection grid so its corner lands on the image border, which silently assumes the subject is entirely inside the frame. When the input is cropped -- an image generator asked for too close a shot, a photo framed tight -- that assumption is wrong and the object gets squeezed into the grid: the visible part comes out compressed and the occluded part is invented at the wrong scale. The reference carries an extend_pixel parameter for exactly this and never exposes it, since MoGe-2 estimates from the whole frame. pixal3d_camera already took it; this only plumbs it through to a flag. At 49.13 degrees and a 512 frame: 0 px gives distance 1.094, 64 px gives 0.875, 128 px gives 0.729. (cherry picked from commit 539c7047a88a0d81191d69ac91bc99437baa9731) (cherry picked from commit 5dc4f6e33ba413cfa7fcb1472a78cee4ab08ecee)
The previous text claimed it rescues an input whose subject is cropped by the frame. It does not, and recommending it that way produces a worse result than leaving it alone. Background removal already reframes: it crops a square around the visible alpha bbox with a 10% margin, so the cutout reaching the camera always has the subject inscribed regardless of the original framing. Pushing the border past that stretches the grid over territory with no pixels behind it -- the projection lands on the border clamp, those cells go unconditioned, and the model fills them with whatever it likes. Observed on a cropped photograph as a tail growing out of the object. No parameter recovers information the file does not contain. Say so, and point at re-framing the source instead. (cherry picked from commit 846cb07b76761e7c87ac042b81a91d3794936908) (cherry picked from commit 7246f6302975afd0ff4cc0cd19c244fc6306f834)
--model already worked as a launch flag, since the server runs the same parser for its defaults and every request copies them. But one model directory now holds both families -- that is what the pixal3d_ prefix bought -- so a single resident server can serve either, and only the request field was missing. The camera fields matter more. fov, mesh_scale and extend_pixel describe the image rather than the run, so pinning them at launch is close to useless on a server taking arbitrary uploads: every image wants its own FOV. (cherry picked from commit 41054689e039806de3d87260bd3e1f08ce4af2f6) (cherry picked from commit 1725a23751d3a4b941ea31f2b40fa689756c90b0)
Proj mode adds proj_linear(proj) into every block's residual, so its activations run larger than TRELLIS.2's. CUDA's tensor-core FA kernels convert BF16 K/V to F16 internally, and those larger values overflow: before the V pre-scaling in sdpa(), the res-1024 cascade decoded 1.19M voxels at a face/vertex ratio of 1.25 while TRELLIS.2 was unaffected at the same token count. With it, and an identical seed, FA gives 3.47M at 2.31 against 3.22M at 2.11 on the exact path, with the sparse-structure stage agreeing to 0.4%. Worth writing down because the failure is silent, resolution dependent, and looks like a bad seed rather than a numerical fault.
Author
vegax87
marked this pull request as ready for review
August 20, 2026 12:38
Owner
|
Yeah, was hoping so :) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
You can choose between trellis or pixal3d with
--modelAdds
--model trellis|pixal3d. The default staystrellisand the TRELLIS.2 path is unchanged.Pixal3D is a fine-tune of TRELLIS.2 rather than a new architecture: same 1.3B DiT, same sampler, same cascade, same decoder checkpoints. Only the conditioning differs. TRELLIS.2 cross-attends over every DINOv3 token; Pixal3D cross-attends over the 5 global tokens and adds a pixel-aligned term, where each DiT token is a grid cell projected into the image and sampled there, then mapped into model space by a per-block
proj_linear. The integration is therefore a conditioning module plus onebranch inside the DiT block.
Models
5 extra GGUF files from https://huggingface.co/vegax87/Pixal3D, dropped into the same
models/folder you already use:The decoders, DINOv3 and BiRefNet are byte identical between the two families and are
shared, so nothing is duplicated.
TRELLIS_FAMILY=pixal3d tools/convert.pyapplies theprefix automatically.
Results
1024 cascade, same image and machine (GB10):
Notes
src/naf.cpp. It stays small because the upsampling itself has no parameters, only theguide encoder does. Verified against a transcription of the reference to 2.3e-6.
--fov, defaulting toPixal3D's own value). The closed-form distance derivation is ported and pinned to golden
values from the reference in
trellis-test-pixal3d.--model pixal3d. On sm_121 it produces ahole-riddled res-1024 surface for this family (face/vertex 1.25 against 2.09 on the exact
path) while TRELLIS.2 is unaffected at the same token count, so the trigger looks like
the activations proj mode adds rather than the shapes.
--faopts back in. TheGGML_FA_DEBUGoutput is in the docs.--res 512writes geometry only.trellis-servertakesmodelper request, plusfov, since one directory holds bothfamilies.
Two pieces are useful beyond this backend:
tools/gguf_probe.pyreads a GGUF header overHTTP range so a conversion can be checked without downloading it, and the loader now
restores tensor shapes from the
orig_shapemetadata that ComfyUI-style converters write.Full documentation in
docs/pixal3d/README.md.