-
-
Notifications
You must be signed in to change notification settings - Fork 672
docs: add contributing guide, model porting guide, and code of conduct #632
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Draft
beshkenadze
wants to merge
6
commits into
Blaizzy:main
Choose a base branch
from
beshkenadze:docs/contributing
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Draft
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit
Hold shift + click to select a range
bdb034c
docs: add CONTRIBUTING.md, ADDING_A_MODEL.md, and CODE_OF_CONDUCT.md
beshkenadze 49706a6
Merge branch 'main' into docs/contributing
Blaizzy c47282c
Merge branch 'main' into docs/contributing
Blaizzy 1f5e36d
Merge branch 'main' into docs/contributing
Blaizzy 1eb4251
Update CONTRIBUTING.md
beshkenadze 983d058
Refine model contribution guides for conversion and STT
beshkenadze File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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) | ||
| ``` | ||
|
|
||
| ## 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. | ||
|
Owner
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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. |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
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
onnxdependency or uses.pttorch files.There was a problem hiding this comment.
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.