mlxrunner: shared audio decoding and chunking for audio models

Audio-capable models need mono PCM at their expected sample rate
before model-specific feature extraction. Like the image decoder set
in base, the supported audio containers are decided once here so
every model accepts the same formats: WAV, covering integer and float
PCM, extensible headers, and multi-channel downmix. Anything else is
rejected as unrecognized; supporting another container later means
one new decoder here, with no model or runner changes.

Input at other sample rates is resampled through a band-limiting
filter, so mismatched rates degrade gracefully instead of aliasing.
Decoded clips are capped at ten minutes: the declared rate comes from
an untrusted header, and the cap is what keeps a small file claiming
an absurdly low rate from resampling into an enormous allocation.

Models whose encoder takes clips only up to a fixed length split
longer clips with Split. Chunks are sized evenly, and each cut moves
to the quietest point within a few seconds of its even share, so a
boundary lands on a pause where the clip has one instead of severing
a word between two independently encoded chunks. The even sizing
bounds the search so no chunk can end up over the limit.

x/mlxrunner/model/audio/audio.go

@@ -0,0 +1,245 @@
+package audio
+
+import (
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"math"
+	"sync"
+)
+
+// maxAudioSeconds caps decoded clip duration. The declared rate is
+// untrusted; this is what bounds Resample's output length.
+const maxAudioSeconds = 600
+
+// Decode decodes a supported audio container into mono PCM samples at
+// the container's declared sample rate, downmixing multiple channels by
+// averaging. The supported containers are decided here so all models accept
+// the same formats; WAV is the only one today. Clips longer than
+// maxAudioSeconds are rejected. Models resample to their own rate with
+// Resample.
+func Decode(data []byte) ([]float32, int, error) {
+	if len(data) < 12 || string(data[0:4]) != "RIFF" || string(data[8:12]) != "WAVE" {
+		return nil, 0, errors.New("unrecognized audio format")
+	}
+	samples, rate, err := decodeWAV(data)
+	if err != nil {
+		return nil, 0, err
+	}
+	if len(samples) > maxAudioSeconds*rate {
+		return nil, 0, fmt.Errorf("audio longer than %d seconds", maxAudioSeconds)
+	}
+	return samples, rate, nil
+}
+
+func decodeWAV(data []byte) ([]float32, int, error) {
+	var format uint16
+	var channels, rate, bits int
+	var pcm []byte
+	foundFmt := false
+
+	offset := 12
+	for offset+8 <= len(data) {
+		id := string(data[offset : offset+4])
+		size := int(binary.LittleEndian.Uint32(data[offset+4 : offset+8]))
+		body := data[offset+8 : min(offset+8+size, len(data))]
+
+		switch id {
+		case "fmt ":
+			if len(body) < 16 {
+				return nil, 0, errors.New("wav: fmt chunk too short")
+			}
+			format = binary.LittleEndian.Uint16(body[0:2])
+			channels = int(binary.LittleEndian.Uint16(body[2:4]))
+			rate = int(binary.LittleEndian.Uint32(body[4:8]))
+			bits = int(binary.LittleEndian.Uint16(body[14:16]))
+			// WAVE_FORMAT_EXTENSIBLE carries the real format code in the
+			// extension's subformat GUID.
+			if format == 0xfffe && len(body) >= 26 {
+				format = binary.LittleEndian.Uint16(body[24:26])
+			}
+			foundFmt = true
+		case "data":
+			pcm = body
+		}
+
+		// Chunks are word-aligned.
+		offset += 8 + size + size%2
+	}
+
+	if !foundFmt {
+		return nil, 0, errors.New("wav: no fmt chunk")
+	}
+	if pcm == nil {
+		return nil, 0, errors.New("wav: no data chunk")
+	}
+	if channels < 1 || rate < 1 {
+		return nil, 0, fmt.Errorf("wav: invalid fmt: %d channels at %d Hz", channels, rate)
+	}
+
+	var sample func([]byte) float64
+	switch {
+	case format == 1 && bits == 8:
+		sample = func(b []byte) float64 { return (float64(b[0]) - 128) / 128 }
+	case format == 1 && bits == 16:
+		sample = func(b []byte) float64 {
+			return float64(int16(binary.LittleEndian.Uint16(b))) / 32768
+		}
+	case format == 1 && bits == 24:
+		sample = func(b []byte) float64 {
+			v := int32(b[0]) | int32(b[1])<<8 | int32(b[2])<<16
+			return float64(v<<8>>8) / 8388608
+		}
+	case format == 1 && bits == 32:
+		sample = func(b []byte) float64 {
+			return float64(int32(binary.LittleEndian.Uint32(b))) / 2147483648
+		}
+	case format == 3 && bits == 32:
+		sample = func(b []byte) float64 {
+			return float64(math.Float32frombits(binary.LittleEndian.Uint32(b)))
+		}
+	default:
+		return nil, 0, fmt.Errorf("wav: unsupported format %d with %d bits", format, bits)
+	}
+
+	stride := bits / 8
+	n := len(pcm) / (stride * channels)
+	mono := make([]float32, n)
+	for i := range n {
+		var sum float64
+		for c := range channels {
+			sum += sample(pcm[(i*channels+c)*stride:])
+		}
+		mono[i] = float32(sum / float64(channels))
+	}
+	return mono, rate, nil
+}
+
+const (
+	resampleLobes   = 16  // sinc zero crossings per side of the kernel
+	resampleRolloff = 0.9 // cutoff as a fraction of the lower Nyquist
+	resampleBeta    = 9.0 // Kaiser shape, roughly 80 dB of stopband
+)
+
+// resampleKernel tabulates the Kaiser-windowed sinc over half its support;
+// weights interpolate linearly between entries.
+var resampleKernel = sync.OnceValue(func() []float64 {
+	const size = 8192
+	i0 := besselI0(resampleBeta)
+	kernel := make([]float64, size+1)
+	for i := range kernel {
+		u := float64(i) / size
+		kernel[i] = sinc(resampleLobes*u) * besselI0(resampleBeta*math.Sqrt(1-u*u)) / i0
+	}
+	return kernel
+})
+
+// Resample converts samples between rates with a Kaiser-windowed sinc
+// kernel evaluated at each output sample's fractional source position. The
+// cutoff sits below the lower of the two Nyquist frequencies, so the same
+// kernel low-passes when downsampling and rejects images when upsampling.
+// Weights are normalized per output sample, which keeps unity gain where
+// the clip edges truncate the kernel.
+func Resample(samples []float32, from, to int) []float32 {
+	if from == to || len(samples) < 2 {
+		return samples
+	}
+
+	ratio := float64(to) / float64(from)
+	scale := resampleRolloff * min(1, ratio)
+	halfWidth := resampleLobes / scale
+	kernel := resampleKernel()
+
+	n := int((int64(len(samples))*int64(to) + int64(from) - 1) / int64(from))
+	out := make([]float32, n)
+	for i := range out {
+		center := float64(i) / ratio
+		lo := max(int(math.Ceil(center-halfWidth)), 0)
+		hi := min(int(math.Floor(center+halfWidth)), len(samples)-1)
+
+		var acc, norm float64
+		for j := lo; j <= hi; j++ {
+			pos := math.Abs(float64(j)-center) / halfWidth * float64(len(kernel)-1)
+			idx := int(pos)
+			if idx >= len(kernel)-1 {
+				continue
+			}
+			w := kernel[idx] + (kernel[idx+1]-kernel[idx])*(pos-float64(idx))
+			acc += w * float64(samples[j])
+			norm += w
+		}
+		if norm != 0 {
+			out[i] = float32(acc / norm)
+		}
+	}
+	return out
+}
+
+// Cuts go at the lowest-energy window of splitWindowMillis found within
+// splitSearchSeconds around each chunk's even share of the clip.
+const (
+	splitSearchSeconds = 4
+	splitWindowMillis  = 100
+)
+
+// Split divides samples into the fewest chunks of at most maxSeconds each,
+// sized evenly, with each cut moved to the quietest point near its even
+// share so a boundary lands on a pause where the clip allows. Chunks are
+// subslices of samples.
+func Split(samples []float32, rate, maxSeconds int) [][]float32 {
+	maxChunk := rate * maxSeconds
+	span := rate * splitSearchSeconds
+	window := max(rate*splitWindowMillis/1000, 1)
+	var chunks [][]float32
+	for k := (len(samples) + maxChunk - 1) / maxChunk; k > 1; k-- {
+		target := len(samples) / k
+		// The rest of the clip must still fit in k-1 chunks.
+		lo := max(target-span/2, len(samples)-(k-1)*maxChunk)
+		hi := min(target+span/2, maxChunk)
+		cut := lo + quietestPoint(samples[lo:hi], window)
+		chunks = append(chunks, samples[:cut])
+		samples = samples[cut:]
+	}
+	return append(chunks, samples)
+}
+
+// quietestPoint returns the center of the lowest-energy window in x,
+// the latest one on ties.
+func quietestPoint(x []float32, window int) int {
+	if len(x) <= window {
+		return len(x) / 2
+	}
+	var energy float64
+	for _, v := range x[:window] {
+		energy += float64(v) * float64(v)
+	}
+	best, bestEnergy := 0, energy
+	for i := window; i < len(x); i++ {
+		energy += float64(x[i])*float64(x[i]) - float64(x[i-window])*float64(x[i-window])
+		if energy <= bestEnergy {
+			best, bestEnergy = i-window+1, energy
+		}
+	}
+	return best + window/2
+}
+
+func sinc(x float64) float64 {
+	if x == 0 {
+		return 1
+	}
+	return math.Sin(math.Pi*x) / (math.Pi * x)
+}
+
+// besselI0 is the zeroth-order modified Bessel function of the first kind,
+// by power series.
+func besselI0(x float64) float64 {
+	sum, term := 1.0, 1.0
+	for k := 1; k < 64; k++ {
+		term *= x * x / (4 * float64(k) * float64(k))
+		sum += term
+		if term < 1e-16*sum {
+			break
+		}
+	}
+	return sum
+}

x/mlxrunner/model/audio/audio_test.go

@@ -0,0 +1,292 @@
+package audio
+
+import (
+	"bytes"
+	"encoding/binary"
+	"math"
+	"strings"
+	"testing"
+)
+
+func wavBytes(format uint16, channels, rate, bits int, pcm []byte) []byte {
+	var b bytes.Buffer
+	b.WriteString("RIFF")
+	binary.Write(&b, binary.LittleEndian, uint32(36+len(pcm)))
+	b.WriteString("WAVE")
+	b.WriteString("fmt ")
+	binary.Write(&b, binary.LittleEndian, uint32(16))
+	binary.Write(&b, binary.LittleEndian, format)
+	binary.Write(&b, binary.LittleEndian, uint16(channels))
+	binary.Write(&b, binary.LittleEndian, uint32(rate))
+	binary.Write(&b, binary.LittleEndian, uint32(rate*channels*bits/8))
+	binary.Write(&b, binary.LittleEndian, uint16(channels*bits/8))
+	binary.Write(&b, binary.LittleEndian, uint16(bits))
+	b.WriteString("data")
+	binary.Write(&b, binary.LittleEndian, uint32(len(pcm)))
+	b.Write(pcm)
+	return b.Bytes()
+}
+
+func requireSamples(t *testing.T, data []byte, wantRate int, want []float32, tol float32) {
+	t.Helper()
+	samples, rate, err := Decode(data)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if rate != wantRate {
+		t.Fatalf("rate %d, want %d", rate, wantRate)
+	}
+	if len(samples) != len(want) {
+		t.Fatalf("%d samples, want %d", len(samples), len(want))
+	}
+	for i := range want {
+		if diff := samples[i] - want[i]; diff > tol || diff < -tol {
+			t.Fatalf("sample %d: %v, want %v", i, samples[i], want[i])
+		}
+	}
+}
+
+func TestDecodeWAVPCM16(t *testing.T) {
+	var pcm bytes.Buffer
+	for _, v := range []int16{0, 16384, -16384, 32767, -32768} {
+		binary.Write(&pcm, binary.LittleEndian, v)
+	}
+	want := []float32{0, 0.5, -0.5, 32767.0 / 32768, -1}
+	requireSamples(t, wavBytes(1, 1, 16000, 16, pcm.Bytes()), 16000, want, 0)
+}
+
+func TestDecodeWAVPCM8(t *testing.T) {
+	pcm := []byte{128, 255, 0, 192}
+	want := []float32{0, 127.0 / 128, -1, 0.5}
+	requireSamples(t, wavBytes(1, 1, 8000, 8, pcm), 8000, want, 0)
+}
+
+func TestDecodeWAVPCM24(t *testing.T) {
+	var pcm []byte
+	for _, v := range []int32{0, 1 << 22, -(1 << 22)} {
+		pcm = append(pcm, byte(v), byte(v>>8), byte(v>>16))
+	}
+	want := []float32{0, 0.5, -0.5}
+	requireSamples(t, wavBytes(1, 1, 16000, 24, pcm), 16000, want, 0)
+}
+
+func TestDecodeWAVPCM32(t *testing.T) {
+	var pcm bytes.Buffer
+	for _, v := range []int32{0, 1 << 30, -(1 << 30)} {
+		binary.Write(&pcm, binary.LittleEndian, v)
+	}
+	want := []float32{0, 0.5, -0.5}
+	requireSamples(t, wavBytes(1, 1, 16000, 32, pcm.Bytes()), 16000, want, 0)
+}
+
+func TestDecodeWAVFloat32(t *testing.T) {
+	var pcm bytes.Buffer
+	want := []float32{0, 0.25, -1, 1}
+	for _, v := range want {
+		binary.Write(&pcm, binary.LittleEndian, v)
+	}
+	requireSamples(t, wavBytes(3, 1, 44100, 32, pcm.Bytes()), 44100, want, 0)
+}
+
+func TestDecodeWAVStereoDownmix(t *testing.T) {
+	var pcm bytes.Buffer
+	for _, v := range []float32{1, 0, -0.5, 0.5} {
+		binary.Write(&pcm, binary.LittleEndian, v)
+	}
+	want := []float32{0.5, 0}
+	requireSamples(t, wavBytes(3, 2, 16000, 32, pcm.Bytes()), 16000, want, 0)
+}
+
+func TestDecodeWAVExtensible(t *testing.T) {
+	var pcm bytes.Buffer
+	binary.Write(&pcm, binary.LittleEndian, int16(16384))
+
+	var b bytes.Buffer
+	b.WriteString("RIFF")
+	binary.Write(&b, binary.LittleEndian, uint32(60+pcm.Len()))
+	b.WriteString("WAVE")
+	b.WriteString("fmt ")
+	binary.Write(&b, binary.LittleEndian, uint32(40))
+	binary.Write(&b, binary.LittleEndian, uint16(0xfffe))
+	binary.Write(&b, binary.LittleEndian, uint16(1))
+	binary.Write(&b, binary.LittleEndian, uint32(16000))
+	binary.Write(&b, binary.LittleEndian, uint32(32000))
+	binary.Write(&b, binary.LittleEndian, uint16(2))
+	binary.Write(&b, binary.LittleEndian, uint16(16))
+	binary.Write(&b, binary.LittleEndian, uint16(22)) // extension size
+	binary.Write(&b, binary.LittleEndian, uint16(16)) // valid bits
+	binary.Write(&b, binary.LittleEndian, uint32(0))  // channel mask
+	binary.Write(&b, binary.LittleEndian, uint16(1))  // subformat: PCM
+	b.Write(make([]byte, 14))                         // rest of subformat GUID
+	b.WriteString("data")
+	binary.Write(&b, binary.LittleEndian, uint32(pcm.Len()))
+	b.Write(pcm.Bytes())
+
+	requireSamples(t, b.Bytes(), 16000, []float32{0.5}, 0)
+}
+
+func TestDecodeErrors(t *testing.T) {
+	adpcm := wavBytes(2, 1, 16000, 4, make([]byte, 8))
+	noData := wavBytes(1, 1, 16000, 16, nil)
+	noData = noData[:len(noData)-8]
+	zeroChannels := wavBytes(1, 0, 16000, 16, make([]byte, 4))
+	// A tiny file whose declared 1 Hz rate makes it run for over 10 minutes.
+	tooLong := wavBytes(1, 1, 1, 8, make([]byte, maxAudioSeconds+1))
+
+	cases := []struct {
+		name string
+		data []byte
+		want string
+	}{
+		{"garbage", []byte("not audio at all"), "unrecognized audio format"},
+		{"empty", nil, "unrecognized audio format"},
+		{"mp3 id3", []byte("ID3\x04\x00rest"), "unrecognized audio format"},
+		{"mp3 sync", []byte{0xff, 0xfb, 0x90, 0x00}, "unrecognized audio format"},
+		{"truncated riff", []byte("RIFF\x00\x00\x00\x00WAVE"), "no fmt chunk"},
+		{"no data chunk", noData, "no data chunk"},
+		{"adpcm", adpcm, "unsupported format 2"},
+		{"zero channels", zeroChannels, "invalid fmt"},
+		{"longer than the duration cap", tooLong, "audio longer"},
+	}
+	for _, tt := range cases {
+		t.Run(tt.name, func(t *testing.T) {
+			_, _, err := Decode(tt.data)
+			if err == nil || !strings.Contains(err.Error(), tt.want) {
+				t.Fatalf("error %v, want %q", err, tt.want)
+			}
+		})
+	}
+}
+
+func TestResample(t *testing.T) {
+	tone := func(n int, freq, rate float64) []float32 {
+		out := make([]float32, n)
+		for i := range out {
+			out[i] = float32(math.Sin(2 * math.Pi * freq * float64(i) / rate))
+		}
+		return out
+	}
+	// The kernel spans ~18 output samples at these ratios; stay well clear
+	// of the edge truncation when comparing against the analytic tone.
+	const margin = 200
+
+	same := tone(100, 1000, 16000)
+	if got := Resample(same, 16000, 16000); &got[0] != &same[0] {
+		t.Fatal("same-rate resample should return the input")
+	}
+
+	// A band-limited tone resampled ideally is the same tone sampled at
+	// the new rate.
+	down := Resample(tone(44100, 1000, 44100), 44100, 16000)
+	if len(down) != 16000 {
+		t.Fatalf("%d samples, want 16000", len(down))
+	}
+	for i := margin; i < len(down)-margin; i++ {
+		want := math.Sin(2 * math.Pi * 1000 * float64(i) / 16000)
+		if diff := math.Abs(float64(down[i]) - want); diff > 5e-3 {
+			t.Fatalf("downsample at %d: %v, want %v", i, down[i], want)
+		}
+	}
+
+	up := Resample(tone(8000, 1000, 8000), 8000, 16000)
+	if len(up) != 16000 {
+		t.Fatalf("%d samples, want 16000", len(up))
+	}
+	for i := margin; i < len(up)-margin; i++ {
+		want := math.Sin(2 * math.Pi * 1000 * float64(i) / 16000)
+		if diff := math.Abs(float64(up[i]) - want); diff > 5e-3 {
+			t.Fatalf("upsample at %d: %v, want %v", i, up[i], want)
+		}
+	}
+
+	// Content above the target Nyquist must be filtered out, not aliased
+	// into the band.
+	alias := Resample(tone(44100, 10000, 44100), 44100, 16000)
+	for i := margin; i < len(alias)-margin; i++ {
+		if math.Abs(float64(alias[i])) > 0.01 {
+			t.Fatalf("aliased content at %d: %v", i, alias[i])
+		}
+	}
+
+	// Per-sample weight normalization keeps DC exact.
+	dc := make([]float32, 1000)
+	for i := range dc {
+		dc[i] = 0.5
+	}
+	for i, v := range Resample(dc, 44100, 16000) {
+		if math.Abs(float64(v)-0.5) > 1e-6 {
+			t.Fatalf("dc at %d: %v", i, v)
+		}
+	}
+}
+
+// requireChunks checks that chunks are consecutive subslices covering
+// samples, each at most maxChunk long, and returns the cut offsets.
+func requireChunks(t *testing.T, samples []float32, chunks [][]float32, maxChunk int) []int {
+	t.Helper()
+	if want := (len(samples) + maxChunk - 1) / maxChunk; len(chunks) != want {
+		t.Fatalf("%d chunks, want %d", len(chunks), want)
+	}
+	var cuts []int
+	off := 0
+	for i, c := range chunks {
+		if len(c) == 0 || len(c) > maxChunk {
+			t.Fatalf("chunk %d has %d samples, limit %d", i, len(c), maxChunk)
+		}
+		if &c[0] != &samples[off] {
+			t.Fatalf("chunk %d does not start at sample %d", i, off)
+		}
+		off += len(c)
+		cuts = append(cuts, off)
+	}
+	if off != len(samples) {
+		t.Fatalf("chunks cover %d of %d samples", off, len(samples))
+	}
+	return cuts[:len(cuts)-1]
+}
+
+func tone(n int) []float32 {
+	s := make([]float32, n)
+	for i := range s {
+		s[i] = float32(0.8 * math.Sin(float64(i)*0.3))
+	}
+	return s
+}
+
+func TestSplitAtPauses(t *testing.T) {
+	const rate, maxSeconds = 16000, 30
+	// 95 s needs four chunks. A 200 ms pause sits 1.3 s past each quarter
+	// mark, inside every cut's search span, and is the only silence.
+	samples := tone(95 * rate)
+	const pause = rate / 5
+	var pauses []int
+	for i := 1; i < 4; i++ {
+		start := i*len(samples)/4 + 13*rate/10
+		clear(samples[start : start+pause])
+		pauses = append(pauses, start)
+	}
+	cuts := requireChunks(t, samples, Split(samples, rate, maxSeconds), maxSeconds*rate)
+	for i, cut := range cuts {
+		if cut < pauses[i] || cut > pauses[i]+pause {
+			t.Fatalf("cut %d at sample %d, pause spans [%d, %d]", i, cut, pauses[i], pauses[i]+pause)
+		}
+	}
+}
+
+func TestSplitLimits(t *testing.T) {
+	const rate, maxSeconds = 16000, 30
+	const maxChunk = maxSeconds * rate
+	whole := tone(maxChunk)
+	if chunks := Split(whole, rate, maxSeconds); len(chunks) != 1 || &chunks[0][0] != &whole[0] || len(chunks[0]) != maxChunk {
+		t.Fatalf("clip at the limit split into %d chunks", len(chunks))
+	}
+	// One sample over, and one short of a whole number of full chunks:
+	// the pause search must not push any chunk past the limit.
+	for _, n := range []int{maxChunk + 1, 2*maxChunk - 1, 3*maxChunk - 1} {
+		s := tone(n)
+		requireChunks(t, s, Split(s, rate, maxSeconds), maxChunk)
+	}
+	// A limit shorter than the search span still yields chunks within it.
+	s := tone(5*rate + 1)
+	requireChunks(t, s, Split(s, rate, 1), rate)
+}