""" Inflect-Micro-v2 TTS Engine — multi-chip, zero PyTorch. Supports: AX650 (full NPU), AX620E (encoder NPU + decoder ONNX), AX637 (encoder NPU + decoder ONNX). """ from __future__ import annotations import sys, os, io, math, re, json from pathlib import Path import numpy as np import onnxruntime as ort PKG = Path(__file__).resolve().parent sys.path.insert(0, str(PKG)) sys.path.insert(0, str(PKG / "runtime")) from inflect_vits_frontend import run_vits_frontend from text.symbols import symbols from text import cleaned_text_to_sequence # ── numpy 工具函数 ────────────────────────────────────────────────────── def _intersperse(lst, item): result = [item] * (len(lst) * 2 + 1) result[1::2] = lst return result def _sequence_mask(length, max_length=None): if max_length is None: max_length = int(length.max()) return np.arange(max_length, dtype=length.dtype)[None, :] < length[:, None] def _generate_path(duration, mask): b, _, t_y, t_x = mask.shape cum_duration = np.cumsum(duration, axis=-1) cum_duration_flat = cum_duration.reshape(b * t_x) path = _sequence_mask(cum_duration_flat, t_y).astype(mask.dtype) path = path.reshape(b, t_x, t_y) padded = np.pad(path[:, :-1], ((0, 0), (1, 0), (0, 0))) path = path - padded path = path[:, np.newaxis, :, :].transpose(0, 1, 3, 2) * mask return path # ── 引擎 ──────────────────────────────────────────────────────────────── class InflectTTSEngine: """Inflect-Micro-v2 TTS — 多芯片，零 PyTorch。""" def __init__(self, model_dir: str | Path | None = None, chip: str = "ax650"): if model_dir: self.root = Path(model_dir) else: # Default: models/{chip}/ self.root = PKG.parent.parent / "models" / chip self.chip = chip # 加载配置 with open(PKG / "config.json") as f: self.hps = json.load(f) self.sample_rate = int(self.hps["data"]["sampling_rate"]) self.hop_length = int(self.hps["data"]["hop_length"]) self.hidden_channels = int(self.hps["model"]["hidden_channels"]) self.inter_channels = int(self.hps["model"]["inter_channels"]) # 加载 chip-specific model_meta meta_path = self.root / "model_meta.json" if meta_path.exists(): with open(meta_path) as f: self.meta = json.load(f) else: self.meta = {} # Embedding (n_vocab, hidden_channels) self._emb = np.load(str(self.root / "emb.npy")) self._emb_scale = math.sqrt(self.hidden_channels) # Duration Predictor ONNX self._dp_sess = ort.InferenceSession( str(self.root / "dp.onnx"), providers=["CPUExecutionProvider"] ) # Encoder — always AXMODEL try: import axengine as axe enc_path = str(self.root / "inflect_encoder.axmodel") self.enc_session = axe.InferenceSession(enc_path) self._has_npu = True except ImportError: self._has_npu = False print("[WARN] axengine not available, encoder will use CPU ONNX") self.enc_session = ort.InferenceSession( str(self.root / "inflect_encoder.onnx"), providers=["CPUExecutionProvider"], ) # Decoder — AXMODEL or ONNX fallback dec_axmodel = self.root / "inflect_decoder.axmodel" if dec_axmodel.exists() and self._has_npu: self.dec_session = axe.InferenceSession(str(dec_axmodel)) self._dec_is_npu = True else: dec_onnx = self.root / "inflect_decoder.onnx" self.dec_session = ort.InferenceSession( str(dec_onnx), providers=["CPUExecutionProvider"] ) self._dec_is_npu = False if self._has_npu: print(f"[INFO] {chip}: decoder running on CPU (ONNX)") def synthesize(self, text: str, speed: float = 1.0, variation: float = 0.667, seed: int = 0): normalized = " ".join(text.split()) if not normalized: raise ValueError("Text must not be empty.") sentences = [p.strip() for p in re.split(r"(?<=[.!?;:])\s+", normalized) if p.strip()] if not sentences: sentences = [normalized] pieces = [] for idx, chunk in enumerate(sentences): if idx > 0: pieces.append(np.zeros(round(self.sample_rate * 0.08), dtype=np.float32)) phonemes = run_vits_frontend(chunk).phoneme_text seq = cleaned_text_to_sequence(phonemes) if self.hps["data"]["add_blank"]: seq = _intersperse(seq, 0) if not seq: continue tokens = np.array(seq, dtype=np.int64) tlen = len(tokens) # ── Embedding ── x_emb = self._emb[tokens] * self._emb_scale x_emb = x_emb.T[np.newaxis, :, :] # ── Encoder ── MAX_TOK = 200 if tlen > MAX_TOK: raise ValueError(f"Text too long: {tlen} tokens > {MAX_TOK}") x_pad = np.zeros((1, self.hidden_channels, MAX_TOK), dtype=np.float32) x_pad[:, :, :tlen] = x_emb enc_out = self.enc_session.run(None, { "x_emb": x_pad, "lengths": np.array([tlen], dtype=np.int32), }) m_p = np.asarray(enc_out[0])[:, :, :tlen] logs_p = np.asarray(enc_out[1])[:, :, :tlen] x_enc = np.asarray(enc_out[2])[:, :, :tlen] x_mask = np.asarray(enc_out[3])[:, :, :tlen] # ── Duration Predictor ── logw = self._dp_sess.run(None, { "x": x_enc.astype(np.float32), "x_mask": x_mask.astype(np.float32), })[0] # ── Duration + Alignment ── rng = np.random.RandomState(seed + idx) w = np.exp(logw) * x_mask * (1.0 / speed) w_ceil = np.ceil(w) y_lengths = np.clip(w_ceil.sum(axis=(1, 2)), 1, None).astype(np.int64) y_mask = _sequence_mask(y_lengths, None)[:, np.newaxis, :].astype(x_mask.dtype) attn_mask = x_mask[:, :, np.newaxis, :] * y_mask[:, :, :, np.newaxis] attn = _generate_path(w_ceil, attn_mask) m_p = (attn.squeeze(1) @ m_p.transpose(0, 2, 1)).transpose(0, 2, 1) logs_p = (attn.squeeze(1) @ logs_p.transpose(0, 2, 1)).transpose(0, 2, 1) z_p = m_p + rng.randn(*m_p.shape).astype(np.float32) * np.exp(logs_p) * variation # ── Decoder ── mel_len = z_p.shape[2] MAX_MEL = 500 if mel_len > MAX_MEL: raise ValueError(f"Audio too long: {mel_len} frames > {MAX_MEL}") if self._dec_is_npu: zp_np = np.zeros((1, self.inter_channels, MAX_MEL), dtype=np.float32) ym_np = np.zeros((1, 1, MAX_MEL), dtype=np.float32) zp_np[:, :, :mel_len] = z_p ym_np[:, :, :mel_len] = y_mask dec_out = self.dec_session.run(None, {"z_p": zp_np, "y_mask": ym_np}) waveform = dec_out[0][0, 0, :mel_len * self.hop_length] else: dec_out = self.dec_session.run(None, { "z_p": z_p.astype(np.float32), "y_mask": y_mask.astype(np.float32), }) waveform = dec_out[0][0, 0, :mel_len * self.hop_length] pieces.append(waveform) waveform = np.clip(np.concatenate(pieces), -1.0, 1.0) return self.sample_rate, waveform def save(self, text: str, output: str | Path, **kwargs): import soundfile as sf dest = Path(output) dest.parent.mkdir(parents=True, exist_ok=True) sr, wav = self.synthesize(text, **kwargs) sf.write(dest, wav, sr) return dest