Initial release
This commit is contained in:
civ
@@ -0,0 +1,293 @@
|
||||
#!/usr/bin/env python
|
||||
"""convert-bs-roformer-gguf.py — BS-RoFormer .ckpt → GGUF for the HOT-Step engine.
|
||||
|
||||
A pure state-dict repack: reads the checkpoint tensor by tensor and writes GGUF.
|
||||
There is NO graph tracing, so peak memory is roughly the size of the checkpoint
|
||||
(~270 MB for Leap Xe) rather than the tens of gigabytes that torch.onnx.export
|
||||
needs to hold a fully unrolled 16-layer axial transformer at T=1722.
|
||||
|
||||
Consumed by engine/src/bs-roformer-ggml.h.
|
||||
|
||||
TENSOR NAMING
|
||||
-------------
|
||||
Flat and index-addressable so the C++ side can build names with snprintf:
|
||||
|
||||
band_split.{b}.norm [dim_in] RMSNorm gamma
|
||||
band_split.{b}.w [dim_in, dim] Linear weight
|
||||
band_split.{b}.b [dim] Linear bias
|
||||
|
||||
blk.{i}.{time|freq}.rope_freqs [dim_head/2] rotary inv-freqs
|
||||
blk.{i}.{time|freq}.attn_norm [dim]
|
||||
blk.{i}.{time|freq}.qkv [dim, 3*dim_inner] fused, no bias
|
||||
blk.{i}.{time|freq}.gates_w [dim, heads]
|
||||
blk.{i}.{time|freq}.gates_b [heads]
|
||||
blk.{i}.{time|freq}.out [dim_inner, dim] no bias
|
||||
blk.{i}.{time|freq}.ff_norm [dim]
|
||||
blk.{i}.{time|freq}.ff1_w [dim, ff_inner]
|
||||
blk.{i}.{time|freq}.ff1_b [ff_inner]
|
||||
blk.{i}.{time|freq}.ff2_w [ff_inner, dim]
|
||||
blk.{i}.{time|freq}.ff2_b [dim]
|
||||
|
||||
final_norm [dim]
|
||||
|
||||
mask.{s}.{b}.w1 [dim, ff_inner] Linear
|
||||
mask.{s}.{b}.b1 [ff_inner]
|
||||
mask.{s}.{b}.w2 [ff_inner, dim_in*2] Linear (GLU halves it)
|
||||
mask.{s}.{b}.b2 [dim_in*2]
|
||||
|
||||
Shapes above are in GGML order (ne[0] innermost). A torch Linear weight is
|
||||
(out, in); gguf reverses numpy dims on write, so passing it through unchanged
|
||||
yields ne = [in, out], which is what ggml_mul_mat wants.
|
||||
|
||||
Note `band_split[b]` input width is freqs_per_bands[b] * 2 (stereo) * 2
|
||||
(complex) — the yaml lists frequency bins, not feature width.
|
||||
|
||||
USAGE
|
||||
-----
|
||||
$py = "d:\\Ace-Step-Latest\\hot-step-9000\\.venv\\Scripts\\python.exe"
|
||||
& $py tools\\convert-bs-roformer-gguf.py `
|
||||
--config models\\supersep-ckpt\\Xe\\leap_xe_config_voc.yaml `
|
||||
--ckpt models\\supersep-ckpt\\Xe\\bs_leap_xe_voc.ckpt `
|
||||
--output models\\supersep\\bs_leap_xe_voc-F32.gguf
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
import gguf
|
||||
|
||||
ARCH = "bs-roformer"
|
||||
|
||||
|
||||
def read_config(path):
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
# freqs_per_bands uses !!python/tuple, so safe_load will not do.
|
||||
return yaml.load(f, Loader=yaml.UnsafeLoader)
|
||||
|
||||
|
||||
def load_state_dict(path):
|
||||
sd = torch.load(path, map_location="cpu", weights_only=False)
|
||||
for key in ("state_dict", "model", "model_state_dict"):
|
||||
if isinstance(sd, dict) and key in sd and isinstance(sd[key], dict):
|
||||
sd = sd[key]
|
||||
break
|
||||
if sd and all(k.startswith("model.") for k in sd):
|
||||
sd = {k[len("model."):]: v for k, v in sd.items()}
|
||||
return sd
|
||||
|
||||
|
||||
class Repacker:
|
||||
"""Pulls named tensors out of the state dict, tracking what was consumed."""
|
||||
|
||||
def __init__(self, sd, writer):
|
||||
self.sd = sd
|
||||
self.writer = writer
|
||||
self.used = set()
|
||||
self.written = 0
|
||||
|
||||
def put(self, dst, src, expect_dims=None):
|
||||
if src not in self.sd:
|
||||
raise SystemExit(f"missing tensor in checkpoint: {src}")
|
||||
t = self.sd[src]
|
||||
if expect_dims is not None and t.dim() != expect_dims:
|
||||
raise SystemExit(
|
||||
f"{src}: expected {expect_dims}-D, got {tuple(t.shape)}")
|
||||
arr = t.detach().to(torch.float32).contiguous().numpy()
|
||||
self.writer.add_tensor(dst, arr)
|
||||
self.used.add(src)
|
||||
self.written += 1
|
||||
|
||||
def report(self):
|
||||
leftover = sorted(set(self.sd) - self.used)
|
||||
if leftover:
|
||||
raise SystemExit(
|
||||
f"{len(leftover)} checkpoint tensors were not written, e.g. "
|
||||
f"{leftover[:6]}\nThe converter does not understand this "
|
||||
"checkpoint layout — refusing to emit a partial model.")
|
||||
print(f"[ok] wrote {self.written} tensors, none left over")
|
||||
|
||||
|
||||
def convert(cfg, sd, out_path):
|
||||
mc = cfg["model"]
|
||||
dim = mc["dim"]
|
||||
depth = mc["depth"]
|
||||
heads = mc["heads"]
|
||||
dim_head = mc["dim_head"]
|
||||
n_stems = mc.get("num_stems", 1)
|
||||
n_fft = mc["stft_n_fft"]
|
||||
hop = mc["stft_hop_length"]
|
||||
win = mc.get("stft_win_length", n_fft)
|
||||
stereo = bool(mc.get("stereo", True))
|
||||
n_ch = 2 if stereo else 1
|
||||
mlp_mult = mc.get("mlp_expansion_factor", 4)
|
||||
chunk = cfg["audio"]["chunk_size"]
|
||||
|
||||
n_freqs = n_fft // 2 + 1
|
||||
|
||||
if "freqs_per_bands" in mc:
|
||||
# BS-RoFormer: contiguous bands that tile the spectrum exactly.
|
||||
arch = "bs"
|
||||
freqs_per_bands = list(mc["freqs_per_bands"])
|
||||
if sum(freqs_per_bands) != n_freqs:
|
||||
raise SystemExit(
|
||||
f"freqs_per_bands sums to {sum(freqs_per_bands)}, expected "
|
||||
f"{n_freqs} (n_fft//2+1)")
|
||||
else:
|
||||
# Mel-Band RoFormer: OVERLAPPING bands from a mel filterbank. The bands
|
||||
# do not tile — the caller gathers freq_indices before the graph and
|
||||
# scatters the mask back afterwards (see mel_band_tables.inc), so the
|
||||
# per-band widths come from how many bins each mel filter touches and
|
||||
# their sum is the *gathered* length, not n_freqs.
|
||||
arch = "mel"
|
||||
from librosa import filters
|
||||
num_bands = mc["num_bands"]
|
||||
fb = filters.mel(sr=mc.get("sample_rate", 44100), n_fft=n_fft, n_mels=num_bands)
|
||||
fb = np.asarray(fb)
|
||||
# Matches MelBandRoformer.__init__: force the first/last bins on so
|
||||
# every frequency is covered by at least one band.
|
||||
fb[0][0] = 1.0
|
||||
fb[-1, -1] = 1.0
|
||||
per_band = (fb > 0)
|
||||
if not per_band.any(axis=0).all():
|
||||
raise SystemExit("mel filterbank leaves some frequencies uncovered")
|
||||
freqs_per_bands = per_band.sum(axis=1).tolist()
|
||||
|
||||
n_bands = len(freqs_per_bands)
|
||||
# Feature width per band: bins * channels * 2 (real/imag).
|
||||
band_widths = [int(f) * n_ch * 2 for f in freqs_per_bands]
|
||||
|
||||
dim_inner = heads * dim_head
|
||||
ff_inner = dim * mlp_mult
|
||||
target = cfg.get("training", {}).get("target_instrument", "unknown")
|
||||
|
||||
print(f"[info] dim={dim} depth={depth} heads={heads} dim_head={dim_head}")
|
||||
print(f"[info] bands={n_bands} stems={n_stems} target={target}")
|
||||
print(f"[info] n_fft={n_fft} hop={hop} chunk={chunk} -> T={chunk // hop + 1}")
|
||||
|
||||
if mc.get("linear_transformer_depth", 0) != 0:
|
||||
raise SystemExit("linear_transformer_depth != 0 is not supported")
|
||||
if mc.get("time_transformer_depth", 1) != 1 or \
|
||||
mc.get("freq_transformer_depth", 1) != 1:
|
||||
raise SystemExit("only time/freq_transformer_depth == 1 is supported")
|
||||
if mc.get("skip_connection", False):
|
||||
raise SystemExit("skip_connection=True is not supported")
|
||||
|
||||
# Structural variations between checkpoints, detected from the state dict
|
||||
# rather than trusted from the yaml (the Mel-Band Karaoke config claims
|
||||
# mask_estimator_depth 2 but its MLPs actually have 3 Linears).
|
||||
has_out_norm = any(re.match(r"layers\.\d+\.\d+\.norm\.", k) for k in sd)
|
||||
has_final_norm = any(k.startswith("final_norm") for k in sd)
|
||||
mask_idx = sorted({int(k.split(".")[5]) for k in sd
|
||||
if k.startswith("mask_estimators.") and k.endswith(".weight")})
|
||||
mask_layers = len(mask_idx)
|
||||
if mask_layers == 0:
|
||||
raise SystemExit("no mask_estimators found in checkpoint")
|
||||
|
||||
print(f"[info] arch={arch} mask_layers={mask_layers} (indices {mask_idx}) "
|
||||
f"out_norm={has_out_norm} final_norm={has_final_norm}")
|
||||
|
||||
w = gguf.GGUFWriter(out_path, ARCH)
|
||||
|
||||
w.add_string("bs_roformer.arch", arch)
|
||||
w.add_uint32("bs_roformer.mask_layers", mask_layers)
|
||||
w.add_bool("bs_roformer.has_out_norm", has_out_norm)
|
||||
w.add_bool("bs_roformer.has_final_norm", has_final_norm)
|
||||
w.add_uint32("bs_roformer.dim", dim)
|
||||
w.add_uint32("bs_roformer.depth", depth)
|
||||
w.add_uint32("bs_roformer.heads", heads)
|
||||
w.add_uint32("bs_roformer.dim_head", dim_head)
|
||||
w.add_uint32("bs_roformer.dim_inner", dim_inner)
|
||||
w.add_uint32("bs_roformer.ff_inner", ff_inner)
|
||||
w.add_uint32("bs_roformer.n_bands", n_bands)
|
||||
w.add_uint32("bs_roformer.n_stems", n_stems)
|
||||
w.add_uint32("bs_roformer.n_channels", n_ch)
|
||||
w.add_uint32("bs_roformer.n_fft", n_fft)
|
||||
w.add_uint32("bs_roformer.hop_length", hop)
|
||||
w.add_uint32("bs_roformer.win_length", win)
|
||||
w.add_uint32("bs_roformer.chunk_size", chunk)
|
||||
w.add_array("bs_roformer.band_widths", band_widths)
|
||||
w.add_string("bs_roformer.target_instrument", str(target))
|
||||
|
||||
r = Repacker(sd, w)
|
||||
|
||||
# ── Band split: 90 × [RMSNorm(dim_in) -> Linear(dim_in, dim)] ──────────
|
||||
for b in range(n_bands):
|
||||
p = f"band_split.to_features.{b}"
|
||||
r.put(f"band_split.{b}.norm", f"{p}.0.gamma", 1)
|
||||
r.put(f"band_split.{b}.w", f"{p}.1.weight", 2)
|
||||
r.put(f"band_split.{b}.b", f"{p}.1.bias", 1)
|
||||
|
||||
# ── Axial transformer body ────────────────────────────────────────────
|
||||
# layers[i][0] = time transformer, layers[i][1] = freq transformer.
|
||||
# Each has depth 1, so exactly one (Attention, FeedForward) pair, and
|
||||
# norm_output is False (no trailing per-Transformer norm — confirmed by
|
||||
# the checkpoint's tensor count reconciling exactly without one).
|
||||
for i in range(depth):
|
||||
for axis_idx, axis in ((0, "time"), (1, "freq")):
|
||||
src = f"layers.{i}.{axis_idx}.layers.0"
|
||||
dst = f"blk.{i}.{axis}"
|
||||
r.put(f"{dst}.rope_freqs", f"{src}.0.rotary_embed.freqs", 1)
|
||||
r.put(f"{dst}.attn_norm", f"{src}.0.norm.gamma", 1)
|
||||
r.put(f"{dst}.qkv", f"{src}.0.to_qkv.weight", 2)
|
||||
r.put(f"{dst}.gates_w", f"{src}.0.to_gates.weight", 2)
|
||||
r.put(f"{dst}.gates_b", f"{src}.0.to_gates.bias", 1)
|
||||
r.put(f"{dst}.out", f"{src}.0.to_out.0.weight", 2)
|
||||
# FeedForward Sequential: 0=RMSNorm 1=Linear 2=GELU 3=Dropout
|
||||
# 4=Linear 5=Dropout
|
||||
r.put(f"{dst}.ff_norm", f"{src}.1.net.0.gamma", 1)
|
||||
r.put(f"{dst}.ff1_w", f"{src}.1.net.1.weight", 2)
|
||||
r.put(f"{dst}.ff1_b", f"{src}.1.net.1.bias", 1)
|
||||
r.put(f"{dst}.ff2_w", f"{src}.1.net.4.weight", 2)
|
||||
r.put(f"{dst}.ff2_b", f"{src}.1.net.4.bias", 1)
|
||||
|
||||
if has_out_norm:
|
||||
r.put(f"{dst}.out_norm", f"layers.{i}.{axis_idx}.norm.gamma", 1)
|
||||
|
||||
if has_final_norm:
|
||||
r.put("final_norm", "final_norm.gamma", 1)
|
||||
|
||||
# ── Mask estimators: per stem, per band ───────────────────────────────
|
||||
# to_freqs[b] = Sequential(MLP, GLU); MLP = Sequential(Linear, Tanh, Linear)
|
||||
for s in range(n_stems):
|
||||
for b in range(n_bands):
|
||||
p = f"mask_estimators.{s}.to_freqs.{b}.0"
|
||||
d = f"mask.{s}.{b}"
|
||||
for n, idx in enumerate(mask_idx):
|
||||
r.put(f"{d}.w{n + 1}", f"{p}.{idx}.weight", 2)
|
||||
r.put(f"{d}.b{n + 1}", f"{p}.{idx}.bias", 1)
|
||||
|
||||
r.report()
|
||||
|
||||
w.write_header_to_file()
|
||||
w.write_kv_data_to_file()
|
||||
w.write_tensors_to_file()
|
||||
w.close()
|
||||
|
||||
size_mb = os.path.getsize(out_path) / (1024 * 1024)
|
||||
print(f"[ok] {out_path} ({size_mb:.1f} MB)")
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--config", required=True, help="Model yaml")
|
||||
ap.add_argument("--ckpt", required=True, help="Checkpoint .ckpt")
|
||||
ap.add_argument("--output", required=True, help="Destination .gguf")
|
||||
args = ap.parse_args()
|
||||
|
||||
cfg = read_config(args.config)
|
||||
sd = load_state_dict(args.ckpt)
|
||||
print(f"[info] checkpoint: {len(sd)} tensors, "
|
||||
f"{sum(v.numel() for v in sd.values()):,} params")
|
||||
|
||||
os.makedirs(os.path.dirname(os.path.abspath(args.output)) or ".", exist_ok=True)
|
||||
convert(cfg, sd, args.output)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user