osobhandClaude Opus 4.6
85b77d49f2
New nn layers: - ConvTranspose1d with stride, padding, output_padding (9 tests) - LSTM/BiLSTM with multi-layer support and hidden state (10 tests) Audio source separation: - Demucs ONNX inference with segmented overlap-add processing - Native HtDemucs architecture (encoder/decoder with BiLSTM bottleneck) - StemType enum: vocals, drums, bass, other, piano, guitar Audio generation: - Stable Audio Open ONNX inference scaffold - GenerationParams (prompt, duration, steps, cfg_scale, seed) ONNX export scripts: - export_demucs_onnx.py — Demucs v4 to ONNX with segment chunking - export_stable_audio_onnx.py — Stable Audio Open components - export_mert_onnx.py — MERT music understanding transformer Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Export Meta's HTDemucs model to ONNX format for use with RustyTorch++.
|
|
|
|
Usage:
|
|
pip install torch demucs onnx
|
|
python export_demucs_onnx.py --stems 4 --output models/htdemucs.onnx
|
|
python export_demucs_onnx.py --stems 6 --output models/htdemucs_6s.onnx
|
|
|
|
The exported model expects:
|
|
Input: "mix" shape [1, 2, segment_length] (stereo audio at 44100 Hz)
|
|
Output: "stems" shape [1, num_stems, 2, segment_length]
|
|
|
|
Segment length is fixed at export time (default: 441000 = 10 seconds at 44.1 kHz).
|
|
The Rust inference code handles longer audio via segmented overlap-add.
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
|
|
import torch
|
|
import torch.onnx
|
|
|
|
|
|
def export_demucs(num_stems: int, output_path: str, segment_length: int = 441000):
|
|
# Import demucs and load the pretrained model
|
|
from demucs.pretrained import get_model
|
|
|
|
model_name = "htdemucs" if num_stems == 4 else "htdemucs_6s"
|
|
print(f"Loading {model_name} ({num_stems} stems)...")
|
|
|
|
model = get_model(model_name)
|
|
model.eval()
|
|
|
|
# Create dummy input: [batch=1, channels=2, samples=segment_length]
|
|
dummy_input = torch.randn(1, 2, segment_length)
|
|
|
|
print(f"Exporting to ONNX: {output_path}")
|
|
print(f" Input shape: [1, 2, {segment_length}]")
|
|
print(f" Output shape: [1, {num_stems}, 2, {segment_length}]")
|
|
|
|
os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
|
|
|
|
torch.onnx.export(
|
|
model,
|
|
dummy_input,
|
|
output_path,
|
|
input_names=["mix"],
|
|
output_names=["stems"],
|
|
opset_version=17,
|
|
do_constant_folding=True,
|
|
dynamic_axes=None, # Fixed shapes for reliable ONNX inference
|
|
)
|
|
|
|
# Verify the exported model
|
|
import onnx
|
|
onnx_model = onnx.load(output_path)
|
|
onnx.checker.check_model(onnx_model)
|
|
|
|
file_size_mb = os.path.getsize(output_path) / (1024 * 1024)
|
|
print(f"Export complete: {output_path} ({file_size_mb:.1f} MB)")
|
|
print(f"Model verified successfully.")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Export HTDemucs to ONNX")
|
|
parser.add_argument("--stems", type=int, default=4, choices=[4, 6],
|
|
help="Number of stems (4 or 6)")
|
|
parser.add_argument("--output", type=str, default=None,
|
|
help="Output ONNX file path")
|
|
parser.add_argument("--segment-length", type=int, default=441000,
|
|
help="Segment length in samples (default: 441000 = 10s at 44.1kHz)")
|
|
args = parser.parse_args()
|
|
|
|
if args.output is None:
|
|
name = "htdemucs" if args.stems == 4 else "htdemucs_6s"
|
|
args.output = f"models/{name}.onnx"
|
|
|
|
export_demucs(args.stems, args.output, args.segment_length)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|