From e23233731b495f4685c3e6679cd866412fe8cbe2 Mon Sep 17 00:00:00 2001 From: John Date: Sun, 2 Aug 2026 11:25:00 +1000 Subject: [PATCH] Pixal3D weight converter + corrected scope 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. --- CLAUDE.md | 10 ++++ README.md | 35 ++++++++++--- pixal3d_mlx/__init__.py | 0 pixal3d_mlx/convert.py | 108 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 146 insertions(+), 7 deletions(-) create mode 100644 pixal3d_mlx/__init__.py create mode 100644 pixal3d_mlx/convert.py diff --git a/CLAUDE.md b/CLAUDE.md index cb5b1f5..9139a7a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -27,3 +27,13 @@ MLX port of TencentARC Pixal3D, on top of the shared `trellis_sparse_mlx` core. the sparse core. - `upstream/Pixal3D` is vendored read-only reference; never edit it. - Weights live at MODELBEAST/vendor/pixal3d-weights (24GB, outside this repo). + +## Weights (measured 2026-08-02) +- Every checkpoint ships a sibling .json with {name, args} — the authoritative config. + Do NOT infer architecture from shapes here (that was needed for LATO.2, not this). +- ss_flow + slat_flow x3 (~20GB, the bulk): ZERO 5-D tensors. Pure transformers, no + sparse conv. They pass through the converter untouched. +- shape_dec / tex_dec: 40 KRSC sparse kernels each. ss_dec: 20 dense Conv3d. + classify_5d() distinguishes them correctly on the real weights. +- `rope_phases` ships as a stored tensor — RoPE phases are precomputed, not derived. +- Converter preserves dtype (fp16 stays fp16); upcasting doubled 24GB for nothing. diff --git a/README.md b/README.md index fbf0776..c00f9b3 100644 --- a/README.md +++ b/README.md @@ -30,9 +30,26 @@ Upstream needs ~24 GB VRAM, which is a wall on consumer Nvidia and a non-issue o | `SparseUpsample(2)` | ❌ **to do** — cache-paired inverse of a downsample | | `SparseSpatial2Channel(2)` | ❌ **to do** — sparse pixel-shuffle, spatial→channel | -So the gap is **two ops**, both of which work through the spatial cache (they pair with a -matching downsample rather than recomputing structure). Everything else is already -covered by work done for LATO.2. +Both of those are now **done** in the shared core. + +### Correction to the earlier scope + +"the gap is two ops" was accurate about `modules/sparse/` — the sparse *primitives*. It +undercounted the **model blocks**, which the configs revealed: + +| Still needed | Where | +|---|---| +| RoPE positional embedding | all 4 flow models (`pe_mode: "rope"`) — but `rope_phases` ships as a **stored tensor**, so phases are precomputed, not derived | +| `qk_rms_norm` on q and k | all 4 flow models | +| AdaLN modulation (`share_mod: true`) | all 4 flow models | +| `image_attn_mode: "proj"` conditioning | all 4 flow models | +| `SparseConvNeXtBlock3d` | shape_dec, tex_dec | +| `SparseResBlockC2S3d` (channel↔spatial) | shape_dec, tex_dec — uses `SparseSpatial2Channel` | + +Offsetting that, a genuine simplification the tensors revealed: **the four flow models +(~20 GB, the bulk of the download) contain ZERO 5-D tensors.** `ss_flow` and the three +`slat_flow` DiTs are pure transformers — they never touch sparse convolution, so they +need none of the sparse core, just DiT blocks. ## Model surface @@ -48,9 +65,13 @@ Weights: 24.04 GB across 19 files (1.3B DiTs at 512/1024 + shape/tex decoders). ## Status -- [x] Scoped against the shared core — gap is two ops -- [x] Weights downloaded -- [ ] `SparseUpsample`, `SparseSpatial2Channel` -- [ ] Weight converter +- [x] Scoped against the shared core +- [x] Weights downloaded (24 GB) — each ships a sibling `.json` with the exact config, + so unlike LATO.2 there is no architecture to infer +- [x] `upsample` (masked) + `downsample(mode=)` landed in the shared core +- [x] **Weight converter** — decoders remap KRSC→`[K³,in,out]`, dtype preserved, + remap verified a pure permutation. Flow models pass through untouched. +- [ ] DiT blocks: RoPE (stored phases), qk_rms_norm, AdaLN modulation +- [ ] `SparseConvNeXtBlock3d`, `SparseResBlockC2S3d`, `SparseSpatial2Channel` - [ ] Model graphs - [ ] End-to-end diff --git a/pixal3d_mlx/__init__.py b/pixal3d_mlx/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pixal3d_mlx/convert.py b/pixal3d_mlx/convert.py new file mode 100644 index 0000000..8f07791 --- /dev/null +++ b/pixal3d_mlx/convert.py @@ -0,0 +1,108 @@ +"""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()