Porting TRELLIS.2 From Python to Rust and Metal: 2D Sprites to 3D Models in One Command

I needed 36 game models for Mazefront, a tower-defence game: towers at three upgrade levels, barriers, creatures, a portal and a core. Each one starts as a single 2D sprite. Microsoft's TRELLIS.2 turns a picture like that into a textured 3D model, and the community port for Apple Silicon works. I used it to make the first set, and it showed me exactly what it costs.

So I ported the pipeline to Rust, on the same Metal runtime that already runs our image models. The result is trellis2-rs. One command, no Python:

trellis2 generate --image cannon.png --seed 42 --output cannon.glb

It takes one to three minutes per model and about 6 GB of memory, where the Python pipeline needs about 33 GB.

From a Sprite to a Model

Every image below was generated by the Rust backend. The left panel is the only input.

The siege cannon. The model invents the back of the tower, which the sprite never shows. The stray chips floating near the base are real: loose fragments that the export does not yet remove.

The boss. Creatures are the hard case: limbs close to the body, and a back the network has to imagine.

The emerald core. Thin geometry like the wreath survives; the back of the gem is the model's guess.

Eight of the set, all from trellis2-rs at seed 42.

This article covers how the port was organised, what each stage came to, the bugs worth retelling, and, honestly, what the port does not do.

---

I – What the Python Pipeline Costs

• Memory. One model peaks near 33 GB. Reusing the pipeline in one process leaked roughly 10 GB per model. After 15 models the process held 194 GB, most of it compressed, swap was full, and the GPU backend ran out of memory. An earlier unguarded remesh took the machine down outright.
• Start-up. Loading the pipeline takes about 75 seconds. Once I ran each model in its own process to contain the leak, I paid that every time.
• Glue. The slowest single step was not a neural network. Turning the decoder output into a triangle mesh took 70 seconds of Python.

The memory guard I had written watched resident memory, which read 22 GB while the real footprint was 194 GB, because compressed and swapped pages do not count as resident. Measure the footprint, not RSS.

II – How the Port Was Organised

The work was done by AI coding agents coordinated from one session. The organisation mattered more than any single piece of code.

A scaffold first. One agent created the project, read the Python source, and wrote the specification everyone else would implement from: each stage's structure, tensor shapes, sampler settings and checkpoint tensor names. It also wrote the shared pieces for real: weight loading that verifies all nine checkpoints against expected tensor lists, the flow sampler, and a bit-exact reproduction of PyTorch's CPU random generator, so Rust starts from the same noise.

Reference tensors second. A script hooked the working Python pipeline and saved the true inputs and outputs of every stage for one model: 20 files, 1.9 GB. A determinism check re-ran the first call of each flow model and got a difference of exactly zero.

Then eight stages at once. Each stage went to its own agent in its own git worktree, so nobody could break anybody else's build. The acceptance test was parity against the reference tensors, with one rule: never loosen a tolerance to get a pass. Report the number instead.

The pipeline, in order:

1. Image preprocessing and the DINOv3 ViT-L encoder
2. Structure flow, a dense transformer over 4,096 tokens
3. Structure decoder, a dense 3D convolutional network that outputs occupied voxels
4. Shape flow at 512, then at 1024 (up to 15,800 sparse tokens for our largest model)
5. Shape decoder, sparse 3D convolutions up to 5.5 million voxels
6. Mesh extraction, about 11 million triangles
7. Texture flow and texture decoder
8. Remesh, UV unwrap, texture bake and the GLB writer

III – What Each Stage Came To

Parity is against PyTorch on the Apple GPU, for one reference model with 15,803 tokens at the 1024 stage and 5.55 million output voxels.

Stage	Runs on	Agreement with PyTorch
Image preprocessing	CPU	Byte-for-byte identical
Image encoder	CPU, full precision	Relative error 5e-6, cosine 1.0
Structure flow	Metal, BF16	Single call: cosine 0.99999
Structure decoder	CPU	Every occupied voxel identical
Shape and texture flows	Metal, BF16	Single call: about 0.5% error, no bias
Shape decoder	Metal, new kernel	With structure forced: all 5,551,390 coordinates identical
Mesh extraction	CPU	10 of 11,253,486 triangles differ
Texture decoder	Metal, new kernel	Coordinates identical; colours within about one 8-bit step
Remesh, UV, bake, GLB	CPU	No tensor reference; fewer non-manifold edges than Python on the test model, holes comparable

IV – Three Bugs Worth Retelling

A more accurate softplus was the wrong softplus

The first mesh comparison was strange. All 5.55 million vertices matched to six parts in a hundred million, the face count matched exactly, and yet 20,548 triangles were different. They turned out to be quads split along the other diagonal.

The split is decided by comparing predicted weights, and those weights pass through softplus. PyTorch on the Apple GPU computes log(1 + exp(x)) in single precision. The Rust code used ln_1p(exp(x)), which is more accurate for small values, and that was the problem. For inputs around −6 to −11 the two differ by up to one part in a thousand, enough to flip the diagonal on 0.18% of quads. Matching PyTorch's less accurate form brought the difference down to 10 triangles, all near-ties.

Rounding the input once too early

The sparse decoders run in half precision in the reference. The first GPU version rounded the incoming latent to half precision too, and the mean error in the subdivision logits doubled. The reference keeps that one tensor in full precision and only the layers after it in half. A CPU run with a deliberately pre-rounded latent reproduced the regression exactly, which confirmed the cause before the fix.

Eight percent of the texture was black

This one was in the export, and it exists in the Python pipeline as well. After remeshing, parts of the new surface sit where the colour volume has no voxels: closed-over gaps, hole caps, interior faces. Sampling there returns zero, so about 8% of texels baked black.

The fix samples a constant-one channel alongside the colours, divides by that coverage, and falls back to the nearest voxel where coverage is too low. Dark area on the test model went from 8.7% to 0.1%.

The same investigation produced a limit I now enforce. Remeshing on the full 1024 grid generates 18.3 million faces and passed 80 GB before I stopped it. A 512 grid closes the surface just as well for a 60,000-face target.

V – The Kernel That Decided It

Our runtime had no sparse 3D convolution. The first correct versions of both decoders ran on the CPU through Accelerate, and they were slow: 181 seconds for the texture decoder, where PyTorch takes 5.6.

The fix was a Metal kernel written inside the project, following the algorithm the Python port uses. Build a neighbour table once per resolution level. Then run a gather-GEMM over it with half-precision inputs and single-precision accumulation, with bias, rounding, activation and the residual add fused into the same pass. It reproduced an isolated reference layer bit for bit.

Step	PyTorch	Rust, first CPU version	Rust, Metal
Texture decoder	5.6 s	181 s	4.8 s
Shape decoder	about 23 s	56 s	2.2–3.0 s
Mesh extraction	70 s	0.16–0.25 s	(stays on CPU)

Mesh extraction needed no GPU at all. A compact hash map and a parallel loop replace 70 seconds of Python with a quarter of a second.

VI – Same Seed, Different Model

With every stage matching, I expected the end-to-end result to match. It does not. Rust and Python, given the same image and seed, produce two different samples of the same object. The first Rust tower was 13% darker than Python's, and I spent a while looking for a wiring bug that was not there.

Same sprite, same seed. Left: PyTorch. Right: trellis2-rs.

A step-by-step comparison settled it. I recorded every step of the 1024 shape flow in Python and fed each of Python's intermediate states to the Rust model.

• The sampler arithmetic is exact. The guidance mix, rescale and update agree to about one part in a million.
• Each model call differs by about 0.5%. That is noise from BF16 arithmetic, with no bias: the gain is 1.0000 at every step, and the errors on the two guidance branches are uncorrelated.
• Guidance amplifies it. At strength 7.5, two independent 0.5% errors become 2 to 5% in the guided prediction, and nine guided steps compound.

Python outputs substituted for the first…	Final difference (1 − cosine)
0 steps (Rust runs freely)	0.0134
3 steps	0.00031
6 steps	0.00012
9 steps	0.00001

The sampler is chaotic at this guidance strength. Perturbing the input noise by a single BF16 unit moves the result as far as switching implementation does, and across five perturbed runs the average surface colour swung between red values of 65 and 80. When I fed Python's structure into Rust, the colours matched to the first decimal. Across eight models, Rust comes out brighter on some and darker on others.

The boss. The Python export has 13 holes and 1,583 non-manifold edges; the Rust one has 11 and 468.

So the claim I can make is narrower than "identical output", and I think more useful: every stage is verified against the reference, and the end-to-end result is a valid sample from the same distribution.

VII – End to End

Stage times for the reference model, the largest of the set. Both columns were measured on a shared GPU, and the Python run carried the tensor-dump hooks, so read the ratio as indicative.

Stage	PyTorch	trellis2-rs
Image encoder	0.9 s	0.7 s on the critical path
Structure flow and decode	20.8 s	13.7 s
Shape flow 512 and upsample	23.8 s	about 11 s
Shape flow 1024	153 s	90 s
Texture flow	86 s	48 s
Shape decode and mesh	93 s	overlapped with the texture flow
Texture decode	5.6 s	2.1 s
Export to GLB	60–90 s	mostly overlapped; 6–9 s alone
Whole command	385 s before export	172–181 s including export

Smaller models finish much faster, because cost follows token count. An amplifier took 52 seconds and the cannon about 74.

A tuning pass at the end found that the flow transformers were already near what the hardware gives, so the gains came from elsewhere: running the heavy encoder pass and the whole geometry path in the background behind GPU stages, and fixing a quadratic loop in the UV unwrapper that cost 8 seconds on some meshes.

VIII – What Is Not Done

• Attention is the ceiling. A flow call runs at about 33 TFLOP/s on a chip that peaks near 62, and self-attention is half of every call. Closing that gap needs a new attention kernel.
• The new kernels live in the project. The sparse convolution should move into the shared runtime so other models can use it.
• Parity was measured on one image. The stage tests are strict, but they cover a single reference model and seed.
• Timings need a quiet machine. Every end-to-end number here was taken with other work on the GPU.
• Geometry is not watertight. Models still carry a handful of pinholes and loose fragments, as the Python exports do.
• Models are static. There is no skeleton, so creatures cannot walk yet. That is the next experiment.

IX – Where the Models Went

All 36 models are in the game now, reduced from 438 MB to 20 MB for play, with a higher-quality tier for the in-game encyclopedia. Four of them were regenerated with the Rust backend after the port, at about a minute each.

The Colossus in Mazefront's encyclopedia, running in the browser.

A few more pairs from the set, sprite on the left and Rust output on the right:

X – Lessons

1. Dump the reference before writing a line of the port. Every argument about correctness ended the moment there was a tensor to compare against. The dumps took one afternoon and paid for themselves in the first hour.

2. Split by stage, and give each stage its own checkout. Eight agents in eight worktrees cannot break each other's build. The integration at the end was mostly dependency conflicts.

3. Match the reference's mistakes. The softplus bug was a case of being more accurate than the thing I was porting. Parity means reproducing the less accurate form.

4. Know what "the same" can mean. With guidance at 7.5 and BF16 arithmetic, two correct implementations do not produce the same model from the same seed. Test the stages strictly, and test the whole by distribution.

5. The neural networks were not the slow part. Mesh extraction, the export and a leaky process cost more than any transformer. Profile the glue first.

TRELLIS.2 is by Microsoft Research and is released under the MIT licence. The Apple Silicon port I started from is the community trellis-mac project.

– Antonio