diff --git a/x/mlxrunner/model/audio/audio.go b/x/mlxrunner/model/audio/audio.go new file mode 100644 index 000000000..9f32b3599 --- /dev/null +++ b/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 +} diff --git a/x/mlxrunner/model/audio/audio_test.go b/x/mlxrunner/model/audio/audio_test.go new file mode 100644 index 000000000..18834c591 --- /dev/null +++ b/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) +}