Add ACE-1.5

README.md

@@ -143,6 +143,12 @@ See the [Sortformer README](mlx_audio/vad/models/sortformer/README.md) for API d
 | **MossFormer2 SE** | Speech enhancement | Noise removal | [starkdmi/MossFormer2_SE_48K_MLX](https://huggingface.co/starkdmi/MossFormer2_SE_48K_MLX) |
 | **DeepFilterNet (1/2/3)** | Speech enhancement | Noise suppression | [mlx-community/DeepFilterNet-mlx](https://huggingface.co/mlx-community/DeepFilterNet-mlx) |
 
+### Music Generation
+
+| Model | Description | Use Case | Repo |
+|-------|-------------|----------|------|
+| **ACE-Step** | Diffusion-based music generation | Text-to-music, covers, track completion | [ACE-Step/Ace-Step1.5](https://huggingface.co/ACE-Step/Ace-Step1.5) |
+
 ## Model Examples
 
 ### Kokoro TTS
@@ -517,6 +523,89 @@ enhanced = model.enhance("noisy_speech.wav")
 save_audio(enhanced, "clean.wav", 48000)
 ```
 
+### ACE-Step (Music Generation)
+
+Generate music from text descriptions with ACE-Step, a diffusion-based music generation model.
+
+**CLI - Basic text-to-music:**
+
+```bash
+mlx_audio.tts.generate \
+    --model ACE-Step/Ace-Step1.5 \
+    --text "A gentle piano melody with soft drums and bass" \
+    --gen-kwargs '{"lyrics": "[Instrumental]", "duration": 30.0, "seed": 42}' \
+    --output_path ./music \
+    --verbose
+```
+
+**CLI - Generate with specific parameters:**
+
+```bash
+mlx_audio.tts.generate \
+    --model ACE-Step/Ace-Step1.5 \
+    --text "Upbeat electronic dance track with heavy bass" \
+    --gen-kwargs '{"lyrics": "[Instrumental]", "duration": 60.0, "num_steps": 8, "shift": 3.0, "seed": 123}' \
+    --play
+```
+
+**Python API - Basic text-to-music:**
+
+```python
+from mlx_audio.tts.models.ace_step import Model
+import soundfile as sf
+import numpy as np
+
+model = Model.from_pretrained("ACE-Step/Ace-Step1.5")
+
+for result in model.generate(
+    text="A gentle piano melody with soft drums and bass",
+    lyrics="[Instrumental]",
+    duration=30.0,
+    seed=42,
+    num_steps=8,
+    shift=3.0,
+    verbose=True,
+):
+    sf.write("music.wav", np.array(result.audio).T, 48000)
+```
+
+**Python API - Complete task (add instruments to existing audio):**
+
+```python
+import mlx.core as mx
+
+# Load source audio
+audio, sr = sf.read("piano.wav")
+source_audio = mx.array(audio.T)  # [channels, samples]
+
+# Add drums to the piano track
+for result in model.generate(
+    text="Piano with energetic drums",
+    lyrics="[Instrumental]",
+    duration=8.0,
+    task_type="complete",
+    complete_track_classes=["drums", "bass"],
+    source_audio=source_audio,
+    source_audio_sample_rate=sr,
+    seed=42,
+):
+    sf.write("piano_with_drums.wav", np.array(result.audio).T, 48000)
+```
+
+**Available task types:**
+
+| Task | Description | Key Parameters |
+|------|-------------|----------------|
+| `text2music` | Generate music from text | `text`, `lyrics` |
+| `complete` | Add instruments to audio | `source_audio`, `complete_track_classes` |
+| `cover` | Generate cover version | `source_audio` |
+| `extract` | Extract track from audio | `source_audio`, `track_name` |
+| `lego` | Generate specific track | `source_audio`, `track_name` |
+| `repaint` | Repaint audio region | `source_audio` |
+
+**Track names for complete/extract/lego:**
+`woodwinds`, `brass`, `fx`, `synth`, `strings`, `percussion`, `keyboard`, `guitar`, `bass`, `drums`, `backing_vocals`, `vocals`
+
 ## Web Interface & API Server
 
 MLX-Audio includes a modern web interface and OpenAI-compatible API.

mlx_audio/tts/generate.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 import os
 import sys
 from typing import Optional, Tuple, Union
@@ -487,6 +488,12 @@ def parse_args():
         default=2.0,
         help="The time interval in seconds for streaming segments",
     )
+    parser.add_argument(
+        "--gen-kwargs",
+        type=str,
+        default=None,
+        help='Additional generation kwargs as JSON string (e.g., \'{"duration": 60.0, "num_steps": 8}\')',
+    )
     parser.add_argument(
         "--save",
         action="store_true",
@@ -505,6 +512,19 @@ def parse_args():
             print("Please enter the text to generate:")
             args.text = input("> ").strip()
 
+    # Parse and merge gen_kwargs if provided
+    if hasattr(args, "gen_kwargs") and args.gen_kwargs:
+        try:
+            gen_kwargs = json.loads(args.gen_kwargs)
+            # Convert args to dict, merge with gen_kwargs, then back to args
+            args_dict = vars(args)
+            args_dict.update(gen_kwargs)
+            # Remove the gen_kwargs key itself
+            args_dict.pop("gen_kwargs", None)
+        except json.JSONDecodeError as e:
+            print(f"Error parsing --gen-kwargs JSON: {e}")
+            sys.exit(1)
+
     return args
 
 

mlx_audio/tts/models/ace_step/README.md

@@ -0,0 +1,162 @@
+# ACE-Step for MLX
+
+MLX implementation of [ACE-Step](https://github.com/ace-step/ACE-Step), a 3.5B parameter flow-matching music generation model that can generate high-quality music with vocals from text prompts and lyrics.
+
+## Features
+
+- **Text-to-Music**: Generate music from text descriptions
+- **Vocals Support**: Generate songs with lyrics in multiple languages
+- **Fast Generation**: Faster than real-time on Apple Silicon (RTF ~0.85x at 20 steps)
+- **High Quality**: 48kHz stereo audio output
+- **4-bit Quantized Variant**: 2.2GB main model for memory-constrained systems
+
+## Quick Start
+
+```python
+from mlx_audio.tts import load
+
+# Load the MLX-converted model (4-bit quantized, 2.2GB)
+model = load("mlx-community/ACE-Step1.5-MLX-4bit")
+
+# Generate instrumental music
+for result in model.generate(
+    text="upbeat electronic dance music with energetic synthesizers",
+    duration=30.0,
+):
+    audio = result.audio  # [samples, 2] stereo audio
+    sample_rate = result.sample_rate  # 48000
+```
+
+Pre-converted weights are available on Hugging Face:
+- [mlx-community/ACE-Step1.5-MLX](https://huggingface.co/mlx-community/ACE-Step1.5-MLX) — fp32 (~9.6GB main model)
+- [mlx-community/ACE-Step1.5-MLX-4bit](https://huggingface.co/mlx-community/ACE-Step1.5-MLX-4bit) — 4-bit quantized (~2.2GB main model)
+
+## Generating Music with Vocals
+
+```python
+from mlx_audio.tts import load
+import scipy.io.wavfile as wavfile
+import numpy as np
+import mlx.core as mx
+
+model = load("mlx-community/ACE-Step1.5-MLX-4bit")
+
+prompt = "upbeat pop song with female vocals, catchy melody, bright synths"
+lyrics = """[Verse 1]
+Dance with me tonight
+Under the neon lights
+Feel the rhythm in your soul
+Let the music take control
+
+[Chorus]
+We're alive, we're on fire
+Dancing higher and higher
+Nothing's gonna stop us now
+We're shining bright somehow
+"""
+
+for result in model.generate(
+    text=prompt,
+    lyrics=lyrics,
+    duration=30.0,
+    vocal_language="en",
+):
+    audio_np = np.array(result.audio.astype(mx.float32))
+    audio_int16 = (np.clip(audio_np, -1, 1) * 32767).astype(np.int16)
+    wavfile.write("song.wav", result.sample_rate, audio_int16)
+```
+
+## How It Works
+
+ACE-Step's turbo model requires a two-stage pipeline:
+
+1. **5Hz Language Model (Planner)**: Takes your text prompt + lyrics and generates an audio code "blueprint" for the song — planning BPM, key signature, structure, and time-aligned audio codes
+2. **Diffusion Transformer (DiT)**: Uses these codes as conditioning to denoise random noise into musical latents
+3. **VAE Decoder**: Decodes the 25Hz latents into 48kHz stereo audio
+
+The LM planner is **required** for the turbo model — without it, the DiT converges back to silence. `use_lm=True` is the default.
+
+## Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `text` | required | Text description / prompt |
+| `lyrics` | `""` | Lyrics (empty for instrumental) |
+| `duration` | `30.0` | Target duration in seconds |
+| `seed` | `None` | Random seed for reproducibility |
+| `num_steps` | `20` | Diffusion steps (20 recommended for vocals, 8 ok for instrumentals) |
+| `shift` | `3.0` | Timestep schedule shift (1.0, 2.0, or 3.0) |
+| `guidance_scale` | `1.0` | CFG scale (1.0 = no guidance; turbo model is distilled without CFG) |
+| `guidance_interval` | `0.5` | Fraction of steps with guidance applied |
+| `cfg_type` | `"apg"` | Guidance type: `"apg"` or `"cfg"` |
+| `vocal_language` | `"unknown"` | Language code: `"en"`, `"zh"`, `"ja"`, etc. |
+| `use_lm` | `True` | Use 5Hz LM planner (required for the turbo model) |
+| `lm_model_size` | `"0.6B"` | LM size: `"0.6B"` (fast) or `"4B"` (higher quality, slower) |
+
+## Lyrics Format
+
+Use section markers to structure your lyrics:
+
+```
+[Verse 1]
+First verse lyrics here
+Line by line
+
+[Chorus]
+Catchy chorus lyrics
+That repeat
+
+[Bridge]
+Bridge section
+Something different
+
+[Outro]
+Final words
+```
+
+## Supported Languages
+
+ACE-Step supports lyrics in multiple languages via the 5Hz LM planner:
+- English (`en`), Chinese (`zh`), Japanese (`ja`), Korean (`ko`)
+- Spanish (`es`), French (`fr`), German (`de`), Italian (`it`)
+- And many more (50+ supported by the LM)
+
+Set `vocal_language` to hint the planner — though the 0.6B LM sometimes picks its own language based on caption/lyrics content. Use the 4B LM for more reliable language adherence.
+
+## Performance
+
+Apple Silicon M4 Max, 4-bit quantized model, 20 diffusion steps:
+
+| Duration | Total Time | Diffusion | LM Planning | VAE Decode | RTF |
+|----------|-----------|-----------|-------------|------------|-----|
+| 30s | ~25s | ~8s | ~12s | ~3.5s | 0.85x |
+| 60s | ~40s | ~17s | ~12s | ~6s | 0.67x |
+
+*RTF (Real-Time Factor) < 1.0 means faster than real-time.*
+
+## Tips
+
+1. **Vocals**: Use `num_steps=20` (default) for clearer vocals; `num_steps=8` is fine for instrumentals
+2. **Cherry-pick seeds**: Generate a few samples with different seeds and pick the best — LM quality varies
+3. **Prompt engineering**: Be specific about genre, instruments, mood, vocal style (e.g., "male vocals", "female choir")
+4. **Section markers** in lyrics help the model structure the song (`[Verse]`, `[Chorus]`, `[Bridge]`, `[Outro]`)
+5. **LM size**: 4B model follows language/caption instructions more reliably but takes ~3x longer to plan
+
+## Known Limitations
+
+- The 0.6B LM occasionally ignores the `vocal_language` parameter and picks its own
+- Vocal clarity varies by seed — use cherry-picking for final output
+
+## Citation
+
+```bibtex
+@article{ace-step,
+  title={ACE-Step: A Step Towards Music Generation Foundation Model},
+  author={ACE-Step Team},
+  year={2024}
+}
+```
+
+## License
+
+This implementation follows the mlx-audio project license. The ACE-Step model weights are subject to their original license from the ACE-Step team.

mlx_audio/tts/models/ace_step/__init__.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2025, Prince Canuma and contributors (https://github.com/Blaizzy/mlx-audio)
+
+from .ace_step import Model
+from .config import (
+    DEFAULT_INSTRUCTION,
+    SFT_GEN_PROMPT,
+    TASK_INSTRUCTIONS,
+    TASK_TYPES,
+    TRACK_NAMES,
+    ModelConfig,
+)
+from .lm import ACEStepLM, LMConfig
+from .vae import AutoencoderOobleck
+
+__all__ = [
+    "Model",
+    "ModelConfig",
+    "AutoencoderOobleck",
+    "ACEStepLM",
+    "LMConfig",
+    "TASK_TYPES",
+    "TASK_INSTRUCTIONS",
+    "DEFAULT_INSTRUCTION",
+    "TRACK_NAMES",
+    "SFT_GEN_PROMPT",
+]