305 lines
10 KiB
Python
305 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
# convert.py: safetensors to GGUF for ACE-Step (LM, DiT, TextEncoder, VAE)
|
|
# Reads from checkpoints/, writes GGUF to models/
|
|
# Each GGUF is self-contained: weights + config + tokenizer + silence_latent
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import struct
|
|
import zipfile
|
|
import numpy as np
|
|
import gguf
|
|
|
|
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
CHECKPOINT_DIR = os.path.join(SCRIPT_DIR, "checkpoints")
|
|
OUTPUT_DIR = os.path.join(SCRIPT_DIR, "models")
|
|
|
|
BF16 = gguf.GGMLQuantizationType.BF16
|
|
|
|
def log(tag, msg):
|
|
print("[%s] %s" % (tag, msg), file=sys.stderr, flush=True)
|
|
|
|
# Safetensors reader
|
|
def read_sf_header(path):
|
|
with open(path, "rb") as f:
|
|
n = struct.unpack("<Q", f.read(8))[0]
|
|
meta = json.loads(f.read(n))
|
|
meta.pop("__metadata__", None)
|
|
return meta, 8 + n
|
|
|
|
def find_sf_files(model_dir):
|
|
"""Return list of safetensors paths (single, sharded, or diffusers VAE)."""
|
|
single = os.path.join(model_dir, "model.safetensors")
|
|
if os.path.exists(single):
|
|
return [single]
|
|
index = os.path.join(model_dir, "model.safetensors.index.json")
|
|
if os.path.exists(index):
|
|
with open(index, "r", encoding="utf-8") as f:
|
|
idx = json.load(f)
|
|
shards = sorted(set(idx["weight_map"].values()))
|
|
return [os.path.join(model_dir, s) for s in shards]
|
|
diffusers = os.path.join(model_dir, "diffusion_pytorch_model.safetensors")
|
|
if os.path.exists(diffusers):
|
|
return [diffusers]
|
|
return []
|
|
|
|
# Model classification
|
|
ARCHS = {
|
|
"lm": "acestep-lm",
|
|
"dit": "acestep-dit",
|
|
"text-enc": "acestep-text-enc",
|
|
"vae": "acestep-vae",
|
|
}
|
|
|
|
def classify(name):
|
|
if name.startswith("acestep-5Hz-lm"):
|
|
return "lm"
|
|
if name.startswith("acestep-v15"):
|
|
return "dit"
|
|
if name.startswith("Qwen3-Embedding"):
|
|
return "text-enc"
|
|
if name == "vae":
|
|
return "vae"
|
|
return None
|
|
|
|
# GGUF metadata from config.json
|
|
def add_metadata(w, cfg, model_type):
|
|
if "num_hidden_layers" in cfg:
|
|
w.add_block_count(cfg["num_hidden_layers"])
|
|
if "hidden_size" in cfg:
|
|
w.add_embedding_length(cfg["hidden_size"])
|
|
if "intermediate_size" in cfg:
|
|
w.add_feed_forward_length(cfg["intermediate_size"])
|
|
if "num_attention_heads" in cfg:
|
|
w.add_head_count(cfg["num_attention_heads"])
|
|
if "num_key_value_heads" in cfg:
|
|
w.add_head_count_kv(cfg["num_key_value_heads"])
|
|
if "head_dim" in cfg:
|
|
w.add_key_length(cfg["head_dim"])
|
|
if "vocab_size" in cfg:
|
|
w.add_vocab_size(cfg["vocab_size"])
|
|
if "max_position_embeddings" in cfg:
|
|
w.add_context_length(cfg["max_position_embeddings"])
|
|
if "rms_norm_eps" in cfg:
|
|
w.add_layer_norm_rms_eps(cfg["rms_norm_eps"])
|
|
rope = cfg.get("rope_theta")
|
|
if rope:
|
|
w.add_rope_freq_base(float(rope))
|
|
|
|
if model_type == "lm":
|
|
if cfg.get("tie_word_embeddings"):
|
|
w.add_bool("acestep.tie_word_embeddings", True)
|
|
|
|
if model_type == "dit":
|
|
for key in [
|
|
"in_channels", "audio_acoustic_hidden_dim", "patch_size",
|
|
"sliding_window", "fsq_dim", "text_hidden_dim", "timbre_hidden_dim",
|
|
"num_lyric_encoder_hidden_layers", "num_timbre_encoder_hidden_layers",
|
|
"num_audio_decoder_hidden_layers", "num_attention_pooler_hidden_layers",
|
|
]:
|
|
if key in cfg:
|
|
w.add_uint32("acestep.%s" % key, cfg[key])
|
|
# XL models have separate encoder dimensions (2B models omit these)
|
|
for key in [
|
|
"encoder_hidden_size", "encoder_intermediate_size",
|
|
"encoder_num_attention_heads", "encoder_num_key_value_heads",
|
|
]:
|
|
if key in cfg:
|
|
w.add_uint32("acestep.%s" % key, cfg[key])
|
|
if cfg.get("is_turbo"):
|
|
w.add_bool("acestep.is_turbo", True)
|
|
levels = cfg.get("fsq_input_levels")
|
|
if levels:
|
|
w.add_array("acestep.fsq_input_levels", levels)
|
|
|
|
w.add_string("acestep.config_json", json.dumps(cfg, separators=(",", ":")))
|
|
|
|
# Tensor packing from safetensors
|
|
def add_tensors_from_sf(w, sf_path, tag, model_type):
|
|
meta, hdr_size = read_sf_header(sf_path)
|
|
names = sorted(meta.keys())
|
|
with open(sf_path, "rb") as f:
|
|
count = 0
|
|
total = 0
|
|
|
|
for name in names:
|
|
info = meta[name]
|
|
|
|
# normalize: some upstream checkpoints omit the "model." prefix
|
|
if model_type == "lm" and not name.startswith("model."):
|
|
name = "model." + name
|
|
|
|
dtype_str = info["dtype"]
|
|
shape = info["shape"]
|
|
off0, off1 = info["data_offsets"]
|
|
nbytes = off1 - off0
|
|
|
|
f.seek(hdr_size + off0)
|
|
raw = f.read(nbytes)
|
|
|
|
if dtype_str == "BF16":
|
|
arr = np.frombuffer(raw, dtype=np.uint16).reshape(shape)
|
|
w.add_tensor(name, arr, raw_dtype=BF16)
|
|
elif dtype_str == "F16":
|
|
arr = np.frombuffer(raw, dtype=np.float16).reshape(shape)
|
|
w.add_tensor(name, arr)
|
|
elif dtype_str == "F32":
|
|
# convert F32 to BF16: truncate lower 16 mantissa bits
|
|
arr = np.frombuffer(raw, dtype=np.uint32).reshape(shape)
|
|
arr = (arr >> 16).astype(np.uint16)
|
|
w.add_tensor(name, arr, raw_dtype=BF16)
|
|
nbytes = nbytes // 2 # actual stored size
|
|
else:
|
|
log(tag, " skip %s: dtype %s" % (name, dtype_str))
|
|
continue
|
|
|
|
count += 1
|
|
total += nbytes
|
|
|
|
return count, total
|
|
|
|
# silence_latent.pt reader (replaces pt2bin C++ tool)
|
|
# PyTorch .pt is a ZIP with entry "*/data/0" containing f32 [64, 15000]
|
|
# We transpose to [15000, 64] (ggml layout: 64 contiguous per frame)
|
|
def read_silence_latent(model_dir):
|
|
pt_path = os.path.join(model_dir, "silence_latent.pt")
|
|
if not os.path.exists(pt_path):
|
|
return None
|
|
with zipfile.ZipFile(pt_path) as z:
|
|
for entry in z.namelist():
|
|
if entry.endswith("/data/0"):
|
|
raw = z.read(entry)
|
|
src = np.frombuffer(raw, dtype=np.float32).reshape(64, 15000)
|
|
return src.T.copy()
|
|
return None
|
|
|
|
# BPE tokenizer embedding (vocab.json + merges.txt -> GGUF KV)
|
|
def add_bpe_tokenizer(w, model_dir, tag):
|
|
vocab_path = os.path.join(model_dir, "vocab.json")
|
|
merges_path = os.path.join(model_dir, "merges.txt")
|
|
if not os.path.exists(vocab_path) or not os.path.exists(merges_path):
|
|
return False
|
|
|
|
with open(vocab_path, "r", encoding="utf-8") as f:
|
|
vocab = json.load(f)
|
|
tokens = [""] * len(vocab)
|
|
for tok_str, tok_id in vocab.items():
|
|
if 0 <= tok_id < len(tokens):
|
|
tokens[tok_id] = tok_str
|
|
|
|
with open(merges_path, "r", encoding="utf-8") as f:
|
|
merges = []
|
|
for line in f:
|
|
line = line.rstrip("\n\r")
|
|
if not line:
|
|
continue
|
|
if line.startswith("#version:"):
|
|
continue
|
|
merges.append(line)
|
|
|
|
w.add_tokenizer_model("gpt2")
|
|
w.add_token_list(tokens)
|
|
w.add_token_merges(merges)
|
|
|
|
log(tag, " tokenizer: %d vocab, %d merges" % (len(tokens), len(merges)))
|
|
return True
|
|
|
|
# Main conversion
|
|
def convert_model(name, model_dir, output_path, model_type):
|
|
tag = "GGUF"
|
|
cfg_path = os.path.join(model_dir, "config.json")
|
|
if not os.path.exists(cfg_path):
|
|
log(tag, "skip %s: no config.json" % name)
|
|
return False
|
|
|
|
with open(cfg_path, "r", encoding="utf-8") as f:
|
|
cfg = json.load(f)
|
|
sf_files = find_sf_files(model_dir)
|
|
if not sf_files:
|
|
log(tag, "skip %s: no safetensors" % name)
|
|
return False
|
|
|
|
arch = ARCHS[model_type]
|
|
log(tag, "%s (%s, %d shard%s) -> %s" % (
|
|
name, arch, len(sf_files), "" if len(sf_files) == 1 else "s",
|
|
os.path.basename(output_path)))
|
|
|
|
w = gguf.GGUFWriter(output_path, arch, use_temp_file=True)
|
|
w.add_name(name)
|
|
add_metadata(w, cfg, model_type)
|
|
|
|
# BPE tokenizer for LM and text encoder
|
|
if model_type in ("lm", "text-enc"):
|
|
add_bpe_tokenizer(w, model_dir, tag)
|
|
|
|
# Model weights
|
|
n_tensors = 0
|
|
n_bytes = 0
|
|
for sf in sf_files:
|
|
c, b = add_tensors_from_sf(w, sf, tag, model_type)
|
|
n_tensors += c
|
|
n_bytes += b
|
|
if len(sf_files) > 1:
|
|
log(tag, " %s: %d tensors" % (os.path.basename(sf), c))
|
|
|
|
# silence_latent for DiT (read .pt, transpose, embed as f32 tensor)
|
|
if model_type == "dit":
|
|
sl = read_silence_latent(model_dir)
|
|
if sl is not None:
|
|
w.add_tensor("silence_latent", sl)
|
|
n_tensors += 1
|
|
n_bytes += sl.nbytes
|
|
log(tag, " silence_latent: [%d, %d] f32 (%.1f MB)" % (
|
|
sl.shape[0], sl.shape[1], sl.nbytes / (1 << 20)))
|
|
else:
|
|
log(tag, " WARNING: no silence_latent.pt found")
|
|
|
|
log(tag, " total: %d tensors, %.1f GB" % (n_tensors, n_bytes / (1 << 30)))
|
|
|
|
w.write_header_to_file()
|
|
w.write_kv_data_to_file()
|
|
w.write_tensors_to_file(progress=True)
|
|
w.close()
|
|
|
|
out_mb = os.path.getsize(output_path) / (1 << 20)
|
|
log(tag, " wrote %.0f MB -> %s" % (out_mb, output_path))
|
|
return True
|
|
|
|
def main():
|
|
if not os.path.isdir(CHECKPOINT_DIR):
|
|
log("GGUF", "checkpoints/ not found, run checkpoints.sh first")
|
|
sys.exit(1)
|
|
|
|
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
|
|
entries = sorted(os.listdir(CHECKPOINT_DIR))
|
|
converted = 0
|
|
skipped = []
|
|
|
|
for name in entries:
|
|
model_dir = os.path.join(CHECKPOINT_DIR, name)
|
|
if not os.path.isdir(model_dir):
|
|
continue
|
|
|
|
model_type = classify(name)
|
|
if model_type is None:
|
|
skipped.append(name)
|
|
continue
|
|
|
|
output_path = os.path.join(OUTPUT_DIR, "%s-BF16.gguf" % name)
|
|
if os.path.exists(output_path):
|
|
log("GGUF", "skip %s: %s exists" % (name, os.path.basename(output_path)))
|
|
converted += 1
|
|
continue
|
|
|
|
if convert_model(name, model_dir, output_path, model_type):
|
|
converted += 1
|
|
|
|
if skipped:
|
|
log("GGUF", "skipped (unknown): %s" % ", ".join(skipped))
|
|
log("GGUF", "done: %d model(s) in %s" % (converted, OUTPUT_DIR))
|
|
|
|
if __name__ == "__main__":
|
|
main()
|