Kokoro-82M on Ascend NPU

1. 简介

本文档记录 Kokoro-82M（StyleTTS2 架构，82M 参数 TTS 模型）在华为昇腾 NPU（Ascend 910 系列）上的适配与验证结果。

Kokoro-82M 是一个轻量级的文本转语音（TTS）模型，采用 StyleTTS2 + ISTFTNet 架构，原生于 PyTorch 实现。本适配工作基于官方 kokoro pip 包源码，针对昇腾 NPU 做了以下关键修改：

• 设备检测：扩展 KPipeline 的设备选择逻辑，支持自动识别 torch.npu
• STFT 兼容：TorchSTFT 在 NPU 上 fallback 到 CustomSTFT（基于 Conv1d 的 STFT/iSTFT 实现），规避 aclnnAbs / aclnnUnfoldGrad 对复数张量的限制
• 张量创建：KModel.forward 中使用 torch.tensor(..., device=device, dtype=torch.long) 替代 torch.LongTensor，确保输入张量直接创建在 NPU 上
• 本地语音缓存：KPipeline.load_single_voice 优先加载本地 voices/ 目录，避免 HuggingFace Hub 不可达时的超时

相关获取地址：

• 权重下载地址（ModelScope）：https://modelscope.cn/models/hexgrad/Kokoro-82M
• 权重下载地址（HuggingFace）：https://huggingface.co/hexgrad/Kokoro-82M
• 原始代码仓库：https://github.com/hexgrad/kokoro

2. 验证环境

• NPU：1 逻辑卡（Ascend910B / Atlas 800 A2）
• 模型路径：/tmp/kokoro-82m-bf16/weights
• 测试架构：aarch64

3. 快速开始

3.1 安装依赖

pip install kokoro soundfile misaki
pip install torch-npu --extra-index-url https://pypi.tuna.tsinghua.edu.cn/simple

注：misaki 为 G2P（Grapheme-to-Phoneme）依赖，英文需要 pip install misaki[en]，中文需要 pip install misaki[zh]。

3.2 下载权重

推荐从 ModelScope 下载（国内网络更稳定）：

modelscope download --model hexgrad/Kokoro-82M --local_dir ./Kokoro-82M

或从 HuggingFace 镜像下载：

export HF_ENDPOINT=https://hf-mirror.com
huggingface-cli download hexgrad/Kokoro-82M --local-dir ./Kokoro-82M

3.3 NPU 推理示例

from kokoro_npu import KPipeline, KModel
import soundfile as sf

# 加载模型到 NPU
pipeline = KPipeline(lang_code='a', model=False)
model = KModel(
    config='./Kokoro-82M/config.json',
    model='./Kokoro-82M/kokoro-v1_0.pth'
).to('npu').eval()

# 生成音频
for result in pipeline(
    'Hello, this is Kokoro running on Ascend NPU.',
    voice='af_heart',
    model=model
):
    audio = result.audio
    sf.write('output.wav', audio.numpy(), 24000)

NPU 建议：在 NPU 上推理时建议显式设置 disable_complex=True，使 STFT/iSTFT 统一走 CustomSTFT 后端，确保与 CPU 对齐的精度表现：

model = KModel(config='...', model='...', disable_complex=True).to('npu').eval()

3.4 运行验证脚本

# 精度验证（CPU vs NPU）
python accuracy_check.py

# 性能基准测试
python perf_benchmark.py

# 单条推理 + 基准
python inference_npu.py \
  --text "Hello world." \
  --voice af_heart \
  --device npu \
  --weights ./Kokoro-82M/kokoro-v1_0.pth \
  --config ./Kokoro-82M/config.json

4. 精度评测

使用对齐后的 CustomSTFT 后端，对比 CPU 与 NPU 输出的一致性。

结论：NPU 输出与 CPU 参考输出高度一致，Cosine Similarity > 0.995，SNR > 20 dB，精度验证通过。

说明：若使用默认的 TorchSTFT（CPU torch.istft）作为参考基准，与 NPU CustomSTFT 的 SNR 约为 5.2 dB，该差异主要来源于 CustomSTFT 卷积近似带来的固有误差，而非 NPU 算子精度问题。实际人耳听觉上难以区分两者差异。

5. 性能参考

测试条件：NPU 单卡，预热 5 轮，正式测试 20 轮，取均值。

汇总：

RTF（Real-Time Factor）< 1.0 表示推理速度快于音频播放速度，可支持实时 TTS 场景。

与 CPU（ARM aarch64）对比：

6. 适配变更摘要

以下文件为相对于上游 kokoro 包的核心改动：

完整变更可通过 git diff 或本仓库源码查看。

7. 注意事项

1. 复数张量限制：当前 CANN 版本下 torch.abs / torch.angle 不支持 torch.complex64 输入，torch.istft 在 NPU 上触发 aclnnUnfoldGrad 内部错误。已在 TorchSTFT 中通过 CustomSTFT fallback 解决。

2. bf16 支持：Kokoro-82M 原生为 fp32 权重。若需 bf16 推理，可在加载后手动 model.bfloat16().to('npu')，但当前未做专门验证。

3. 语音文件缓存：首次运行需下载 voices/*.pt，建议预先通过 modelscope download 或 huggingface-cli 下载完整仓库到本地，避免运行时网络超时。

4. 长文本分块：KPipeline 对英文自动按 510 音素分块，对非英语按 400 字符分块。长文本会生成多段音频，需在外层拼接。

5. ESpeak 依赖：非英语语音（如法语、西班牙语）依赖 espeak-ng，需提前安装：apt-get install espeak-ng。