e23233731b
Converter is light because Pixal3D ships safetensors with sibling .json configs, so there is no architecture to infer and no pickle to unpack. Only rank-5 tensors are rearranged; dtype is preserved (upcasting fp16->fp32 doubled 24GB for no benefit) and the KRSC remap is verified a pure permutation of values. Measured: the four flow models (~20GB) have ZERO 5-D tensors - pure transformers that never touch sparse conv. shape_dec/tex_dec carry 40 KRSC kernels each, ss_dec 20 dense Conv3d, and classify_5d separates them correctly on the real weights. Corrects the earlier 'gap is two ops' claim: that was right about modules/sparse but undercounted the model blocks. The configs show all four flow models need RoPE, qk_rms_norm and AdaLN modulation, and the decoders need SparseConvNeXtBlock3d and SparseResBlockC2S3d.
109 lines
3.8 KiB
Python
109 lines
3.8 KiB
Python
"""Convert Pixal3D checkpoints to MLX-ready safetensors.
|
|
|
|
Much lighter than LATO.2's converter: Pixal3D ships safetensors already (not pickled
|
|
.pt), and each checkpoint has a sibling .json naming the model class and its exact
|
|
constructor args — so there is no architecture to infer from tensor shapes.
|
|
|
|
The only real work is the rank-5 layout question, handled by `trellis_sparse_mlx.convert`
|
|
(spconv KRSC vs dense torch Conv3d — indistinguishable by rank, silently wrong if
|
|
guessed). Measured against the released weights:
|
|
|
|
ss_flow / slat_flow x3 ZERO 5-D tensors — pure transformers, pass straight through
|
|
shape_dec / tex_dec 40 KRSC sparse kernels each
|
|
ss_dec dense Conv3d
|
|
|
|
So the ~20GB of flow models need no conversion at all; only the decoders do.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
|
|
from trellis_sparse_mlx.convert import convert_tensor
|
|
|
|
|
|
def load_config(ckpt: Path) -> dict:
|
|
"""Each checkpoint has a sibling .json with {name, args} — the authoritative config."""
|
|
cfg = ckpt.with_suffix(".json")
|
|
if not cfg.exists():
|
|
raise FileNotFoundError(f"no config beside {ckpt.name}")
|
|
return json.loads(cfg.read_text())
|
|
|
|
|
|
def convert_checkpoint(src: Path, dst: Path, flip_kernel: bool = False) -> dict:
|
|
w = mx.load(str(src))
|
|
out, counts = {}, {"dense": 0, "krsc": 0, "conv3d": 0}
|
|
for k, v in w.items():
|
|
dt = v.dtype
|
|
# Only rank-5 tensors are ever rearranged, and the rearrangement is a pure
|
|
# transpose/reshape — so round-trip through float32 for numpy only where needed
|
|
# and restore the original dtype. Upcasting everything to fp32 would double
|
|
# 24GB of checkpoints on disk for no benefit.
|
|
if v.ndim != 5:
|
|
out[k] = (v, "dense")
|
|
counts["dense"] += 1
|
|
continue
|
|
arr = np.asarray(v.astype(mx.float32))
|
|
arr, kind = convert_tensor(k, arr, flip_kernel)
|
|
counts[kind] += 1
|
|
out[k] = (mx.array(arr).astype(dt), kind)
|
|
dst.parent.mkdir(parents=True, exist_ok=True)
|
|
mx.save_safetensors(str(dst), {k: v for k, (v, _) in out.items()})
|
|
cfg = load_config(src)
|
|
return {
|
|
"model": cfg.get("name"),
|
|
"tensors": len(out),
|
|
**counts,
|
|
"gb": dst.stat().st_size / 1e9,
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser(description="Pixal3D -> MLX weight converter")
|
|
ap.add_argument(
|
|
"--ckpt",
|
|
default="/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts",
|
|
)
|
|
ap.add_argument("--out", default="weights")
|
|
ap.add_argument("--only", default=None, help="substring filter, e.g. 'dec'")
|
|
ap.add_argument(
|
|
"--flip-kernel",
|
|
action="store_true",
|
|
help="mirror sparse kernel offsets (orientation experiment; see shared core)",
|
|
)
|
|
ap.add_argument(
|
|
"--skip-passthrough",
|
|
action="store_true",
|
|
help="skip checkpoints with no 5-D tensors — they need no conversion",
|
|
)
|
|
a = ap.parse_args()
|
|
|
|
src_dir, out_dir = Path(a.ckpt), Path(a.out)
|
|
files = sorted(src_dir.glob("*.safetensors"))
|
|
if a.only:
|
|
files = [f for f in files if a.only in f.name]
|
|
if not files:
|
|
raise SystemExit(f"no .safetensors in {src_dir}")
|
|
|
|
for f in files:
|
|
if a.skip_passthrough:
|
|
w = mx.load(str(f))
|
|
if not any(v.ndim == 5 for v in w.values()):
|
|
print(f" {f.stem[:44]:46s} skipped (no 5-D tensors)")
|
|
continue
|
|
info = convert_checkpoint(f, out_dir / f.name, a.flip_kernel)
|
|
print(
|
|
f" {f.stem[:44]:46s} {info['model'][:26]:28s} "
|
|
f"{info['tensors']:4d}t krsc {info['krsc']:3d} conv3d {info['conv3d']:3d} "
|
|
f"{info['gb']:5.2f} GB"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|