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()
|
||||
@@ -0,0 +1,248 @@
|
||||
#!/usr/bin/env python
|
||||
"""convert-mdx23c-gguf.py — MDX23C (TFC-TDF v3) .ckpt → GGUF for HOT-Step.
|
||||
|
||||
A pure state-dict repack, like convert-bs-roformer-gguf.py: no tracing, peak
|
||||
memory ~= the checkpoint (~440 MB for DrumSep).
|
||||
|
||||
Consumed by engine/src/mdx23c-ggml.h.
|
||||
|
||||
ARCHITECTURE (models_without_stft/mdx23c_tfc_tdf_v3_no_stft.py)
|
||||
---------------------------------------------------------------
|
||||
A 5-scale conv U-Net over the complex spectrogram, with subband folding:
|
||||
|
||||
x = cac2cws(x) fold num_subbands into channels
|
||||
mix = x
|
||||
first_conv_out = x = first_conv(x) Conv2d(dim_c -> c, 1x1)
|
||||
x = x.transpose(-1, -2) [b, c, f, t] -> [b, c, t, f]
|
||||
for s in scales: x = tfc_tdf(x); skip.append(x); x = downscale(x)
|
||||
x = bottleneck(x)
|
||||
for s in scales: x = upscale(x); x = cat(x, skip.pop()); x = tfc_tdf(x)
|
||||
x = x.transpose(-1, -2)
|
||||
x = x * first_conv_out artifact reduction
|
||||
x = final_conv(cat(mix, x))
|
||||
x = cws2cac(x) unfold subbands
|
||||
|
||||
TFC_TDF sub-block (residual):
|
||||
s = shortcut(x) Conv2d 1x1
|
||||
x = tfc1(x) norm -> act -> Conv2d 3x3
|
||||
x = x + tdf(x) norm->act->Linear->norm->act->Linear
|
||||
x = tfc2(x) norm -> act -> Conv2d 3x3
|
||||
x = x + s
|
||||
|
||||
`norm` is InstanceNorm2d(affine=True) — normalise each (sample, channel) over
|
||||
its spatial extent, i.e. GroupNorm with n_groups == n_channels, then an affine
|
||||
weight/bias per channel. `act` is GELU.
|
||||
|
||||
TENSOR NAMING / LAYOUT
|
||||
----------------------
|
||||
gguf reverses numpy dims on write, so a torch Conv2d weight (OC, IC, KH, KW)
|
||||
lands as ne = [KW, KH, IC, OC] — exactly ggml_conv_2d's kernel layout. A torch
|
||||
ConvTranspose2d weight is (IC, OC, KH, KW) and lands as [KW, KH, OC, IC],
|
||||
which is ggml_conv_transpose_2d_p0's layout. Both pass through unchanged.
|
||||
|
||||
first_conv [1,1,dim_c,c]
|
||||
enc.{s}.blk.{b}.* per TFC_TDF sub-block (see below)
|
||||
enc.{s}.down_norm_w / _b [c]
|
||||
enc.{s}.down_conv [sw,sh,c,c+g]
|
||||
bot.blk.{b}.*
|
||||
dec.{s}.up_norm_w / _b [c]
|
||||
dec.{s}.up_conv [sw,sh,c-g,c] (transposed)
|
||||
dec.{s}.blk.{b}.*
|
||||
final1 [1,1,c+dim_c,c]
|
||||
final2 [1,1,c,n_inst*dim_c]
|
||||
|
||||
per TFC_TDF sub-block:
|
||||
tfc1_norm_w/_b, tfc1_conv [3,3,in_c,c]
|
||||
tdf_n1_w/_b, tdf_l1 [f, f//bn]
|
||||
tdf_n2_w/_b, tdf_l2 [f//bn, f]
|
||||
tfc2_norm_w/_b, tfc2_conv [3,3,c,c]
|
||||
shortcut [1,1,in_c,c]
|
||||
|
||||
USAGE
|
||||
-----
|
||||
$py = "d:\\Ace-Step-Latest\\hot-step-9000\\.venv\\Scripts\\python.exe"
|
||||
& $py scripts\\convert-mdx23c-gguf.py `
|
||||
--config <SuperSep>\\models\\config_drumsep_mdx23c.yaml `
|
||||
--ckpt <SuperSep>\\models\\MDX23C-DrumSep-aufr33-jarredou.ckpt `
|
||||
--output models\\supersep\\mdx23c_drumsep-F32.gguf
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
import gguf
|
||||
|
||||
ARCH = "mdx23c"
|
||||
|
||||
|
||||
def read_config(path):
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
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:
|
||||
def __init__(self, sd, writer):
|
||||
self.sd, self.writer = sd, writer
|
||||
self.used, self.written = set(), 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)}")
|
||||
self.writer.add_tensor(dst, t.detach().to(torch.float32).contiguous().numpy())
|
||||
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]}\nRefusing to emit a partial model.")
|
||||
print(f"[ok] wrote {self.written} tensors, none left over")
|
||||
|
||||
|
||||
def put_tfc_tdf(r, dst, src, n_blocks):
|
||||
"""One TFC_TDF stack: n_blocks residual sub-blocks.
|
||||
|
||||
Sequential indices in the checkpoint:
|
||||
tfc1 / tfc2 : 0=norm, 1=act, 2=Conv2d
|
||||
tdf : 0=norm, 1=act, 2=Linear, 3=norm, 4=act, 5=Linear
|
||||
"""
|
||||
for b in range(n_blocks):
|
||||
s = f"{src}.blocks.{b}"
|
||||
d = f"{dst}.blk.{b}"
|
||||
r.put(f"{d}.tfc1_norm_w", f"{s}.tfc1.0.weight", 1)
|
||||
r.put(f"{d}.tfc1_norm_b", f"{s}.tfc1.0.bias", 1)
|
||||
r.put(f"{d}.tfc1_conv", f"{s}.tfc1.2.weight", 4)
|
||||
|
||||
r.put(f"{d}.tdf_n1_w", f"{s}.tdf.0.weight", 1)
|
||||
r.put(f"{d}.tdf_n1_b", f"{s}.tdf.0.bias", 1)
|
||||
r.put(f"{d}.tdf_l1", f"{s}.tdf.2.weight", 2)
|
||||
r.put(f"{d}.tdf_n2_w", f"{s}.tdf.3.weight", 1)
|
||||
r.put(f"{d}.tdf_n2_b", f"{s}.tdf.3.bias", 1)
|
||||
r.put(f"{d}.tdf_l2", f"{s}.tdf.5.weight", 2)
|
||||
|
||||
r.put(f"{d}.tfc2_norm_w", f"{s}.tfc2.0.weight", 1)
|
||||
r.put(f"{d}.tfc2_norm_b", f"{s}.tfc2.0.bias", 1)
|
||||
r.put(f"{d}.tfc2_conv", f"{s}.tfc2.2.weight", 4)
|
||||
|
||||
r.put(f"{d}.shortcut", f"{s}.shortcut.weight", 4)
|
||||
|
||||
|
||||
def convert(cfg, sd, out_path):
|
||||
mc, ac = cfg["model"], cfg["audio"]
|
||||
n_scales = mc["num_scales"]
|
||||
n_blocks = mc["num_blocks_per_scale"]
|
||||
c0 = mc["num_channels"]
|
||||
growth = mc["growth"]
|
||||
bn = mc["bottleneck_factor"]
|
||||
n_subbands = mc["num_subbands"]
|
||||
scale = list(mc["scale"])
|
||||
norm_type = mc.get("norm", "InstanceNorm")
|
||||
act_type = mc.get("act", "gelu")
|
||||
|
||||
n_audio_ch = ac["num_channels"]
|
||||
dim_c = n_subbands * n_audio_ch * 2
|
||||
dim_f = ac["dim_f"]
|
||||
n_fft = ac["n_fft"]
|
||||
hop = ac["hop_length"]
|
||||
chunk = ac["chunk_size"]
|
||||
|
||||
instruments = cfg.get("training", {}).get("instruments", [])
|
||||
target = cfg.get("training", {}).get("target_instrument", None)
|
||||
n_inst = 1 if target else len(instruments)
|
||||
|
||||
if norm_type != "InstanceNorm":
|
||||
raise SystemExit(f"only InstanceNorm is supported, got {norm_type}")
|
||||
if act_type != "gelu":
|
||||
raise SystemExit(f"only gelu is supported, got {act_type}")
|
||||
|
||||
print(f"[info] scales={n_scales} blocks/scale={n_blocks} c={c0} growth={growth}")
|
||||
print(f"[info] subbands={n_subbands} dim_c={dim_c} dim_f={dim_f} bn={bn}")
|
||||
print(f"[info] instruments={n_inst} {instruments}")
|
||||
print(f"[info] n_fft={n_fft} hop={hop} chunk={chunk} scale={scale}")
|
||||
|
||||
w = gguf.GGUFWriter(out_path, ARCH)
|
||||
w.add_uint32("mdx23c.num_scales", n_scales)
|
||||
w.add_uint32("mdx23c.blocks_per_scale", n_blocks)
|
||||
w.add_uint32("mdx23c.num_channels", c0)
|
||||
w.add_uint32("mdx23c.growth", growth)
|
||||
w.add_uint32("mdx23c.bottleneck_factor", bn)
|
||||
w.add_uint32("mdx23c.num_subbands", n_subbands)
|
||||
w.add_uint32("mdx23c.dim_c", dim_c)
|
||||
w.add_uint32("mdx23c.dim_f", dim_f)
|
||||
w.add_uint32("mdx23c.n_fft", n_fft)
|
||||
w.add_uint32("mdx23c.hop_length", hop)
|
||||
w.add_uint32("mdx23c.chunk_size", chunk)
|
||||
w.add_uint32("mdx23c.n_instruments", n_inst)
|
||||
w.add_uint32("mdx23c.n_audio_channels", n_audio_ch)
|
||||
w.add_array("mdx23c.scale", scale)
|
||||
w.add_string("mdx23c.instruments", ",".join(instruments))
|
||||
|
||||
r = Repacker(sd, w)
|
||||
r.put("first_conv", "first_conv.weight", 4)
|
||||
|
||||
c = c0
|
||||
for s in range(n_scales):
|
||||
put_tfc_tdf(r, f"enc.{s}", f"encoder_blocks.{s}.tfc_tdf", n_blocks)
|
||||
r.put(f"enc.{s}.down_norm_w", f"encoder_blocks.{s}.downscale.conv.0.weight", 1)
|
||||
r.put(f"enc.{s}.down_norm_b", f"encoder_blocks.{s}.downscale.conv.0.bias", 1)
|
||||
r.put(f"enc.{s}.down_conv", f"encoder_blocks.{s}.downscale.conv.2.weight", 4)
|
||||
c += growth
|
||||
|
||||
put_tfc_tdf(r, "bot", "bottleneck_block", n_blocks)
|
||||
|
||||
for s in range(n_scales):
|
||||
r.put(f"dec.{s}.up_norm_w", f"decoder_blocks.{s}.upscale.conv.0.weight", 1)
|
||||
r.put(f"dec.{s}.up_norm_b", f"decoder_blocks.{s}.upscale.conv.0.bias", 1)
|
||||
r.put(f"dec.{s}.up_conv", f"decoder_blocks.{s}.upscale.conv.2.weight", 4)
|
||||
put_tfc_tdf(r, f"dec.{s}", f"decoder_blocks.{s}.tfc_tdf", n_blocks)
|
||||
c -= growth
|
||||
|
||||
r.put("final1", "final_conv.0.weight", 4)
|
||||
r.put("final2", "final_conv.2.weight", 4)
|
||||
|
||||
r.report()
|
||||
|
||||
w.write_header_to_file()
|
||||
w.write_kv_data_to_file()
|
||||
w.write_tensors_to_file()
|
||||
w.close()
|
||||
print(f"[ok] {out_path} ({os.path.getsize(out_path)/1024/1024:.1f} MB)")
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--config", required=True)
|
||||
ap.add_argument("--ckpt", required=True)
|
||||
ap.add_argument("--output", required=True)
|
||||
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()
|
||||
@@ -0,0 +1,200 @@
|
||||
#!/usr/bin/env python
|
||||
"""dump_bs_roformer_goldens.py — PyTorch reference activations for the GGML port.
|
||||
|
||||
Runs ZFTurbo's no-STFT BS-RoFormer under torch.no_grad at a SHORT time length
|
||||
and dumps intermediate tensors, so engine/src/bs-roformer-ggml.h can be
|
||||
validated stage by stage instead of only at the output.
|
||||
|
||||
MEMORY
|
||||
------
|
||||
Deliberately runs at T=256 (~3 s of audio) rather than the model's trained
|
||||
T=1722. A no_grad forward frees activations as it goes, so peak is bounded by
|
||||
the largest few tensors (~23 MB each at T=256) plus the 268 MB of weights.
|
||||
Numerics do not depend on sequence length, so a short T validates the graph
|
||||
just as well. Do NOT raise --time-steps to 1722 "to be thorough" — that is not
|
||||
what broke before (tracing was), but there is no reason to pay for it either.
|
||||
|
||||
WHAT IT DUMPS
|
||||
-------------
|
||||
Full tensors (for elementwise comparison):
|
||||
input [1, T, F*C*2]
|
||||
band_split [1, T, n_bands, dim]
|
||||
layer_00 [1, T, n_bands, dim] after layers[0] (time + freq)
|
||||
layer_01 [1, T, n_bands, dim]
|
||||
layer_last [1, T, n_bands, dim]
|
||||
final_norm [1, T, n_bands, dim]
|
||||
mask [1, S, F*C, T, 2] the graph's actual output
|
||||
|
||||
Per-layer summary stats for all `depth` layers (mean/std/absmax), enough to
|
||||
localise which layer a divergence starts in without storing 16 full tensors.
|
||||
|
||||
USAGE
|
||||
-----
|
||||
$py = "d:\\Ace-Step-Latest\\hot-step-9000\\.venv\\Scripts\\python.exe"
|
||||
& $py tools\\dump_bs_roformer_goldens.py `
|
||||
--config models\\supersep-ckpt\\Xe\\leap_xe_config_voc.yaml `
|
||||
--ckpt models\\supersep-ckpt\\Xe\\bs_leap_xe_voc.ckpt `
|
||||
--output models\\supersep-ckpt\\goldens_voc.npz
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
DEFAULT_MSS_REPO = os.path.normpath(os.path.join(
|
||||
os.path.dirname(os.path.abspath(__file__)), "..", "..", "MSS_ONNX_TensorRT"))
|
||||
|
||||
|
||||
def load_no_stft_class(mss_repo, arch):
|
||||
"""Import the amputated model class for `arch` ('bs' or 'mel')."""
|
||||
fname, cls = (("bs_roformer_no_stft.py", "BSRoformer") if arch == "bs"
|
||||
else ("mel_band_roformer_no_stft.py", "MelBandRoformer"))
|
||||
mod_path = os.path.join(mss_repo, "models_without_stft", fname)
|
||||
if not os.path.isfile(mod_path):
|
||||
raise SystemExit(f"Could not find {mod_path}\nPass --mss-repo <checkout>.")
|
||||
if mss_repo not in sys.path:
|
||||
sys.path.insert(0, mss_repo)
|
||||
spec = importlib.util.spec_from_file_location(fname[:-3], mod_path)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
return getattr(mod, cls)
|
||||
|
||||
|
||||
def read_config(path):
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
return yaml.load(f, Loader=yaml.UnsafeLoader)
|
||||
|
||||
|
||||
def build_model(BSRoformer, cfg, ckpt_path):
|
||||
kwargs = dict(cfg["model"])
|
||||
kwargs["flash_attn"] = False # match what the GGML graph computes
|
||||
kwargs["use_torch_checkpoint"] = False
|
||||
model = BSRoformer(**kwargs)
|
||||
|
||||
sd = torch.load(ckpt_path, map_location="cpu", weights_only=False)
|
||||
for k in ("state_dict", "model", "model_state_dict"):
|
||||
if isinstance(sd, dict) and k in sd and isinstance(sd[k], dict):
|
||||
sd = sd[k]
|
||||
break
|
||||
if sd and all(k.startswith("model.") for k in sd):
|
||||
sd = {k[len("model."):]: v for k, v in sd.items()}
|
||||
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False)
|
||||
if missing:
|
||||
raise SystemExit(f"missing weights: {missing[:6]}")
|
||||
if unexpected:
|
||||
print(f"[warn] unexpected: {unexpected[:6]}")
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--config", required=True)
|
||||
ap.add_argument("--ckpt", required=True)
|
||||
ap.add_argument("--output", required=True)
|
||||
ap.add_argument("--mss-repo", default=DEFAULT_MSS_REPO)
|
||||
ap.add_argument("--time-steps", type=int, default=256,
|
||||
help="T to run at (default 256; see MEMORY in the docstring)")
|
||||
ap.add_argument("--seed", type=int, default=1234)
|
||||
args = ap.parse_args()
|
||||
|
||||
cfg = read_config(args.config)
|
||||
mc = cfg["model"]
|
||||
arch = "bs" if "freqs_per_bands" in mc else "mel"
|
||||
ModelCls = load_no_stft_class(os.path.abspath(args.mss_repo), arch)
|
||||
model = build_model(ModelCls, cfg, args.ckpt)
|
||||
|
||||
n_ch = 2 if mc.get("stereo", True) else 1
|
||||
depth = mc["depth"]
|
||||
T = args.time_steps
|
||||
|
||||
torch.manual_seed(args.seed)
|
||||
# Scaled to roughly the magnitude of a real STFT bin so activations land in
|
||||
# a representative range rather than an artificially tiny one.
|
||||
if arch == "bs":
|
||||
# BSRoformer.forward takes the flattened [b, t, (f c)] directly.
|
||||
in_dim = (mc["stft_n_fft"] // 2 + 1) * n_ch * 2
|
||||
x = torch.randn(1, T, in_dim) * 0.05
|
||||
model_in = x
|
||||
else:
|
||||
# MelBandRoformer.forward takes [b, f, t, c] and flattens it itself,
|
||||
# where f is the GATHERED index count (bands overlap), so derive it
|
||||
# from the model's own freq_indices buffer rather than from n_fft.
|
||||
n_gathered = int(model.freq_indices.numel())
|
||||
in_dim = n_gathered * 2
|
||||
model_in = torch.randn(1, n_gathered, T, 2) * 0.05
|
||||
# Flattened view is what the C++ side feeds: [t, (f c)]
|
||||
x = model_in.permute(0, 2, 1, 3).reshape(1, T, in_dim).contiguous()
|
||||
print(f"[info] arch={arch} in_dim={in_dim}")
|
||||
|
||||
out = {}
|
||||
stats = []
|
||||
|
||||
def store(name, t):
|
||||
out[name] = t.detach().to(torch.float32).contiguous().numpy()
|
||||
|
||||
# Forward hooks capture intermediates without touching the model source.
|
||||
captured = {}
|
||||
|
||||
def hook(tag):
|
||||
def fn(_mod, _inp, output):
|
||||
captured[tag] = output.detach()
|
||||
return fn
|
||||
|
||||
handles = [model.band_split.register_forward_hook(hook("band_split"))]
|
||||
# Mel-Band Karaoke has norm_output=True on each Transformer and NO
|
||||
# final_norm; hooking the last layer stands in for it there.
|
||||
has_final_norm = not isinstance(model.final_norm, torch.nn.Identity) \
|
||||
if hasattr(model, "final_norm") else False
|
||||
if has_final_norm:
|
||||
handles.append(model.final_norm.register_forward_hook(hook("final_norm")))
|
||||
for i, block in enumerate(model.layers):
|
||||
# block is ModuleList([time_transformer, freq_transformer]); hooking the
|
||||
# freq transformer captures the state after the full layer.
|
||||
handles.append(block[-1].register_forward_hook(hook(f"L{i}")))
|
||||
|
||||
print(f"[info] running T={T}, in_dim={in_dim}, depth={depth}")
|
||||
with torch.no_grad():
|
||||
mask = model(model_in)
|
||||
|
||||
for h in handles:
|
||||
h.remove()
|
||||
|
||||
store("input", x)
|
||||
store("mask", mask)
|
||||
store("band_split", captured["band_split"])
|
||||
# With no final_norm the C++ stage depth+1 is a no-op passthrough, so the
|
||||
# last layer's output is exactly what it should produce.
|
||||
store("final_norm", captured["final_norm"] if has_final_norm
|
||||
else captured[f"L{depth - 1}"])
|
||||
store("layer_00", captured["L0"])
|
||||
store("layer_01", captured["L1"])
|
||||
store("layer_last", captured[f"L{depth - 1}"])
|
||||
|
||||
for i in range(depth):
|
||||
t = captured[f"L{i}"].float()
|
||||
stats.append([t.mean().item(), t.std().item(), t.abs().max().item()])
|
||||
out["layer_stats"] = np.asarray(stats, dtype=np.float32) # [depth, 3]
|
||||
|
||||
out["meta"] = np.asarray([T, in_dim, depth, mask.shape[1], mask.shape[2]],
|
||||
dtype=np.int64)
|
||||
|
||||
os.makedirs(os.path.dirname(os.path.abspath(args.output)) or ".", exist_ok=True)
|
||||
np.savez(args.output, **out)
|
||||
|
||||
print(f"[ok] {args.output} ({os.path.getsize(args.output)/1024/1024:.1f} MB)")
|
||||
print(f"[ok] mask shape {tuple(mask.shape)} "
|
||||
f"range [{mask.min():.4f}, {mask.max():.4f}]")
|
||||
print("[ok] per-layer absmax: " +
|
||||
" ".join(f"{s[2]:.2f}" for s in stats))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,120 @@
|
||||
#!/usr/bin/env python
|
||||
"""dump_mdx23c_goldens.py — PyTorch reference output for the MDX23C GGML port.
|
||||
|
||||
Runs ZFTurbo's no-STFT TFC_TDF_net under torch.no_grad and writes a flat binary
|
||||
that engine/tools/mdx23c-test.cpp compares against.
|
||||
|
||||
magic "MDXG"
|
||||
int32 T, dim_f, cin, n_inst
|
||||
f32 input [T * dim_f * cin] torch [b, cin, dim_f, T], t fastest
|
||||
f32 output [n_inst * cin * dim_f * T]
|
||||
|
||||
Memory is bounded: a no_grad forward at the trained T=256 frees activations as
|
||||
it goes. No tracing anywhere.
|
||||
|
||||
The upstream module does `from utils import prefer_target_instrument`, and that
|
||||
utils.py drags in the whole MSS model zoo (demucs, scnet, omegaconf). We stub
|
||||
the one function instead — it is three lines — so this needs only torch.
|
||||
|
||||
USAGE
|
||||
$py = "d:\\Ace-Step-Latest\\hot-step-9000\\.venv\\Scripts\\python.exe"
|
||||
& $py scripts\\dump_mdx23c_goldens.py `
|
||||
--config <SuperSep>\\models\\config_drumsep_mdx23c.yaml `
|
||||
--ckpt <SuperSep>\\models\\MDX23C-DrumSep-aufr33-jarredou.ckpt `
|
||||
--output models\\supersep-ckpt\\goldens_mdx23c.bin
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import importlib.util
|
||||
import os
|
||||
import struct
|
||||
import sys
|
||||
import types
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
DEFAULT_MSS_REPO = os.path.normpath(os.path.join(
|
||||
os.path.dirname(os.path.abspath(__file__)), "..", "..", "MSS_ONNX_TensorRT"))
|
||||
|
||||
|
||||
def prefer_target_instrument(config):
|
||||
"""Mirror of MSS utils.prefer_target_instrument (avoids importing its utils)."""
|
||||
if getattr(config.training, "target_instrument", None):
|
||||
return [config.training.target_instrument]
|
||||
return config.training.instruments
|
||||
|
||||
|
||||
def load_net_class(mss_repo):
|
||||
mod_path = os.path.join(mss_repo, "models_without_stft",
|
||||
"mdx23c_tfc_tdf_v3_no_stft.py")
|
||||
if not os.path.isfile(mod_path):
|
||||
raise SystemExit(f"Could not find {mod_path}\nPass --mss-repo <checkout>.")
|
||||
if mss_repo not in sys.path:
|
||||
sys.path.insert(0, mss_repo)
|
||||
# Stub `utils` so the module's single import does not pull in the zoo.
|
||||
stub = types.ModuleType("utils")
|
||||
stub.prefer_target_instrument = prefer_target_instrument
|
||||
sys.modules.setdefault("utils", stub)
|
||||
spec = importlib.util.spec_from_file_location("mdx23c_no_stft", mod_path)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
return mod.TFC_TDF_net
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--config", required=True)
|
||||
ap.add_argument("--ckpt", required=True)
|
||||
ap.add_argument("--output", required=True)
|
||||
ap.add_argument("--mss-repo", default=DEFAULT_MSS_REPO)
|
||||
ap.add_argument("--seed", type=int, default=1234)
|
||||
args = ap.parse_args()
|
||||
|
||||
from ml_collections import ConfigDict
|
||||
with open(args.config, "r", encoding="utf-8") as f:
|
||||
cfg = ConfigDict(yaml.load(f, Loader=yaml.UnsafeLoader))
|
||||
|
||||
Net = load_net_class(os.path.abspath(args.mss_repo))
|
||||
model = Net(cfg)
|
||||
|
||||
sd = torch.load(args.ckpt, map_location="cpu", weights_only=False)
|
||||
for k in ("state_dict", "model", "model_state_dict"):
|
||||
if isinstance(sd, dict) and k in sd and isinstance(sd[k], dict):
|
||||
sd = sd[k]
|
||||
break
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False)
|
||||
if missing:
|
||||
raise SystemExit(f"missing weights: {missing[:6]}")
|
||||
if unexpected:
|
||||
print(f"[warn] unexpected: {unexpected[:6]}")
|
||||
model.eval()
|
||||
|
||||
dim_f = cfg.audio.dim_f
|
||||
cin = cfg.audio.num_channels * 2
|
||||
T = cfg.audio.chunk_size // cfg.audio.hop_length + 1
|
||||
n_inst = len(prefer_target_instrument(cfg))
|
||||
|
||||
print(f"[info] T={T} dim_f={dim_f} cin={cin} instruments={n_inst}")
|
||||
|
||||
torch.manual_seed(args.seed)
|
||||
x = torch.randn(1, cin, dim_f, T) * 0.05
|
||||
|
||||
with torch.no_grad():
|
||||
y = model(x)
|
||||
print(f"[info] output shape {tuple(y.shape)} range "
|
||||
f"[{y.min():.4f}, {y.max():.4f}]")
|
||||
|
||||
os.makedirs(os.path.dirname(os.path.abspath(args.output)) or ".", exist_ok=True)
|
||||
with open(args.output, "wb") as f:
|
||||
f.write(b"MDXG")
|
||||
f.write(struct.pack("<4i", T, dim_f, cin, n_inst))
|
||||
f.write(np.ascontiguousarray(x[0], dtype=np.float32).tobytes())
|
||||
f.write(np.ascontiguousarray(y[0], dtype=np.float32).tobytes())
|
||||
|
||||
print(f"[ok] {args.output} ({os.path.getsize(args.output)/1024/1024:.1f} MB)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,387 @@
|
||||
"""Export AceStep DiT decoder from safetensors to ONNX for native TRT.
|
||||
|
||||
Self-contained script: doesn't require the 'acestep' package. Loads the model
|
||||
definition directly from the model directory's modeling file.
|
||||
|
||||
The exported ONNX matches HOT-Step's dit-trt.h I/O layout:
|
||||
- input_latents: [B, T, 192] fp32 (context[128] + xt[64], pre-concatenated)
|
||||
- enc_hidden: [B, S, 2048] fp32
|
||||
- t: [B] fp32
|
||||
- t_r: [B] fp32
|
||||
- velocity: [B, T, 64] fp32 (output)
|
||||
|
||||
Usage:
|
||||
python scripts/export_dit_onnx.py ^
|
||||
--model-dir models/acestep-v15-merge-sft-turbo-xl-ta-0.7 ^
|
||||
--output-dir models/onnx/dit-stream
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import importlib.util
|
||||
import shutil
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
# ── Traceable Lambda replacement ────────────────────────────────────────────
|
||||
class _Transpose12(nn.Module):
|
||||
"""Drop-in for Lambda(lambda x: x.transpose(1, 2))."""
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return x.transpose(1, 2)
|
||||
|
||||
|
||||
# ── Export wrapper ──────────────────────────────────────────────────────────
|
||||
class DiTForTRTExport(nn.Module):
|
||||
"""Wraps AceStepDiTModel for ONNX export matching dit-trt.h I/O layout.
|
||||
|
||||
Takes pre-concatenated input_latents [B, T, 192] and exposes t/t_r
|
||||
as separate inputs, matching the C++ engine's tensor layout.
|
||||
"""
|
||||
|
||||
def __init__(self, decoder: nn.Module):
|
||||
super().__init__()
|
||||
self.decoder = decoder
|
||||
self._replace_lambdas()
|
||||
self.decoder.config._attn_implementation = "sdpa"
|
||||
self._patch_decoder_for_trace()
|
||||
|
||||
def _replace_lambdas(self) -> None:
|
||||
for seq in (self.decoder.proj_in, self.decoder.proj_out):
|
||||
for i, mod in enumerate(seq):
|
||||
if type(mod).__name__ == "Lambda":
|
||||
seq[i] = _Transpose12()
|
||||
|
||||
def _patch_decoder_for_trace(self) -> None:
|
||||
"""Monkey-patch decoder forward for ONNX traceability."""
|
||||
try:
|
||||
import transformers.integrations.sdpa_attention as _sdpa_mod
|
||||
_sdpa_mod.use_gqa_in_sdpa = lambda *args, **kwargs: False
|
||||
except (ImportError, AttributeError):
|
||||
pass
|
||||
|
||||
decoder = self.decoder
|
||||
sliding_window = decoder.config.sliding_window
|
||||
layer_types = decoder.config.layer_types
|
||||
|
||||
time_embed_dim = decoder.time_embed.time_proj.out_features // 6
|
||||
|
||||
def _patched_time_embed_forward(self_te, t):
|
||||
t_freq = self_te.timestep_embedding(t, self_te.in_channels)
|
||||
temb = self_te.linear_1(t_freq.to(t.dtype))
|
||||
temb = self_te.act1(temb)
|
||||
temb = self_te.linear_2(temb)
|
||||
timestep_proj = self_te.time_proj(self_te.act2(temb)).reshape(-1, 6, time_embed_dim)
|
||||
return temb, timestep_proj
|
||||
|
||||
decoder.time_embed.forward = types.MethodType(
|
||||
_patched_time_embed_forward, decoder.time_embed)
|
||||
decoder.time_embed_r.forward = types.MethodType(
|
||||
_patched_time_embed_forward, decoder.time_embed_r)
|
||||
|
||||
def _export_forward(
|
||||
self_dec, hidden_states, timestep, timestep_r,
|
||||
attention_mask, encoder_hidden_states, encoder_attention_mask,
|
||||
context_latents, use_cache=None, past_key_values=None,
|
||||
cache_position=None, position_ids=None, output_attentions=False,
|
||||
return_hidden_states=None, custom_layers_config=None,
|
||||
enable_early_exit=False, **flash_attn_kwargs,
|
||||
):
|
||||
temb_t, timestep_proj_t = self_dec.time_embed(timestep)
|
||||
temb_r, timestep_proj_r = self_dec.time_embed_r(timestep - timestep_r)
|
||||
temb = temb_t + temb_r
|
||||
timestep_proj = timestep_proj_t + timestep_proj_r
|
||||
|
||||
hidden_states = torch.cat([context_latents, hidden_states], dim=-1)
|
||||
hidden_states = self_dec.proj_in(hidden_states)
|
||||
encoder_hidden_states = self_dec.condition_embedder(encoder_hidden_states)
|
||||
|
||||
seq_len_pat = hidden_states.shape[1]
|
||||
cache_position = torch.arange(seq_len_pat, device=hidden_states.device)
|
||||
position_ids = cache_position.unsqueeze(0)
|
||||
position_embeddings = self_dec.rotary_emb(hidden_states, position_ids)
|
||||
|
||||
indices = cache_position
|
||||
diff = indices.unsqueeze(0) - indices.unsqueeze(1)
|
||||
sw_mask = torch.where(
|
||||
torch.abs(diff) <= sliding_window,
|
||||
torch.zeros(1, device=hidden_states.device, dtype=hidden_states.dtype),
|
||||
torch.full((1,), torch.finfo(hidden_states.dtype).min,
|
||||
device=hidden_states.device, dtype=hidden_states.dtype),
|
||||
)
|
||||
sw_mask = sw_mask.unsqueeze(0).unsqueeze(0)
|
||||
|
||||
for i, layer_module in enumerate(self_dec.layers):
|
||||
attn_mask = sw_mask if layer_types[i] == "sliding_attention" else None
|
||||
layer_outputs = layer_module(
|
||||
hidden_states, position_embeddings, timestep_proj, attn_mask,
|
||||
position_ids, None, False, False, cache_position,
|
||||
encoder_hidden_states, None,
|
||||
)
|
||||
hidden_states = layer_outputs[0]
|
||||
|
||||
shift, scale = (self_dec.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)
|
||||
hidden_states = (self_dec.norm_out(hidden_states) * (1 + scale) + shift).type_as(hidden_states)
|
||||
hidden_states = self_dec.proj_out(hidden_states)
|
||||
return (hidden_states, None)
|
||||
|
||||
decoder.forward = types.MethodType(_export_forward, decoder)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_latents: torch.Tensor, # [B, T, 192]
|
||||
enc_hidden: torch.Tensor, # [B, S, 2048]
|
||||
t: torch.Tensor, # [B]
|
||||
t_r: torch.Tensor, # [B]
|
||||
) -> torch.Tensor:
|
||||
context_latents = input_latents[..., :128]
|
||||
hidden_states = input_latents[..., 128:]
|
||||
outputs = self.decoder(
|
||||
hidden_states=hidden_states, timestep=t, timestep_r=t_r,
|
||||
attention_mask=None, encoder_hidden_states=enc_hidden,
|
||||
encoder_attention_mask=None, context_latents=context_latents,
|
||||
use_cache=False, past_key_values=None, output_attentions=False,
|
||||
)
|
||||
return outputs[0]
|
||||
|
||||
|
||||
# ── Load model from directory ────────────────────────────────────────────────
|
||||
|
||||
def _load_module_from_file(name: str, filepath: str):
|
||||
"""Load a Python module from a file path."""
|
||||
spec = importlib.util.spec_from_file_location(name, filepath)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
sys.modules[name] = mod
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
|
||||
|
||||
def load_model(model_dir: str, device: str = "cpu"):
|
||||
"""Load AceStepConditionGenerationModel from a safetensors directory.
|
||||
|
||||
Handles the 'acestep' module dependency by loading the config class
|
||||
from the DEMON project or from a standalone copy.
|
||||
"""
|
||||
model_dir = Path(model_dir).resolve()
|
||||
|
||||
# Load config.json
|
||||
config_path = model_dir / "config.json"
|
||||
with open(config_path) as f:
|
||||
config_dict = json.load(f)
|
||||
|
||||
# ── Bootstrap the 'acestep' package from DEMON project ──────────────
|
||||
# The model dir's shim files (configuration_acestep_v15.py, apg_guidance.py)
|
||||
# import from 'acestep.models.common.*', which maps to DEMON's 'acestep.models.*'.
|
||||
# We add DEMON to sys.path AND create a 'common' alias pointing to 'models'.
|
||||
demon_root = Path("d:/Ace-Step-Latest/Demon")
|
||||
if not demon_root.exists():
|
||||
raise RuntimeError(
|
||||
"DEMON project not found at d:/Ace-Step-Latest/Demon — "
|
||||
"needed for model class definitions")
|
||||
|
||||
if str(demon_root) not in sys.path:
|
||||
sys.path.insert(0, str(demon_root))
|
||||
print(f"[Export] Using DEMON project for acestep package: {demon_root}")
|
||||
|
||||
# Import the real acestep package so its __init__.py runs
|
||||
import acestep
|
||||
import acestep.models
|
||||
|
||||
# The model-dir shims import from 'acestep.models.common.*' but DEMON
|
||||
# puts everything directly in 'acestep.models.*'. Create a 'common'
|
||||
# alias that points to the real 'acestep.models' module.
|
||||
sys.modules["acestep.models.common"] = sys.modules["acestep.models"]
|
||||
|
||||
# Now load the config
|
||||
from acestep.models.configuration_acestep_v15 import AceStepConfig
|
||||
config = AceStepConfig(**config_dict)
|
||||
print(f"[Export] Config: hidden_size={config.hidden_size}, layers={config.num_hidden_layers}, "
|
||||
f"heads={config.num_attention_heads}, kv_heads={config.num_key_value_heads}")
|
||||
|
||||
# Register the config module under its short name too (the modeling file
|
||||
# tries 'from configuration_acestep_v15 import AceStepConfig' as fallback)
|
||||
config_shim = types.ModuleType("configuration_acestep_v15")
|
||||
config_shim.AceStepConfig = AceStepConfig
|
||||
sys.modules["configuration_acestep_v15"] = config_shim
|
||||
|
||||
# Register apg_guidance similarly (modeling file does 'from apg_guidance import ...')
|
||||
import acestep.models.apg_guidance as _real_apg
|
||||
sys.modules["apg_guidance"] = _real_apg
|
||||
# Also alias 'acestep.models.common.apg_guidance' → the real module
|
||||
sys.modules["acestep.models.common.apg_guidance"] = _real_apg
|
||||
sys.modules["acestep.models.common.configuration_acestep_v15"] = sys.modules[
|
||||
"acestep.models.configuration_acestep_v15"]
|
||||
|
||||
# Load modeling file
|
||||
modeling_file = model_dir / "modeling_acestep_v15_xl_base.py"
|
||||
if not modeling_file.exists():
|
||||
raise FileNotFoundError(f"Modeling file not found: {modeling_file}")
|
||||
|
||||
print(f"[Export] Loading model definition from {modeling_file.name}...")
|
||||
|
||||
# Add model_dir to path for local imports
|
||||
if str(model_dir) not in sys.path:
|
||||
sys.path.insert(0, str(model_dir))
|
||||
|
||||
model_mod = _load_module_from_file("_acestep_modeling", str(modeling_file))
|
||||
model_cls = model_mod.AceStepConditionGenerationModel
|
||||
|
||||
# Create model on CPU (meta device doesn't work with ResidualFSQ)
|
||||
print(f"[Export] Creating model on CPU (this uses ~10GB RAM temporarily)...")
|
||||
model = model_cls(config)
|
||||
|
||||
# Load safetensors weights
|
||||
from safetensors.torch import load_file
|
||||
|
||||
safetensors_path = model_dir / "model.safetensors"
|
||||
if not safetensors_path.exists():
|
||||
index_path = model_dir / "model.safetensors.index.json"
|
||||
if index_path.exists():
|
||||
with open(index_path) as f:
|
||||
index = json.load(f)
|
||||
weight_files = set(index["weight_map"].values())
|
||||
state_dict = {}
|
||||
for wf in sorted(weight_files):
|
||||
print(f"[Export] Loading shard: {wf}")
|
||||
shard = load_file(str(model_dir / wf), device="cpu")
|
||||
state_dict.update(shard)
|
||||
else:
|
||||
raise FileNotFoundError(f"No model.safetensors or index file in {model_dir}")
|
||||
else:
|
||||
fsize = safetensors_path.stat().st_size / (1 << 20)
|
||||
print(f"[Export] Loading weights: {safetensors_path.name} ({fsize:.0f} MB)...")
|
||||
state_dict = load_file(str(safetensors_path), device="cpu")
|
||||
|
||||
print(f"[Export] Loading {len(state_dict)} tensors into model...")
|
||||
missing, unexpected = model.load_state_dict(state_dict, strict=False)
|
||||
if missing:
|
||||
print(f"[Export] WARNING: {len(missing)} missing keys (first 5): {missing[:5]}")
|
||||
if unexpected:
|
||||
print(f"[Export] WARNING: {len(unexpected)} unexpected keys (first 5): {unexpected[:5]}")
|
||||
del state_dict # Free RAM
|
||||
|
||||
# Only move the decoder to GPU (we don't need encoders/tokenizer for export)
|
||||
model.decoder = model.decoder.to(device)
|
||||
model.eval()
|
||||
|
||||
return model, config
|
||||
|
||||
|
||||
def export_onnx(
|
||||
model_dir: str,
|
||||
output_dir: str,
|
||||
device: str = "cuda",
|
||||
batch_size: int = 1,
|
||||
seq_len: int = 750,
|
||||
enc_len: int = 200,
|
||||
opset: int = 17,
|
||||
):
|
||||
"""Export DiT decoder to ONNX in FP32."""
|
||||
output_dir = Path(output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"[Export] Loading model from {model_dir}...")
|
||||
model, config = load_model(model_dir, device=device)
|
||||
|
||||
# Extract the decoder (DiT model)
|
||||
decoder = model.decoder
|
||||
param_count = sum(p.numel() for p in decoder.parameters()) / 1e6
|
||||
print(f"[Export] Decoder: {config.num_hidden_layers} layers, "
|
||||
f"hidden={config.hidden_size}, {param_count:.0f}M params")
|
||||
|
||||
# Wrap for export
|
||||
wrapper = DiTForTRTExport(decoder).float().to(device).eval()
|
||||
print("[Export] Exporting in FP32 precision")
|
||||
|
||||
B, T, S = batch_size, seq_len, enc_len
|
||||
print(f"[Export] Trace shapes: B={B}, T={T}, S={S}")
|
||||
|
||||
example_inputs = (
|
||||
torch.randn(B, T, 192, device=device, dtype=torch.float32),
|
||||
torch.randn(B, S, 2048, device=device, dtype=torch.float32),
|
||||
torch.full((B,), 0.5, device=device, dtype=torch.float32),
|
||||
torch.full((B,), 0.5, device=device, dtype=torch.float32),
|
||||
)
|
||||
|
||||
input_names = ["input_latents", "enc_hidden", "t", "t_r"]
|
||||
output_names = ["velocity"]
|
||||
|
||||
dynamic_axes = {
|
||||
"input_latents": {0: "batch", 1: "seq_len"},
|
||||
"enc_hidden": {0: "batch", 1: "enc_len"},
|
||||
"t": {0: "batch"},
|
||||
"t_r": {0: "batch"},
|
||||
"velocity": {0: "batch", 1: "seq_len"},
|
||||
}
|
||||
|
||||
onnx_path = output_dir / "dit.onnx"
|
||||
|
||||
with torch.no_grad():
|
||||
print("[Export] Running test forward pass...")
|
||||
test_out = wrapper(*example_inputs)
|
||||
print(f"[Export] Test output shape: {test_out.shape}, "
|
||||
f"mean={test_out.float().mean():.6f}, "
|
||||
f"std={test_out.float().std():.6f}")
|
||||
|
||||
if torch.isnan(test_out).any():
|
||||
print("[Export] ERROR: test output contains NaN! Aborting.")
|
||||
return None
|
||||
|
||||
print("[Export] Exporting to ONNX (this may take a few minutes)...")
|
||||
torch.onnx.export(
|
||||
wrapper,
|
||||
example_inputs,
|
||||
str(onnx_path),
|
||||
input_names=input_names,
|
||||
output_names=output_names,
|
||||
dynamic_axes=dynamic_axes,
|
||||
opset_version=opset,
|
||||
do_constant_folding=True,
|
||||
dynamo=False,
|
||||
)
|
||||
|
||||
size_mb = onnx_path.stat().st_size / (1 << 20)
|
||||
print(f"[Export] ONNX saved to {onnx_path} ({size_mb:.1f} MB)")
|
||||
|
||||
# Copy config
|
||||
model_dir_path = Path(model_dir)
|
||||
for fname in ["config.json", "silence_latent.pt"]:
|
||||
src = model_dir_path / fname
|
||||
if src.exists():
|
||||
shutil.copy2(str(src), str(output_dir / fname))
|
||||
print(f"[Export] Copied {fname}")
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"[Export] SUCCESS! ONNX model at: {onnx_path}")
|
||||
print(f"[Export] Size: {size_mb:.1f} MB")
|
||||
print(f"{'='*60}")
|
||||
|
||||
return onnx_path
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Export AceStep DiT to ONNX")
|
||||
parser.add_argument("--model-dir", required=True)
|
||||
parser.add_argument("--output-dir", required=True)
|
||||
parser.add_argument("--device", default="cuda")
|
||||
parser.add_argument("--batch-size", type=int, default=1)
|
||||
parser.add_argument("--seq-len", type=int, default=750)
|
||||
parser.add_argument("--enc-len", type=int, default=200)
|
||||
parser.add_argument("--opset", type=int, default=17)
|
||||
|
||||
args = parser.parse_args()
|
||||
export_onnx(
|
||||
model_dir=args.model_dir,
|
||||
output_dir=args.output_dir,
|
||||
device=args.device,
|
||||
batch_size=args.batch_size,
|
||||
seq_len=args.seq_len,
|
||||
enc_len=args.enc_len,
|
||||
opset=args.opset,
|
||||
)
|
||||
@@ -0,0 +1,86 @@
|
||||
#!/usr/bin/env python
|
||||
"""goldens_to_bin.py — flatten dump_bs_roformer_goldens.py output for C++.
|
||||
|
||||
engine/tools/bs-roformer-test.cpp reads a flat little-endian binary rather than
|
||||
an .npz, because parsing zip+npy in C++ buys nothing here.
|
||||
|
||||
Layout:
|
||||
magic "BSRG" 4 bytes
|
||||
int32 T, in_dim, depth, dim, n_bands, n_stems
|
||||
f32 input [in_dim * T]
|
||||
f32 band_split [dim * n_bands * T]
|
||||
f32 layer_00 [dim * n_bands * T]
|
||||
f32 layer_01 [dim * n_bands * T]
|
||||
f32 layer_last [dim * n_bands * T]
|
||||
f32 final_norm [dim * n_bands * T]
|
||||
f32 mask [n_stems * in_dim * T]
|
||||
|
||||
Hidden-state arrays arrive from torch as [1, T, n_bands, dim] and are written
|
||||
in GGML memory order (dim fastest, then n_bands, then T) so the C++ side can
|
||||
compare against its debug tensor elementwise without reindexing.
|
||||
|
||||
The mask arrives as [1, S, fs, T, 2] and is already in the engine's order.
|
||||
|
||||
USAGE
|
||||
$py = "d:\\Ace-Step-Latest\\hot-step-9000\\.venv\\Scripts\\python.exe"
|
||||
& $py tools\\goldens_to_bin.py `
|
||||
--input models\\supersep-ckpt\\goldens_voc.npz `
|
||||
--output models\\supersep-ckpt\\goldens_voc.bin
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--input", required=True)
|
||||
ap.add_argument("--output", required=True)
|
||||
args = ap.parse_args()
|
||||
|
||||
z = np.load(args.input)
|
||||
T, in_dim, depth, n_stems, fs = (int(v) for v in z["meta"])
|
||||
|
||||
band_split = z["band_split"] # [1, T, n_bands, dim]
|
||||
_, t_chk, n_bands, dim = band_split.shape
|
||||
assert t_chk == T, f"T mismatch {t_chk} vs {T}"
|
||||
|
||||
print(f"T={T} in_dim={in_dim} depth={depth} dim={dim} "
|
||||
f"bands={n_bands} stems={n_stems} fs={fs}")
|
||||
|
||||
expect = T * n_bands * dim
|
||||
|
||||
def hidden(name):
|
||||
# Rank varies: band_split/final_norm are hooked outside the pack and
|
||||
# come back [1, T, n_bands, dim]; the per-layer hooks sit on the freq
|
||||
# transformer, which sees the packed [(b t), n_bands, dim]. Batch is 1,
|
||||
# so both flatten to the same GGML order (dim fastest, n_bands, T) —
|
||||
# just flatten whatever rank arrived and check the count.
|
||||
a = np.ascontiguousarray(z[name], dtype=np.float32).reshape(-1)
|
||||
if a.size != expect:
|
||||
raise SystemExit(
|
||||
f"{name}: {a.size} elements, expected {expect} "
|
||||
f"(shape {z[name].shape})")
|
||||
return a
|
||||
|
||||
with open(args.output, "wb") as f:
|
||||
f.write(b"BSRG")
|
||||
f.write(struct.pack("<6i", T, in_dim, depth, dim, n_bands, n_stems))
|
||||
|
||||
# input is [1, T, in_dim]; GGML wants in_dim fastest then T — same order.
|
||||
f.write(np.ascontiguousarray(z["input"][0], dtype=np.float32).tobytes())
|
||||
|
||||
for name in ("band_split", "layer_00", "layer_01", "layer_last", "final_norm"):
|
||||
f.write(hidden(name).tobytes())
|
||||
|
||||
# mask [1, S, fs, T, 2] -> linear [s][fs][t][2], already correct
|
||||
f.write(np.ascontiguousarray(z["mask"][0], dtype=np.float32).tobytes())
|
||||
|
||||
import os
|
||||
print(f"[ok] {args.output} ({os.path.getsize(args.output)/1024/1024:.1f} MB)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user