Add Supertonic FP16 MLX export and native runner

[seyeong-han]

Summary

Add an ExecuTorch MLX implementation of Supertonic 3, Supertone's on-device multilingual TTS model.

Published ONNX models
  -> export-only Python code
  -> one FP16 ExecuTorch PTE
  -> C++ runner with MLX
  -> 44.1 kHz PCM16 WAV

The Python model files translate the published ONNX architecture and weights into an ExecuTorch PTE. They are needed only when creating the PTE; the C++ runner does not use Python or the source ONNX model files.

This PR supports one-shot, batch-1, five-step synthesis on macOS arm64. It does not include model weights, generated PTEs, WAVs, a persistent server, benchmarks, or application integration.

What to review

The 49 files do not need equal review depth.

1. Create the PTE (Python, export time only)

• examples/models/supertonic/loaders/checkpoint_loader.py
• examples/models/supertonic/model/duration_predictor.py
• examples/models/supertonic/model/text_encoder.py
• examples/models/supertonic/model/vector_estimator.py
• examples/models/supertonic/model/vocoder.py
• examples/models/supertonic/export/common.py
• examples/models/supertonic/export/export_supertonic.py
• examples/models/supertonic/source_transformations/mlx.py

These files define the four Supertonic stages in PyTorch, copy and convert the ONNX weights, apply MLX-compatible graph rewrites, and export four PTE methods: duration_predictor, text_encoder, vector_estimator, and vocoder.

Review the layer/weight mapping, equivalent outputs, dynamic bounds, FP16 conversion, metadata, and MLX delegation. The highest-signal tests are test_checkpoint_loader.py, test_stage_parity.py, and test_mlx_pipeline.py.

2. Run the PTE (C++)

• examples/models/supertonic/runtime/supertonic_runner.cpp
• examples/models/supertonic/runtime/text_processor.cpp
• examples/models/supertonic/runtime/style_loader.cpp
• examples/models/supertonic/runtime/wav_writer.cpp
• examples/models/supertonic/runtime/main.cpp

These files load the PTE, call its four methods, preprocess and chunk text, run five generation steps, trim the waveform, and write a WAV file. Review the PTE compatibility checks, Unicode handling, tensor flow, and output validation.

The highest-signal tests are examples/models/supertonic/runtime/tests/supertonic_runtime_test.cpp and examples/models/supertonic/runtime/tests/run_integration.cmake.

3. Build and CI

• examples/models/supertonic/CMakeLists.txt
• examples/models/supertonic/CMakePresets.json
• .github/workflows/mlx.yml
• examples/models/supertonic/README.md
• examples/models/supertonic/NOTICE

Review the macOS arm64 restriction, MLX runtime files, pinned-asset CI flow, usage, and licensing. The remaining unit tests, package files, and root README/Makefile entries are support changes that can be skimmed.

Runtime files

After export, synthesis needs only supertonic_runner, mlx.metallib, the exported PTE, unicode_indexer.json, and one voice-style JSON file.

Validation

• 151 Python tests and 2/2 native CTests passed, including real PTE synthesis.
• Local Apple M5 Max benchmark (batch-1 FP16, five steps): short-input median 45.9 ms, RTF 0.0244 (40.9x real time).
• ONNX FP32 vs MLX FP16 waveform cosine/Pearson: 0.999603; SQNR: 30.95 dB.
• GitHub CI: 201 successful, 30 skipped, 1 neutral, 0 failed.

Provenance

• Source: supertone-inc/supertonic at 7e2804f96016a7028cb1ed627353c61c1e9dd281 (MIT).
• Model assets: Supertone/supertonic-3 at 3cadd1ee6394adea1bd021217a0e650ede09a323 (BigScience Open RAIL-M).
• Assets are downloaded separately and are not redistributed. Exported PTEs remain subject to the model asset license.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22063

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 39c750c with merge base cff6f4d ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 4m 41s —— View job

---

Code Review — Supertonic FP16 MLX export and native runner

I reviewed the export pipeline (Python), the native runner (C++), and the build/CI wiring. Overall this is a high-quality, defensively-written contribution: the validation density in supertonic_runner.cpp, common.py, and style_loader.cpp is excellent, integer-overflow checks are consistent, and the metadata contract between exporter and runner is well enforced. Notes below, ordered by significance. I couldn't build/run locally (no origin/main merge base in this shallow checkout, and the runner is macOS-arm64/MLX only), so this is a static review.

---

Correctness

1. Dead latent-padding mask in synthesize_chunk — supertonic_runner.cpp:727-739
valid_latents is computed with the same ceil formula as latent_length in latent_layout (supertonic_runner.cpp:362-363):

const int64_t valid_latents =
    (layout.waveform_samples + base_chunk_size*chunk_compress_factor - 1) /
    (base_chunk_size*chunk_compress_factor);   // == layout.latent_length, always

Since layout.waveform_samples == samples and latent_length = ceil(samples/chunk_size), valid_latents is always equal to layout.latent_length. The loop for (position = valid_latents; position < layout.latent_length; ...) therefore never executes — the latent_mask stays all-ones and no latent positions are ever zeroed. This is functionally fine for the current one-latent-per-ceil layout (there is no padding to mask), but the block is unreachable and misleading. Either remove it, or if it's a placeholder for a future padded/batched layout, add a comment saying so. Fix this →

2. Internal chunk padding retained on multi-chunk synthesis — supertonic_runner.cpp:416-448
Each chunk's vocoder waveform has latent_length * chunk_size samples, which rounds up past the chunk's true duration * sample_rate by up to chunk_size - 1 samples. combine_vocoder_chunks only trims the combined waveform to the summed target_duration at the very end, so the trailing overrun of every non-final chunk is baked into the middle of the output rather than trimmed. For the single-chunk one-shot path (the headline use case) this is correct, but multi-chunk output will carry small audible tails between chunks. Consider trimming each chunk to its own duration before concatenating, or documenting this as a known limitation.

3. accumulate_chunk_durations checks finiteness after use — supertonic_runner.cpp:407-411

const float increment = durations[index] + inter_chunk_silence;
result += increment;
if (!std::isfinite(increment) || !std::isfinite(result)) { ... }

The result += increment happens before the finiteness guard. It still throws on the same iteration so the result is never returned, but reordering the check before the accumulation would be cleaner and clearly intentional.

---

Parity between Python and C++ preprocessing

4. preprocess_text — preprocessing.py vs text_processor.cpp
The two implementations must stay in lockstep since the Python path is used for parity tests and the C++ path runs in production. A couple of small divergences worth confirming are intentional:

• Python has a while "``" in text collapse loop (preprocessing.py:111-112) with no C++ equivalent. In both, backtick is already mapped to ' earlier, so this is effectively dead in Python too — but the asymmetry invites drift.
• normalize_nfkd uses CoreFoundation kCFStringNormalizationFormKD (text_processor.cpp:87-112) while Python uses unicodedata.normalize("NFKD"). These are generally equivalent but can differ on rare codepoints; the parity tests (test_stage_parity.py) only exercise the Python path, so the C++ normalizer is not directly parity-checked against Python. Worth noting the C++ preprocessing is only validated by supertonic_runtime_test.cpp, not cross-checked against the reference.

---

Licensing / attribution

5. BSD/Meta headers on adapted MIT code
All new files carry the standard Copyright (c) Meta Platforms, Inc. ... BSD-style header, but the model definitions (model/*.py, source_transformations/mlx.py) and preprocessing are described in the PR as translations of the upstream MIT-licensed Supertonic repo. The NOTICE file preserves the upstream copyright, which is good, but you may want a maintainer to confirm whether the directly-derived files should additionally carry an upstream attribution line rather than only the Meta BSD header. This is a policy call for the ExecuTorch maintainers, not a blocker.

---

Minor / nits

• main.cpp:63-71: require_single_voice_style_path and validate_language are called, then require_file(style, ...) — the style path existence check happens after language validation. Harmless, but grouping the file-existence checks together reads more clearly.
• text_processor.cpp:256: index = next == 0 ? 0 : next - 1; — next is always > index >= 0 here so it can never be 0 at this point; the ternary guard is defensive but unreachable. Minor.
• CMakeLists.txt:58: gflags_DIR is hardcoded to ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/gflags, coupling this project to the exact mlx-release install layout. It matches the documented make supertonic-mlx flow, so it's fine, but a stray direct cmake invocation without the prior install will fail with a less obvious error than the explicit FATAL_ERROR guards you added elsewhere.
• .github/workflows/mlx.yml: the Supertonic job downloads a pinned HF revision and runs a real end-to-end synthesis in CI — nicely thorough. Note I can't modify workflow files, so this is observation only.

---

Things done well

• Exhaustive input validation with overflow-safe arithmetic (element_count, checked_sample_count, checked_add_size, validate_wav_layout).
• The exporter's initializer-map verification (load_onnx_initializers rejects unknown/unmapped/duplicate/missing/unused initializers) makes weight-mapping regressions loud at export time.
• Strong metadata contract enforced on both sides (runtime_metadata ↔ validate_metadata_contract), including the latent_channels == latent_dim * chunk_compress_factor consistency check.
• The SUPERTONIC_PURE_HELPERS_ONLY split lets the pure numeric helpers be unit-tested without the MLX/ExecuTorch runtime — clean design.

None of the above are hard blockers for the one-shot single-chunk path this PR targets; items 1 and 2 are the ones I'd most want addressed or explicitly documented.
· branch supertonic/01-runtime