Initial release
This commit is contained in:
civ
@@ -0,0 +1,216 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user