217 lines
8.6 KiB
Python
217 lines
8.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Convert the TRELLIS.2 sparse-structure DECODER checkpoint
|
|
(ss_dec_conv3d_16l8_fp16.safetensors) to a GGUF file for trellis2.cpp.
|
|
|
|
This is the stage-1 decoder D_S: a dense 3D-conv ResNet that turns the
|
|
sparse-structure latent z_s ([8, 16, 16, 16]) into an occupancy logit grid at
|
|
64^3 ([1, 64, 64, 64]). It upsamples 16 -> 32 -> 64 with two pixel-shuffle
|
|
blocks. Architecture (config ss_dec_conv3d_16l8_fp16.json):
|
|
|
|
SparseStructureDecoder(out_channels=1, latent_channels=8,
|
|
num_res_blocks=2, num_res_blocks_middle=2,
|
|
channels=[512, 128, 32], norm_type="layer")
|
|
|
|
Like convert_ss_flow_to_gguf.py this is self-contained (safetensors + numpy +
|
|
torch) and writes a standard GGUF v3 file read back by ggml's
|
|
gguf_init_from_file(). Hyperparameters travel as KV metadata under the
|
|
`trellis2.ss_dec.*` namespace; tensors keep their original checkpoint names.
|
|
|
|
Conv3d weights are 5-D in PyTorch ([OC, IC, kD, kH, kW]). ggml tensors are 4-D,
|
|
and ggml_conv_3d_direct wants the kernel as ne = [kW, kH, kD, IC*OC] with the
|
|
merged channel index packed oc*IC + ic. A C-contiguous [OC, IC, kD, kH, kW]
|
|
array reshaped to [OC*IC, kD, kH, kW] is exactly that packing (and identical
|
|
bytes), so we just reshape before writing — no permute, no data movement.
|
|
|
|
Usage:
|
|
python convert_ss_dec_to_gguf.py --output ss_dec.gguf --ftype 0
|
|
# --model/--config default to the microsoft/TRELLIS-image-large HF snapshot.
|
|
|
|
ftype: 0 = f32 (lossless upcast from the fp16 checkpoint; use for validation),
|
|
1 = f16 (default; conv weight matrices f16, norms/biases f32).
|
|
"""
|
|
|
|
import argparse
|
|
import glob
|
|
import json
|
|
import os
|
|
import struct
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
# ── GGUF / GGML constants (must match the bundled ggml) ──────────────────────
|
|
GGUF_MAGIC = b"GGUF"
|
|
GGUF_VERSION = 3
|
|
GGUF_ALIGNMENT = 32
|
|
|
|
GGML_TYPE_F32 = 0
|
|
GGML_TYPE_F16 = 1
|
|
|
|
GGUF_VT_UINT32 = 4
|
|
GGUF_VT_INT32 = 5
|
|
GGUF_VT_FLOAT32 = 6
|
|
GGUF_VT_BOOL = 7
|
|
GGUF_VT_STRING = 8
|
|
|
|
ARCH = "trellis2-ss-dec"
|
|
KV_PREFIX = "trellis2.ss_dec."
|
|
|
|
DEFAULT_SNAPSHOT = os.path.expanduser(
|
|
"~/.cache/huggingface/hub/models--microsoft--TRELLIS-image-large/snapshots/*/ckpts"
|
|
)
|
|
CKPT_STEM = "ss_dec_conv3d_16l8_fp16"
|
|
|
|
|
|
# ── GGUF writer (minimal, v3) ────────────────────────────────────────────────
|
|
def _gguf_str(s: str) -> bytes:
|
|
b = s.encode("utf-8")
|
|
return struct.pack("<Q", len(b)) + b
|
|
|
|
|
|
def _kv(key: str, vtype: int, payload: bytes) -> bytes:
|
|
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_bool(key, v): return _kv(key, GGUF_VT_BOOL, struct.pack("<?", bool(v)))
|
|
def kv_str(key, v): return _kv(key, GGUF_VT_STRING, _gguf_str(str(v)))
|
|
|
|
|
|
def _align(n: int, a: int = GGUF_ALIGNMENT) -> int:
|
|
return (n + a - 1) // a * a
|
|
|
|
|
|
def choose_type(shape, ftype: int) -> int:
|
|
"""f32 always for ftype 0; for ftype 1 the big conv weight matrices are f16
|
|
while everything 1-D (biases, norm weight/bias) stays f32."""
|
|
if ftype == 0:
|
|
return GGML_TYPE_F32
|
|
return GGML_TYPE_F16 if len(shape) >= 2 else GGML_TYPE_F32
|
|
|
|
|
|
def to_bytes(arr_f32: np.ndarray, ggml_type: int) -> bytes:
|
|
if ggml_type == GGML_TYPE_F32:
|
|
return arr_f32.astype("<f4", copy=False).tobytes()
|
|
if ggml_type == GGML_TYPE_F16:
|
|
return arr_f32.astype("<f2", copy=False).tobytes()
|
|
raise ValueError(f"unhandled ggml type {ggml_type}")
|
|
|
|
|
|
def resolve_paths(args):
|
|
model = args.model
|
|
if model is None:
|
|
hits = sorted(glob.glob(os.path.join(DEFAULT_SNAPSHOT, CKPT_STEM + ".safetensors")))
|
|
if not hits:
|
|
sys.exit("error: --model not given and no TRELLIS-image-large snapshot in HF cache")
|
|
model = hits[-1]
|
|
cfg = args.config or (os.path.splitext(model)[0] + ".json")
|
|
out = args.output or (CKPT_STEM + ".gguf")
|
|
return model, cfg, out
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="Convert TRELLIS.2 SS decoder to GGUF")
|
|
ap.add_argument("--model", default=None, help="path to ...ss_dec...safetensors (default: HF cache)")
|
|
ap.add_argument("--config", default=None, help="path to matching .json (default: alongside model)")
|
|
ap.add_argument("--output", default=None, help="output .gguf (default: <stem>.gguf)")
|
|
ap.add_argument("--ftype", type=int, default=1, choices=[0, 1],
|
|
help="0=f32 (lossless, for validation), 1=f16 (default)")
|
|
args = ap.parse_args()
|
|
|
|
from safetensors.numpy import load_file
|
|
|
|
model_path, cfg_path, out_path = resolve_paths(args)
|
|
print(f"model : {model_path}")
|
|
print(f"config: {cfg_path}")
|
|
print(f"output: {out_path} (ftype={args.ftype})")
|
|
|
|
with open(cfg_path) as f:
|
|
cfg = json.load(f)
|
|
a = cfg["args"]
|
|
channels = a["channels"]
|
|
print(f"arch : {cfg.get('name')} channels={channels}, "
|
|
f"res_blocks={a['num_res_blocks']}, mid={a['num_res_blocks_middle']}, "
|
|
f"latent={a['latent_channels']}, out={a['out_channels']}")
|
|
|
|
# ── KV metadata ──────────────────────────────────────────────────────────
|
|
metadata = [
|
|
kv_str("general.architecture", ARCH),
|
|
kv_str("general.name", CKPT_STEM),
|
|
kv_u32("general.file_type", args.ftype),
|
|
kv_u32("general.alignment", GGUF_ALIGNMENT),
|
|
kv_u32(KV_PREFIX + "out_channels", a["out_channels"]),
|
|
kv_u32(KV_PREFIX + "latent_channels", a["latent_channels"]),
|
|
kv_u32(KV_PREFIX + "num_res_blocks", a["num_res_blocks"]),
|
|
kv_u32(KV_PREFIX + "num_res_blocks_middle", a["num_res_blocks_middle"]),
|
|
kv_u32(KV_PREFIX + "n_levels", len(channels)),
|
|
kv_str(KV_PREFIX + "norm_type", a.get("norm_type", "layer")),
|
|
kv_f32(KV_PREFIX + "norm_eps", 1e-5),
|
|
]
|
|
for i, ch in enumerate(channels):
|
|
metadata.append(kv_u32(KV_PREFIX + f"channels.{i}", ch))
|
|
|
|
# ── tensors ──────────────────────────────────────────────────────────────
|
|
print("loading state_dict...")
|
|
sd = load_file(model_path) # numpy arrays (fp16/fp32 as stored)
|
|
|
|
tensors = [] # (name, ggml_type, dims_ggml_order, raw_bytes)
|
|
counts = {GGML_TYPE_F32: 0, GGML_TYPE_F16: 0}
|
|
for name in sorted(sd.keys()):
|
|
arr = np.asarray(sd[name]).astype(np.float32)
|
|
shape = tuple(arr.shape)
|
|
if arr.ndim == 5: # Conv3d weight [OC, IC, kD, kH, kW] -> [OC*IC, kD, kH, kW]
|
|
OC, IC, kD, kH, kW = shape
|
|
arr = np.ascontiguousarray(arr).reshape(OC * IC, kD, kH, kW)
|
|
shape = arr.shape
|
|
gtype = choose_type(shape, args.ftype)
|
|
raw = to_bytes(np.ascontiguousarray(arr), gtype)
|
|
dims = list(reversed(shape)) if len(shape) > 0 else [1] # ggml ne[] order
|
|
tensors.append((name, gtype, dims, raw))
|
|
counts[gtype] += 1
|
|
|
|
print(f"tensors: {len(tensors)} (f32={counts[GGML_TYPE_F32]}, f16={counts[GGML_TYPE_F16]})")
|
|
|
|
# ── assemble header + infos, compute aligned data offsets ────────────────
|
|
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(out_path, "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 {out_path} ({os.path.getsize(out_path):,} bytes)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|