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

@@ -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.

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

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()