Merge pull request #13 from shivammehta25/dev

Merging dev to main | adding ONNX support

README.md

@@ -36,7 +36,6 @@ Check out our [demo page](https://shivammehta25.github.io/Matcha-TTS) and read [
 
 [![Watch the video](https://img.youtube.com/vi/xmvJkz3bqw0/hqdefault.jpg)](https://youtu.be/xmvJkz3bqw0)
 
-
 ## Installation
 
 1. Create an environment (suggested but optional)
@@ -46,7 +45,7 @@ conda create -n matcha-tts python=3.10 -y
 conda activate matcha-tts
 ```
 
-2. Install Matcha TTS using pip  or from source
+2. Install Matcha TTS using pip or from source
 
 ```bash
 pip install matcha-tts
@@ -189,6 +188,70 @@ python matcha/train.py experiment=ljspeech trainer.devices=[0,1]
 matcha-tts --text "<INPUT TEXT>" --checkpoint_path <PATH TO CHECKPOINT>
 ```
 
+## ONNX support
+
+> Special thanks to [@mush42](https://github.com/mush42) for implementing ONNX export and inference support.
+
+It is possible to export Matcha checkpoints to [ONNX](https://onnx.ai/), and run inference on the exported ONNX graph.
+
+### ONNX export
+
+To export a checkpoint to ONNX, first install ONNX with
+
+```bash
+pip install onnx
+```
+
+then run the following:
+
+```bash
+python3 -m matcha.onnx.export matcha.ckpt model.onnx --n-timesteps 5
+```
+
+Optionally, the ONNX exporter accepts **vocoder-name** and **vocoder-checkpoint** arguments. This enables you to embed the vocoder in the exported graph and generate waveforms in a single run (similar to end-to-end TTS systems).
+
+**Note** that `n_timesteps` is treated as a hyper-parameter rather than a model input. This means you should specify it during export (not during inference). If not specified, `n_timesteps` is set to **5**.
+
+**Important**: for now, torch>=2.1.0 is needed for export since the `scaled_product_attention` operator is not exportable in older versions. Until the final version is released, those who want to export their models must install torch>=2.1.0 manually as a pre-release.
+
+### ONNX Inference
+
+To run inference on the exported model, first install `onnxruntime` using
+
+```bash
+pip install onnxruntime
+pip install onnxruntime-gpu # for GPU inference
+```
+
+then use the following:
+
+```bash
+python3 -m matcha.onnx.infer model.onnx --text "hey" --output-dir ./outputs
+```
+
+You can also control synthesis parameters:
+
+```bash
+python3 -m matcha.onnx.infer model.onnx --text "hey" --output-dir ./outputs --temperature 0.4 --speaking_rate 0.9 --spk 0
+```
+
+To run inference on **GPU**, make sure to install **onnxruntime-gpu** package, and then pass `--gpu` to the inference command:
+
+```bash
+python3 -m matcha.onnx.infer model.onnx --text "hey" --output-dir ./outputs --gpu
+```
+
+If you exported only Matcha to ONNX, this will write mel-spectrogram as graphs and `numpy` arrays to the output directory.
+If you embedded the vocoder in the exported graph, this will write `.wav` audio files to the output directory.
+
+If you exported only Matcha to ONNX, and you want to run a full TTS pipeline, you can pass a path to a vocoder model in `ONNX` format:
+
+```bash
+python3 -m matcha.onnx.infer model.onnx --text "hey" --output-dir ./outputs --vocoder hifigan.small.onnx
+```
+
+This will write `.wav` audio files to the output directory.
+
 ## Citation information
 
 If you use our code or otherwise find this work useful, please cite our paper:

matcha/VERSION

@@ -1 +1 @@
-0.0.3
+0.0.4

matcha/models/matcha_tts.py

@@ -116,7 +116,7 @@ def synthesise(self, x, x_lengths, n_timesteps, temperature=1.0, spks=None, leng
  w = torch.exp(logw) * x_mask
  w_ceil = torch.ceil(w) * length_scale
  y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
- y_max_length = int(y_lengths.max())
+ y_max_length = y_lengths.max()
  y_max_length_ = fix_len_compatibility(y_max_length)
 
  # Using obtained durations `w` construct alignment map `attn`

matcha/onnx/export.py

@@ -0,0 +1,181 @@
+import argparse
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+from lightning import LightningModule
+
+from matcha.cli import VOCODER_URLS, load_matcha, load_vocoder
+
+DEFAULT_OPSET = 15
+
+SEED = 1234
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed(SEED)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+
+class MatchaWithVocoder(LightningModule):
+ def __init__(self, matcha, vocoder):
+ super().__init__()
+ self.matcha = matcha
+ self.vocoder = vocoder
+
+ def forward(self, x, x_lengths, scales, spks=None):
+ mel, mel_lengths = self.matcha(x, x_lengths, scales, spks)
+ wavs = self.vocoder(mel).clamp(-1, 1)
+ lengths = mel_lengths * 256
+ return wavs.squeeze(1), lengths
+
+
+def get_exportable_module(matcha, vocoder, n_timesteps):
+ """
+ Return an appropriate `LighteningModule` and output-node names
+ based on whether the vocoder is embedded in the final graph
+ """
+
+ def onnx_forward_func(x, x_lengths, scales, spks=None):
+ """
+ Custom forward function for accepting
+ scaler parameters as tensors
+ """
+ # Extract scaler parameters from tensors
+ temperature = scales[0]
+ length_scale = scales[1]
+ output = matcha.synthesise(x, x_lengths, n_timesteps, temperature, spks, length_scale)
+ return output["mel"], output["mel_lengths"]
+
+ # Monkey-patch Matcha's forward function
+ matcha.forward = onnx_forward_func
+
+ if vocoder is None:
+ model, output_names = matcha, ["mel", "mel_lengths"]
+ else:
+ model = MatchaWithVocoder(matcha, vocoder)
+ output_names = ["wav", "wav_lengths"]
+ return model, output_names
+
+
+def get_inputs(is_multi_speaker):
+ """
+ Create dummy inputs for tracing
+ """
+ dummy_input_length = 50
+ x = torch.randint(low=0, high=20, size=(1, dummy_input_length), dtype=torch.long)
+ x_lengths = torch.LongTensor([dummy_input_length])
+
+ # Scales
+ temperature = 0.667
+ length_scale = 1.0
+ scales = torch.Tensor([temperature, length_scale])
+
+ model_inputs = [x, x_lengths, scales]
+ input_names = [
+ "x",
+ "x_lengths",
+ "scales",
+ ]
+
+ if is_multi_speaker:
+ spks = torch.LongTensor([1])
+ model_inputs.append(spks)
+ input_names.append("spks")
+
+ return tuple(model_inputs), input_names
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Export 🍵 Matcha-TTS to ONNX")
+
+ parser.add_argument(
+ "checkpoint_path",
+ type=str,
+ help="Path to the model checkpoint",
+ )
+ parser.add_argument("output", type=str, help="Path to output `.onnx` file")
+ parser.add_argument(
+ "--n-timesteps", type=int, default=5, help="Number of steps to use for reverse diffusion in decoder (default 5)"
+ )
+ parser.add_argument(
+ "--vocoder-name",
+ type=str,
+ choices=list(VOCODER_URLS.keys()),
+ default=None,
+ help="Name of the vocoder to embed in the ONNX graph",
+ )
+ parser.add_argument(
+ "--vocoder-checkpoint-path",
+ type=str,
+ default=None,
+ help="Vocoder checkpoint to embed in the ONNX graph for an `e2e` like experience",
+ )
+ parser.add_argument("--opset", type=int, default=DEFAULT_OPSET, help="ONNX opset version to use (default 15")
+
+ args = parser.parse_args()
+
+ print(f"[🍵] Loading Matcha checkpoint from {args.checkpoint_path}")
+ print(f"Setting n_timesteps to {args.n_timesteps}")
+
+ checkpoint_path = Path(args.checkpoint_path)
+ matcha = load_matcha(checkpoint_path.stem, checkpoint_path, "cpu")
+
+ if args.vocoder_name or args.vocoder_checkpoint_path:
+ assert (
+ args.vocoder_name and args.vocoder_checkpoint_path
+ ), "Both vocoder_name and vocoder-checkpoint are required when embedding the vocoder in the ONNX graph."
+ vocoder, _ = load_vocoder(args.vocoder_name, args.vocoder_checkpoint_path, "cpu")
+ else:
+ vocoder = None
+
+ is_multi_speaker = matcha.n_spks > 1
+
+ dummy_input, input_names = get_inputs(is_multi_speaker)
+ model, output_names = get_exportable_module(matcha, vocoder, args.n_timesteps)
+
+ # Set dynamic shape for inputs/outputs
+ dynamic_axes = {
+ "x": {0: "batch_size", 1: "time"},
+ "x_lengths": {0: "batch_size"},
+ }
+
+ if vocoder is None:
+ dynamic_axes.update(
+ {
+ "mel": {0: "batch_size", 2: "time"},
+ "mel_lengths": {0: "batch_size"},
+ }
+ )
+ else:
+ print("Embedding the vocoder in the ONNX graph")
+ dynamic_axes.update(
+ {
+ "wav": {0: "batch_size", 1: "time"},
+ "wav_lengths": {0: "batch_size"},
+ }
+ )
+
+ if is_multi_speaker:
+ dynamic_axes["spks"] = {0: "batch_size"}
+
+ # Create the output directory (if not exists)
+ Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+
+ model.to_onnx(
+ args.output,
+ dummy_input,
+ input_names=input_names,
+ output_names=output_names,
+ dynamic_axes=dynamic_axes,
+ opset_version=args.opset,
+ export_params=True,
+ do_constant_folding=True,
+ )
+ print(f"[🍵] ONNX model exported to {args.output}")
+
+
+if __name__ == "__main__":
+ main()