mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-16 13:20:39 -07:00
Enhance MLX Audio Documentation and Testing Framework
- Updated MLX_AUDIO.md to reflect validated status and included detailed validation results, model mapping, and API usage examples. - Added a demo script (demo.py) for testing audio generation speed and functionality. - Introduced a test script (test_tts.py) to validate MLX audio model loading and generation, ensuring robust testing for future developments. - Created a .gitignore file in the mlx-test directory to exclude unnecessary files from version control.
This commit is contained in:
Jamie Pine
+87
-4
@@ -1,8 +1,58 @@
|
||||
# MLX Audio Integration
|
||||
|
||||
**Status:** Planned
|
||||
**Status:** Validated ✅
|
||||
**Context:** [mlx-audio v0.3.1 release](https://github.com/Blaizzy/mlx-audio)
|
||||
|
||||
## Validation Results
|
||||
|
||||
We validated mlx-audio in an isolated environment (`mlx-test/`). Key findings:
|
||||
|
||||
| Metric | Result |
|
||||
|--------|--------|
|
||||
| MLX Version | 0.30.4 |
|
||||
| Model Load Time | ~1s (after initial download) |
|
||||
| Generation RTF | **0.5-0.6x** (1.7-2x faster than real-time) |
|
||||
| Test Hardware | Apple Silicon Mac |
|
||||
|
||||
### Model Mapping
|
||||
|
||||
| voicebox (PyTorch) | mlx-audio (MLX) |
|
||||
|--------------------|-----------------|
|
||||
| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | `mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16` |
|
||||
| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | (not yet converted) |
|
||||
|
||||
### mlx-audio API
|
||||
|
||||
The API uses a **generator-based streaming pattern**:
|
||||
|
||||
```python
|
||||
from mlx_audio.tts import load
|
||||
|
||||
model = load("mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16")
|
||||
|
||||
# generate() yields GenerationResult objects
|
||||
for result in model.generate("Hello world"):
|
||||
audio = result.audio # numpy array of samples
|
||||
sample_rate = result.sample_rate # 24000
|
||||
rtf = result.real_time_factor # e.g., 0.55
|
||||
```
|
||||
|
||||
### Known Warnings (harmless)
|
||||
|
||||
```
|
||||
You are using a model of type qwen3_tts to instantiate a model of type .
|
||||
The tokenizer you are loading... with an incorrect regex pattern...
|
||||
```
|
||||
|
||||
These warnings appear but don't affect functionality or output quality.
|
||||
|
||||
### Demo Script
|
||||
|
||||
Run `mlx-test/demo.py` to test:
|
||||
```bash
|
||||
cd mlx-test && source venv/bin/activate && python demo.py "Your text here"
|
||||
```
|
||||
|
||||
## Problem
|
||||
|
||||
Apple Silicon users are stuck on CPU inference while Windows and Linux users get CUDA acceleration. The current PyTorch MPS backend has stability issues (lines 34-36 in `backend/tts.py` and `backend/transcribe.py`), forcing a CPU fallback that makes voicebox significantly slower on M1/M2/M3 Macs.
|
||||
@@ -105,6 +155,33 @@ class STTBackend(Protocol):
|
||||
def unload_model(self) -> None: ...
|
||||
```
|
||||
|
||||
**MLX backend implementation notes:**
|
||||
|
||||
mlx-audio's `generate()` returns a generator by default (streaming is built-in):
|
||||
|
||||
```python
|
||||
# MLX backend wrapper
|
||||
from mlx_audio.tts import load
|
||||
|
||||
class MLXTTSBackend:
|
||||
def __init__(self):
|
||||
self.model = None
|
||||
|
||||
async def load_model(self, model_size: str) -> None:
|
||||
model_map = {
|
||||
"1.7B": "mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16",
|
||||
# "0.6B": needs conversion to mlx format
|
||||
}
|
||||
self.model = load(model_map[model_size])
|
||||
|
||||
async def generate(self, text: str, voice_prompt: dict, **kwargs) -> Tuple[np.ndarray, int]:
|
||||
# Collect all chunks from generator
|
||||
chunks = []
|
||||
for result in self.model.generate(text): # TODO: add voice_prompt support
|
||||
chunks.append(np.array(result.audio))
|
||||
return np.concatenate(chunks), 24000
|
||||
```
|
||||
|
||||
**MLX-specific features to expose:**
|
||||
- Streaming TTS (new endpoint: `/api/generate/stream`)
|
||||
- Memory-optimized model loading
|
||||
@@ -261,7 +338,13 @@ Nothing needs migrating, macos users will just notice a speed-boost in inference
|
||||
|
||||
## Performance Expectations
|
||||
|
||||
Based on mlx-audio benchmarks and community reports:
|
||||
### Measured Results (from validation)
|
||||
|
||||
| Metric | MLX (measured) | PyTorch CPU (estimated) |
|
||||
|--------|----------------|-------------------------|
|
||||
| **6s audio generation** | ~3-4s | ~10-15s |
|
||||
| **Real-time factor** | 0.5-0.6x | 2-3x |
|
||||
| **Model load (cached)** | ~1s | ~3-5s |
|
||||
|
||||
### TTS Generation (1.7B model, ~20s output)
|
||||
- **PyTorch CPU (M2 Max):** ~45-60s (slower than real-time)
|
||||
@@ -278,7 +361,7 @@ Based on mlx-audio benchmarks and community reports:
|
||||
- **MLX:** ~4-6GB (unified memory, better optimization)
|
||||
- **Improvement:** ~40% less RAM
|
||||
|
||||
These are estimates. Actual benchmarks will be in `docs/overview/performance.md` after Phase 6.
|
||||
Full benchmarks will be in `docs/overview/performance.md` after Phase 6.
|
||||
|
||||
## Open Questions
|
||||
|
||||
@@ -304,7 +387,7 @@ How we'll know this worked:
|
||||
|
||||
## Next Steps
|
||||
|
||||
1. Validate mlx-audio can load Qwen3-TTS models (quick test)
|
||||
1. ~~Validate mlx-audio can load Qwen3-TTS models (quick test)~~ ✅ Done - see `mlx-test/`
|
||||
2. Get approval on dual-backend architecture
|
||||
3. Start Phase 1 (platform detection)
|
||||
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
# Virtual environment
|
||||
venv/
|
||||
|
||||
# Generated test files
|
||||
*.wav
|
||||
|
||||
# Python cache
|
||||
__pycache__/
|
||||
*.pyc
|
||||
Executable
+64
@@ -0,0 +1,64 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Quick demo script to test MLX audio generation speed.
|
||||
|
||||
Usage:
|
||||
python demo.py # Use default text
|
||||
python demo.py "Your custom text" # Use custom text
|
||||
"""
|
||||
import sys
|
||||
import time
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from mlx_audio.tts import load
|
||||
|
||||
# Default demo text
|
||||
DEFAULT_TEXT = "Hello! This is MLX audio running natively on Apple Silicon. It's incredibly fast!"
|
||||
|
||||
def main():
|
||||
text = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_TEXT
|
||||
|
||||
print(f"\n🎙️ MLX Audio Demo")
|
||||
print(f"{'=' * 50}")
|
||||
print(f"Text: \"{text}\"\n")
|
||||
|
||||
# Load model
|
||||
print("Loading model...", end=" ", flush=True)
|
||||
start = time.time()
|
||||
model = load("mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16")
|
||||
print(f"done ({time.time() - start:.1f}s)\n")
|
||||
|
||||
# Generate
|
||||
print("Generating audio...")
|
||||
start = time.time()
|
||||
|
||||
for result in model.generate(text):
|
||||
# Calculate duration from audio samples
|
||||
audio = np.array(result.audio)
|
||||
sample_rate = result.sample_rate
|
||||
duration = len(audio) / sample_rate
|
||||
gen_time = float(result.processing_time_seconds)
|
||||
rtf = gen_time / duration if duration > 0 else 0
|
||||
|
||||
print(f" Audio duration: {duration:.2f}s")
|
||||
print(f" Generation time: {gen_time:.2f}s")
|
||||
print(f" Real-time factor: {rtf:.2f}x", end="")
|
||||
|
||||
if rtf < 1.0:
|
||||
print(f" ⚡ ({1/rtf:.1f}x faster than real-time)")
|
||||
else:
|
||||
print()
|
||||
|
||||
# Save audio
|
||||
sf.write("test_output.wav", audio, sample_rate)
|
||||
|
||||
print(f"\n✅ Saved to test_output.wav")
|
||||
print(f"{'=' * 50}")
|
||||
|
||||
# Play audio
|
||||
print("\n🔊 Playing audio...\n")
|
||||
import subprocess
|
||||
subprocess.run(["afplay", "test_output.wav"])
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,207 @@
|
||||
"""
|
||||
Test script to validate mlx-audio can load and run Qwen3-TTS models.
|
||||
"""
|
||||
import sys
|
||||
import time
|
||||
|
||||
def test_mlx_available():
|
||||
"""Step 1: Verify MLX is available and working."""
|
||||
print("=" * 60)
|
||||
print("Step 1: Testing MLX availability")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
import mlx.core as mx
|
||||
print(f"✓ MLX imported successfully")
|
||||
print(f" Version: {mx.__version__ if hasattr(mx, '__version__') else 'unknown'}")
|
||||
|
||||
# Quick compute test
|
||||
a = mx.array([1.0, 2.0, 3.0])
|
||||
b = mx.array([4.0, 5.0, 6.0])
|
||||
c = a + b
|
||||
print(f" Compute test: {a.tolist()} + {b.tolist()} = {c.tolist()}")
|
||||
print("✓ MLX compute working\n")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"✗ MLX error: {e}\n")
|
||||
return False
|
||||
|
||||
|
||||
def test_mlx_audio_import():
|
||||
"""Step 2: Verify mlx-audio modules can be imported."""
|
||||
print("=" * 60)
|
||||
print("Step 2: Testing mlx-audio imports")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
import mlx_audio
|
||||
print(f"✓ mlx_audio imported")
|
||||
|
||||
from mlx_audio.tts import load
|
||||
print(f"✓ mlx_audio.tts.load imported")
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"✗ Import error: {e}\n")
|
||||
return False
|
||||
|
||||
|
||||
def test_model_loading():
|
||||
"""Step 3: Load Qwen3-TTS model (1.7B - same as voicebox uses)."""
|
||||
print("=" * 60)
|
||||
print("Step 3: Loading Qwen3-TTS model (1.7B)")
|
||||
print("=" * 60)
|
||||
print("(This will download the model on first run, ~3.4GB)")
|
||||
print()
|
||||
|
||||
# Model mapping - same as backend/tts.py but for MLX
|
||||
# PyTorch: Qwen/Qwen3-TTS-12Hz-1.7B-Base
|
||||
# MLX: mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16
|
||||
|
||||
try:
|
||||
from mlx_audio.tts import load
|
||||
|
||||
start = time.time()
|
||||
# Load the MLX-converted version of the same model voicebox uses
|
||||
model = load("mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16")
|
||||
load_time = time.time() - start
|
||||
|
||||
print(f"✓ Model loaded in {load_time:.1f}s\n")
|
||||
return model
|
||||
except Exception as e:
|
||||
print(f"✗ Model loading error: {e}\n")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return None
|
||||
|
||||
|
||||
def test_generation(model):
|
||||
"""Step 4: Generate a short audio clip."""
|
||||
print("=" * 60)
|
||||
print("Step 4: Generating test audio")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
test_text = "Hello, this is a test of MLX audio generation."
|
||||
print(f" Text: \"{test_text}\"")
|
||||
print(f" Model type: {type(model).__name__}")
|
||||
|
||||
start = time.time()
|
||||
|
||||
# mlx-audio generate() returns a generator yielding GenerationResult objects
|
||||
# Each result has: audio, sample_rate, real_time_factor, etc.
|
||||
audio_chunks = []
|
||||
sample_rate = 24000
|
||||
|
||||
for result in model.generate(test_text):
|
||||
# result is a GenerationResult with audio and metadata
|
||||
audio_chunks.append(result.audio)
|
||||
sample_rate = result.sample_rate
|
||||
|
||||
# Print streaming progress info
|
||||
if hasattr(result, 'real_time_factor') and result.real_time_factor:
|
||||
print(f" Chunk: {result.audio.shape[0]} samples, RTF: {result.real_time_factor:.2f}x")
|
||||
|
||||
gen_time = time.time() - start
|
||||
|
||||
# Concatenate all audio chunks
|
||||
import numpy as np
|
||||
audio = np.concatenate([np.array(chunk) for chunk in audio_chunks])
|
||||
|
||||
samples = len(audio)
|
||||
duration = samples / sample_rate
|
||||
rtf = gen_time / duration if duration > 0 else float('inf')
|
||||
|
||||
print(f"✓ Audio generated:")
|
||||
print(f" Samples: {samples}")
|
||||
print(f" Sample rate: {sample_rate} Hz")
|
||||
print(f" Duration: {duration:.2f}s")
|
||||
print(f" Generation time: {gen_time:.2f}s")
|
||||
print(f" Real-time factor: {rtf:.2f}x (lower is faster)")
|
||||
|
||||
if rtf < 1.0:
|
||||
print(f" → Faster than real-time!")
|
||||
|
||||
return audio, sample_rate
|
||||
except Exception as e:
|
||||
print(f"✗ Generation error: {e}\n")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return None, None
|
||||
|
||||
|
||||
def test_save_audio(audio, sample_rate):
|
||||
"""Step 5: Save the generated audio to a file."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Step 5: Saving audio file")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
# Audio should already be a numpy array from test_generation
|
||||
audio_np = np.asarray(audio, dtype=np.float32)
|
||||
|
||||
# Ensure 1D
|
||||
if len(audio_np.shape) > 1:
|
||||
audio_np = audio_np.squeeze()
|
||||
|
||||
output_path = "test_output.wav"
|
||||
sf.write(output_path, audio_np, sample_rate)
|
||||
print(f"✓ Saved to: {output_path}")
|
||||
|
||||
# Get file size
|
||||
import os
|
||||
size_kb = os.path.getsize(output_path) / 1024
|
||||
print(f" File size: {size_kb:.1f} KB\n")
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"✗ Save error: {e}\n")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
print("\n" + "=" * 60)
|
||||
print("MLX Audio Validation Test")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
# Step 1: MLX
|
||||
if not test_mlx_available():
|
||||
print("FAILED: MLX not available")
|
||||
sys.exit(1)
|
||||
|
||||
# Step 2: Imports
|
||||
if not test_mlx_audio_import():
|
||||
print("FAILED: mlx-audio import failed")
|
||||
sys.exit(1)
|
||||
|
||||
# Step 3: Model loading
|
||||
tts = test_model_loading()
|
||||
if tts is None:
|
||||
print("FAILED: Model loading failed")
|
||||
sys.exit(1)
|
||||
|
||||
# Step 4: Generation
|
||||
audio, sr = test_generation(tts)
|
||||
if audio is None:
|
||||
print("FAILED: Audio generation failed")
|
||||
sys.exit(1)
|
||||
|
||||
# Step 5: Save
|
||||
if not test_save_audio(audio, sr):
|
||||
print("FAILED: Could not save audio")
|
||||
sys.exit(1)
|
||||
|
||||
print("=" * 60)
|
||||
print("ALL TESTS PASSED ✓")
|
||||
print("=" * 60)
|
||||
print("\nMLX Audio is working correctly on this system.")
|
||||
print("You can play the generated audio with: afplay test_output.wav\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user