diff --git a/ADDING_A_MODEL.md b/ADDING_A_MODEL.md new file mode 100644 index 000000000..e85a8e3cb --- /dev/null +++ b/ADDING_A_MODEL.md @@ -0,0 +1,272 @@ +# Adding a New Model + +## Directory layout + +Place the model under the appropriate category: + +``` +mlx_audio/ +├── tts/models// # Text-to-speech +├── stt/models// # Speech-to-text +├── sts/models// # Speech-to-speech / enhancement +└── codec/models// # 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 +``` + +Add `convert.py` only when the conversion needs model-specific logic that is +worth keeping reproducible in-tree, for example starting from `.pt` or ONNX +checkpoints, remapping unusual keys, or exporting auxiliary files. + +## 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 + +The exact fields depend on the category. TTS models usually include +`sample_rate`; STT models are more likely to focus on tokenizer, frontend, or +decoder settings. + +```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 + +The example below is TTS-oriented because it shows the full +`GenerationResult` shape. For STT models, the same package/export pattern +applies, but inference is usually audio-in / text-out with structured results +such as `.text`, `.segments`, or timestamps. + +```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 + +This section is TTS-specific. STT models typically accept audio input and may +return structured transcription metadata instead of waveform chunks. + +`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`, optional) + +Add `convert.py` when the conversion path needs model-specific logic or a +reproducible custom pipeline. Even upstream `safetensors` may still need +remapping or auxiliary exports before the model is ready to load. + +- If the upstream release already provides `safetensors`, load those directly + and skip `torch`. +- Prefer the repository-wide converter when it is enough. +- Add a model-specific `convert.py` when the conversion has custom steps that + are worth preserving for future contributors. +- If conversion is required, document the exact input files and command so + another contributor can reproduce the same `model.safetensors` output. +- Keep key renaming and shape fixes in `Model.sanitize()` — the loader calls it + automatically after reading the safetensors file. + +## Acoustic codecs / tokenizers + +If your model needs an audio codec (encode waveform → tokens or decode tokens +→ waveform), add it under `mlx_audio/codec/models//` 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 the matching category directory, for example +`mlx_audio/tts/tests/test_.py` or `mlx_audio/stt/tests/test_.py`. +Tests should match the model output type — waveform assertions for TTS, +transcription/text/segment assertions for STT — and should avoid requiring a +real checkpoint when random MLX weights are enough. + +```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/[-]-- +``` + +| 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 diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 000000000..123cc9641 --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -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. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 000000000..5dbbd865d --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,121 @@ +# Contributing to mlx-audio + +Thanks for contributing to mlx-audio. + +## Scope + +We welcome: +- New model ports (TTS, STT, STS, SFX, VAD) +- Bug fixes and performance improvements +- Documentation and example improvements + +For large API changes or new features, **open an issue first** to discuss the +approach before starting implementation. This avoids wasted effort on PRs that +won't be accepted. + +## Reporting Bugs + +Search [existing issues](https://github.com/Blaizzy/mlx-audio/issues) before +opening a new one. When reporting a bug, include: + +- OS and Apple Silicon chip (e.g., macOS 15.3, M3 Pro) +- Python and MLX versions (`python --version`, `python -c "import mlx; print(mlx.__version__)"`) +- Full traceback +- Minimal reproducible snippet + +## Security Vulnerabilities + +Do **not** open a public issue for security vulnerabilities. Use +[GitHub Security Advisories](https://github.com/Blaizzy/mlx-audio/security/advisories/new) +to report privately. + +## Development Setup + +```bash +# Install in editable mode with dev dependencies +uv pip install -e ".[all,dev]" + +# Install pre-commit hooks (required — CI uses pinned formatter versions) +pre-commit install +``` + +## Pull Requests + +- Open pull requests against `Blaizzy/mlx-audio:main`. +- If you are contributing from a fork, make sure the base repository is + `Blaizzy/mlx-audio` and the base branch is `main`. +- Keep pull requests focused. Include tests and documentation updates when + behavior changes. +- Keep PRs atomic and touch the smallest possible amount of code. This helps + reviewers evaluate and merge changes faster and with higher confidence. +- A checklist is pre-filled when you open a PR — fill it out before requesting review. + +Run local checks before opening a PR: + +```bash +# Formatting (always run via pre-commit — CI uses pinned 24.x version) +pre-commit run black --files +pre-commit run isort --files + +# Core tests +pytest -s mlx_audio/tests/ +``` + +## Keeping the repository clean + +Do not commit personal or temporary files. Before opening a PR, make sure your +diff does not include: + +- Planning docs, notes, or spec files (`docs/plans/`, `TODO.md`, etc.) +- Test scripts, scratch notebooks, or one-off debug files +- Local model weights, audio samples, or large binaries +- Changes to `.gitignore` that only cover your personal setup + +**Use a global gitignore for personal patterns** so you never have to touch the +project's `.gitignore`: + +```bash +# Create (or append to) your global gitignore +echo "*.local.*" >> ~/.gitignore_global +echo ".env.local" >> ~/.gitignore_global + +# Register it with git (one-time setup) +git config --global core.excludesFile ~/.gitignore_global +``` + +The `*.local.*` pattern (e.g., `TODO.local.md`, `config.local.json`) is a +useful convention for files that should always stay local. + +## Adding a New Model + +See [ADDING_A_MODEL.md](ADDING_A_MODEL.md) for the full guide — directory +layout, required class interfaces, `generate()` kwargs, weight conversion, +codec integration, tests, and PR checklist. + +## Good First Issues + +Issues labeled [`good first issue`](https://github.com/Blaizzy/mlx-audio/contribute) +are a good starting point if you are new to the codebase. + +## Code of Conduct + +This project follows the [Contributor Covenant Code of Conduct](CODE_OF_CONDUCT.md). +By participating, you agree to uphold it. + +## Commit Signing and Account Security + +To improve commit provenance and reduce supply chain risk, please sign commits +submitted to this repository. This is a one-time setup on your machine. + +- Any GitHub-supported signing method is fine: GPG, SSH, or S/MIME. +- Enable GitHub vigilant mode so commits and tags always show a verification + status. +- Enable two-factor authentication on your GitHub account. Passkeys are + preferred when available. + +## References + +- [About commit signature verification](https://docs.github.com/en/authentication/managing-commit-signature-verification/about-commit-signature-verification) +- [Displaying verification statuses for all of your commits](https://docs.github.com/en/authentication/managing-commit-signature-verification/displaying-verification-statuses-for-all-of-your-commits) +- [Enable vigilant mode](https://docs.github.com/en/authentication/managing-commit-signature-verification/displaying-verification-statuses-for-all-of-your-commits#enabling-vigilant-mode) +- [GPG setup walkthrough](https://docs.github.com/en/authentication/managing-commit-signature-verification/telling-git-about-your-signing-key) diff --git a/README.md b/README.md index 2a9eb2f93..e3f27074d 100644 --- a/README.md +++ b/README.md @@ -630,6 +630,11 @@ WAV format works without ffmpeg. } ``` +## Contributing + +See [CONTRIBUTING.md](CONTRIBUTING.md) for the PR workflow, local checks, and +commit-signing expectations for contributions to `main`. + ## Acknowledgements - [Apple MLX Team](https://github.com/ml-explore/mlx) for the MLX framework diff --git a/docs/contributing/adding-a-model.md b/docs/contributing/adding-a-model.md index 9dcafccb7..9cb816bbc 100644 --- a/docs/contributing/adding-a-model.md +++ b/docs/contributing/adding-a-model.md @@ -40,15 +40,27 @@ mlx_audio/tts/models/my_model/ ### Base Classes -TTS models use the base classes from `mlx_audio/tts/models/base.py`: +TTS and STT models share the same high-level pattern — a config dataclass plus +an `mlx.nn.Module` exported as `Model` — but their runtime interfaces differ. + +- **TTS** models typically use `mlx_audio/tts/models/base.py` and yield + `GenerationResult` objects containing waveform data. +- **STT** models usually return structured transcription results such as + `.text`, `.segments`, `.sentences`, timestamps, or streaming chunks. + +For shared config parsing, TTS models commonly use `BaseModelArgs`: - **`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. +- **`GenerationResult`** -- TTS dataclass returned by `generate()`. Contains `audio`, `sample_rate`, `token_count`, timing information, and streaming flags. +- **`BatchGenerationResult`** -- TTS dataclass for batch generation results. ### Model Configuration -Create a dataclass for your model's config that extends `BaseModelArgs`: +Create a dataclass for your model's config. Use fields that match your domain: +TTS models often include `sample_rate`, while STT models usually focus on audio +frontend, tokenizer, or decoder settings. + +For example, a TTS config can extend `BaseModelArgs`: ```python from dataclasses import dataclass @@ -65,7 +77,12 @@ class MyModelConfig(BaseModelArgs): ### Model Class -Your model should be an `mlx.nn.Module` with a `generate()` method that yields `GenerationResult` objects: +Your model should be an `mlx.nn.Module`, but the public inference method should +match the domain. + +#### TTS example + +TTS models usually implement `generate()` and yield `GenerationResult` objects: ```python import mlx.nn as nn @@ -112,6 +129,26 @@ class MyModel(nn.Module): ) ``` +#### STT example + +STT models usually accept audio input and return structured transcription +results: + +```python +import mlx.nn as nn + +class MySTTModel(nn.Module): + def __init__(self, config: MyModelConfig): + super().__init__() + self.config = config + self.model_type = config.model_type + + def generate(self, audio: str, **kwargs): + """Transcribe audio and return a structured result.""" + # Your transcription logic here + return self._transcribe(audio, **kwargs) +``` + ### `__init__.py` Export the model class and config so the loader can find them: @@ -166,7 +203,13 @@ If the `model_type` in `config.json` differs from your directory name, add an en ### Convert Weights -If your model's original weights are in PyTorch format, use the conversion script: +If the upstream model already ships loadable `.safetensors` weights, use those +directly and skip `torch`. + +If the original weights are only available as `.pt`, ONNX, or another +non-MLX format, you will need a conversion step to produce the final MLX +`model.safetensors` layout. When the source model is already supported by the +repository-wide conversion entry point, a typical command looks like this: ```bash python -m mlx_audio.convert \ @@ -175,6 +218,11 @@ python -m mlx_audio.convert \ --dtype bfloat16 ``` +Otherwise, add a model-specific `convert.py` when the model needs custom +remapping, extra exported artifacts, or other one-off steps that should stay +reproducible for future contributors. Document the exact input files and +conversion steps so someone else can regenerate the same output. + ### Publish to Hugging Face If you plan to share the converted model, prefer publishing it on the @@ -189,7 +237,9 @@ the `mlx-community` org when possible. ### Test -Write a basic test: +Write a basic test that matches the model category. + +#### TTS test ```python from mlx_audio.tts.utils import load_model @@ -201,6 +251,18 @@ def test_my_model(): assert results[0].audio.shape[0] > 0 ``` +#### STT test + +```python +from mlx_audio.stt.utils import load_model + +def test_my_stt_model(): + model = load_model("path/to/my-stt-model") + result = model.generate("audio.wav") + assert isinstance(result.text, str) + assert result.text +``` + Run it: ```bash