修改程序说明

doc.md

@@ -0,0 +1,731 @@
+# 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 输入构造方式
+- 缓存与流式输出机制