lato.2_mrp_mlx

An MLX port of LATO.2 — factorised 3D mesh generation (vertex flow, then connectivity flow) — so it runs natively on Apple Silicon.

Why

LATO.2 generates meshes to a controllable vertex budget, which is the interesting part: it sidesteps the generate-dense-then-decimate loop that TRELLIS-style pipelines force on you. But upstream inherits TRELLIS.2's setup.sh and hard-requires CUDA:

# upstream modules/sparse/__init__.py
BACKEND = 'spconv'    # accepts only ['spconv', 'torchsparse']  — both CUDA-only
ATTN = 'flash_attn'   # accepts only ['xformers', 'flash_attn'] — both CUDA-only

No SDPA fallback, no MPS path. Neither sparse backend has a Metal build.

The scope, once measured

The CUDA surface is far smaller than setup.sh --all implies. Of the seven released checkpoints, five are entirely dense. Sparse code touches only vertex_autoencoder and vertex_structured_flow, and every SparseConv3d in the model is constructed with the defaults stride=1, padding=None — which upstream dispatches to SubMConv3d.

So the whole blocker is one operation: submanifold 3×3×3 convolution.

Upstream	Here
spconv.SubMConv3d	lato_mlx/sparse/conv.py — gather/scatter over a sorted-key indice map
SparseInverseConv3d	never instantiated upstream; not needed
strided sparse conv	never used; upsampling is SparseSubdivide (coord expansion ×8)
nn.Conv3d (voxel encoder)	dense, maps to mlx.nn.Conv3d
flash_attn / xformers	attn_mode="full" everywhere → plain SDPA

Status

• SparseTensor container + subdivide
• SubMConv3d (k=3 and k=1) — 7/7 correctness tests pass, max err 3e-7 vs an independent naive reference
• Weight converter, all 7 checkpoints → MLX safetensors (3.3 GB)
• Remaining sparse ops: SparseLinear, SparseGroupNorm32, activations, attention
• Model graphs: V-VAE, V-Flow, T-VAE, T-Flow, encoders
• End-to-end inference
• Fleet benchmark (m1max / m2max / m4pro / m1ultra / m3ultra)

Correctness

spconv cannot be installed here — that is the reason this port exists — so there is no numerical diff against upstream. Instead tests/test_sparse.py checks the vectorised implementation against a deliberately naive one written straight from the definition (dict lookup, per-voxel loop). They share no indexing code, so an off-by-one cannot hide in both. Tests also cover batch isolation, isolated voxels, and indice-map hit rates.

One assumption remains unverified: whether spconv gathers feats[c+d] (cross-correlation — the deep-learning convention, and what this implements) or feats[c-d]. A flipped kernel is numerically silent. It gets settled end-to-end: the V-VAE is an autoencoder, so a clean reconstruction confirms the orientation. Run the converter with --flip-kernel to test the alternative without touching code.

Use

uv venv --python 3.12 .venv
VIRTUAL_ENV=.venv uv pip install mlx numpy torch trimesh

hf download 0x4c48/LATO.2 --local-dir ckpt      # 3.3 GB upstream weights
.venv/bin/python -m lato_mlx.convert --ckpt ckpt --out weights
.venv/bin/python tests/test_sparse.py

Upstream source is vendored read-only under upstream/LATO.2 for reference.

Licence

Upstream LATO.2 is MIT (Copyright the LATO.2 authors); its sparse module carries Microsoft and VAST-AI-Research copyright, also MIT. This port is MIT on the same terms.