# ONNX TTS 当前实现说明

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

## 1. 文件定位

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

- [speech_tts_onnx_opt.py](/home/tts-server/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. 调用：

```python
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

逻辑：

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

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

### 5.4 文本转音素

逻辑：

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

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

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

逻辑：

```python
batched_phonemes = engine._split_phonemes(phonemes)
```

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

### 5.6 取 voice style

逻辑：

```python
voice_style = engine.get_voice_style(voice)
```

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

### 5.7 逐批构造 ONNX 输入

对每一个 `phoneme_batch`：

1. 音素转 token：

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

2. 跳过空 token

3. 构造 `feeds`：

```python
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 推理

逻辑：

```python
outputs = session.run(None, feeds)
```

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

返回结果后：

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

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

### 5.9 拼接所有音频片段

如果有多个 batch，就做：

```python
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
- 支持过期清理
- 限制最大条目数
- 限制总字节数

缓存值结构大致为：

```python
{
    "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)`

生成方式：

```python
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`

每一行数据结构大致是：

```json
{
  "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 输入构造方式
- 缓存与流式输出机制