ONNX & Web Inference

What is ONNX?

ONNX (Open Neural Network Exchange) is an open format for representing machine learning models. It allows models trained in one framework (PyTorch, TensorFlow) to run in another.

Benefits:

• Portability: Train in Python, run in browser
• Optimization: ONNX Runtime optimizes for different hardware
• Standardization: One format, many platforms

ONNX Runtime Web

ONNX Runtime Web runs ONNX models in the browser using WebAssembly and WebGL.

// Load ONNX Runtime
import * as ort from 'onnxruntime-web';

// Configure execution providers
ort.env.wasm.numThreads = 4; // Use multiple threads

// Load model
const session = await ort.InferenceSession.create('./model.onnx', {
    executionProviders: ['wasm'] // 'webgl' also available
});

// Prepare input
const inputData = new Float32Array([/* ... */]);
const inputTensor = new ort.Tensor('float32', inputData, [1, 1, 512, 512]);

// Run inference
const feeds = { input: inputTensor };
const results = await session.run(feeds);

// Get output
const output = results.output.data; // Float32Array

WebAssembly (WASM)

WebAssembly is a low-level binary format that runs at near-native speed in the browser.

• Compiled from C/C++/Rust
• Runs in a sandbox
• 3-5x slower than native, but 10-100x faster than JavaScript for math

ONNX Runtime uses WASM to run neural networks efficiently in the browser.

// WASM loading in MixMind
async function loadWasm() {
    // Load WASM binaries from CDN
    const wasmUrl = 'https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort-wasm-simd.wasm';
    
    // Configure ONNX Runtime
    ort.env.wasm.wasmPaths = {
        'ort-wasm-simd.wasm': wasmUrl
    };
    
    // Now we can create inference sessions
    const session = await ort.InferenceSession.create('model.onnx');
}

Performance Considerations

Running ML in the browser has constraints:

• Model size: Large models = slow download, high memory
• Computation: Limited by CPU/GPU power
• Memory: Browser memory limits
• Threading: Use Web Workers for background processing

// Optimization strategies
// 1. Model quantization: reduce precision (float32 → int8)
// 2. Pruning: remove unimportant weights
// 3. Chunking: process audio in small pieces
// 4. Streaming: show progress, don't block UI

// Example: chunked processing
async function processInChunks(audioBuffer, chunkSize = 5) {
    const duration = audioBuffer.duration;
    const stems = { vocals: [], drums: [], bass: [], other: [] };
    
    for (let start = 0; start < duration; start += chunkSize) {
        const end = Math.min(start + chunkSize, duration);
        const chunk = extractChunk(audioBuffer, start, end);
        
        const chunkStems = await model.separate(chunk);
        
        stems.vocals.push(chunkStems.vocals);
        stems.drums.push(chunkStems.drums);
        stems.bass.push(chunkStems.bass);
        stems.other.push(chunkStems.other);
        
        // Update progress
        updateProgress(end / duration * 100);
    }
    
    // Concatenate chunks
    return {
        vocals: concatenateAudio(stems.vocals),
        drums: concatenateAudio(stems.drums),
        bass: concatenateAudio(stems.bass),
        other: concatenateAudio(stems.other)
    };
}

// MixMind's ONNX integration
class OnnxEngine {
    constructor() {
        this.session = null;
        this.initialized = false;
    }
    
    async initialize() {
        if (this.initialized) return;
        
        try {
            // Configure WASM
            ort.env.wasm.numThreads = navigator.hardwareConcurrency || 4;
            
            // Load model
            this.session = await ort.InferenceSession.create('./models/spleeter.onnx', {
                executionProviders: ['wasm'],
                graphOptimizationLevel: 'all'
            });
            
            this.initialized = true;
            console.log('✅ ONNX model loaded');
        } catch (error) {
            console.error('Failed to load ONNX model:', error);
            throw error;
        }
    }
    
    async runInference(inputTensor) {
        if (!this.initialized) {
            await this.initialize();
        }
        
        const feeds = { 'input_1': inputTensor };
        const results = await this.session.run(feeds);
        
        return results;
    }
}

// Audio preprocessing
function prepareInput(audioBuffer) {
    // 1. Convert to mono
    const mono = toMono(audioBuffer);
    
    // 2. Convert to spectrogram
    const spectrogram = audioToSpectrogram(mono);
    
    // 3. Normalize
    const normalized = normalize(spectrogram);
    
    // 4. Create tensor
    const tensor = new ort.Tensor(
        'float32',
        normalized.data,
        [1, 1, normalized.height, normalized.width]
    );
    
    return tensor;
}

// Full pipeline
async function separateStems(audioBuffer) {
    const engine = new OnnxEngine();
    await engine.initialize();
    
    // Process in chunks to avoid memory issues
    const chunkDuration = 5; // seconds
    const chunks = splitIntoChunks(audioBuffer, chunkDuration);
    
    const results = [];
    
    for (let i = 0; i < chunks.length; i++) {
        const chunk = chunks[i];
        const input = prepareInput(chunk);
        const output = await engine.runInference(input);
        
        results.push({
            vocals: tensorToAudio(output.vocals),
            drums: tensorToAudio(output.drums),
            bass: tensorToAudio(output.bass),
            other: tensorToAudio(output.other)
        });
        
        // Update progress
        updateProgress((i + 1) / chunks.length * 100);
    }
    
    // Concatenate chunks
    return mergeChunks(results);
}