Kokoro-82M-bf16 on Ascend NPU

1. 简介

本文档记录 Kokoro-82M-bf16 文本转语音（TTS）模型在华为昇腾 NPU 上的适配与验证结果。

Kokoro-82M 是基于 StyleTTS2 架构的轻量级 TTS 模型（82M 参数），支持多说话人、多语言语音合成。模型使用 PlBERT 文本编码器 + 韵律预测器 + iSTFTNet 解码器的流水线，输出 24kHz 音频。

适配工作主要解决以下问题：

• STFT 算子兼容性：NPU 不支持 torch.stft 返回复数张量，改用基于 Conv1d 的 CustomSTFT 实现
• 权重格式转换：支持 safetensors 格式权重的直接加载
• 设备迁移：完整的 CPU -> NPU 推理链路适配

相关地址：

• 模型地址（AtomGit）：https://gitcode.com/hf_mirrors/mlx-community/Kokoro-82M-bf16
• 原始模型（HuggingFace）：https://huggingface.co/hexgrad/Kokoro-82M
• 源码仓库：https://github.com/hexgrad/kokoro

2. 验证环境

• NPU：1 逻辑卡（Atlas 800 A2/A3）
• 模型路径：./kokoro-v1_0.safetensors
• 模型大小：327 MB（bf16 精度）

3. 环境准备

3.1 安装依赖

pip install torch torch-npu transformers safetensors soundfile numpy

3.2 下载模型

# 方式一：使用 huggingface-cli（推荐）
HF_ENDPOINT=https://hf-mirror.com huggingface-cli download mlx-community/Kokoro-82M-bf16 --local-dir ./Kokoro-82M-bf16

# 方式二：使用 ModelScope
modelscope download --model mlx-community/Kokoro-82M-bf16 --local_dir ./Kokoro-82M-bf16

3.3 NPU 环境变量

export ASCEND_RT_VISIBLE_DEVICES=0
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True

4. 推理脚本

4.1 基础推理

python inference.py \
  --text "Hello world, this is a test of the Kokoro text to speech system." \
  --voice af_heart \
  --output output.wav \
  --device npu

4.2 参数说明

4.3 可用声音

模型支持 54 种说话人声音（voices/ 目录），命名规则：

• af_*：女性声音（American Female）
• am_*：男性声音（American Male）
• bf_*：女性声音（British Female）
• bm_*：男性声音（British Male）

常用声音：af_heart, af_bella, af_nicole, am_adam, bf_emma

5. 性能参考

测试条件：单卡 NPU，重复 10 次取平均值。

关键指标：

• 平均 RTF：0.03 ~ 0.07（14~33 倍实时速度）
• 首包延迟：~0.11s（短文本）
• 吞吐量：105 ~ 496 phonemes/s

CPU vs NPU 对比

6. 精度验证

使用 5 条英文文本进行推理验证，所有文本均成功生成音频。

所有推理结果音频已保存至 eval_output/ 目录。

7. 适配要点

7.1 STFT 算子适配

NPU 不支持 torch.stft 返回复数张量（DT_COMPLEX64），报错：

RuntimeError: abs: NPU function error: call aclnnAbs failed
Tensor out expected dtype is DT_FLOAT but found DT_COMPLEX64

解决方案：在 Generator 中自动检测 NPU 设备，使用基于 Conv1d 的 CustomSTFT 替代 TorchSTFT。CustomSTFT 通过预计算 DFT 矩阵并使用 Conv1d/ConvTranspose1d 实现 STFT/iSTFT，完全避免复数运算。

# kokoro_istftnet.py 中的关键修改
use_custom = disable_complex or (hasattr(torch, 'npu') and torch.npu.is_available())
self.stft = CustomSTFT(...) if use_custom else TorchSTFT(...)

7.2 权重加载

原始代码使用 torch.load() 加载 .pth 权重，而本适配版本使用 .safetensors 格式。修改模型加载逻辑以自动检测文件格式：

if model.endswith('.safetensors'):
    from safetensors import safe_open
    f = safe_open(model, framework='pt')
    # 按前缀分组为子模块 state_dict
    grouped = defaultdict(dict)
    for k in f.keys():
        prefix = k.split('.')[0]
        grouped[prefix][k[len(prefix)+1:]] = f.get_tensor(k)
    state_dict = dict(grouped)
else:
    state_dict = torch.load(model, map_location='cpu', weights_only=True)

7.3 设备迁移

模型整体迁移至 NPU：

model = model.npu().eval()

推理时需确保所有张量在同一设备：

ref_s = pack[phoneme_len - 1].npu()
input_ids = torch.LongTensor([[0, *input_ids, 0]]).npu()

8. 文件说明

9. 注意事项

1. G2P 依赖：完整 G2P 需要安装 misaki[en]（pip install misaki[en]）。当前脚本内置了简单的字母到音素映射作为 fallback，安装 misaki 后可获得更好的音素转换质量。

2. 长文本处理：模型最大支持 510 个音素的输入序列。对于长文本，建议按句子分段合成后拼接。

3. 声音文件格式：模型同时提供 .pt 和 .safetensors 两种格式的声音文件，两者等价。

4. 首次推理延迟：首次推理包含 NPU 图编译开销，后续推理会显著加速。建议使用 --warmup 1 进行预热。

Ascend NPU 精度评测

NPU vs CPU 精度对比（CPU 为基线，NPU 为验证目标）：

精度评测源代码和日志详见 eval/ 目录。