176 lines
7.9 KiB
Python
176 lines
7.9 KiB
Python
#!/usr/bin/env python3
|
|
# convert-sa3.py: Stable Audio 3 checkpoints -> GGUF for the StableStep GGML backend.
|
|
#
|
|
# Produces four GGUFs into models/:
|
|
# sa3-dit-BF16.gguf arch "sa3-dit" from stable-audio-3-medium (model.* keys)
|
|
# sa3-same-enc-F16.gguf arch "sa3-same-enc" from pretransform.model.* (encoder side)
|
|
# sa3-same-dec-F16.gguf arch "sa3-same-dec" from pretransform.model.* (decoder side)
|
|
# sa3-text-enc-BF16.gguf arch "sa3-t5gemma" from the t5gemma-b-b-ul2 subfolder
|
|
#
|
|
# The SAME pair is stored F16 (not BF16): the decoder's sinusoidal FF blocks
|
|
# amplify per-weight rounding noise across its 12 layers (see write_sa3_gguf).
|
|
#
|
|
# Tensor policy: >=2D weights -> BF16; 1D tensors (norms, biases, scales) -> F32
|
|
# (precision finding from the ONNX leg: this model's norm/timestep paths are
|
|
# fp32-sensitive — measured cosine 0.966 with blanket fp16 vs 0.9995 scoped).
|
|
# Tensor names are the source names minus the strip prefix; the C++ graph
|
|
# builders consume them as-is. The full model_config.json is embedded verbatim
|
|
# under metadata key "sa3.config_json" (the C++ side parses what it needs).
|
|
#
|
|
# Runs in the StableAudio3 uv venv:
|
|
# cd d:/Ace-Step-Latest/StableAudio3
|
|
# uv run --with gguf python d:/Ace-Step-Latest/hot-step-cpp/engine/convert-sa3.py
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
import numpy as np
|
|
import gguf
|
|
|
|
sys.path.insert(0, r"d:/Ace-Step-Latest/StableAudio3")
|
|
from huggingface_hub import hf_hub_download
|
|
from safetensors import safe_open
|
|
|
|
OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
|
|
REPO = "stabilityai/stable-audio-3-medium"
|
|
|
|
|
|
def log(msg):
|
|
print(f"[convert-sa3] {msg}", file=sys.stderr, flush=True)
|
|
|
|
|
|
def to_np(t):
|
|
import torch
|
|
if t.dtype == torch.bfloat16:
|
|
return t.float().numpy()
|
|
return t.numpy()
|
|
|
|
|
|
def write_sa3_gguf(out_path, arch, tensors, config_json, extra_meta=None, half="bf16"):
|
|
"""tensors: list of (name, np.float32 array). >=2D stored in `half`
|
|
("bf16" or "f16"), 1D tensors (norms, biases, scales) stored F32.
|
|
|
|
half="f16" is used for the SAME autoencoder halves: the decoder's
|
|
sinusoidal FF layers amplify weight rounding noise layer over layer
|
|
(bf16 ~0.4% rel error -> parity cosine 0.9987 < 0.999; f16 ~0.05%
|
|
passes). Same file size either way; weight magnitudes are far inside
|
|
f16 range."""
|
|
w = gguf.GGUFWriter(out_path, arch)
|
|
w.add_string("sa3.config_json", config_json)
|
|
for k, v in (extra_meta or {}).items():
|
|
w.add_string(k, v)
|
|
import torch
|
|
n_half = n_f32 = 0
|
|
for name, arr in tensors:
|
|
arr = np.ascontiguousarray(arr, dtype=np.float32)
|
|
if arr.size == 0:
|
|
# e.g. bottleneck.noise_scaling_factor (1, 0, 1) when
|
|
# noise_augment_dim == 0. ggml's gguf reader hits an integer
|
|
# divide-by-zero on ne==0 tensors, and the C++ side never reads
|
|
# them — drop.
|
|
log(f" skipping zero-element tensor {name} {arr.shape}")
|
|
continue
|
|
if arr.ndim >= 2:
|
|
# raw_dtype does NOT convert — it labels. Convert to 16-bit bytes
|
|
# explicitly (uint16 view keeps the logical shape).
|
|
if half == "f16":
|
|
h = torch.from_numpy(arr).to(torch.float16).view(torch.uint16).numpy()
|
|
w.add_tensor(name, h, raw_dtype=gguf.GGMLQuantizationType.F16)
|
|
else:
|
|
h = torch.from_numpy(arr).to(torch.bfloat16).view(torch.uint16).numpy()
|
|
w.add_tensor(name, h, raw_dtype=gguf.GGMLQuantizationType.BF16)
|
|
n_half += 1
|
|
else:
|
|
w.add_tensor(name, arr) # F32
|
|
n_f32 += 1
|
|
w.write_header_to_file()
|
|
w.write_kv_data_to_file()
|
|
w.write_tensors_to_file()
|
|
w.close()
|
|
size = os.path.getsize(out_path) / 1e9
|
|
log(f"{os.path.basename(out_path)}: {n_half} {half.upper()} + {n_f32} F32 tensors, {size:.2f} GB")
|
|
|
|
|
|
def main():
|
|
# Local-checkpoint support (e.g. LoRA-merged models): --ckpt/--config
|
|
# override the HF download; --out-dir redirects output; --dit-only skips
|
|
# the SAME + T5Gemma GGUFs (unchanged when only the DiT was fine-tuned).
|
|
import argparse
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--ckpt", default=None, help="Local model.safetensors (default: HF download)")
|
|
ap.add_argument("--config", default=None, help="Local model_config.json (default: HF download)")
|
|
ap.add_argument("--out-dir", default=OUTPUT_DIR)
|
|
ap.add_argument("--dit-only", action="store_true")
|
|
args = ap.parse_args()
|
|
out_dir = args.out_dir
|
|
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
cfg_path = args.config or hf_hub_download(REPO, "model_config.json")
|
|
ckpt_path = args.ckpt or hf_hub_download(REPO, "model.safetensors")
|
|
with open(cfg_path) as f:
|
|
config_json = f.read()
|
|
|
|
# ── Split the combined checkpoint by prefix ─────────────────────────
|
|
dit_tensors, enc_tensors, dec_tensors = [], [], []
|
|
with safe_open(ckpt_path, framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
if key.startswith("model."):
|
|
dit_tensors.append((key[len("model."):], to_np(f.get_tensor(key))))
|
|
elif key.startswith("pretransform.model."):
|
|
sub = key[len("pretransform.model."):]
|
|
# AudioAutoencoder members: encoder.*, decoder.*, bottleneck.*,
|
|
# pretransform.* (patched — no weights). Bottleneck params go to BOTH
|
|
# (encoder applies scale/bias+running_std, decoder inverts).
|
|
if sub.startswith("encoder."):
|
|
enc_tensors.append((sub, to_np(f.get_tensor(key))))
|
|
elif sub.startswith("decoder."):
|
|
dec_tensors.append((sub, to_np(f.get_tensor(key))))
|
|
elif sub.startswith("bottleneck."):
|
|
t = to_np(f.get_tensor(key))
|
|
enc_tensors.append((sub, t))
|
|
dec_tensors.append((sub, t))
|
|
# conditioner.* (learned padding, seconds embedder) rides with the DiT
|
|
# gguf — small and needed by the same backend module.
|
|
elif key.startswith("conditioner."):
|
|
dit_tensors.append((key, to_np(f.get_tensor(key))))
|
|
|
|
write_sa3_gguf(os.path.join(out_dir, "sa3-dit-BF16.gguf"),
|
|
"sa3-dit", dit_tensors, config_json)
|
|
if args.dit_only:
|
|
log("Done (dit-only).")
|
|
return
|
|
write_sa3_gguf(os.path.join(out_dir, "sa3-same-enc-F16.gguf"),
|
|
"sa3-same-enc", enc_tensors, config_json, half="f16")
|
|
write_sa3_gguf(os.path.join(out_dir, "sa3-same-dec-F16.gguf"),
|
|
"sa3-same-dec", dec_tensors, config_json, half="f16")
|
|
|
|
# ── T5Gemma encoder (separate HF model in the repo subfolder) ───────
|
|
t5_cfg = hf_hub_download(REPO, "config.json", subfolder="t5gemma-b-b-ul2")
|
|
t5_ckpt = hf_hub_download(REPO, "model.safetensors", subfolder="t5gemma-b-b-ul2")
|
|
with open(t5_cfg) as f:
|
|
t5_config_json = f.read()
|
|
t5_tensors = []
|
|
with safe_open(t5_ckpt, framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
# Encoder-only: drop the decoder half (never used by SA3)
|
|
if key.startswith("decoder."):
|
|
continue
|
|
t5_tensors.append((key, to_np(f.get_tensor(key))))
|
|
# The SA3 conditioner's learned padding embedding is applied to the text
|
|
# encoder's output (padded positions replaced) — it belongs to this module,
|
|
# so duplicate it here (it also rides in the DiT gguf with the rest of
|
|
# conditioner.*).
|
|
with safe_open(ckpt_path, framework="pt", device="cpu") as f:
|
|
key = "conditioner.conditioners.prompt.padding_embedding"
|
|
t5_tensors.append((key, to_np(f.get_tensor(key))))
|
|
write_sa3_gguf(os.path.join(out_dir, "sa3-text-enc-BF16.gguf"),
|
|
"sa3-t5gemma", t5_tensors, t5_config_json,
|
|
extra_meta={"sa3.parent_config_json": config_json})
|
|
|
|
log("Done.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|