python/openai_server.py

#!/usr/bin/env python3
"""OpenAI-Compatible 板端 TTS 服务示例（POST /v1/audio/speech，torch-free）。

仅依赖 numpy + pyaxengine（AX 芯片）：
  input(speech tokens) ──S3Gen AXMODEL(axengine)──> mel ──numpy Griffin-Lim──> wav

说明：
  - 本服务输出的是 S3Gen 的"语音 token -> 音频"能力；若需要 text -> audio，
    在宿主用 T3（torch，官方 chatterbox）把文本转成 speech tokens 后再调用本服务
    （或直接把文本换成语义 token 序列）。板端本身不依赖 torch。
  - mel_basis_24k.npy（S3Gen mel 前端滤波矩阵）与 default_embedding.npy（内置音色）
    随包提供；voice 参数可传自定义 192 维 xvector。
  - response_format 支持 wav（默认）/ mel（JSON 调试）。纯标准库写 WAV，不依赖 ffmpeg。

用法（板端）：
  python3 openai_server.py --model models/model.axmodel --port 8000
  python3 openai_client.py --tokens "12,34,56" --out out.wav
"""

from __future__ import annotations

import argparse
import io
import json
import math
import sys
import wave
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import numpy as np

from hift_vocoder import wav_bytes

SR = 24000
N_FFT = 1920
HOP = 480
WINDOW = np.hanning(N_FFT).astype(np.float32)


def _load_pkg_assets(here: Path):
    mel_basis = np.load(here / "mel_basis_24k.npy")  # (80, 961)
    emb = np.load(here / "default_embedding.npy")  # (1, 192)
    # 最小二乘伪逆：mel 幅度 -> 线性谱幅度
    inv = np.linalg.pinv(mel_basis.astype(np.float64)).astype(np.float32)
    return mel_basis, inv, emb


def _stft(x: np.ndarray) -> np.ndarray:
    """numpy STFT（center=False，hop=480, n_fft=1920）-> (961, T) 复数谱"""
    n = (len(x) - N_FFT) // HOP + 1
    frames = np.stack([x[i * HOP:i * HOP + N_FFT] for i in range(n)])  # (T, 1920)
    return np.fft.rfft(frames * WINDOW, axis=1).T  # (961, T)


def _istft(spec: np.ndarray) -> np.ndarray:
    """numpy ISTFT（overlap-add + COLA 归一化）"""
    T = spec.shape[1]
    n = (T - 1) * HOP + N_FFT
    out = np.zeros(n, dtype=np.float64)
    wsum = np.zeros(n, dtype=np.float64)
    frames = np.fft.irfft(spec.T, n=N_FFT, axis=1)  # (T, 1920)
    for i in range(T):
        s = i * HOP
        out[s:s + N_FFT] += frames[i] * WINDOW
        wsum[s:s + N_FFT] += WINDOW ** 2
    eps = 1e-8
    return (out / np.maximum(wsum, eps)).astype(np.float32)


def griffin_lim(mel_log10: np.ndarray, mel_inv: np.ndarray, iters: int = 32) -> np.ndarray:
    """log10 mel(1,80,T) -> 24kHz 波形（纯 numpy，无 torch）"""
    M = np.power(10.0, mel_log10[0].astype(np.float64))  # (80,T) 线性 mel 幅度
    V = np.maximum(mel_inv @ M, 0.0)  # (961,T) 线性谱幅度
    phase = np.random.rand(*V.shape) * 2.0 * math.pi
    spec = V * np.exp(1j * phase)
    for _ in range(iters):
        x = _istft(spec)
        X = _stft(x)
        spec = V * np.exp(1j * np.angle(X))
    x = _istft(spec)
    peak = np.max(np.abs(x)) + 1e-8
    return (x / peak).astype(np.float32)


def mel_to_wav(mel_log10: np.ndarray, mel_inv: np.ndarray) -> bytes:
    x = griffin_lim(mel_log10, mel_inv)
    pcm = (x * 32767).clip(-32768, 32767).astype(np.int16)
    buf = io.BytesIO()
    with wave.open(buf, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(SR)
        w.writeframes(pcm.tobytes())
    return buf.getvalue()


def _load_hift_vocoder(args) -> "HiftVocoder | None":
    """指定 --f0-model/--decode-model 时启用 HiFT 神经声码器；缺省回退 Griffin-Lim。"""
    f0 = getattr(args, "f0_model", "") or ""
    dec = getattr(args, "decode_model", "") or ""
    if not (f0 and dec):
        return None
    here = Path(__file__).resolve().parent
    lw = np.load(here / "hift_linear_w.npy")
    lb = np.load(here / "hift_linear_b.npy")
    from hift_vocoder import HiftVocoder
    return HiftVocoder(f0, dec, lw, lb)


class TTSApp:
    def __init__(self, args):
        sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
        from chatterbox_s3gen_onestep_sdk.inference import ModelSession
        from chatterbox_s3gen_onestep_sdk.preprocess import preprocess

        self.session = ModelSession(args.model)
        self.clone_session = ModelSession(args.clone_model) if getattr(args, "clone_model", "") else None
        self.clone_z_ensemble = getattr(args, "clone_z_ensemble", 4)
        self.preprocess = preprocess
        here = Path(__file__).resolve().parent
        self.mel_basis, self.mel_inv, self.default_emb = _load_pkg_assets(here)
        self.hift = _load_hift_vocoder(args)
        if self.hift is not None:
            print("[openai-server] vocoder: HiFT NPU (f0+decode axmodel), 无 GL 兜底")
        else:
            print("[openai-server] vocoder: Griffin-Lim（未指定 --f0-model/--decode-model）")
        # clone 模型默认 prompt（可用 extract_voice_embedding.py 生成后放在同目录）
        self.default_prompt = None
        if self.clone_session:
            pt_path, pf_path = here / "ref_prompt_token.npy", here / "ref_prompt_feat.npy"
            if pt_path.exists() and pf_path.exists():
                self.default_prompt = (np.load(pt_path), np.load(pf_path))

    def synthesize(self, tokens: list[int], voice=None, fmt="wav") -> bytes:
        tokens = np.clip(np.asarray(tokens, dtype=np.int32).reshape(1, -1), 0, 6560)
        tlen = np.asarray([tokens.shape[1]], dtype=np.int32)
        if tokens.shape[1] < 256:
            pad = np.zeros((1, 256 - tokens.shape[1]), dtype=np.int32)
            tokens = np.concatenate([tokens, pad], axis=1)
        embedding = None if voice in (None, "default") else (voice.get("embedding") if isinstance(voice, dict) else voice)
        emb = self.default_emb if embedding is None else np.asarray(embedding, dtype=np.float32).reshape(1, -1)
        use_clone = self.clone_session and isinstance(voice, dict) and bool(voice.get("prompt_token")) and bool(voice.get("prompt_feat"))
        if use_clone:
            # 完整克隆路径（5 输入）：tokens_all = prompt_token + 生成 token（宿主侧拼接，
            # 避免图内 Concat——AX650 NPU 对 AxConcat 有 wdma bug）+ embedding + z + prompt_feat
            pt, pf = self.default_prompt if self.default_prompt is not None else (None, None)
            if isinstance(voice, dict):
                pt = np.asarray(voice.get("prompt_token"), dtype=np.int32).reshape(1, -1) if voice.get("prompt_token") else pt
                pf = np.asarray(voice.get("prompt_feat"), dtype=np.float32).reshape(1, -1, 80) if voice.get("prompt_feat") else pf
            if pt is None or pf is None:
                raise ValueError("clone 模型需要 prompt：请用 extract_voice_embedding.py 生成 ref_prompt_token/ref_prompt_feat.npy 放同目录，或请求 voice.prompt_token/prompt_feat")
            gen = np.asarray(tokens, dtype=np.int32).reshape(-1)[:99]
            if len(gen) < 99:
                gen = np.pad(gen, (0, 99 - len(gen)), mode="edge")
            pt_pad = np.zeros(157, dtype=np.int32)
            pt_pad[: min(pt.shape[1], 157)] = pt[:, :157].reshape(-1)
            tokens_all = np.concatenate([pt_pad, gen]).reshape(1, 256)
            gen_len = min(int(tlen[0]), 99)
            token_len_all = np.asarray([157 + gen_len], dtype=np.int32)
            pf_pad = np.zeros((1, 314, 80), dtype=np.float32)
            pf_pad[:, : min(pf.shape[1], 314)] = pf[:, :314, :]
            # 单步学生对 z 敏感（个别 z 会静音）：多 z 平均增强稳定性
            z_ensemble = self.clone_z_ensemble
            mels = []
            for _ in range(max(1, z_ensemble)):
                z = np.random.randn(1, 80, 512).astype(np.float32)
                feeds = [tokens_all, token_len_all, emb, z, pf_pad]
                mels.append(np.asarray(self.clone_session.run_named(feeds)[0], dtype=np.float32))
            raw = [np.mean(mels, axis=0)]
        else:
            feeds = self.preprocess(tokens, tlen, emb, None)
            raw = self.session.run_named(feeds)
        mel = np.asarray(raw[0], dtype=np.float32)[:, :, :int(tlen[0]) * 2]  # (1,80,T)
        if fmt == "mel":
            return json.dumps({"mel": mel[0].tolist()}).encode("utf-8")
        if self.hift is not None:
            return wav_bytes(self.hift.synth(mel, int(tlen[0]) * 2))
        wav = mel_to_wav(mel, self.mel_inv)
        return wav


def make_handler(app: TTSApp):
    class Handler(BaseHTTPRequestHandler):
        def log_message(self, fmt, *args):
            sys.stderr.write("[openai-server] %s\n" % (fmt % args))

        def _json(self, code, obj):
            body = json.dumps(obj, ensure_ascii=False).encode("utf-8")
            self.send_response(code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self):
            if self.path.rstrip("/") == "/v1/models":
                self._json(200, {"object": "list", "data": [{
                    "id": "chatterbox-onestep", "object": "model", "created": 0, "owned_by": "axera",
                }]})
            else:
                self._json(404, {"error": {"message": "not found", "type": "invalid_request_error", "code": None}})

        def do_POST(self):
            if self.path.rstrip("/") != "/v1/audio/speech":
                return self._json(404, {"error": {"message": "not found", "type": "invalid_request_error", "code": None}})
            try:
                req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0)) or 0))
            except Exception:
                return self._json(400, {"error": {"message": "invalid JSON body", "type": "invalid_request_error", "code": None}})
            inp = req.get("input")
            fmt = str(req.get("response_format", "wav"))
            if isinstance(inp, str):
                inp = [int(t) for t in inp.split(",") if t.strip()]
            if not isinstance(inp, list) or not inp or not all(isinstance(t, int) for t in inp):
                return self._json(400, {"error": {"message": "input must be a non-empty list of S3 token ids (or comma-separated string)", "type": "invalid_request_error", "code": None}})
            if fmt not in ("wav", "mel"):
                return self._json(400, {"error": {"message": f"response_format '{fmt}' not supported (wav|mel)", "type": "invalid_request_error", "code": None}})
            voice = req.get("voice")
            try:
                audio = app.synthesize(inp, voice=voice, fmt=fmt)
            except Exception as e:  # noqa: BLE001
                return self._json(500, {"error": {"message": f"synthesis failed: {e}", "type": "server_error", "code": None}})
            ctype = {"wav": "audio/wav", "mel": "application/json"}[fmt]
            self.send_response(200)
            self.send_header("Content-Type", ctype)
            self.send_header("Content-Length", str(len(audio)))
            self.end_headers()
            self.wfile.write(audio)

    return Handler


def main():
    p = argparse.ArgumentParser(description="OpenAI-Compatible 板端 S3Gen TTS（torch-free）")
    p.add_argument("--model", default="models/model.axmodel")
    p.add_argument("--clone-model", default="", help="完整克隆模型 model_clone.axmodel（可选）")
    p.add_argument("--clone-z-ensemble", type=int, default=4, help="克隆路径多 z 平均次数（缓解单步 z 敏感）")
    p.add_argument("--f0-model", default="models/hifift_f0.axmodel", help="HiFT f0 声码器模型（缺省 models/hifift_f0.axmodel）")
    p.add_argument("--decode-model", default="models/hifift_decode.axmodel", help="HiFT decode 声码器模型（缺省 models/hifift_decode.axmodel）")
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=8000)
    args = p.parse_args()
    app = TTSApp(args)
    server = ThreadingHTTPServer((args.host, args.port), make_handler(app))
    print(f"[openai-server] listening on http://{args.host}:{args.port} (torch-free, axengine)")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()