Skip to content
Draft
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
261 changes: 261 additions & 0 deletions ADDING_A_MODEL.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,261 @@
# Adding a New Model

## Directory layout

Place the model under the appropriate category:

```
mlx_audio/
├── tts/models/<model_name>/ # Text-to-speech
├── stt/models/<model_name>/ # Speech-to-text
├── sts/models/<model_name>/ # Speech-to-speech / enhancement
└── codec/models/<model_name>/ # Standalone audio codecs / tokenizers
```

Minimum required files:

```
mlx_audio/tts/models/my_model/
├── __init__.py # must export Model and ModelConfig
├── my_model.py # Model class + ModelConfig dataclass
└── convert.py # weight conversion script (PyTorch → safetensors)

@Blaizzy Blaizzy Apr 14, 2026

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is optional only if the model has onnx dependency or uses .pt torch files.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Blaizzy my key point is that the result must be reproducible and independently verifiable.

```

## Auto-discovery

**No registration needed.** The loader resolves models dynamically:

1. Reads `"model_type"` from `config.json` in the model directory.
2. Imports `mlx_audio.{tts|stt|sts}.models.{model_type}`.
3. Instantiates `module.ModelConfig.from_dict(config)` then `module.Model(config)`.

So `model_type` in `config.json` **must match the directory name exactly**.

## ModelConfig

```python
from dataclasses import dataclass
from mlx_audio.tts.models.base import BaseModelArgs

@dataclass
class ModelConfig(BaseModelArgs):
model_type: str = "my_model" # must match directory name
sample_rate: int = 24000 # output audio sample rate

# model-specific fields …

@classmethod
def from_dict(cls, config: dict) -> "ModelConfig":
return cls(
model_type=config.get("model_type", "my_model"),
sample_rate=config.get("sample_rate", 24000),
# …
)
```

## Model class

```python
import mlx.nn as nn
import mlx.core as mx
from typing import Generator, Optional

class Model(nn.Module):
def __init__(self, config: ModelConfig):
super().__init__()
self.config = config
# define layers …

@property
def sample_rate(self) -> int: # required
return self.config.sample_rate

@property
def model_type(self) -> str: # recommended
return self.config.model_type

def sanitize(self, weights: dict) -> dict:
"""Rename / reshape PyTorch keys to match MLX attribute paths.

Common transforms:
- Strip "model." prefix
- Conv1d weights: PyTorch (out, in, k) — MLX loads as (out, k, in),
so no manual transpose is needed; mlx.load_weights handles it.
- LayerNorm: .gamma → .weight, .beta → .bias
- Drop unused keys (position_ids, etc.)
"""
result = {}
for k, v in weights.items():
if k.startswith("model."):
k = k[4:]
result[k] = v
return result

def generate(
self,
text: str,
voice: Optional[str] = None,
speed: float = 1.0,
lang_code: str = "en",
temperature: float = 0.7,
max_tokens: int = 1200,
**kwargs, # absorb unused args from generate_audio()
) -> Generator["GenerationResult", None, None]:
import time
from mlx_audio.tts.models.base import GenerationResult

start = time.time()
audio = self._run_inference(text) # mx.array [samples]
elapsed = time.time() - start

n = int(audio.shape[0])
dur = n / self.sample_rate
d = int(dur)

yield GenerationResult(
audio=audio,
samples=n,
sample_rate=self.sample_rate,
segment_idx=0,
token_count=0,
audio_duration=f"{d//3600:02d}:{(d%3600)//60:02d}:{d%60:02d}.{int((dur%1)*1000):03d}",
real_time_factor=dur / elapsed if elapsed > 0 else 0.0,
prompt={"tokens": 0, "tokens-per-sec": 0},
audio_samples={"samples": n, "samples-per-sec": round(n / elapsed, 2)},
processing_time_seconds=elapsed,
peak_memory_usage=mx.get_peak_memory() / 1e9,
)

@staticmethod
def post_load_hook(model: "Model", model_path) -> "Model":
"""Optional. Called by the loader after weights are applied.
Use to load auxiliary tokenizers or preprocessors."""
return model
```

## `generate()` — kwargs passed by the CLI

`generate_audio()` always passes these kwargs; use `**kwargs` to absorb any
you don't need:

| kwarg | type | description |
|---|---|---|
| `text` | `str` | input text |
| `voice` | `str \| None` | speaker / voice ID |
| `speed` | `float` | playback speed multiplier |
| `lang_code` | `str` | BCP-47 language code |
| `temperature` | `float` | sampling temperature |
| `max_tokens` | `int` | token budget |
| `ref_audio` | `mx.array \| None` | reference waveform for voice cloning |
| `ref_text` | `str \| None` | transcript of reference audio |
| `cfg_scale` | `float \| None` | classifier-free guidance strength |
| `instruct` | `str \| None` | style / emotion instruction |
| `stream` | `bool` | streaming mode |

## Weight conversion (`convert.py`)

```python
# convert.py — minimal pattern
from pathlib import Path
import numpy as np
from safetensors.numpy import save_file
import torch

def convert(pt_path: str, output_dir: str, dtype: str = "bfloat16"):
mlx_dtype = {"float16": np.float16, "bfloat16": np.float32}[dtype]
weights = torch.load(pt_path, map_location="cpu")
out = {}
for k, v in weights.items():
arr = v.detach().cpu().numpy().astype(mlx_dtype)
out[k] = arr
Path(output_dir).mkdir(parents=True, exist_ok=True)
save_file(out, f"{output_dir}/model.safetensors")
```

For key renaming and shape fixes, implement `Model.sanitize()` — the loader
calls it automatically after reading the safetensors file.

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would remove this because if you can use safetensors to open the file without using torch


## Acoustic codecs / tokenizers

If your model needs an audio codec (encode waveform → tokens or decode tokens
→ waveform), add it under `mlx_audio/codec/models/<codec_name>/` and export
from `mlx_audio/codec/__init__.py`. Reference it from your TTS/STT model by
import — do not bundle codec weights inside the TTS model directory.

## Tests

Add tests under `mlx_audio/tts/tests/test_<model_name>.py` (or the equivalent
category). Tests should use random MLX weights — no real checkpoint required:

```python
import unittest
import mlx.core as mx
from mlx_audio.tts.models.my_model import Model, ModelConfig
from mlx_audio.tts.models.base import GenerationResult

class TestMyModel(unittest.TestCase):
def setUp(self):
self.model = Model(ModelConfig())

def test_sample_rate(self):
self.assertEqual(self.model.sample_rate, 24000)

def test_generate_yields_result(self):
results = list(self.model.generate("Hello"))
self.assertIsInstance(results[0], GenerationResult)
self.assertGreater(results[0].samples, 0)
```

## Publishing weights to mlx-community

Model weights must be published to the
[mlx-community](https://huggingface.co/mlx-community) HuggingFace organization,
not bundled in this repository. The only exception is when an existing
mlx-community model needs to be updated and the PR is waiting for approval —
in that case a temporary personal fork is acceptable.

### Naming convention

```
mlx-community/<ModelName>[-<Variant>]-<ParameterCount>-<Dtype>
```

| part | description | examples |
|---|---|---|
| `ModelName` | base model name, preserve original casing | `Kokoro`, `Qwen3-TTS`, `Voxtral` |
| `Variant` | optional variant tag | `Base`, `VoiceDesign`, `Realtime` |
| `ParameterCount` | size indicator | `82M`, `0.6B`, `1B`, `4B` |
| `Dtype` | precision / quantization level | `bf16`, `fp16`, `8bit`, `4bit`, `6bit` |

Real examples from mlx-community:

```
mlx-community/Kokoro-82M-bf16
mlx-community/Kokoro-82M-4bit
mlx-community/OuteTTS-1.0-0.6B-fp16
mlx-community/Qwen3-TTS-12Hz-1.7B-VoiceDesign-bf16
mlx-community/Voxtral-4B-TTS-2603-mlx-bf16
mlx-community/chatterbox-fp16
mlx-community/LongCat-AudioDiT-1B-bf16
mlx-community/whisper-large-v3-turbo-asr-fp16
mlx-community/parakeet-tdt-0.6b-v3
```

**Notes:**
- Prefer `bf16` as the primary upload; add quantized variants (`4bit`, `8bit`) if
the model is large enough to benefit.
- Include the parameter count when the model family has multiple sizes.
- Do not add an `-mlx` suffix unless the upstream name already contains it.
- Link the mlx-community repo in your PR description so reviewers can verify
the weights are accessible.

## PR checklist

- [ ] `config.json` has `"model_type"` matching the directory name
- [ ] `__init__.py` exports `Model` and `ModelConfig`
- [ ] `ModelConfig` is a `@dataclass` with `from_dict()` and `sample_rate`
- [ ] `Model.generate()` yields `GenerationResult` and accepts `**kwargs`
- [ ] `Model.sanitize()` covers all key renames / shape fixes
- [ ] `convert.py` produces a loadable `model.safetensors`
- [ ] Tests pass with random weights (no real checkpoint needed)
- [ ] Model listed in `README.md` model table
56 changes: 56 additions & 0 deletions CODE_OF_CONDUCT.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
# Contributor Covenant Code of Conduct

## Our Pledge

We as members, contributors, and leaders pledge to make participation in our
community a harassment-free experience for everyone, regardless of age, body
size, visible or invisible disability, ethnicity, sex characteristics, gender
identity and expression, level of experience, education, socio-economic status,
nationality, personal appearance, race, caste, color, religion, or sexual
identity and orientation.

We pledge to act and interact in ways that contribute to an open, welcoming,
diverse, inclusive, and healthy community.

## Our Standards

Examples of behavior that contributes to a positive environment:

- Demonstrating empathy and kindness toward other people
- Being respectful of differing opinions, viewpoints, and experiences
- Giving and gracefully accepting constructive feedback
- Accepting responsibility and apologizing to those affected by our mistakes
- Focusing on what is best not just for us as individuals, but for the overall community

Examples of unacceptable behavior:

- The use of sexualized language or imagery, and sexual attention or advances of any kind
- Trolling, insulting or derogatory comments, and personal or political attacks
- Public or private harassment
- Publishing others' private information without their explicit permission
- Other conduct which could reasonably be considered inappropriate in a professional setting

## Enforcement Responsibilities

Community leaders are responsible for clarifying and enforcing our standards of
acceptable behavior and will take appropriate and fair corrective action in
response to any behavior that they deem inappropriate, threatening, offensive,
or harmful.

## Scope

This Code of Conduct applies within all community spaces, and also applies when
an individual is officially representing the community in public spaces.

## Enforcement

Instances of abusive, harassing, or otherwise unacceptable behavior may be
reported to the community leaders responsible for enforcement via
[GitHub Security Advisories](https://github.com/Blaizzy/mlx-audio/security/advisories/new).
All complaints will be reviewed and investigated promptly and fairly.

## Attribution

This Code of Conduct is adapted from the [Contributor Covenant](https://www.contributor-covenant.org),
version 2.1, available at
https://www.contributor-covenant.org/version/2/1/code_of_conduct.html.
Loading