mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-16 21:30: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:
+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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user