Rebase pc/add-ace on upstream main#739
Closed
shreyaskarnik wants to merge 173 commits into
Closed
Conversation
No need to continue generating for clients that have hung up. This reduces latency processing the next request. Co-authored-by: Prince Canuma <prince.gdt@gmail.com>
…ons (Blaizzy#624) Quantized Kokoro checkpoints (8bit/6bit/4bit) store conv-like weights in MLX layout, while bf16 checkpoints use PyTorch layout. The sanitize() functions were unconditionally transposing these weights, breaking quantized models with shape mismatches. - Detect packed quantized checkpoints via .scales/.biases key suffixes - Skip 3D conv transposition for already-converted quantized weights - Guard duration path with nan_to_num and cap max frames per phoneme - Return silence instead of crashing on empty concatenation Fixes Blaizzy#623 Co-authored-by: Prince Canuma <prince.gdt@gmail.com>
…1 dB
The parity test was permanently skipped, so several load-bearing bugs in
the MLX port went undetected. Wiring up the reference inference surfaced
them; this commit fixes them and verifies numerical parity.
Config / conversion
- Add zfturbo_vocals_v1 preset (dim=192, depth=8, hop=512,
mask_estimator_depth=1) for ZFTurbo's v1.0.0 MIT-licensed vocals
release asset. None of the existing three presets matched it.
Port bugs fixed in model.py
- RoPE convention: MLX used the halved split (x[:half] / x[half:]);
rotary_embedding_torch (what ZFTurbo imports) uses interleaved pairs
(x[2i] / x[2i+1]) with a doubled-frequency layout. Same math family,
different rotation plane — attention output was uncorrelated with
the reference until this was fixed. Single biggest jump in SDR.
- RMSNorm eps: mlx.nn.RMSNorm uses additive eps=1e-5; ZFTurbo uses
F.normalize(x, dim=-1) which is effectively max(||x||, 1e-12). The
two agree for typical magnitudes but diverge by up to 70% at the
small-magnitude high-frequency STFT bins that the band-split's
per-band RMSNorm ingests. Replaced with a local RMSNorm matching
F.normalize semantics.
- to_gates: declared bias=False, but ZFTurbo's Linear(dim, heads) has
the default bias=True and ships a to_gates.bias tensor.
- MaskEstimator: hardcoded to 3 linears (mask_estimator_depth=2 in
effect). Now honors config.mask_estimator_depth, so depth-1
checkpoints (like this one) load with the correct number of linears.
- MelFilterbank: custom Slaney triangular filter did not quite match
librosa's. Swapped for librosa.filters.mel + ZFTurbo's explicit
fb[0,0]=1 / fb[-1,-1]=1 force-assigns so every freq bin is covered.
- STFT: zero-padded the audio before framing, but torch.stft(center=True)
pad_mode defaults to reflect. Built reflect pad via slice+concat
(mx.pad has no reflect mode).
- STFT post-processing: .real() / .imag() as method calls; these are
properties on mlx.core.array.
- BandSplit: freq_indices stored as np.ndarray, but MLX fancy-indexing
rejects raw numpy — convert once at MelFilterbank init.
sanitize() remaps added
- Drop rotary_embed.freqs — MLX computes RoPE freqs on the fly.
- Mask-estimator per-band Sequential -> list index: PyTorch stores
linears at positions {0, 2} (depth=1) or {0, 2, 4} (depth=2); MLX
list positions are consecutive {0, 1, 2}.
- to_out.0.weight -> to_out.weight — PyTorch wraps in Sequential(Linear,
Dropout), MLX uses a bare Linear.
- .gamma -> .weight — PyTorch RMSNorm names its scale "gamma";
mlx.nn.RMSNorm names it "weight".
Test wiring in tests/sts/test_mel_roformer_parity.py
- Replace the pytest.skip stub in test_sdr_parity with the real
single-chunk SDR comparison: dynamically imports a user-supplied
torch_infer.py (see PARITY_TESTING.md) via importlib, runs both
models on the same chunk of audio, asserts SDR > target (default
40 dB).
- Add fixtures for MEL_ROFORMER_TORCH_INFER, MEL_ROFORMER_TORCH_CONFIG,
MEL_ROFORMER_CHUNK_SAMPLES.
Verified
- All 25 existing unit tests in tests/sts/test_mel_roformer.py still pass.
- test_qkv_split_preserves_weights: pass.
- test_sdr_parity on the v1.0.0 ZFTurbo MIT vocals ckpt +
ZFTurbo config + a deterministic synthetic 30s stereo signal:
SDR = 58.11 dB (target > 40 dB).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Register MelRoFormer, MelRoFormerConfig, and MelRoFormerResult in the parent package so users can import via mlx_audio.sts.models without reaching into the subpackage. Rename our result dataclass from SeparationResult to MelRoFormerResult to avoid a name collision with sam_audio.SeparationResult, which is already upstream and has a different streaming-oriented shape. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Pre-commit auto-fixes from running hooks before PR submission. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…el_type Two related STT fixes: 1. granite_speech sanitize(): The Conv1d weight transposition was applied unconditionally to all non-quantized models. Models with weights already saved in MLX-native layout (out, kernel, in) — such as bf16 safetensors converted directly from HuggingFace — were being incorrectly transposed, producing garbled or empty transcriptions. Fix: only transpose when the weight is actually in PyTorch layout (out, in, kernel), detected by checking shape[-1] > shape[-2] (true for depthwise convs where kernel > 1). 2. parakeet ModelConfig.from_dict(): NeMo-format config.json files do not include a top-level model_type field. Downstream tools that rely on this field (e.g. omlx model discovery) would misidentify the model as an LLM, causing a KeyError at inference time. Fix: inject model_type: 'parakeet' in from_dict() when the field is absent.
…point Adds an end-to-end streaming ingestion path for Voxtral Realtime so audio can be consumed incrementally and transcription deltas emitted as soon as the encoder is caught up, without changing the batch generate() path. Exposes it over a new OpenAI Realtime-compatible WebSocket endpoint on the built-in server, and closes the burst-mode perf gap with voxmlx. Streaming core (mlx_audio/stt/models/voxtral_realtime/streaming.py): - StreamingAudioSource: thread-safe queue bridging any producer (thread or asyncio task) to the MLX consumer, with coalescing, close signal, blocking read, and dtype coercion. - StreamingMel / StreamingCausalConv1d / StreamingConvStem: per-chunk equivalents of the batch mel + conv stem, bit-exact for mel and the conv stem via explicit edge state. StreamingConvStem casts mel to the conv weight dtype so bf16 weights don't promote the whole graph to fp32. - StreamingEncoder: transformer layers with RotatingKVCache per layer and a running RoPE position; forces an explicit array mask so SDPA stays correct when Q_len < K_len. - StreamingDownsampler: buffers encoder output to emit adapter frames in downsample_factor-aligned groups. - VoxtralStreamingSession: stateful session where feed() is a cheap thread-safe push and step(max_decode_tokens=N) runs a bounded unit of MLX work. Splitting feed and step lets callers release the MLX executor between steps so multiple streams can interleave on a single-worker executor. Model surface (voxtral_realtime.py): - create_streaming_session(...) -> VoxtralStreamingSession - generate_streaming(source, ...) as a thin convenience wrapper Perf (c2c0f0c): port decoder/encoder to mx.fast.rope + RotatingKVCache, quantize tok_embeddings, normalize checkpoint dtype, vectorize mel window extraction, and cast mel -> conv weight dtype at the conv stem. Burst-mode wall is within 1% of voxmlx on 10s audio (3.594s vs 3.559s). Config cleanup: SAMPLE_RATE / FRAME_RATE / HOP_LENGTH / RAW_AUDIO_LENGTH_PER_TOK / AUDIO_LENGTH_PER_TOK and the _num_audio_tokens / _num_delay_tokens helpers move to config.py so streaming.py imports them at module scope. Expanded session docstring documents transcription_delay_ms (incl. the 2400 ms preset that matches offline-mode WER) and clarifies that max_decode_tokens is a cooperative yield granularity, not a speedup. Server (mlx_audio/server.py): - New /v1/realtime WebSocket speaking a subset of the OpenAI Realtime protocol: session.created / session.update(d), input_audio_buffer.append (base64 PCM16), and input_audio_buffer.commit {final: true} -> streamed response.audio_transcript.delta plus a final response.audio_transcript.done. - Model picked via ?model=<repo> (default: $MLX_AUDIO_REALTIME_MODEL or mlx-community/Voxtral-Mini-4B-Realtime-2602-4bit) and reused via the existing ModelProvider cache. - Thin wrapper over create_streaming_session; session.step runs in asyncio.to_thread so the loop keeps receiving audio while MLX decodes. Drains between appends for low-latency deltas and fully flushes on final commit. After done we start a fresh session so the same socket handles multiple utterances. - Existing VAD-based /v1/audio/transcriptions/realtime is left in place; the new endpoint is purely additive. - Weight loading: support Mistral-native repos and improve default model resolution. Docs: new Live Input Streaming section in docs/models/stt/voxtral-realtime.md with StreamingAudioSource + generate_streaming() example and a lower-level create_streaming_session() example for async servers and multi-stream executors. Tests (mlx_audio/stt/tests/test_voxtral_realtime_streaming.py): - StreamingAudioSource threading behavior (coalescing, close signal, blocking read, dtype coercion). - StreamingMel parity vs compute_mel_spectrogram across chunk sizes. - StreamingCausalConv1d parity on a freshly-initialized CausalConv1d (no model weights required). - Encoder + session parity behind a weight-gated TestCase that skips when VOXTRAL_REALTIME_MODEL is not available.
Port of Boson AI's Higgs Audio v2 to MLX. Llama-3.2-3B backbone with a
dual-FFN decoder layer (text + audio paths share self-attention; LN + MLP
are per-path, routed by audio_out_mask) and delay-pattern audio emission
(codebook i lags by i frames).
Framework interface
-------------------
Exposes `Model` and `ModelConfig` in the shape expected by
`mlx_audio.tts.utils.load()` and `python -m mlx_audio.tts.generate`:
python -m mlx_audio.tts.generate \
--model mlx-community/higgs-audio-v2-3B-mlx-q8 \
--text "Hello" --ref_audio voice.wav --ref_text "..."
`Model` subclasses the underlying `HiggsAudioModel` so safetensors keys
land unchanged. Tokenizer + codec attach via `post_load_hook`. The
generate() signature matches the standard TTS convention (text, voice,
ref_audio, ref_text, ...) and yields a mlx_audio.tts.models.base
GenerationResult.
`HiggsAudioServer` is kept as an additional Python entrypoint for
Higgs-specific kwargs (max_new_frames, ras_win_len, fade_in_ms, etc.).
Load-bearing details
--------------------
The first generated audio frame must be a synthetic all-stream_bos
frame — sampling from the model's audio_logits at the <|audio_out_bos|>
text position collapses to stream-EOS on half the codebooks (those
positions were never trained for direct audio prediction). See
Model.generate and HiggsAudioModel._generate_raw_frames for the full
AUDIO_INIT + K-frame ramp-in + EOS ramp-out state machine.
Quantization
------------
MLX native 4/6/8-bit on the Llama backbone. `Model.model_quant_predicate`
protects `audio_codebook_embeddings` and `audio_decoder_proj.audio_lm_head`
at bf16 — quantizing them introduces voice-character drift (pitch
register shifts at q6, trajectory instability at q4). RTF on M5 Max:
bf16 0.60× / q8 0.36× / q6 0.33×.
Bundled assets
--------------
Three drop-in reference voices in examples/voice_prompts/ (en_woman,
en_man, en_man_deep), generated via smart-voice mode so they're
license-clean. Example: examples/higgs_audio_clone_demo.py.
Tests
-----
16 unit tests in mlx_audio/tts/tests/test_higgs_audio.py — delay-pattern
round-trip, audio-embedding lookup, sampling equivalence at T=0, tiny
config model forward, selective-quant predicate verification, and the
framework-interface contract (Model subclass, ModelConfig.from_dict,
sample_rate, model_quant_predicate, generate-before-load guard).
References
----------
Original: https://github.com/boson-ai/higgs-audio
HF (bf16): https://huggingface.co/bosonai/higgs-audio-v2-generation-3B-base
HF (q8): https://huggingface.co/mlx-community/higgs-audio-v2-3B-mlx-q8
HF (q6): https://huggingface.co/mlx-community/higgs-audio-v2-3B-mlx-q6
HF codec: https://huggingface.co/mlx-community/higgs-audio-v2-tokenizer
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Per @lucasnewman review comment on PR Blaizzy#654 — the final PR should not contain scratchpad / planning .md files. PARITY_TESTING.md was an internal M5 Pro setup guide used during parity verification; the content has no use for downstream pip-install users of mlx-audio. The parity test itself (`tests/sts/test_mel_roformer_parity.py`) remains and is self-documenting via its docstrings and fixture parameters. Users who want to run parity can follow the @pytest.mark.requires_torch flag and the env-var fixtures directly. README.md is unchanged — it's user-facing architecture + license + usage documentation, not a planning doc. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Removed the workflow that checks for documentation updates on user-facing changes. Signed-off-by: Prince Canuma <prince.gdt@gmail.com>
Matches the pre-commit config (.pre-commit-config.yaml — black 24.2.0, isort 5.13.2 profile=black) enforced by the PR CI style job. No code changes — pure formatting + import-order + uv.lock reconciliation with the pydub dep already declared in pyproject.toml.
…-sanitize-and-parakeet-model-type fix(stt): correct granite_speech Conv1d weight sanitization and add parakeet model_type
Addresses @lucasnewman's review request. Covers usage (CLI + standard Python API + Higgs-specific HiggsAudioServer), voice cloning, parameter table, available models / quantizations / RTF / memory, conversion, architecture, and license. Links to the full port docs at docs/models/tts/higgs_audio.md for depth.
…inference fix(cohere): restore quantized inference for 8-bit and 4-bit checkpoints
…improve audio resampling
…DACVAE decode - Support Irodori-TTS v2 (latent_dim=32, Semantic-DACVAE) alongside v1 - Add VoiceDesign caption conditioning (use_caption_condition=True) - Accept instruct kwarg as alias for caption in model.generate() - Use chunked DACVAE decoding (chunk_size=50) to stay within 16GB memory limits Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: tian-sweetaylor <153503632+tian-sweetaylor@users.noreply.github.com>
Signed-off-by: tian-sweetaylor <153503632+tian-sweetaylor@users.noreply.github.com>
Signed-off-by: tian-sweetaylor <153503632+tian-sweetaylor@users.noreply.github.com>
Signed-off-by: tian-sweetaylor <153503632+tian-sweetaylor@users.noreply.github.com>
Signed-off-by: tian-sweetaylor <153503632+tian-sweetaylor@users.noreply.github.com>
style: run pre-commit formatting Signed-off-by: tian-sweetaylor <153503632+tian-sweetaylor@users.noreply.github.com>
Signed-off-by: tian-sweetaylor <153503632+tian-sweetaylor@users.noreply.github.com>
feat(vad): Add FSMN-VAD model support
…ad-for-speech Pre-flight model load on streaming audio routes
- Added methods for encoding audio to latent representations and normalizing audio inputs. - Implemented instruction generation based on task types, improving flexibility for various audio tasks. - Introduced new task type constants and instruction templates in the configuration file. - Updated imports and module exports to include new constants and configurations. These changes enhance the model's capabilities for audio processing and task management.
- Introduced a new section in the README for music generation, detailing the ACE-Step model for text-to-music tasks. - Added CLI and Python API examples for generating music from text descriptions and enhancing existing audio tracks. - Updated the `generate.py` script to accept additional generation parameters via a JSON string. - Enhanced the ACE-Step model loading logic to check for weights in both turbo and root directories. These changes expand the functionality of the MLX-Audio framework to include music generation features.
- Converter: export all weight prefixes (tokenizer, detokenizer, null_condition_emb) not just decoder/encoder - Converter: keep decoder. prefix instead of renaming to dit. - Converter: use underscore naming for decoder conv params (proj_in_weight) to match DiTModel bare parameters - Sanitize: add dit. -> decoder. mapping for legacy weights - Sanitize: scope conv param renaming to decoder only to avoid breaking detokenizer.proj_out (nn.Linear) - Add "ace" to MODEL_REMAPPING for HF repo name resolution
The turbo model requires 5Hz LM hints to generate music. Without LM-generated audio codes guiding the DiT, the diffusion simply denoises back to silence. Previously, LM hints were only generated for cover tasks.
The turbo model produces silence without LM hints, so enable them by default for a working out-of-the-box experience.
Add explicit language section to the LM prompt template so the 5Hz planner model better respects the vocal_language parameter. The turbo HF model uses Qwen3 word embeddings for lyrics (not BPE phoneme tokens), which limits vocal synthesis to the model's learned capabilities.
8 steps produces good instrumentals but vocals need more diffusion steps to resolve clearly. 20 steps still runs well under real-time (~8s diffusion for 30s of audio).
- Cache sliding window mask in DiT instead of recreating per-layer (240x)
- Pass None instead of all-zeros no-op self-attention mask
- Remove dead: test_model(), omega_scale param, silence_latent param,
unused nested_weights, TASK_TYPES_TURBO/BASE, VAEConfig, TextEncoderConfig
- Remove redundant hidden_size attrs stored but never read (3 classes)
- Fix fragile .replace(".pt", ".npy") path handling
- Strip narrating comments that restate what code already expresses
22 tests covering: config, KVCache, attention masks, RMSNorm, TimestepEmbedding, RotaryEmbedding, Attention (self + cross with cache), MLP, DiTLayer, DiTModel (forward, cache, conv_transpose), and weight sanitization (dit prefix remap, proj_in scoping, scale_shift_table reshape). Tests caught a stale cross_attn_mask variable left from the simplify pass — fixed to None.
- Document use_lm=True default (required for turbo model) - Document num_steps=20 default (needed for clear vocals) - Add "How It Works" section explaining the LM planner + DiT pipeline - Add Hugging Face repo links (fp32 and 4-bit quantized) - Update performance table with actual M4 Max measurements - Document LM size trade-offs (0.6B vs 4B) - List known limitations (language detection, seed variance)
Draft
Contributor
Author
|
Closing — rebased branch carries too much unrelated diff from main. Will open a focused PR instead. |
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
19 participants
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.
Summary
pc/add-aceon currentmain(v0.4.3) to resolve merge conflictsmlx_audio/tts/utils.py(new model entries from both sides merged)mlx_audio/tts/generate.py(imports + argparse args merged)Ref: #499
Context
Per Prince's comment on #499 — driving the rebase so the branch is mergeable.