File size: 4,165 Bytes
f444284 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """
Kokoro TTS Engine for AX650 NPU.
Pipeline: AXERA Encoder U16 → F0N Mixed FP32 → Decoder S16 → HAR + ISTFT
"""
import os, json, time, numpy as np
import onnxruntime as ort
MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "models")
MAX_LEN = 96
T_PAD = 192
class KokoroEngine:
def __init__(self, model_dir=None):
if model_dir is None:
model_dir = MODEL_DIR
self.model_dir = model_dir
# Load axengine
from axengine import InferenceSession
# Load vocabulary
with open(os.path.join(model_dir, "vocab.json")) as f:
self.vocab = json.load(f)
# Load voice
self.voice = np.load(os.path.join(model_dir, "voice.npy")).astype(np.float32)
self.style_head = self.voice[:, :128]
self.style_tail = self.voice[:, 128:]
# Load models
self.enc = InferenceSession(os.path.join(model_dir, "kokoro_enc_axera.axmodel"))
self.f0n = InferenceSession(os.path.join(model_dir, "kokoro_f0n.axmodel"))
self.dec = InferenceSession(os.path.join(model_dir, "kokoro_dec.axmodel"))
self.har = ort.InferenceSession(os.path.join(model_dir, "kokoro_har_noup.onnx"),
providers=["CPUExecutionProvider"])
self.istft = ort.InferenceSession(os.path.join(model_dir, "kokoro_istft.onnx"),
providers=["CPUExecutionProvider"])
# Sigmoid via frompyfunc (avoids numpy float64 issue on ARM)
self._sigmoid = np.frompyfunc(lambda x: 1.0 / (1.0 + np.exp(-float(x))), 1, 1)
def tokenize(self, phonemes):
"""Convert phoneme string to input_ids."""
ids = [self.vocab[p] for p in phonemes if p in self.vocab][:MAX_LEN - 2]
input_ids = np.zeros((1, MAX_LEN), dtype=np.int32)
input_ids[0, 0] = 0
for i, t in enumerate(ids):
if i + 1 < MAX_LEN - 1:
input_ids[0, i + 1] = t
return input_ids
def synthesize(self, phonemes, speed=1.0):
"""Generate audio from phoneme string. Returns (audio, sample_rate)."""
input_ids = self.tokenize(phonemes)
# 1. Encoder
d, t_en, dur_logits = self.enc.run(
["d", "t_en", "duration_logits"],
{"input_ids": input_ids, "style_tail": self.style_tail}
)
# 2. Duration + Alignment (CPU)
dt = d.transpose(0, 2, 1) # [1, 96, 640]
duration = self._sigmoid(dur_logits).sum(axis=-1).astype(np.float64).squeeze()
pred_dur = np.round(duration / speed).clip(1).astype(np.int64)
Ta_raw = int(pred_dur.sum())
Ta = int(Ta_raw * 224 / 151) # duration correction factor
indices = np.repeat(np.arange(MAX_LEN), pred_dur)
T_out = indices.shape[0]
alignment = np.zeros((1, MAX_LEN, T_out), dtype=np.float32)
for j, i in enumerate(indices):
alignment[0, i, j] = 1.0
en = (dt @ alignment).astype(np.float32) # [1, 640, T_out]
asr = (t_en @ alignment).astype(np.float32) # [1, 512, T_out]
# Pad to fixed length
if en.shape[-1] < T_PAD:
en = np.pad(en, ((0, 0), (0, 0), (0, T_PAD - en.shape[-1])))
asr = np.pad(asr, ((0, 0), (0, 0), (0, T_PAD - asr.shape[-1])))
else:
en = en[:, :, :T_PAD]
asr = asr[:, :, :T_PAD]
# 3. F0N
F0, N = self.f0n.run(["F0", "N"],
{"en": en.astype(np.float32), "style_tail": self.style_tail})
# 4. HAR (CPU)
f0_up = np.repeat(F0, 300, axis=-1).astype(np.float32)
har = self.har.run(["har"], {"f0_up": f0_up})[0]
# 5. Decoder
raw_x = self.dec.run(
["raw_x"],
{"asr": asr.astype(np.float32), "F0_pred": F0.astype(np.float32),
"N_pred": N.astype(np.float32), "ref_s_head": self.style_head,
"har": har.astype(np.float32)}
)[0]
# 6. ISTFT (CPU)
waveform = self.istft.run(["waveform"], {"raw_x": raw_x.astype(np.float32)})[0]
audio = waveform.squeeze()[:int(Ta * 300)]
return audio, 24000
|