tts-server/doc.md

ONNX TTS 当前实现说明

本文档说明当前 speech_tts_onnx_opt.py 的完整实现，包括整体架构、模型导入流程、TTS 推理流程、缓存机制、接口行为以及关键函数职责。

1. 文件定位

当前 ONNX 版本服务主文件是：

• speech_tts_onnx_opt.py

它是一个基于 FastAPI 的 TTS 服务，依赖：

• onnxruntime：执行 ONNX 模型推理
• kokoro-onnx：封装 Kokoro ONNX 的 tokenizer、voice style 和模型适配
• numpy：构造推理输入、拼接输出音频
• soundfile：把 float 音频写成 WAV
• aiofiles：异步读写缓存元数据

2. 整体架构

当前服务可以分成 6 层：

1. API 层
   负责对外暴露 HTTP 接口：/tts、/generate、/clear-cache、/cache-info。

2. 请求控制层
   负责并发限制、请求打断、按客户端跟踪当前流式请求。

3. 文本切分层
   负责把整段文本拆成多个句子或片段，降低单次推理长度。

4. 模型与音色加载层
   负责加载 ONNX 模型、初始化 kokoro_onnx.Kokoro 引擎、准备 voice 文件。

5. 推理层
   负责文本转音素、音素转 token、构造 ONNX 输入、调用 session.run() 得到音频。

6. 缓存层
   包括内存缓存、磁盘缓存和同 key 合并计算，避免重复句子反复推理。

可以概括成下面这条链路：

HTTP 请求 -> 参数校验 -> 文本切分 -> 查缓存 -> ONNX 推理 -> WAV 编码 -> HTTP 返回

对于 /generate：

HTTP 请求 -> 参数校验 -> 文本切分 -> 逐句查缓存/推理 -> Base64 编码 -> NDJSON 流式返回

3. 启动和初始化

3.1 环境变量

文件开头定义了运行时配置，核心参数如下：

• TTS_ONNX_MODEL_NAME：默认模型名，默认 model_uint8.onnx
• TTS_ONNX_MODEL_DIR：模型目录，默认 /home/tts-server/onnx
• TTS_ONNX_CONFIG_PATH：配置文件路径
• TTS_ONNX_VOICES_DIR：voice .bin 文件目录
• TTS_ONNX_VOICES_V1_PATH：备用 voice 打包文件
• CACHE_DIR：磁盘缓存目录
• MAX_CONCURRENT_REQUESTS：流式接口最大并发数
• TTS_SAMPLE_RATE：采样率，默认 24000
• ORT_INTRA_OP_THREADS：ORT 单次推理内部线程数
• ORT_INTER_OP_THREADS：ORT 算子间线程数
• ORT_ENABLE_CPU_MEM_ARENA：是否开启 ORT CPU 内存池
• ORT_ENABLE_MEM_PATTERN：是否开启 ORT 内存模式

这些参数在模块导入时读取，因此进程启动前设置环境变量即可生效。

3.2 全局对象

模块初始化时创建了几个关键全局对象：

• model_lock：保护模型加载，避免并发重复初始化
• request_semaphore：限制 /generate 并发量
• current_requests：记录当前客户端请求状态，支持打断
• memory_cache：内存音频缓存
• executor：线程池，主要用于文件 I/O
• inflight_tasks：同一个缓存 key 的“进行中任务”复用表
• model_session：全局 ONNX Runtime Session
• _KOKORO_ONNX_ENGINE：全局 Kokoro ONNX 引擎

3.3 FastAPI 生命周期

应用使用了 lifespan：

• lifespan() 在服务启动时调用 load_model()
• 也就是说模型会在服务启动阶段预加载，而不是等第一个请求才懒加载

对应函数：

• lifespan(app)
• load_model(force_reload=False, name=None)

4. 模型和音色是如何导入的

这一部分是当前实现的核心。

4.1 模型文件定位

函数：

• resolve_model_path(name: str) -> str

逻辑：

1. 优先拼接 MODEL_DIR / name
2. 如果存在，直接返回
3. 如果 name 是绝对路径且文件存在，也允许直接使用
4. 否则抛出 FileNotFoundError

默认情况下，会去找：

• /home/tts-server/onnx/model_uint8.onnx

4.2 音色文件定位和打包

函数：

• resolve_voices_path() -> str

逻辑分两种：

1. 如果 VOICES_DIR 存在并且里面有多个 *.bin

   • 遍历每个 voice 文件
   • 用 np.fromfile(..., dtype=np.float32) 读取原始数据
   • 校验长度必须是 510 * 1 * 256
   • reshape 成 510 x 1 x 256
   • 最终保存成 _voices.generated.npz
   • 返回这个 .npz 路径

2. 如果 VOICES_DIR 不可用

   • 尝试回退到 VOICES_V1_PATH

这样做的目的，是把多个独立 voice .bin 文件打成一个 npz，方便 kokoro_onnx.Kokoro 统一读取。

4.3 预留的 voice 下载函数

函数：

• _download_voice_file(voice_path: Path) -> Path

它会从 Hugging Face 下载 voice 文件到本地，但当前主调用链并没有自动调用这个函数。
也就是说，当前代码里下载能力存在，但主流程实际上主要依赖本地已有 voice 文件。

4.4 ONNX Runtime Session 的创建

函数：

• load_onnx_session(name: str)

逻辑：

1. 导入 onnxruntime as ort
2. 调 resolve_model_path(name) 拿到模型路径
3. 创建 ort.SessionOptions()
4. 设置：
   • intra_op_num_threads
   • inter_op_num_threads
   • enable_cpu_mem_arena
   • enable_mem_pattern
5. 创建：
   • ort.InferenceSession(model_path, sess_options=sess_options, providers=["CPUExecutionProvider"])

这里已经明确指定只使用 CPU provider。

4.5 Kokoro 引擎的创建

函数：

• load_kokoro_engine(name: str)

逻辑：

1. 导入 from kokoro_onnx import Kokoro
2. 解析模型路径
3. 调用：

Kokoro(
    model_path=model_path,
    voices_path=resolve_voices_path(),
    vocab_config=CONFIG_PATH if Path(CONFIG_PATH).exists() else None,
)

这个对象负责：

• tokenizer
• phonemize
• voice style 获取
• 音素批切分

4.6 全局模型加载入口

函数：

• load_model(force_reload=False, name=None)

逻辑：

1. 用 model_lock 加锁
2. 判断是否需要重载：
   • force_reload=True
   • model_session is None
   • 传入的 target != model_name
3. 如果需要重载：
   • model_session = load_onnx_session(target)
   • _KOKORO_ONNX_ENGINE = load_kokoro_engine(target)
   • model_name = target

这意味着当前实现会同时维护两套模型相关对象：

• 原生 onnxruntime.InferenceSession
• kokoro_onnx.Kokoro

原因是：

• session 真正执行推理
• engine 负责 tokenizer、voice style、文本到模型输入的辅助逻辑

4.7 获取当前引擎

函数：

• get_kokoro_engine(name=None)

如果引擎为空，或者请求的模型名和当前不一致，就触发 load_model()。

5. TTS 是如何实现的

TTS 核心在 synthesize_audio()。

函数：

• synthesize_audio(text, voice, speed, model_name=None) -> np.ndarray

它的完整处理流程如下。

5.1 文本合法性检查

先判断：

• text.strip() 是否为空

为空直接报 HTTPException(400)。

5.2 获取模型对象

函数内部先拿两个对象：

• engine = get_kokoro_engine(name=model_name)
• session = load_model(name=model_name)

其中：

• engine 用于文本前处理
• session 用于真正推理

5.3 校验 voice

逻辑：

if voice not in set(engine.get_voices()):
    raise HTTPException(...)

也就是说 voice 必须存在于当前 Kokoro 引擎已加载的 voice 列表中。

5.4 文本转音素

逻辑：

phonemes = engine.tokenizer.phonemize(text, "en-us")

这里把输入文本转成音素串。当前实现写死了 "en-us"，因此这条 ONNX 逻辑本质上按英文音素流程在走。

5.5 按模型可接受长度切分音素

逻辑：

batched_phonemes = engine._split_phonemes(phonemes)

这是对长文本的第二次切分。
前面的 iter_text_parts() 是句子级切分，这里是音素级切分，目的是避免单次 token 太长。

5.6 取 voice style

逻辑：

voice_style = engine.get_voice_style(voice)

voice style 是当前 voice 对应的风格张量集合，后续会根据 token 长度取其中一个切片。

5.7 逐批构造 ONNX 输入

对每一个 phoneme_batch：

1. 音素转 token：

tokens = np.array(engine.tokenizer.tokenize(phoneme_batch), dtype=np.int64)

1. 跳过空 token

2. 构造 feeds：

feeds = {
    "input_ids": np.asarray([[0, *tokens.tolist(), 0]], dtype=np.int64),
    "style": np.asarray(voice_style[len(tokens)], dtype=np.float32),
    "speed": np.asarray([speed], dtype=np.float32),
}

三个输入含义如下：

• input_ids 模型文本输入 token，前后补 0

• style 根据 token 长度索引 voice style，说明 voice style 不是固定一份，而是按长度取对应条目

• speed 语速控制参数

5.8 执行 ONNX 推理

逻辑：

outputs = session.run(None, feeds)

这里就是实际调用 onnxruntime 执行模型。

返回结果后：

audio_segments.append(to_mono_numpy(outputs[0]))

第一路输出被当成音频，随后通过 to_mono_numpy() 归一成一维 float32 单声道数组。

5.9 拼接所有音频片段

如果有多个 batch，就做：

audio = np.concatenate(audio_segments, axis=0)

最终返回一个完整的一维 numpy.ndarray 音频数组。

5.10 输出校验

推理完成后还会检查：

• 是否有输出
• 输出长度是否为 0
• 是否包含 NaN / Inf

不合法就抛 500 错误。

6. 文本切分机制

当前实现有两级切分。

6.1 句子级切分

函数：

• split_sentences(text)
• iter_text_parts(text, split_pattern)

逻辑：

1. iter_text_parts() 先按调用方传入的 split_pattern 切成 block
2. 再对每个 block 调 split_sentences()
3. split_sentences() 再按标点和换行拆句
4. 如果句子太短，例如长度小于 3，会尽量并回上一句

作用：

• 降低单次推理长度
• 方便缓存
• 方便流式逐句返回

6.2 音素级切分

函数：

• engine._split_phonemes(phonemes)

这是模型输入长度控制，属于更底层的切分。

7. 音频编码和格式转换

7.1 单声道归一

函数：

• to_mono_numpy(audio)

作用：

• 把各种形状的音频输出转成一维 float32
• 如果是二维多声道，按规则压成单声道

7.2 编码成 WAV 字节

函数：

• encode_wav_bytes(audio, sr) -> bytes

作用：

• 把 numpy 音频写入 BytesIO
• 格式固定为 WAV PCM_16

7.3 直接返回整段 WAV

函数：

• synthesize_wav_bytes(text, voice, speed, split_pattern, model_name=None) -> io.BytesIO

逻辑：

1. 按句切分文本
2. 对每一段：
   • 先查内存缓存
   • 有缓存则直接读取 WAV 并解码成 float32
   • 没缓存则调用 synthesize_audio()
3. 拼接所有段
4. 最终写成一个完整 WAV 返回

注意：

• /tts 这条路径只直接使用内存缓存
• 它不会像 /generate 那样主动走“内存缓存 + 磁盘缓存 + inflight 合并”完整链路

8. 缓存机制

当前实现有 3 层缓存/去重能力。

8.1 内存缓存

类：

• MemoryAudioCache

这是一个带 TTL 和容量控制的内存缓存。

特点：

• 基于 OrderedDict 维护近似 LRU
• 支持过期清理
• 限制最大条目数
• 限制总字节数

缓存值结构大致为：

{
    "sentence": "...",
    "sample_rate": 24000,
    "audio_bytes": b"..."
}

关键方法：

• get(key)
• set(key, value)
• clear()
• info()

8.2 磁盘缓存

函数：

• sentence_cache_path(key)
• meta_cache_path(key)
• save_sentence_to_disk(...)
• load_sentence_from_disk(key)
• clean_disk_cache()

存储方式：

• 音频保存为 CACHE_DIR/<key>.wav
• 元数据保存为 CACHE_DIR/<key>.json

元数据里保存：

• sentence
• sample_rate

清理策略：

• 按 .wav 文件修改时间排序
• 超过 DISK_CACHE_SIZE 后删除最旧文件
• 同时删除对应 .json

8.3 进行中任务合并

函数：

• get_or_create_sentence_cache_item(sentence, voice, speed, model)

这里用 inflight_tasks 做了一个很实用的优化：

• 如果两个请求同时要同一句文本、同一个 voice、同一个 speed、同一个 model
• 第一个请求会创建一个异步任务 _compute_sentence_item(...)
• 第二个请求不会重复推理，而是直接 await 同一个 future

这样能避免热点句子被并发重复算多次。

8.4 缓存 key

函数：

• sentence_cache_key(sentence, voice, speed, model)

生成方式：

raw = f"{sentence}|{voice}|{speed:.4f}|{model}"
md5(raw.encode("utf-8"))

说明：

• 句子内容、音色、语速、模型名任一变化，都会生成新 key

9. HTTP 接口说明

9.1 /tts POST

函数：

• tts_post(req: TTSRequest)

输入模型：

• text
• voice
• speed
• split_pattern
• model_name

行为：

1. 调 synthesize_wav_bytes(...)
2. 返回 audio/wav

适合：

• 直接拿完整 WAV 文件

9.2 /tts GET

函数：

• tts_get(...)

行为和 POST 基本一致，只是参数来自 query string。

9.3 /generate POST

函数：

• generate_audio_stream(data: Dict = Body(...))

这是更复杂的一条链路。

逻辑：

1. 进入 request_semaphore，限制整体并发
2. 读取：
   • text
   • voice
   • speed
   • model_name
   • split_pattern
   • client_id
3. 按句拆分文本
4. 如果同一个 client_id 已经有活跃请求：
   • 把旧请求标记为 interrupt=True
   • 稍等 0.05s
5. 新建当前请求状态
6. 定义异步生成器 stream()
7. 对每一段文本：
   • 检查是否被打断
   • 调 get_or_create_sentence_cache_item(...)
   • 转成 Base64
   • 以一行 JSON 输出
8. 返回 application/x-ndjson

每一行数据结构大致是：

{
  "index": 0,
  "sentence": "...",
  "sample_rate": 24000,
  "audio": "base64..."
}

适合：

• 前端逐句播放
• 长文本流式处理
• 边生成边下发

9.4 /clear-cache

函数：

• clear_cache()

行为：

• 清空内存缓存
• 删除 CACHE_DIR 目录下所有缓存文件

9.5 /cache-info

函数：

• get_cache_info()

返回：

• 内存缓存条数
• 内存缓存总字节数
• 磁盘缓存文件数

10. 关键函数关系图

10.1 启动链路

FastAPI lifespan -> load_model() -> load_onnx_session() + load_kokoro_engine()

10.2 /tts 链路

tts_post/tts_get -> synthesize_wav_bytes() -> iter_text_parts() -> memory_cache.get() or synthesize_audio() -> session.run() -> sf.write() -> StreamingResponse

10.3 /generate 链路

generate_audio_stream() -> iter_text_parts() -> get_or_create_sentence_cache_item() -> memory_cache/disk_cache/inflight -> _compute_sentence_item() -> synthesize_audio() -> encode_wav_bytes() -> NDJSON stream

11. 当前实现的设计特点

11.1 优点

• 模型启动时预加载，避免首请求冷启动
• ONNX 推理和 tokenizer/voice style 职责分离
• 有内存缓存、磁盘缓存、并发去重
• 支持整段返回和逐句流式返回两种模式
• 文本切分策略比较务实，适合长文本

11.2 当前代码中值得注意的点

1. executor = ThreadPoolExecutor(max_workers=8) 主要用于文件 I/O
   但真正推理主要是通过 asyncio.to_thread() 和同步函数组合完成。

2. /tts 路径没有走完整的磁盘缓存和 inflight 合并逻辑
   它更偏向直接同步生成完整 WAV。

3. phonemize(text, "en-us") 是写死的
   这说明当前 ONNX TTS 主流程是按英文音素化逻辑设计的。

4. _download_voice_file() 当前没有接入主路径
   voice 自动下载并不是当前服务的主行为。

5. 同时保留 model_session 和 _KOKORO_ONNX_ENGINE 是必要设计
   因为 kokoro_onnx 并没有完全替代手动 session.run() 这条调用链。

12. 维护时最常关注的函数

如果你后续要改这个服务，最值得优先看的是下面这些函数：

• load_model()
  模型重载入口。

• load_onnx_session()
  ORT 线程、provider、session 配置都在这里。

• load_kokoro_engine()
  Kokoro ONNX 引擎初始化在这里。

• resolve_voices_path()
  voice 文件处理逻辑在这里。

• synthesize_audio()
  文本到音频数组的核心推理函数。

• synthesize_wav_bytes()
/tts 直接返回整段音频的关键函数。

• get_or_create_sentence_cache_item()
/generate 的缓存和并发合并核心。

• generate_audio_stream()
  流式接口主控制逻辑。

13. 一句话总结

当前这套 ONNX TTS 的实现，本质上是：

• 用 kokoro_onnx 负责文本前处理、voice style 和 tokenizer
• 用 onnxruntime.InferenceSession.run() 真正执行 ONNX 推理
• 用 FastAPI 暴露同步整段返回和异步逐句流式返回两套接口
• 用内存缓存、磁盘缓存和 inflight 去重减少重复计算

如果后续你要重构成 Go，最需要完整复刻的不是 HTTP 层，而是下面这几块：

• 文本切分策略
• voice 文件加载和 style 索引逻辑
• phonemize/tokenize 流程
• ONNX 输入构造方式
• 缓存与流式输出机制