Adding a Model¶

This guide explains how to add a new TTS or STT model to MLX Audio. The library uses a consistent pattern across all models, so once you understand the structure, adding a new one is straightforward.

Directory Structure¶

Every model lives in its own package under the appropriate category directory:

mlx_audio/
├── tts/models/
│   ├── base.py              # Shared base classes
│   ├── __init__.py
│   ├── kokoro/              # Example: Kokoro TTS
│   │   ├── __init__.py      # Exports Model, ModelConfig, Pipeline
│   │   ├── kokoro.py        # Model implementation
│   │   └── pipeline.py      # Generation pipeline
│   ├── your_model/          # Your new model
│   │   ├── __init__.py
│   │   ├── your_model.py
│   │   └── ...
├── stt/models/
│   ├── whisper/
│   ├── parakeet/
│   └── your_model/

Step 1: Create the Model Package¶

Create a new directory under mlx_audio/tts/models/ (or stt/models/ for STT):

mlx_audio/tts/models/my_model/
├── __init__.py
├── my_model.py
└── README.md          # Optional but recommended

Step 2: Implement the Model¶

Base Classes¶

TTS models use the base classes from mlx_audio/tts/models/base.py:

BaseModelArgs -- Dataclass for model configuration. Includes a from_dict() class method that filters unknown keys automatically.
GenerationResult -- Dataclass returned by generate(). Contains audio, sample_rate, token_count, timing information, and streaming flags.
BatchGenerationResult -- Dataclass for batch generation results.

Transformer Components¶

Many speech models wrap a language-model backbone. Import that machinery from mlx_audio.lm, not from mlx-lm -- mlx-audio vendors it so the package has no runtime dependency on mlx-lm, and a test fails the build if mlx_lm is imported anywhere outside sts/voice_pipeline.py.

Need	Import from
KV caches (`KVCache`, `RotatingKVCache`, `BatchKVCache`, `make_prompt_cache`)	`mlx_audio.lm.models.cache`
Attention masks, `scaled_dot_product_attention`	`mlx_audio.lm.models.base`
Samplers, logits processors	`mlx_audio.lm.sample_utils`
`generate_step`, `stream_generate`	`mlx_audio.lm.generate`
Backbones (llama, qwen2, qwen3, qwen3_5, gpt2, granite, lfm2, bailing_moe, gemma3)	`mlx_audio.lm.models.<name>`
Quantize / save helpers	`mlx_audio.lm.convert`

If your model needs a backbone that is not vendored yet, copy it from mlx-lm into mlx_audio/lm/models/ verbatim and add the provenance header used by the other files there (upstream path, version, commit).

Model Configuration¶

Create a dataclass for your model's config that extends BaseModelArgs:

from dataclasses import dataclass
from ..base import BaseModelArgs

@dataclass
class MyModelConfig(BaseModelArgs):
    hidden_size: int = 512
    num_layers: int = 6
    vocab_size: int = 32000
    sample_rate: int = 24000
    model_type: str = "my_model"

Model Class¶

Your model should be an mlx.nn.Module with a generate() method that yields GenerationResult objects:

import mlx.nn as nn
from ..base import GenerationResult

class MyModel(nn.Module):
    def __init__(self, config: MyModelConfig):
        super().__init__()
        self.config = config
        self.sample_rate = config.sample_rate
        self.model_type = config.model_type
        # Build layers ...

    def generate(
        self,
        text: str,
        voice: str = "default",
        speed: float = 1.0,
        lang_code: str = "en",
        stream: bool = False,
        streaming_interval: float = 2.0,
        **kwargs,
    ):
        """Generate speech from text.

        Yields:
            GenerationResult for each audio segment.
        """
        # Your generation logic here
        audio = self._synthesize(text)

        yield GenerationResult(
            audio=audio,
            samples=audio.shape[0],
            sample_rate=self.sample_rate,
            segment_idx=0,
            token_count=0,
            audio_duration="0.0s",
            real_time_factor=0.0,
            prompt={},
            audio_samples={},
            processing_time_seconds=0.0,
            peak_memory_usage=0.0,
        )

`init.py`¶

Export the model class and config so the loader can find them:

from .my_model import MyModel as Model, MyModelConfig as ModelConfig

__all__ = ["Model", "ModelConfig"]

Naming convention

The loader expects Model and ModelConfig (or ModelArgs) to be importable from the package. Use these exact names in your exports.

Step 3: Register the Model¶

Config File¶

The model type is detected from the config.json file in the HuggingFace repo. The model_type field must match your directory name (or a key in the remapping dict).

If the model_type in config.json differs from your directory name, add an entry to the MODEL_REMAPPING dict in the appropriate utils file:

TTSSTT

# mlx_audio/tts/utils.py
MODEL_REMAPPING = {
    # ... existing entries ...
    "my_model": "my_model",  # config.json model_type -> directory name
}

# mlx_audio/stt/utils.py
MODEL_REMAPPING = {
    # ... existing entries ...
    "my_model": "my_model",
}

How Model Loading Works¶

load() downloads the model from HuggingFace (or uses a local path).
It reads config.json to determine model_type.
It looks up the model type in MODEL_REMAPPING (or uses the type directly).
It dynamically imports the corresponding package from mlx_audio/tts/models/ (or stt/models/).
It instantiates ModelConfig from the config dict and Model from the config.
It loads the .safetensors weights into the model.

For checkpoints that require unpacking or format adaptation, the model package can export prepare_model_path(model_path, *, source=None, revision=None, force_download=False) -> Path. The shared loader calls this optional hook before constructing ModelConfig or reading weights. Return the input path for an ordinary checkpoint, or a prepared directory containing the runtime config and weights. When the path changes, the loader reads the prepared config and resolves its model type again. Keep archive downloading, verification, and unpacking inside the model package.

Step 4: Convert and Test¶

Convert Weights¶

If your model's original weights are in PyTorch format, use the conversion script:

python -m mlx_audio.convert \
    --hf-path original/my-model \
    --mlx-path ./my-model-bf16 \
    --dtype bfloat16

Publish to Hugging Face¶

If you plan to share the converted model, prefer publishing it on the mlx-community organization on Hugging Face. It is the shared home for ready-to-use MLX weights across projects like mlx-lm, mlx-vlm, mlx-swift-examples, and mlx-audio, so publishing there keeps MLX-native checkpoints discoverable in one place.

If you cannot publish directly to mlx-community, use your own namespace first and link it from the docs. We should still encourage new model contributors to be part of the mlx-community org when possible.

Test¶

Write a basic test:

from mlx_audio.tts.utils import load_model

def test_my_model():
    model = load_model("path/to/my-model-bf16")
    results = list(model.generate("Hello, world!"))
    assert len(results) > 0
    assert results[0].audio.shape[0] > 0

Run it:

pytest mlx_audio/tts/tests/test_my_model.py

Step 5: Add Documentation¶

Add a model page in docs/models/tts/ (or stt/).
Add the page to the nav section in mkdocs.yml.
Optionally add a README.md inside your model directory for model-specific details.
If the model is published on Hugging Face, prefer an mlx-community/... repo when available and link that repo from the docs.
Make sure your PR includes the docs change. The docs workflow will fail if model files change without a matching docs update.

Checklist¶

[ ] Model package created under mlx_audio/{tts,stt}/models/
[ ] __init__.py exports Model and ModelConfig
[ ] generate() method yields GenerationResult objects
[ ] Model type registered in MODEL_REMAPPING (if needed)
[ ] Weights converted to MLX .safetensors format
[ ] Hugging Face repo chosen and linked in docs (mlx-community/... preferred)
[ ] Basic test written and passing
[ ] Documentation page added
[ ] PR submitted with a clear description