python/inflect_tts_sdk/tts_engine.py

"""
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