lato.2_mrp_mlx/lato_mlx/convert.py

"""Convert LATO.2 PyTorch checkpoints to MLX safetensors.

Run with any torch-bearing interpreter; torch is needed only to unpickle the .pt files
and is not a runtime dependency of the port itself.

    python -m lato_mlx.convert --ckpt ckpt --out weights

Weight layouts
--------------
spconv 2.x stores SubMConv3d weights in KRSC order, verified empirically against the
released vvae.pt (e.g. `encoder.downsample.0.conv1.conv.weight` is (64,3,3,3,32) =
[out, kz, ky, kx, in]). The MLX SubMConv3d here wants [K^3, in, out], so:

    [O, kz, ky, kx, I] --transpose--> [kz, ky, kx, I, O] --reshape--> [K^3, I, O]

The C-order flatten of (kz,ky,kx) matches `_kernel_offsets`, which meshgrids
`arange(k) - k//2` with indexing="ij" — so kernel index 13 is the centre tap at k=3.

UNVERIFIED ASSUMPTION — kernel orientation
------------------------------------------
Whether spconv gathers `feats[c + d]` (cross-correlation, the deep-learning convention
and what this port implements) or `feats[c - d]` (true convolution) cannot be checked
numerically on this machine, because spconv has no Metal build — which is the whole
reason this port exists. A flipped kernel is silent: shapes and norms look right and the
output is subtly wrong.

It is settled end-to-end rather than by assertion: the V-VAE is an autoencoder, so if
encode→decode reconstructs the input mesh, the orientation is right; a flip yields
visibly broken geometry. `--flip-kernel` writes the mirrored variant so the alternative
is a one-flag experiment rather than a code change.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np

# Checkpoints built on the sparse module tree. Measured against the released weights:
# only vvae actually carries SubMConv3d kernels (18 of them); vflow is sparse in the
# structural sense (SparseTensor + SparseLinear + sparse attention) but has no conv.
SPARSE_CKPTS = {"vvae.pt", "vflow.pt"}


def _load_state_dict(path: Path):
    import torch

    obj = torch.load(str(path), map_location="cpu", weights_only=False)
    if not isinstance(obj, dict):
        obj = obj.state_dict()
    for key in ("state_dict", "model", "module"):
        inner = obj.get(key)
        if isinstance(inner, dict):
            obj = inner
            break
    return obj


def classify_5d(name: str, shape: tuple) -> str:
    """'krsc' (spconv sparse) vs 'conv3d' (dense torch) — they collide at 5-D.

    spconv SubMConv3d : [out, kz, ky, kx, in]   -> dims 1..3 are the cubic kernel
    torch nn.Conv3d   : [out, in, kz, ky, kx]   -> dims 2..4 are the cubic kernel

    LATO.2 contains both (the V-VAE uses sparse convs; the voxel encoder uses dense
    nn.Conv3d), so guessing by rank alone silently mangles one of them.
    """
    krsc = shape[1] == shape[2] == shape[3]
    conv3d = shape[2] == shape[3] == shape[4]
    if krsc and not conv3d:
        return "krsc"
    if conv3d and not krsc:
        return "conv3d"
    if krsc and conv3d:
        # e.g. (O,3,3,3,3): genuinely ambiguous by shape. spconv weights always sit
        # under the wrapper's `.conv.weight`; fail loudly rather than coin-flip.
        if name.endswith(".conv.weight"):
            return "krsc"
        raise ValueError(
            f"{name}: ambiguous 5-D layout {shape} — cannot tell KRSC from Conv3d"
        )
    raise ValueError(f"{name}: unrecognised 5-D layout {shape}")


def convert_tensor(name: str, arr: np.ndarray, flip_kernel: bool) -> tuple:
    """Returns (array, kind). Only spconv KRSC kernels are reshaped."""
    if arr.ndim != 5:
        return arr, "dense"
    kind = classify_5d(name, arr.shape)
    if kind == "conv3d":
        # Dense conv is handled by mlx.nn.Conv3d, which wants [out, kz, ky, kx, in]
        # (channels-last), so move the input-channel axis to the end.
        return arr.transpose(0, 2, 3, 4, 1).copy(), "conv3d"
    o, kz, ky, kx, i = arr.shape
    w = arr.transpose(1, 2, 3, 4, 0).reshape(kz * ky * kx, i, o)
    if flip_kernel:
        w = w[::-1].copy()  # mirror all offsets: correlation <-> convolution
    return w, "krsc"


def convert_checkpoint(src: Path, dst: Path, flip_kernel: bool = False) -> dict:
    import mlx.core as mx

    sd = _load_state_dict(src)
    out, n_krsc, n_conv3d, skipped = {}, 0, 0, 0
    for k, v in sd.items():
        if not hasattr(v, "detach"):
            skipped += 1
            continue
        arr = v.detach().to("cpu").float().numpy()
        arr, kind = convert_tensor(k, arr, flip_kernel)
        n_krsc += kind == "krsc"
        n_conv3d += kind == "conv3d"
        out[k] = arr

    dst.parent.mkdir(parents=True, exist_ok=True)
    mx.save_safetensors(str(dst), {k: mx.array(v) for k, v in out.items()})
    return {
        "tensors": len(out),
        "sparse_kernels": n_krsc,
        "dense_conv3d": n_conv3d,
        "skipped": skipped,
        "mb": dst.stat().st_size / 1e6,
    }


def main() -> None:
    ap = argparse.ArgumentParser(description="LATO.2 -> MLX weight converter")
    ap.add_argument("--ckpt", default="ckpt", help="dir of upstream .pt files")
    ap.add_argument("--out", default="weights", help="output dir for .safetensors")
    ap.add_argument(
        "--flip-kernel",
        action="store_true",
        help="mirror sparse kernel offsets (see module docstring: orientation experiment)",
    )
    ap.add_argument("--only", default=None, help="convert just this file, e.g. vvae.pt")
    a = ap.parse_args()

    src_dir, out_dir = Path(a.ckpt), Path(a.out)
    files = sorted(src_dir.glob("*.pt"))
    if a.only:
        files = [f for f in files if f.name == a.only]
    if not files:
        raise SystemExit(f"no .pt files in {src_dir}")

    for f in files:
        dst = out_dir / (f.stem + ".safetensors")
        info = convert_checkpoint(f, dst, a.flip_kernel)
        tag = " [sparse]" if f.name in SPARSE_CKPTS else ""
        print(
            f"  {f.name:20s} -> {dst.name:26s} "
            f"{info['tensors']:4d} tensors, {info['sparse_kernels']:2d} sparse, {info['dense_conv3d']:2d} conv3d, "
            f"{info['mb']:7.1f} MB{tag}"
        )


if __name__ == "__main__":
    main()