# lato.2_mrp_mlx

An MLX port of [LATO.2](https://github.com/LoHhhha/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:

```python
# 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

- [x] `SparseTensor` container + `subdivide`
- [x] `SubMConv3d` (k=3 and k=1) — **7/7 correctness tests pass**, max err 3e-7 vs an
      independent naive reference
- [x] 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

```bash
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.