diff --git a/docs/plans/MLX_AUDIO.md b/docs/plans/MLX_AUDIO.md index c42f4064..9c2a100a 100644 --- a/docs/plans/MLX_AUDIO.md +++ b/docs/plans/MLX_AUDIO.md @@ -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) diff --git a/mlx-test/.gitignore b/mlx-test/.gitignore new file mode 100644 index 00000000..420f8dc1 --- /dev/null +++ b/mlx-test/.gitignore @@ -0,0 +1,9 @@ +# Virtual environment +venv/ + +# Generated test files +*.wav + +# Python cache +__pycache__/ +*.pyc diff --git a/mlx-test/demo.py b/mlx-test/demo.py new file mode 100755 index 00000000..006da72c --- /dev/null +++ b/mlx-test/demo.py @@ -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() diff --git a/mlx-test/test_tts.py b/mlx-test/test_tts.py new file mode 100644 index 00000000..eda0ad67 --- /dev/null +++ b/mlx-test/test_tts.py @@ -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()