174 lines
7.9 KiB
Python
174 lines
7.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Convert a TRELLIS.2 texture-SLAT flow DiT checkpoint
|
|
(slat_flow_imgshape2tex_dit_1_3B_{512,1024}_bf16.safetensors) to GGUF.
|
|
|
|
Architecturally identical to the shape SLAT flow (same 1536-d / 30-block /
|
|
12-head sparse DiT, adaLN-Zero share_mod, self+cross attention to DINOv3 cond,
|
|
3D RoPE, QK-RMSNorm) with ONE difference: it is conditioned on the shape via
|
|
`concat_cond`. So in_channels = 64 = 32 tex-noise + 32 shape-SLAT, out = 32.
|
|
At sample time the (normalized) shape SLAT is concatenated onto the noise before
|
|
the input layer each step.
|
|
|
|
Emits the same `trellis2.slat_flow.*` KV as the shape-flow converter so the C++
|
|
loader is shared, plus:
|
|
* trellis2.slat_flow.concat_cond_channels = 32 (0 on the shape flow)
|
|
* norm_mean/std = tex_slat_normalization (output de-norm, 32)
|
|
* concat_norm_mean/std = shape_slat_normalization (concat_cond in-norm, 32)
|
|
both read from texturing_pipeline.json.
|
|
|
|
Usage:
|
|
python convert_tex_flow_to_gguf.py \
|
|
--model models/TRELLIS.2-4B/ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16.safetensors \
|
|
--texturing-json models/TRELLIS.2-4B/texturing_pipeline.json \
|
|
--output ggufs/tex_slat_flow_512_f16.gguf --ftype 1
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import struct
|
|
|
|
import numpy as np
|
|
|
|
GGUF_MAGIC = b"GGUF"
|
|
GGUF_VERSION = 3
|
|
GGUF_ALIGNMENT = 32
|
|
GGML_TYPE_F32 = 0
|
|
GGML_TYPE_F16 = 1
|
|
GGML_TYPE_BF16 = 30
|
|
GGUF_VT_UINT32 = 4
|
|
GGUF_VT_FLOAT32 = 6
|
|
GGUF_VT_BOOL = 7
|
|
GGUF_VT_STRING = 8
|
|
|
|
ARCH = "trellis2-slat-flow" # shared loader with the shape SLAT flow
|
|
KV_PREFIX = "trellis2.slat_flow."
|
|
|
|
|
|
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_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, a=GGUF_ALIGNMENT): return (n + a - 1) // a * a
|
|
|
|
|
|
def choose_type(name, shape, ftype):
|
|
if ftype == 0: return GGML_TYPE_F32
|
|
if ftype == 2: return GGML_TYPE_BF16
|
|
keep_f32 = ("gamma" in name) or ("modulation" in name) or ("norm" in name)
|
|
return GGML_TYPE_F16 if (len(shape) >= 2 and not keep_f32) else GGML_TYPE_F32
|
|
|
|
|
|
def to_bytes(t, ggml_type):
|
|
t = t.detach().cpu().contiguous()
|
|
if ggml_type == GGML_TYPE_F32: return t.float().numpy().astype("<f4", copy=False).tobytes()
|
|
if ggml_type == GGML_TYPE_F16: return t.float().numpy().astype("<f2", copy=False).tobytes()
|
|
if ggml_type == GGML_TYPE_BF16:
|
|
u32 = t.float().numpy().view(np.uint32)
|
|
return (((u32 + 0x7FFF + ((u32 >> 16) & 1)) >> 16).astype("<u2", copy=False)).tobytes()
|
|
raise ValueError(ggml_type)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="Convert TRELLIS.2 texture SLAT flow DiT to GGUF")
|
|
ap.add_argument("--model", required=True)
|
|
ap.add_argument("--config", default=None)
|
|
ap.add_argument("--texturing-json", default=None,
|
|
help="texturing_pipeline.json (tex + shape slat normalization)")
|
|
ap.add_argument("--output", required=True)
|
|
ap.add_argument("--ftype", type=int, default=1, choices=[0, 1, 2])
|
|
args = ap.parse_args()
|
|
|
|
from safetensors.torch import load_file
|
|
|
|
cfg_path = args.config or (os.path.splitext(args.model)[0] + ".json")
|
|
with open(cfg_path) as f:
|
|
a = json.load(f)["args"]
|
|
print(f"model : {args.model}")
|
|
print(f"output: {args.output} (ftype={args.ftype})")
|
|
print(f"arch : in={a['in_channels']} out={a['out_channels']} cond={a['cond_channels']} "
|
|
f"d={a['model_channels']} blks={a['num_blocks']} heads={a['num_heads']}")
|
|
|
|
concat_cond_channels = a["in_channels"] - a["out_channels"]
|
|
assert concat_cond_channels > 0, "tex flow must have concat_cond (in > out)"
|
|
|
|
tj = args.texturing_json or os.path.join(os.path.dirname(os.path.dirname(args.model)),
|
|
"texturing_pipeline.json")
|
|
with open(tj) as f:
|
|
tpa = json.load(f)["args"]
|
|
tex_norm = tpa["tex_slat_normalization"]
|
|
shape_norm = tpa["shape_slat_normalization"]
|
|
assert len(tex_norm["mean"]) == a["out_channels"], "tex_slat_norm must match out_channels"
|
|
assert len(shape_norm["mean"]) == concat_cond_channels, "shape_slat_norm must match concat_cond"
|
|
|
|
rope_freq = a.get("rope_freq", (1.0, 10000.0))
|
|
metadata = [
|
|
kv_str("general.architecture", ARCH),
|
|
kv_str("general.name", os.path.splitext(os.path.basename(args.model))[0]),
|
|
kv_u32("general.file_type", args.ftype),
|
|
kv_u32("general.alignment", GGUF_ALIGNMENT),
|
|
kv_u32(KV_PREFIX + "resolution", a["resolution"]),
|
|
kv_u32(KV_PREFIX + "in_channels", a["in_channels"]),
|
|
kv_u32(KV_PREFIX + "out_channels", a["out_channels"]),
|
|
kv_u32(KV_PREFIX + "concat_cond_channels", concat_cond_channels),
|
|
kv_u32(KV_PREFIX + "model_channels", a["model_channels"]),
|
|
kv_u32(KV_PREFIX + "cond_channels", a["cond_channels"]),
|
|
kv_u32(KV_PREFIX + "num_blocks", a["num_blocks"]),
|
|
kv_u32(KV_PREFIX + "num_heads", a["num_heads"]),
|
|
kv_f32(KV_PREFIX + "mlp_ratio", a["mlp_ratio"]),
|
|
kv_str(KV_PREFIX + "pe_mode", a.get("pe_mode", "rope")),
|
|
kv_bool(KV_PREFIX + "share_mod", a.get("share_mod", False)),
|
|
kv_bool(KV_PREFIX + "qk_rms_norm", a.get("qk_rms_norm", False)),
|
|
kv_bool(KV_PREFIX + "qk_rms_norm_cross", a.get("qk_rms_norm_cross", False)),
|
|
kv_f32(KV_PREFIX + "rope_freq_min", float(rope_freq[0])),
|
|
kv_f32(KV_PREFIX + "rope_freq_base", float(rope_freq[1])),
|
|
]
|
|
# output de-normalization (tex slat) and concat_cond in-normalization (shape slat)
|
|
for i, v in enumerate(tex_norm["mean"]): metadata.append(kv_f32(KV_PREFIX + f"norm_mean.{i}", v))
|
|
for i, v in enumerate(tex_norm["std"]): metadata.append(kv_f32(KV_PREFIX + f"norm_std.{i}", v))
|
|
for i, v in enumerate(shape_norm["mean"]): metadata.append(kv_f32(KV_PREFIX + f"concat_norm_mean.{i}", v))
|
|
for i, v in enumerate(shape_norm["std"]): metadata.append(kv_f32(KV_PREFIX + f"concat_norm_std.{i}", v))
|
|
|
|
print("loading state_dict...")
|
|
sd = load_file(args.model)
|
|
tensors, counts = [], {GGML_TYPE_F32: 0, GGML_TYPE_F16: 0, GGML_TYPE_BF16: 0}
|
|
for name in sorted(sd.keys()):
|
|
t = sd[name]; shape = tuple(t.shape)
|
|
gtype = choose_type(name, shape, args.ftype)
|
|
raw = to_bytes(t, gtype)
|
|
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]}, bf16={counts[GGML_TYPE_BF16]})")
|
|
|
|
header = bytearray(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, offsets = 0, []
|
|
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()
|