dance-ai

Real-time musical bar-phase estimation from live audio streams. Runs the dance models via ONNX in the browser (onnxruntime-web) or Node.js (onnxruntime-node).

This package provides:

1. On-the-fly Mel Log-Frontend: A streaming-aligned Mel log-spectrogram generator (MelFrontend) matching the PyTorch reference implementation exactly.
2. Causal Level Normalizer: Keeps the audio input volume normalized without lookahead (RunningPeakNormalizer), preventing volume differences from throwing off the estimator.
3. ONNX Inference Wrapper: Manages model input/output state (DancePhaseEstimator), feeding mono audio sequences and maintaining hidden recurrent states between frames.
4. Phase Extrapolator: A smooth, monotonic phase clock (PhaseExtrapolator) that estimates and continues the tempo forward between model inferences to avoid phase jitter.

Installation

npm install dance-ai

[!NOTE] onnxruntime-web (or onnxruntime-node for server-side environments) is required as a peer dependency.

Core Concepts

For each audio frame (≈16.7 ms), the model outputs a FrameEstimate:

• phase: Position inside the current musical bar, in [0, 1) (where 0 is the exact start of the bar/downbeat).
• barDurationS: The predicted duration of the current bar in seconds (directly related to tempo, i.e., 60 * beats_per_bar / bpm).
• expectedPhaseError: Calibrated uncertainty of the phase prediction in bars, in [0, 0.5].

No Anticipation Input

To ensure zero lookahead delay, the model has no future anticipation input. Under a constant tempo, the phase can be extrapolated forward by your application using the predicted barDurationS (or by passing the outputs to the included PhaseExtrapolator).

Usage Guide

1. Browser (Web Audio API & AudioWorklet)

To run in the browser, you need to capture mono audio at the model's required samplerate (typically 24 kHz) and chunk it into the model's expected frame size (typically 400 samples). dance-ai provides helper functions to build a web-audio worklet node easily.

import * as ort from 'onnxruntime-web';
import {
  DancePhaseEstimator,
  PhaseExtrapolator,
  createAudioProcessorUrl,
  AUDIO_PROCESSOR_NAME,
} from 'dance-ai';

// 1. Initialize the estimator with the ONNX model and metadata
const estimator = await DancePhaseEstimator.create({
  ort,
  model: '/model.onnx',
  meta: '/model.meta.json', // served alongside your model
});

// 2. Initialize the smooth phase extrapolator clock
const clock = new PhaseExtrapolator({ smoothing: 0.5 });

// 3. Setup the Web Audio context matching the model's samplerate
const audioContext = new AudioContext({ sampleRate: estimator.meta.samplerate });
await audioContext.audioWorklet.addModule(createAudioProcessorUrl());

const node = new AudioWorkletNode(audioContext, AUDIO_PROCESSOR_NAME, {
  processorOptions: { frameSize: estimator.meta.frame_size },
});

// 4. Stream audio frames from the worklet to the estimator
node.port.onmessage = async (event) => {
  const samples = event.data as Float32Array;
  const estimates = await estimator.feed(samples);
  for (const e of estimates) {
    clock.update(e.phase, e.barDurationS);
  }
};

// Connect microphone/audio source to the worklet node
const source = audioContext.createMediaStreamSource(micStream);
node.connect(audioContext.destination);
source.connect(node);

// 5. Render loop (runs at screen refresh rate, e.g., 60fps)
function render() {
  const smoothPhase = clock.phaseAt(); // Smoothly extrapolated value in [0, 1)
  // Update your visuals, animations, or UI using the smooth phase...
  
  requestAnimationFrame(render);
}

2. Node.js (Server / CLI)

To run on the server, use onnxruntime-node and load the files from the filesystem.

import * as ort from 'onnxruntime-node';
import { DancePhaseEstimator } from 'dance-ai';
import * as fs from 'fs';

// Load files
const modelBytes = fs.readFileSync('./model.onnx');
const meta = JSON.parse(fs.readFileSync('./model.meta.json', 'utf8'));

// Initialize estimator
const estimator = await DancePhaseEstimator.create({
  ort,
  model: modelBytes,
  meta,
});

// Feed mono audio samples at the required sample rate
const audioSamples = new Float32Array(/* ... read from wav / decoded audio ... */);
const estimates = await estimator.feed(audioSamples);

for (const e of estimates) {
  console.log(`Phase: ${e.phase.toFixed(3)}, Temp: ${(60 * 4 / e.barDurationS).toFixed(1)} BPM`);
}

Numerical Alignment with PyTorch

For deep-learning inference to be accurate, the preprocessing frontend and recurrent context of the inference engine must match the training script exactly:

1. Exact-Size DFT Frontend: The MelFrontend performs a real DFT of size $N_{fft} = \text{fftFrames} \times \text{frameSize}$ (commonly $2 \times 400 = 800$). Standard FFT libraries pad inputs to power-of-2 sizes (e.g., 1024), which changes the frequency grid spacing and leads to severe prediction mismatch. This library implements a precise DFT matching the PyTorch implementation without power-of-2 padding.
2. Context Continuity: The Mel spectrogram requires past audio samples to compute log-mel features for the current frame. MelFrontend buffers the last (fftFrames - 1) * frameSize audio samples to use as context for the next frame. Call estimator.reset() when switching audio streams to clear this recurrent context.
3. Causal Level Normalizer: Training uses a global peak normalization (scaling the audio to a max peak of 1.0 over the entire sample). Since a live stream cannot see the future, RunningPeakNormalizer implements an exponential decay running peak tracking mechanism to causally scale incoming frames.

API Reference

DancePhaseEstimator

Manages ONNX inference, audio feature preprocessing, and recurrent state.

• static async create(options: EstimatorOptions): Promise<DancePhaseEstimator> Creates an instance of the estimator. Fetches or parses the metadata and loads the ONNX session.
• feed(samples: Float32Array): Promise<FrameEstimate[]> Feeds mono audio samples at meta.samplerate. Partial frames are buffered. Returns estimates for any completed frames. Handles re-entrant calls sequentially.
• reset(): void Resets the recurrent model states, Mel frontend audio context, normalizer history, and the sample buffer. Call this when switching songs or restarting the stream.
• dispose(): void Releases ONNX inference session resources and model tensors.

EstimatorOptions

| Property | Type | Description | | :--- | :--- | :--- | | ort | OrtNamespaceLike | The imported ONNX Runtime namespace (onnxruntime-web or onnxruntime-node). | | model | string \| ArrayBufferLike \| Uint8Array | The path/URL to the .onnx file, or raw model bytes. | | meta | ModelMeta \| string | Parsed metadata object, or path/URL to the .meta.json file. | | normalize | boolean | Enable causal level normalization (default: true). | | normalizerHalfLifeS | number | Exponential half-life in seconds for the running peak tracker (default: 30). | | normalizerFloor | number | Minimum peak floor to avoid boosting noise/silence (default: 0.01). | | sessionOptions | object | Native options passed to the ONNX InferenceSession. |

PhaseExtrapolator

Converts discrete, per-frame model predictions into a smooth, continuous phase clock.

• update(phase: number, barDurationS: number, nowMs?: number): void Feeds a new observation from the model.
• phaseAt(nowMs?: number): number Returns the current extrapolated, wrapped phase in [0, 1).
• unwrappedPhaseAt(nowMs?: number): number Returns the current monotonically increasing unwrapped phase (total bars since start).
• getRate(): number Returns the current estimated bar-rate in bars/second.
• setSmoothing(value: number): void Dynamically changes the smoothing factor.
• reset(): void Resets the clock rate, aligning the next phase to the next integer boundary to prevent jump-back.

PhaseExtrapolatorOptions

| Property | Type | Description | | :--- | :--- | :--- | | smoothing | number | 0 follows raw model output instantly (jittery, low latency); 1 is pure tempo extrapolation (smooth, slow to correct). Default: 0.5. | | rateAlpha | number | EMA smoothing factor for the bar rate (default: 0.15 ≈ 100 ms at 60 fps). |

Audio Worklet Helpers

Exposes helper utilities to capture frame-aligned audio samples inside an browser AudioWorklet.

• audioProcessorSource: string The raw source code of the AudioWorkletProcessor class.
• AUDIO_PROCESSOR_NAME: string The name the processor is registered under ('dance-audio-processor').
• createAudioProcessorUrl(): string Creates a self-referencing Blob URL pointing to the worklet source code that can be immediately registered using audioContext.audioWorklet.addModule(url).

Exporting Models

To export a model checkpoint from the parent dance repository into the format required by dance-ai:

# Inside the dance parent repository:
venv/bin/python export_onnx.py checkpoints/<experiment>/<tag>/<epoch>.pt

This generates:

1. <epoch>.onnx: The ONNX model file.
2. <epoch>.meta.json: The model metadata containing neural network config and samplerate.

Development

Install dependencies:

npm install

Build the TypeScript files:

npm run build

Testing

To test the log-mel frontend and normalizer against PyTorch reference vectors:

# 1. Regenerate golden vectors if you modified the Python frontend:
(cd .. && venv/bin/python dance-ai/scripts/gen_test_vectors.py)

# 2. Run the tests:
npm test

License

MIT License. Designed and developed by Felix Niemeyer.