Initial release

convert_shape_dec_to_gguf.py

@@ -0,0 +1,167 @@
+#!/usr/bin/env python3
+"""
+Convert the TRELLIS.2 shape-SLAT VAE decoder checkpoint
+(shape_dec_next_dc_f16c32_fp16.safetensors) to a GGUF file for trellis2.cpp.
+
+This is FlexiDualGridVaeDecoder (a SparseUnetVaeDecoder): a sparse ConvNeXt
+U-Net decoder that turns the 32-channel structured latent on active voxels
+into 7 channels per voxel at 16x the input resolution (dual-vertex offset,
+per-axis intersection flags, quad split weight) for flexible-dual-grid mesh
+extraction. Architecture (config shape_dec_next_dc_f16c32_fp16.json):
+
+    model_channels [1024, 512, 256, 128, 64], num_blocks [4, 16, 8, 4, 0],
+    SparseConvNeXtBlock3d blocks, SparseResBlockC2S3d up-blocks, out 7.
+
+Checkpoint conv weights are stored in FlexGEMM layout [Co, kD, kH, kW, Ci]
+(permuted at module init, then saved). ggml tensors are 4-D, so the kernel
+axes merge: we reshape to [Co, kD*kH*kW, Ci] (identical bytes) and the C++
+side slices per kernel offset into [Ci, Co] GEMM operands.
+
+Usage:
+    python convert_shape_dec_to_gguf.py --output shape_dec.gguf --ftype 1
+
+ftype: 0 = f32 (lossless upcast, validation), 1 = f16 (default).
+"""
+
+import argparse
+import json
+import os
+import struct
+import sys
+
+import numpy as np
+
+GGUF_MAGIC = b"GGUF"
+GGUF_VERSION = 3
+GGUF_ALIGNMENT = 32
+
+GGML_TYPE_F32 = 0
+GGML_TYPE_F16 = 1
+
+GGUF_VT_UINT32 = 4
+GGUF_VT_FLOAT32 = 6
+GGUF_VT_STRING = 8
+
+ARCH = "trellis2-shape-dec"
+KV_PREFIX = "trellis2.shape_dec."
+
+
+def _gguf_str(s):
+    b = s.encode("utf-8")
+    return struct.pack("<Q", len(b)) + b
+
+
+def _kv(key, vtype, payload):
+    return _gguf_str(key) + struct.pack("<I", vtype) + payload
+
+
+def kv_u32(key, v):  return _kv(key, GGUF_VT_UINT32,  struct.pack("<I", int(v)))
+def kv_f32(key, v):  return _kv(key, GGUF_VT_FLOAT32, struct.pack("<f", float(v)))
+def kv_str(key, v):  return _kv(key, GGUF_VT_STRING,  _gguf_str(str(v)))
+
+
+def _align(n, a=GGUF_ALIGNMENT):
+    return (n + a - 1) // a * a
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Convert TRELLIS.2 shape decoder to GGUF")
+    ap.add_argument("--model", default=os.path.join(os.path.dirname(__file__),
+                    "models", "TRELLIS.2-4B", "ckpts", "shape_dec_next_dc_f16c32_fp16.safetensors"))
+    ap.add_argument("--config", default=None)
+    ap.add_argument("--output", default="shape_dec.gguf")
+    ap.add_argument("--ftype", type=int, default=1, choices=[0, 1])
+    args = ap.parse_args()
+
+    import torch
+    from safetensors.torch import load_file
+
+    cfg_path = args.config or (os.path.splitext(args.model)[0] + ".json")
+    with open(cfg_path) as f:
+        cfg = json.load(f)
+    a = cfg["args"]
+    channels = a["model_channels"]
+    nblocks = a["num_blocks"]
+    print(f"model : {args.model}")
+    print(f"output: {args.output}  (ftype={args.ftype})")
+    print(f"arch  : channels={channels} blocks={nblocks} latent={a['latent_channels']}")
+
+    metadata = [
+        kv_str("general.architecture", ARCH),
+        kv_str("general.name", "shape_dec_next_dc_f16c32_fp16"),
+        kv_u32("general.file_type", args.ftype),
+        kv_u32("general.alignment", GGUF_ALIGNMENT),
+        kv_u32(KV_PREFIX + "latent_channels", a["latent_channels"]),
+        kv_u32(KV_PREFIX + "out_channels",    7),
+        kv_u32(KV_PREFIX + "n_levels",        len(channels)),
+        kv_f32(KV_PREFIX + "norm_eps",        1e-6),
+        kv_f32(KV_PREFIX + "voxel_margin",    0.5),
+    ]
+    for i, ch in enumerate(channels):
+        metadata.append(kv_u32(KV_PREFIX + f"channels.{i}", ch))
+    for i, nb in enumerate(nblocks):
+        metadata.append(kv_u32(KV_PREFIX + f"num_blocks.{i}", nb))
+
+    print("loading state_dict...")
+    sd = load_file(args.model)
+
+    tensors = []
+    counts = {GGML_TYPE_F32: 0, GGML_TYPE_F16: 0}
+    for name in sorted(sd.keys()):
+        t = sd[name].float()
+        arr = t.numpy().astype(np.float32)
+        shape = tuple(arr.shape)
+        if arr.ndim == 5:  # FlexGEMM conv weight [Co, kD, kH, kW, Ci] -> [Co, 27, Ci]
+            Co, kD, kH, kW, Ci = shape
+            arr = np.ascontiguousarray(arr).reshape(Co, kD * kH * kW, Ci)
+            shape = arr.shape
+        gtype = GGML_TYPE_F32
+        if args.ftype == 1 and len(shape) >= 2:
+            gtype = GGML_TYPE_F16
+        raw = (arr.astype("<f2") if gtype == GGML_TYPE_F16 else arr.astype("<f4")).tobytes()
+        dims = list(reversed(shape)) if len(shape) > 0 else [1]
+        tensors.append((name, gtype, dims, raw))
+        counts[gtype] += 1
+
+    print(f"tensors: {len(tensors)}  (f32={counts[GGML_TYPE_F32]}, f16={counts[GGML_TYPE_F16]})")
+
+    header = bytearray()
+    header += GGUF_MAGIC
+    header += struct.pack("<I", GGUF_VERSION)
+    header += struct.pack("<Q", len(tensors))
+    header += struct.pack("<Q", len(metadata))
+    for m in metadata:
+        header += m
+
+    infos = bytearray()
+    offset = 0
+    offsets = []
+    for name, gtype, dims, raw in tensors:
+        offsets.append(offset)
+        offset = _align(offset + len(raw))
+    for (name, gtype, dims, raw), off in zip(tensors, offsets):
+        infos += _gguf_str(name)
+        infos += struct.pack("<I", len(dims))
+        for d in dims:
+            infos += struct.pack("<Q", int(d))
+        infos += struct.pack("<I", gtype)
+        infos += struct.pack("<Q", off)
+
+    pre_data = len(header) + len(infos)
+    pad0 = _align(pre_data) - pre_data
+
+    with open(args.output, "wb") as fout:
+        fout.write(header)
+        fout.write(infos)
+        fout.write(b"\x00" * pad0)
+        for (name, gtype, dims, raw), off in zip(tensors, offsets):
+            fout.write(raw)
+            pad = _align(len(raw)) - len(raw)
+            if pad:
+                fout.write(b"\x00" * pad)
+
+    print(f"wrote {args.output}  ({os.path.getsize(args.output):,} bytes)")
+
+
+if __name__ == "__main__":
+    main()