ONNX Model Export
Relevant source files
The following files were used as context for generating this wiki page:
The ONNX Model Export system is a critical component of the Piper project, responsible for converting trained PyTorch Lightning checkpoints into optimized Open Neural Network Exchange (ONNX) format. This transition enables the neural text-to-speech (TTS) models to run efficiently on a variety of hardware platforms using the onnxruntime library, which is integrated via both Python and C++ APIs.
The export process involves taking a VITS-based generator model, stripping training-specific logic (such as weight normalization), and defining a standardized set of inputs and outputs for inference. This ensures that downstream consumers, such as the libpiper C++ shared library or the PiperVoice Python class, can interact with the model using a predictable interface.
Sources: src/piper/train/export_onnx.py, libpiper/src/piper.cpp
Export Workflow
The conversion from a PyTorch checkpoint to an ONNX model follows a specific pipeline. It begins with loading the VitsModel from a checkpoint file, followed by preparing the generator for inference-only operations.
Model Preparation
Before export, the generator model (model_g) is switched to evaluation mode. A key optimization step involves removing weight normalization from the decoder using model_g.dec.remove_weight_norm(). The model's forward method is then wrapped in a custom inference function that specifically handles inputs like phoneme sequences, lengths, scales, and optional speaker IDs.
Sources: src/piper/train/export_onnx.py:53-75, src/piper/train/export_generator.py:46-52
Export Logic Flow
The following diagram illustrates the sequence of operations performed during the export process.
The export uses an opset_version of 15 to ensure compatibility with modern features of the ONNX format.
Sources: src/piper/train/export_onnx.py:14, tests/generate_test_voice.py:27
Input and Output Specifications
To ensure interoperability, all exported Piper models adhere to a strict signature for their input and output tensors.
Input Tensors
| Input Name | Data Type | Description |
|---|---|---|
input | int64 | Long tensor of phoneme IDs. Shape is [batch_size, phonemes]. |
input_lengths | int64 | Length of each phoneme sequence in the batch. |
scales | float32 | A tensor containing [noise_scale, length_scale, noise_w_scale]. |
sid | int64 | (Optional) Speaker ID for multi-speaker models. |
Sources: src/piper/train/export_onnx.py:95-102, libpiper/src/piper.cpp:322-348
Output Tensors
The primary output is a tensor named output (or audio) containing the synthesized waveform. However, models can be patched to include a second output for phoneme alignments.
| Output Name | Data Type | Description |
|---|---|---|
output | float32 | Synthesized audio samples. Shape is [batch_size, 1, time]. |
(Alignment) | float32 | (Optional) Phoneme width/alignment data, typically detected from "Ceil" nodes. |
Sources: src/piper/train/export_onnx.py:103-108, src/piper/patch_voice_with_alignment.py:38-60
Alignment Patching
Piper supports an experimental feature for phoneme-to-audio alignments. This is achieved by "patching" an existing ONNX model to expose internal tensors as outputs.
The patch_voice_with_alignment.py utility scans the ONNX graph for Ceil operation nodes. These nodes often represent the calculated phoneme widths during the inference process. By marking these internal tensors as graph outputs, the ONNX runtime can return alignment data alongside the audio.
Implementation in Downstream Consumers
The exported ONNX model is consumed by both the Python and C++ implementations of the Piper engine.
C++ (libpiper)
In libpiper, the piper_create function initializes an Ort::Session using the model path. During synthesis, the piper_synthesize_next function prepares the input tensors (including the scales tensor) and runs the ONNX session to retrieve audio and alignment data.
Sources: libpiper/src/piper.cpp:103-107, libpiper/src/piper.cpp:322-370
Python (PiperVoice)
The PiperVoice class uses onnxruntime.InferenceSession to load the model. The phoneme_ids_to_audio method constructs the args dictionary containing the "input", "input_lengths", and "scales" tensors before executing session.run.
Sources: src/piper/voice.py:133-145, src/piper/voice.py:483-500
Summary
The ONNX Model Export system provides the bridge between the research-oriented training environment (PyTorch) and the production-oriented inference environments (C++/Python). By standardizing inputs like phoneme IDs and scales, and providing mechanisms for alignment patching, the project ensures that high-quality neural TTS can be deployed across diverse platforms with minimal overhead.