Repackage frontend as third-party dependency
This commit is contained in:
Owen Song
system
@@ -0,0 +1,199 @@
|
||||
"""
|
||||
ONNX Runtime inference engine for TinyTTS.
|
||||
|
||||
Replaces the PyTorch VoiceSynthesizer.infer() with equivalent
|
||||
ONNX Runtime sessions + NumPy ops for the non-exported parts
|
||||
(alignment path computation).
|
||||
"""
|
||||
import os
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
|
||||
from tiny_tts.text import phonemes_to_ids
|
||||
from tiny_tts.nn import commons
|
||||
from tiny_tts.utils.config import (
|
||||
SAMPLING_RATE, ADD_BLANK, SPK2ID,
|
||||
)
|
||||
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
except ImportError:
|
||||
raise ImportError("onnxruntime is required. Run: pip install onnxruntime")
|
||||
|
||||
|
||||
def _build_session(path: str, use_gpu: bool = False):
|
||||
"""Create an ORT InferenceSession with optional GPU support."""
|
||||
providers = (
|
||||
["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
if use_gpu else
|
||||
["CPUExecutionProvider"]
|
||||
)
|
||||
opts = ort.SessionOptions()
|
||||
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
opts.intra_op_num_threads = os.cpu_count() or 4
|
||||
return ort.InferenceSession(path, sess_options=opts, providers=providers)
|
||||
|
||||
|
||||
def _create_length_mask_np(lengths, max_len=None):
|
||||
"""NumPy equivalent of commons.create_length_mask."""
|
||||
if max_len is None:
|
||||
max_len = int(lengths.max())
|
||||
ids = np.arange(max_len, dtype=np.float32) # [T]
|
||||
mask = (ids[None, :] < lengths[:, None]).astype(np.float32) # [B, T]
|
||||
return mask
|
||||
|
||||
|
||||
def _compute_alignment_path_np(w_ceil, attn_mask):
|
||||
"""
|
||||
Monotonic alignment path - vectorized via cumsum (much faster than Python loops).
|
||||
w_ceil: [B, 1, T_x] — integer duration per phone
|
||||
attn_mask: [B, 1, T_y, T_x] — joint mask
|
||||
Returns attn: [B, 1, T_y, T_x]
|
||||
"""
|
||||
B, _, T_x = w_ceil.shape
|
||||
T_y = attn_mask.shape[2]
|
||||
|
||||
# Build duration matrix: for each phone column expand the duration
|
||||
# cumulative sum of durations gives us the end frame index for each phone
|
||||
dur = w_ceil[:, 0, :] # [B, T_x]
|
||||
cum_dur = np.cumsum(dur, axis=1) # [B, T_x] — end frame (1-indexed)
|
||||
cum_dur_prev = np.pad(cum_dur[:, :-1], ((0,0),(1,0))) # [B, T_x] — start frame
|
||||
|
||||
# Frame indices: [1, T_y, 1]
|
||||
frame_idx = np.arange(T_y, dtype=np.float32)[None, :, None] # [1, T_y, 1]
|
||||
# For each phone, mark frames [start, end)
|
||||
# cum_dur_prev: [B,1,T_x], cum_dur: [B,1,T_x]
|
||||
start = cum_dur_prev[:, None, :] # [B, 1, T_x]
|
||||
end = cum_dur[:, None, :] # [B, 1, T_x]
|
||||
attn = ((frame_idx >= start) & (frame_idx < end)).astype(np.float32) # [B, T_y, T_x]
|
||||
attn = attn[:, None, :, :] # [B, 1, T_y, T_x]
|
||||
return attn * attn_mask
|
||||
|
||||
|
||||
class OnnxTinyTTS:
|
||||
"""
|
||||
Inference using ONNX Runtime.
|
||||
|
||||
Args:
|
||||
onnx_dir: directory containing the 4 .onnx files
|
||||
use_gpu: if True, try CUDAExecutionProvider
|
||||
"""
|
||||
|
||||
def __init__(self, onnx_dir: str = "onnx", use_gpu: bool = False):
|
||||
onnx_dir = os.path.abspath(onnx_dir)
|
||||
print(f"Loading ONNX sessions from: {onnx_dir}")
|
||||
|
||||
self._enc = _build_session(os.path.join(onnx_dir, "text_encoder.onnx"), use_gpu)
|
||||
self._dp = _build_session(os.path.join(onnx_dir, "duration_predictor.onnx"), use_gpu)
|
||||
self._flow = _build_session(os.path.join(onnx_dir, "flow.onnx"), use_gpu)
|
||||
self._dec = _build_session(os.path.join(onnx_dir, "decoder.onnx"), use_gpu)
|
||||
|
||||
print("ONNX sessions ready ✅")
|
||||
|
||||
def _text_to_ids(self, text: str):
|
||||
normalized = normalize_text(text)
|
||||
phones, tones, _ = grapheme_to_phoneme(normalized)
|
||||
phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
|
||||
|
||||
if ADD_BLANK:
|
||||
phone_ids = commons.insert_blanks(phone_ids, 0)
|
||||
tone_ids = commons.insert_blanks(tone_ids, 0)
|
||||
lang_ids = commons.insert_blanks(lang_ids, 0)
|
||||
|
||||
return phone_ids, tone_ids, lang_ids
|
||||
|
||||
def speak(
|
||||
self,
|
||||
text: str,
|
||||
output_path: str = "onnx_output.wav",
|
||||
speaker: str = "female",
|
||||
noise_scale: float = 0.667,
|
||||
noise_scale_w: float = 0.8,
|
||||
length_scale: float = 1.0,
|
||||
output_sr: int = None,
|
||||
) -> np.ndarray:
|
||||
"""Synthesize speech and save to output_path.
|
||||
|
||||
Args:
|
||||
output_sr: If set (e.g. 22050), resample the output from 44100 Hz.
|
||||
Useful to reduce file size while keeping quality.
|
||||
"""
|
||||
print(f"[ONNX] Synthesizing: {text}")
|
||||
|
||||
phone_ids, tone_ids, lang_ids = self._text_to_ids(text)
|
||||
T = len(phone_ids)
|
||||
|
||||
# Prepare inputs as float32 / int64 arrays
|
||||
x = np.array(phone_ids, dtype=np.int64)[None, :] # [1, T]
|
||||
x_len = np.array([T], dtype=np.int64) # [1]
|
||||
tone = np.array(tone_ids, dtype=np.int64)[None, :] # [1, T]
|
||||
lang = np.array(lang_ids, dtype=np.int64)[None, :] # [1, T]
|
||||
bert = np.zeros((1, 1024, T), dtype=np.float32)
|
||||
ja_bert = np.zeros((1, 768, T), dtype=np.float32)
|
||||
sid_val = SPK2ID.get(speaker, 0)
|
||||
sid = np.array([sid_val], dtype=np.int64) # [1]
|
||||
|
||||
# ── 1. Text Encoder ──────────────────────────────────────────────
|
||||
x_enc, m_p, logs_p, x_mask, g = self._enc.run(
|
||||
None,
|
||||
{
|
||||
"phone_ids": x,
|
||||
"phone_lengths":x_len,
|
||||
"tone_ids": tone,
|
||||
"language_ids": lang,
|
||||
"bert": bert,
|
||||
"ja_bert": ja_bert,
|
||||
"speaker_id": sid,
|
||||
},
|
||||
)
|
||||
|
||||
# ── 2. Duration Predictor ─────────────────────────────────────────
|
||||
logw = self._dp.run(None, {"x": x_enc, "x_mask": x_mask, "g": g})[0]
|
||||
|
||||
# ── 3. Alignment Path (NumPy) ─────────────────────────────────────
|
||||
w = np.exp(logw) * x_mask * length_scale # [1, 1, T]
|
||||
w_ceil = np.ceil(w) # [1, 1, T]
|
||||
y_len = max(1, int(w_ceil.sum()))
|
||||
y_lens = np.array([y_len], dtype=np.int64)
|
||||
|
||||
y_mask = _create_length_mask_np(y_lens, y_len) # [1, T_y]
|
||||
y_mask = y_mask[:, None, :] # [1, 1, T_y]
|
||||
# attn_mask: [1, 1, T_y, T_x] (outer product of frame mask and phone mask)
|
||||
attn_mask = y_mask[:, :, :, None] * x_mask[:, :, None, :] # [1,1,T_y,T_x]
|
||||
attn = _compute_alignment_path_np(w_ceil, attn_mask) # [1, 1, T_y, T_x]
|
||||
|
||||
# Expand prior stats via alignment
|
||||
m_p_exp = np.matmul(attn[:, 0], m_p.transpose(0, 2, 1)).transpose(0, 2, 1)
|
||||
logs_p_exp = np.matmul(attn[:, 0], logs_p.transpose(0, 2, 1)).transpose(0, 2, 1)
|
||||
|
||||
# ── 4. Sample z_p ─────────────────────────────────────────────────
|
||||
z_p = m_p_exp + np.random.randn(*m_p_exp.shape).astype(np.float32) * \
|
||||
np.exp(logs_p_exp) * noise_scale
|
||||
|
||||
# ── 5. Flow (reverse) ─────────────────────────────────────────────
|
||||
z = self._flow.run(
|
||||
None,
|
||||
{"z_p": z_p, "y_mask": y_mask.astype(np.float32), "g": g},
|
||||
)[0]
|
||||
|
||||
# ── 6. Decoder ────────────────────────────────────────────────────
|
||||
z_masked = (z * y_mask).astype(np.float32)
|
||||
audio = self._dec.run(None, {"z": z_masked, "g": g})[0] # [1, 1, samples]
|
||||
|
||||
audio_np = audio[0, 0]
|
||||
save_sr = SAMPLING_RATE
|
||||
if output_sr is not None and output_sr != SAMPLING_RATE:
|
||||
try:
|
||||
import torchaudio
|
||||
import torch
|
||||
wav_t = torch.from_numpy(audio_np).unsqueeze(0)
|
||||
resampler = torchaudio.transforms.Resample(SAMPLING_RATE, output_sr)
|
||||
audio_np = resampler(wav_t).squeeze(0).numpy()
|
||||
save_sr = output_sr
|
||||
except Exception as e:
|
||||
print(f"[ONNX] Resampling failed ({e}), saving at {SAMPLING_RATE}Hz")
|
||||
|
||||
sf.write(output_path, audio_np, save_sr)
|
||||
print(f"[ONNX] Saved: {output_path} ({save_sr}Hz)")
|
||||
return audio_np
|
||||
Reference in New Issue
Block a user