61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for audio utility functions."""
|
|
|
|
import io
|
|
import wave
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
|
|
from omlx.engine.audio_utils import audio_to_wav_bytes
|
|
|
|
|
|
def _read_wav(data: bytes):
|
|
"""Helper: parse WAV bytes and return (samples_int16, sample_rate)."""
|
|
buf = io.BytesIO(data)
|
|
with wave.open(buf, "rb") as wf:
|
|
sr = wf.getframerate()
|
|
frames = wf.readframes(wf.getnframes())
|
|
return np.frombuffer(frames, dtype=np.int16), sr
|
|
|
|
|
|
class TestAudioToWavBytes:
|
|
def test_numpy_float32(self):
|
|
audio = np.array([0.0, 0.5, -0.5, 1.0, -1.0], dtype=np.float32)
|
|
wav = audio_to_wav_bytes(audio, 24000)
|
|
samples, sr = _read_wav(wav)
|
|
assert sr == 24000
|
|
assert len(samples) == 5
|
|
|
|
def test_mlx_float32(self):
|
|
audio = mx.array([0.0, 0.5, -0.5], dtype=mx.float32)
|
|
wav = audio_to_wav_bytes(audio, 16000)
|
|
samples, sr = _read_wav(wav)
|
|
assert sr == 16000
|
|
assert len(samples) == 3
|
|
|
|
def test_mlx_bfloat16(self):
|
|
"""bfloat16 arrays (e.g. Kokoro-82M-bf16) must not crash."""
|
|
audio = mx.array([0.0, 0.5, -0.5, 1.0, -1.0], dtype=mx.bfloat16)
|
|
wav = audio_to_wav_bytes(audio, 24000)
|
|
samples, sr = _read_wav(wav)
|
|
assert sr == 24000
|
|
assert len(samples) == 5
|
|
# Verify values are reasonable after bf16 → f32 → int16 conversion
|
|
assert samples[0] == 0
|
|
assert samples[3] == 32767 # 1.0 clipped
|
|
assert samples[4] == -32767 # -1.0 clipped
|
|
|
|
def test_mlx_float16(self):
|
|
audio = mx.array([0.0, 0.25], dtype=mx.float16)
|
|
wav = audio_to_wav_bytes(audio, 24000)
|
|
samples, _ = _read_wav(wav)
|
|
assert len(samples) == 2
|
|
|
|
def test_clipping(self):
|
|
audio = np.array([2.0, -2.0], dtype=np.float32)
|
|
wav = audio_to_wav_bytes(audio, 24000)
|
|
samples, _ = _read_wav(wav)
|
|
assert samples[0] == 32767
|
|
assert samples[1] == -32767
|