WebAssembly Audio Pipeline: High-Performance MIR in the Browser

1// Complete WASM Audio Pipeline Implementation
2class WASMAudioPipeline {
3  constructor() {
4    this.audioContext = null;
5    this.decoder = null;
6    this.essentia = null;
7    this.onnxSession = null;
8    this.initialized = false;
9  }
10  
11  async initialize() {
12    // Initialize Web Audio API
13    this.audioContext = new (window.AudioContext || window.webkitAudioContext)();
14    
15    // Load WASM modules in parallel
16    const [decoder, essentia, onnx] = await Promise.all([
17      this.initializeDecoder(),
18      this.initializeEssentia(),
19      this.initializeONNX()
20    ]);
21    
22    this.decoder = decoder;
23    this.essentia = essentia;
24    this.onnxSession = onnx;
25    this.initialized = true;
26  }
27  
28  async initializeDecoder() {
29    const { FLACDecoder } = await import('@wasm-audio-decoders/flac');
30    const decoder = new FLACDecoder();
31    await decoder.ready;
32    return decoder;
33  }
34  
35  async initializeEssentia() {
36    const { Essentia, EssentiaWASM } = await import('essentia.js');
37    const essentia = new Essentia(EssentiaWASM);
38    return essentia;
39  }
40  
41  async initializeONNX() {
42    const ort = await import('onnxruntime-web');
43    
44    // Configure WASM backend with optimizations
45    ort.env.wasm.numThreads = navigator.hardwareConcurrency || 4;
46    ort.env.wasm.simd = true;
47    
48    // Load model (e.g., Demucs)
49    const session = await ort.InferenceSession.create(
50      '/models/demucs_quantized.onnx',
51      {
52        executionProviders: ['wasm'],
53        graphOptimizationLevel: 'all'
54      }
55    );
56    
57    return session;
58  }
59  
60  async processAudioFile(file) {
61    if (!this.initialized) {
62      await this.initialize();
63    }
64    
65    // Step 1: Decode audio file
66    const pcmData = await this.decodeAudio(file);
67    
68    // Step 2: Extract features
69    const features = await this.extractFeatures(pcmData);
70    
71    // Step 3: Run AI model for source separation
72    const separatedSources = await this.separateSources(features.spectrogram);
73    
74    // Step 4: Post-process and prepare for playback
75    const processedAudio = await this.postProcess(separatedSources);
76    
77    return {
78      pcmData,
79      features,
80      separatedSources,
81      processedAudio
82    };
83  }
84  
85  async decodeAudio(file) {
86    const arrayBuffer = await file.arrayBuffer();
87    const uint8Array = new Uint8Array(arrayBuffer);
88    
89    // Decode based on file type
90    let decoder;
91    if (file.name.endsWith('.flac')) {
92      const { FLACDecoder } = await import('@wasm-audio-decoders/flac');
93      decoder = new FLACDecoder();
94    } else if (file.name.endsWith('.mp3')) {
95      const { MPEGDecoder } = await import('@wasm-audio-decoders/mpeg');
96      decoder = new MPEGDecoder();
97    }
98    
99    await decoder.ready;
100    
101    // Decode file
102    const { channelData, sampleRate, numberOfChannels } = 
103      await decoder.decodeFile(uint8Array);
104    
105    // Convert to mono if needed
106    let monoData;
107    if (numberOfChannels > 1) {
108      monoData = new Float32Array(channelData[0].length);
109      for (let i = 0; i < channelData[0].length; i++) {
110        monoData[i] = (channelData[0][i] + channelData[1][i]) / 2;
111      }
112    } else {
113      monoData = channelData[0];
114    }
115    
116    // Free decoder resources
117    decoder.free();
118    
119    return {
120      data: monoData,
121      sampleRate,
122      duration: monoData.length / sampleRate
123    };
124  }
125  
126  async extractFeatures(pcmData) {
127    // Convert to Essentia vector
128    const audioVector = this.essentia.arrayToVector(pcmData.data);
129    
130    // Compute various features
131    const features = {};
132    
133    // Mel Spectrogram
134    features.melSpectrogram = this.computeMelSpectrogram(audioVector);
135    
136    // MFCC
137    features.mfcc = this.computeMFCC(audioVector);
138    
139    // Chroma
140    features.chroma = this.computeChroma(audioVector);
141    
142    // Beat tracking
143    features.beats = this.detectBeats(audioVector);
144    
145    // Key detection
146    features.key = this.detectKey(audioVector);
147    
148    // Clean up
149    this.essentia.delete(audioVector);
150    
151    return features;
152  }
153  
154  computeMelSpectrogram(audioVector) {
155    const frameSize = 2048;
156    const hopSize = 512;
157    const melBands = 128;
158    
159    // Frame generator
160    const frames = this.essentia.FrameGenerator(
161      audioVector, frameSize, hopSize
162    );
163    
164    const melSpectrogram = [];
165    
166    for (let i = 0; i < frames.size(); i++) {
167      const frame = frames.get(i);
168      
169      // Window the frame
170      const windowed = this.essentia.Windowing(
171        frame, true, frameSize, 'hann'
172      );
173      
174      // Compute spectrum
175      const spectrum = this.essentia.Spectrum(
176        windowed.frame, frameSize
177      );
178      
179      // Compute mel bands
180      const mels = this.essentia.MelBands(
181        spectrum.spectrum,
182        melBands,
183        44100,
184        0,
185        22050
186      );
187      
188      melSpectrogram.push(
189        this.essentia.vectorToArray(mels.bands)
190      );
191      
192      // Clean up frame memory
193      this.essentia.delete(frame);
194      this.essentia.delete(windowed.frame);
195      this.essentia.delete(spectrum.spectrum);
196      this.essentia.delete(mels.bands);
197    }
198    
199    this.essentia.delete(frames);
200    
201    return melSpectrogram;
202  }
203  
204  async separateSources(spectrogram) {
205    // Prepare input tensor
206    const inputData = Float32Array.from(spectrogram.flat());
207    const inputTensor = new ort.Tensor(
208      'float32',
209      inputData,
210      [1, 1, spectrogram.length, spectrogram[0].length]
211    );
212    
213    // Run inference
214    const feeds = { 'input': inputTensor };
215    const results = await this.onnxSession.run(feeds);
216    
217    // Extract separated sources
218    const sources = {
219      vocals: results.vocals.data,
220      drums: results.drums.data,
221      bass: results.bass.data,
222      other: results.other.data
223    };
224    
225    return sources;
226  }
227}

1// Memory-efficient WASM operations
2class WASMMemoryManager {
3  constructor() {
4    this.allocations = new Map();
5  }
6  
7  // Track allocations
8  allocate(module, size, id) {
9    const ptr = module._malloc(size);
10    this.allocations.set(id, { module, ptr, size });
11    return ptr;
12  }
13  
14  // Automatic cleanup with try-finally
15  async withMemory(module, size, callback) {
16    const ptr = module._malloc(size);
17    const heap = new Float32Array(
18      module.HEAPF32.buffer,
19      ptr,
20      size / 4
21    );
22    
23    try {
24      return await callback(heap, ptr);
25    } finally {
26      module._free(ptr);
27    }
28  }
29  
30  // Batch processing to minimize allocations
31  processBatch(data, chunkSize = 8192) {
32    const results = [];
33    
34    for (let i = 0; i < data.length; i += chunkSize) {
35      const chunk = data.slice(i, i + chunkSize);
36      
37      this.withMemory(module, chunk.length * 4, (heap) => {
38        // Copy data to WASM heap
39        heap.set(chunk);
40        
41        // Process in WASM
42        const result = module.processAudio(heap.byteOffset, chunk.length);
43        
44        // Copy result back
45        const output = new Float32Array(chunk.length);
46        output.set(heap.subarray(0, chunk.length));
47        
48        results.push(output);
49      });
50    }
51    
52    return results;
53  }
54  
55  // Clean up all allocations
56  cleanup() {
57    for (const [id, { module, ptr }] of this.allocations) {
58      module._free(ptr);
59    }
60    this.allocations.clear();
61  }
62}
63
64// Web Worker for heavy processing
65class AudioProcessorWorker {
66  constructor() {
67    this.worker = new Worker('/audio-processor.worker.js');
68    this.pending = new Map();
69    this.nextId = 0;
70    
71    this.worker.onmessage = (e) => {
72      const { id, result, error } = e.data;
73      const { resolve, reject } = this.pending.get(id);
74      
75      if (error) {
76        reject(new Error(error));
77      } else {
78        resolve(result);
79      }
80      
81      this.pending.delete(id);
82    };
83  }
84  
85  async process(data, operation) {
86    const id = this.nextId++;
87    
88    return new Promise((resolve, reject) => {
89      this.pending.set(id, { resolve, reject });
90      
91      // Transfer data to worker (zero-copy)
92      this.worker.postMessage(
93        { id, operation, data },
94        [data.buffer]
95      );
96    });
97  }
98  
99  terminate() {
100    this.worker.terminate();
101  }
102}

1// audio-processor.worklet.js
2class WASMProcessor extends AudioWorkletProcessor {
3  constructor() {
4    super();
5    this.initialized = false;
6    this.wasmModule = null;
7    
8    // Initialize WASM in worklet
9    this.port.onmessage = async (e) => {
10      if (e.data.type === 'init') {
11        await this.initialize(e.data.wasmUrl);
12      }
13    };
14  }
15  
16  async initialize(wasmUrl) {
17    // Load WASM module
18    const response = await fetch(wasmUrl);
19    const wasmBytes = await response.arrayBuffer();
20    
21    const { instance } = await WebAssembly.instantiate(wasmBytes, {
22      env: {
23        memory: new WebAssembly.Memory({ initial: 256 })
24      }
25    });
26    
27    this.wasmModule = instance.exports;
28    this.initialized = true;
29    
30    // Allocate processing buffers
31    this.inputPtr = this.wasmModule.allocate(128 * 4);
32    this.outputPtr = this.wasmModule.allocate(128 * 4);
33  }
34  
35  process(inputs, outputs, parameters) {
36    if (!this.initialized || inputs[0].length === 0) {
37      return true;
38    }
39    
40    const input = inputs[0][0];
41    const output = outputs[0][0];
42    
43    // Copy input to WASM memory
44    const inputHeap = new Float32Array(
45      this.wasmModule.memory.buffer,
46      this.inputPtr,
47      128
48    );
49    inputHeap.set(input);
50    
51    // Process in WASM
52    this.wasmModule.processBlock(
53      this.inputPtr,
54      this.outputPtr,
55      128
56    );
57    
58    // Copy output from WASM memory
59    const outputHeap = new Float32Array(
60      this.wasmModule.memory.buffer,
61      this.outputPtr,
62      128
63    );
64    output.set(outputHeap);
65    
66    return true;
67  }
68}
69
70registerProcessor('wasm-processor', WASMProcessor);
71
72// Main thread code
73async function setupRealtimeProcessing() {
74  const audioContext = new AudioContext();
75  
76  // Register worklet
77  await audioContext.audioWorklet.addModule(
78    '/audio-processor.worklet.js'
79  );
80  
81  // Create processor node
82  const processorNode = new AudioWorkletNode(
83    audioContext,
84    'wasm-processor'
85  );
86  
87  // Initialize WASM in worklet
88  processorNode.port.postMessage({
89    type: 'init',
90    wasmUrl: '/processing.wasm'
91  });
92  
93  // Connect audio graph
94  const source = audioContext.createMediaStreamSource(stream);
95  source.connect(processorNode);
96  processorNode.connect(audioContext.destination);
97}

1// Lazy load WASM modules based on user needs
2class LazyWASMLoader {
3  constructor() {
4    this.modules = new Map();
5    this.loading = new Map();
6  }
7  
8  async loadModule(name) {
9    // Return cached module
10    if (this.modules.has(name)) {
11      return this.modules.get(name);
12    }
13    
14    // Wait for ongoing load
15    if (this.loading.has(name)) {
16      return this.loading.get(name);
17    }
18    
19    // Start loading
20    const loadPromise = this._loadModule(name);
21    this.loading.set(name, loadPromise);
22    
23    try {
24      const module = await loadPromise;
25      this.modules.set(name, module);
26      this.loading.delete(name);
27      return module;
28    } catch (error) {
29      this.loading.delete(name);
30      throw error;
31    }
32  }
33  
34  async _loadModule(name) {
35    switch (name) {
36      case 'decoder-mp3':
37        return import('@wasm-audio-decoders/mpeg');
38        
39      case 'decoder-flac':
40        return import('@wasm-audio-decoders/flac');
41        
42      case 'essentia-core':
43        // Load only core features
44        return import('essentia.js/dist/essentia-wasm.core.js');
45        
46      case 'essentia-full':
47        // Load full library when needed
48        return import('essentia.js/dist/essentia-wasm.js');
49        
50      case 'onnx-runtime':
51        return import('onnxruntime-web');
52        
53      default:
54        throw new Error(`Unknown module: ${name}`);
55    }
56  }
57  
58  // Preload modules in the background
59  preload(modules) {
60    modules.forEach(name => {
61      requestIdleCallback(() => this.loadModule(name));
62    });
63  }
64}
65
66// Usage
67const loader = new LazyWASMLoader();
68
69// Preload common modules
70loader.preload(['decoder-mp3', 'essentia-core']);
71
72// Load on demand
73async function processFile(file) {
74  const extension = file.name.split('.').pop();
75  
76  // Load appropriate decoder
77  const decoderModule = await loader.loadModule(`decoder-${extension}`);
78  
79  // Process file...
80}