The ONNX Revolution: Deploying AI Audio Models in Production

# Step 1: Install conversion tools
pip install torch torchvision onnx onnxruntime
pip install spleeter-pytorch  # Community port

# Step 2: Load PyTorch version
import torch
from spleeter_pytorch import Separator

model = Separator('2stems')
dummy_input = torch.randn(1, 2, 44100 * 10)  # 10 seconds stereo

# Step 3: Export to ONNX
torch.onnx.export(
    model,
    dummy_input,
    "spleeter_2stems.onnx",
    input_names=['mixture'],
    output_names=['vocals', 'accompaniment'],
    dynamic_axes={
        'mixture': {2: 'sequence_length'},
        'vocals': {2: 'sequence_length'},
        'accompaniment': {2: 'sequence_length'}
    }
)

# Modified Demucs forward pass for ONNX
class DemucsONNX(nn.Module):
    def __init__(self, original_demucs):
        super().__init__()
        # Extract only the neural network components
        self.encoder = original_demucs.encoder
        self.decoder = original_demucs.decoder
        self.lstm = original_demucs.lstm
        
    def forward(self, mix_spec):
        # Input: pre-computed spectrogram
        # Skip internal STFT operations
        encoded = self.encoder(mix_spec)
        processed = self.lstm(encoded)
        decoded = self.decoder(processed)
        return decoded  # Output: separated spectrograms

# Application-level pipeline
def separate_audio(audio_tensor):
    # External STFT
    mix_spec = torch.stft(audio_tensor, ...)
    
    # ONNX model inference
    separated_specs = onnx_session.run(None, {'input': mix_spec})
    
    # External iSTFT
    separated_audio = [torch.istft(spec, ...) for spec in separated_specs]
    return separated_audio

// C++ production server example
#include <onnxruntime_cxx_api.h>
#include <vector>
#include <memory>

class AudioSeparationService {
private:
    std::unique_ptr<Ort::Session> session_;
    std::vector<const char*> input_names_;
    std::vector<const char*> output_names_;

public:
    AudioSeparationService(const std::string& model_path) {
        Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "AudioSeparation");
        
        Ort::SessionOptions options;
        options.SetIntraOpNumThreads(8);
        options.SetGraphOptimizationLevel(
            GraphOptimizationLevel::ORT_ENABLE_ALL);
        
        // Use CUDA if available
        Ort::ThrowOnError(
            OrtSessionOptionsAppendExecutionProvider_CUDA(options, 0));
        
        session_ = std::make_unique<Ort::Session>(
            env, model_path.c_str(), options);
    }
    
    std::vector<std::vector<float>> separate(
        const std::vector<float>& audio_data) {
        
        // Create input tensor
        std::vector<int64_t> input_shape = {1, 2, audio_data.size() / 2};
        Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(
            OrtArenaAllocator, OrtMemTypeDefault);
        
        Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
            memory_info, const_cast<float*>(audio_data.data()), 
            audio_data.size(), input_shape.data(), input_shape.size());
        
        // Run inference
        auto output_tensors = session_->Run(
            Ort::RunOptions{nullptr}, 
            input_names_.data(), &input_tensor, 1,
            output_names_.data(), output_names_.size());
        
        // Extract results
        std::vector<std::vector<float>> results;
        for (auto& tensor : output_tensors) {
            float* tensor_data = tensor.GetTensorMutableData<float>();
            size_t tensor_size = tensor.GetTensorTypeAndShapeInfo().GetElementCount();
            results.emplace_back(tensor_data, tensor_data + tensor_size);
        }
        
        return results;
    }
};

# Quantization for mobile deployment
import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

# Dynamic quantization - reduces model size by ~75%
quantized_model = quantize_dynamic(
    model_input="demucs_f32.onnx",
    model_output="demucs_int8.onnx",
    weight_type=QuantType.QInt8
)

# For even smaller models, use static quantization
from onnxruntime.quantization import quantize_static, CalibrationDataReader

class AudioCalibrationDataReader(CalibrationDataReader):
    def __init__(self, calibration_dataset):
        self.dataset = calibration_dataset
        self.iter = iter(calibration_dataset)
    
    def get_next(self):
        return next(self.iter, None)

quantize_static(
    model_input="demucs_f32.onnx",
    model_output="demucs_static_int8.onnx",
    calibration_data_reader=AudioCalibrationDataReader(calib_data)
)

# Enable detailed profiling
session_options = onnxruntime.SessionOptions()
session_options.enable_profiling = True

session = onnxruntime.InferenceSession(
    "model.onnx", 
    sess_options=session_options,
    providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
)

# Run inference
outputs = session.run(None, inputs)

# Analyze performance profile
prof_file = session.end_profiling()
print(f"Profile saved to: {prof_file}")

# Memory usage tracking
import tracemalloc
tracemalloc.start()

# Your inference code here
outputs = session.run(None, inputs)

current, peak = tracemalloc.get_traced_memory()
print(f"Current memory usage: {current / 1024 / 1024:.1f} MB")
print(f"Peak memory usage: {peak / 1024 / 1024:.1f} MB")