Kokoro-82M 昇腾 NPU 部署实战课程

┌──────────────────────────────────────────────────────┐
│              Uvicorn API 服务层                       │
│  ┌──────────────────────────────────────────────┐    │
│  │  FastAPI (api/main.py)                       │    │
│  │  ├─ /health          → 健康检查              │    │
│  │  └─ /v1/audio/speech → TTS 推理接口          │    │
│  └──────────────────┬───────────────────────────┘    │
│                     │                                 │
│  ┌──────────────────▼───────────────────────────┐    │
│  │  KokoroModelLoader (api/model_loader.py)     │    │
│  │  ├─ _patch_hf_download() → 离线/镜像加载     │    │
│  │  ├─ KPipeline(lang_code='z') → 中文TTS管道   │    │
│  │  └─ voice_exists() → 语音包验证              │    │
│  └──────────────────┬───────────────────────────┘    │
│                     │                                 │
├─────────────────────┼────────────────────────────────┤
│       CANN 8.5 异构计算栈          │                  │
│  ┌──────────────┐  ┌─▼──────────────┐                │
│  │  GE 图编译   │  │ CUBE+Vector    │                │
│  │  (算子融合)  │  │  混合执行      │                │
│  └──────────────┘  └────────────────┘                │
├──────────────────────────────────────────────────────┤
│              昇腾 NPU 硬件层                          │
│  ┌──────────┐ ┌──────────┐ ┌──────────────┐          │
│  │ Scalar   │ │ Vector   │ │ CUBE 矩阵乘  │          │
│  │ 控制流   │ │ 上采样   │ │ Transformer  │          │
│  └──────────┘ │ Conv1D   │ │ Encoder      │          │
│               └──────────┘ └──────────────┘          │
└──────────────────────────────────────────────────────┘

# [系统依赖] 安装 espeak-ng 和其他系统级工具
!sudo apt-get update
!sudo apt-get install -y espeak-ng git curl wget libsndfile1 build-essential

# [目录准备] 创建模型存储和工作目录
!mkdir -p /opt/atomgit/models
!mkdir -p /opt/atomgit/models/Kokoro-82M/work/outputs
!mkdir -p /opt/atomgit/models/Kokoro-82M/work/temp
!mkdir -p /opt/atomgit/models/Kokoro-82M/work/logs

# [代码获取] 从 AtomGit 克隆 Kokoro-82M 服务代码
!mkdir -p /opt/atomgit
!git clone https://atomgit.com/atomgit-ascend/Kokoro-82M.git /opt/atomgit/Kokoro-82M

# [依赖安装] 安装 atomgit_hub 工具
!pip install -U atomgit_hub -i https://mirrors.huaweicloud.com/repository/pypi/simple

# [模型下载] 从 AtomGit 平台下载 Kokoro-82M 模型
from atomgit_hub import snapshot_download

# 下载模型到本地
model_repo = "hf_mirrors/AI-ModelScope/Kokoro-82M"
local_dir = "/opt/atomgit/models/Kokoro-82M"

snapshot_download(model_repo, local_dir=local_dir)
print(f"模型已下载到: {local_dir}")

# [依赖安装] 安装项目依赖和中文 TTS 扩展包
!pip install -r /opt/atomgit/Kokoro-82M/requirements.txt -i https://mirrors.huaweicloud.com/repository/pypi/simple
!pip install ordered_set pypinyin cn2an jieba -i https://mirrors.huaweicloud.com/repository/pypi/simple

# [昇腾环境] 加载 CANN 工具链环境变量
!source /usr/local/Ascend/ascend-toolkit/set_env.sh

# [设备配置] 设置 NPU 可见设备，禁用 CUDA 回退
!export ASCEND_RT_VISIBLE_DEVICES=0
!export CUDA_VISIBLE_DEVICES=

# [核心适配] model_loader.py 完整替换内容
# 写入路径：/opt/atomgit/Kokoro-82M/api/model_loader.py

"""
Kokoro-82M 模型加载器
使用 kokoro 库进行 TTS 推理
"""
import os
import logging
import yaml
from typing import Optional, Dict, List

from huggingface_hub import hf_hub_download
from kokoro import KPipeline
import kokoro.model as km
import kokoro.pipeline as kp

logger = logging.getLogger(__name__)

# 仓库配置
YOUR_REPO_ID = "hf_mirrors-hexgrad/Kokoro-82M"
ORIGINAL_REPO_ID = "hexgrad/Kokoro-82M"
MODEL_FILENAME = "kokoro-v1_0.pth"

# 默认语音目录路径
DEFAULT_VOICES_DIR = "/root/.cache/huggingface/hub/models--hf_mirrors-hexgrad--Kokoro-82M/snapshots/f3ff3571791e39611d31c381e3a41a3af07b4987/voices"

# 离线 / atomgit 本机目录：含 config.json、kokoro-v1_0.pth、voices/ 等
_ENV_LOCAL_MODEL = "KOKORO_LOCAL_MODEL"


def _kokoro_repo_ids():
    return frozenset({ORIGINAL_REPO_ID, YOUR_REPO_ID})


def _local_model_root() -> str:
    return os.environ.get(_ENV_LOCAL_MODEL, "").strip()


def _resolve_under_local_root(local_root: str, filename: str) -> Optional[str]:
    """若 local_root/filename 为真实文件则返回绝对路径，否则 None（防目录穿越）。"""
    if not local_root or not filename:
        return None
    base = os.path.realpath(local_root)
    if not os.path.isdir(base):
        return None
    joined = os.path.realpath(os.path.join(base, filename))
    if not joined.startswith(base + os.sep) and joined != base:
        logger.warning("拒绝非安全路径: root=%s filename=%s", local_root, filename)
        return None
    if os.path.isfile(joined):
        return joined
    return None


def _patch_hf_download():
    """
    Patch 下载函数：
    1) 若设置 KOKORO_LOCAL_MODEL 且文件存在，直接返回本地路径（支持 HF_HUB_OFFLINE）
    2) 否则将 hexgrad/Kokoro-82M 重定向到 hf_mirrors-hexgrad/Kokoro-82M
    """
    km.KModel.MODEL_NAMES[ORIGINAL_REPO_ID] = MODEL_FILENAME
    kokoro_ids = _kokoro_repo_ids()
    upstream = hf_hub_download

    def patched_download(repo_id, filename, *args, **kwargs):
        local_root = _local_model_root()
        if local_root and repo_id in kokoro_ids:
            local_path = _resolve_under_local_root(local_root, filename)
            if local_path is not None:
                logger.info("✓ %s 本地命中: %s", _ENV_LOCAL_MODEL, local_path)
                return local_path
            if os.environ.get("HF_HUB_OFFLINE", "").lower() in ("1", "true", "yes"):
                raise FileNotFoundError(
                    "HF_HUB_OFFLINE 已开启且未在 {} 找到文件 {!r}（repo_id={!r}）。"
                    "请确认 atomgit 已完整下载到该目录，或关闭 HF_HUB_OFFLINE。".format(
                        os.path.realpath(local_root), filename, repo_id
                    )
                )

        hub_repo_id = YOUR_REPO_ID if repo_id == ORIGINAL_REPO_ID else repo_id
        return upstream(repo_id=hub_repo_id, filename=filename, *args, **kwargs)

    # kokoro.model 和 kokoro.pipeline 各自引用了 hf_hub_download，需分别 patch
    km.hf_hub_download = patched_download
    if hasattr(kp, "hf_hub_download"):
        kp.hf_hub_download = patched_download
    else:
        logger.warning("kokoro.pipeline 无 hf_hub_download，语音加载可能仍走未 patch 的路径；请检查 kokoro 版本")
    logger.info(
        "✓ 下载已 patch（kokoro.model + kokoro.pipeline）: %s -> %s；%s=%s",
        ORIGINAL_REPO_ID, YOUR_REPO_ID,
        _ENV_LOCAL_MODEL, _local_model_root() or "(未设置)",
    )


class KokoroModelLoader:
    """Kokoro-82M 模型加载器"""

    def __init__(self, config_path: Optional[str] = None):
        if config_path is None:
            config_path = os.path.join(os.path.dirname(__file__), "../config/config.yaml")

        with open(config_path, 'r', encoding='utf-8') as f:
            self.config = yaml.safe_load(f)

        model_config = self.config.get('model', {})
        self.model_id = YOUR_REPO_ID
        self.original_repo_id = ORIGINAL_REPO_ID
        self.voices_dir = os.getenv('VOICES_DIR', model_config.get('voices_dir', DEFAULT_VOICES_DIR))
        self.lang_code = os.getenv('LANG_CODE', model_config.get('lang_code', 'z'))
        self.default_voice = os.getenv('DEFAULT_VOICE', model_config.get('default_voice', 'af_heart'))
        self.sample_rate = int(os.getenv('SAMPLE_RATE', model_config.get('sample_rate', 24000)))
        self.device = model_config.get('device', None)
        self.device_id = model_config.get('device_id', 0)

        self.pipeline = None
        self.is_loaded = False
        self._available_voices = None

    def load_model(self):
        """加载 Kokoro-82M 模型"""
        try:
            logger.info(f"开始加载 Kokoro-82M 模型: {self.model_id}")
            logger.info(f"语音目录: {self.voices_dir}")
            logger.info(f"语言代码: {self.lang_code}")
            logger.info(f"默认语音: {self.default_voice}")

            # 应用下载 patch
            _patch_hf_download()

            # 验证语音目录
            if not os.path.exists(self.voices_dir):
                logger.warning(f"语音目录不存在: {self.voices_dir}，将在首次下载后创建")
            else:
                self._scan_available_voices()

            # 验证默认语音
            if not self.voice_exists(self.default_voice):
                available = self.get_available_voices()
                raise ValueError(f"默认语音 '{self.default_voice}' 不存在于 {self.voices_dir}，可用语音: {available}")

            # 尝试导入 NPU 支持
            try:
                import torch
                import torch_npu
                logger.info("✓ torch_npu 已加载")
                npu_device_id = self.device_id
                torch.npu.set_device(npu_device_id)
                logger.info(f"✓ NPU设备设置为: {npu_device_id}")
            except ImportError:
                logger.info("torch_npu 未安装，将使用 CPU")

            # 创建 Kokoro Pipeline
            logger.info("正在初始化 KPipeline...")
            self.pipeline = KPipeline(lang_code=self.lang_code, repo_id=self.original_repo_id)
            logger.info(f"✓ Kokoro Pipeline 已初始化")

            # 重新扫描语音
            self._scan_available_voices()

            # 打印配置
            logger.info("=" * 60)
            logger.info("Kokoro TTS 配置:")
            logger.info(f"  模型ID: {self.model_id}")
            logger.info(f"  语音目录: {self.voices_dir}")
            logger.info(f"  语言代码: {self.lang_code}")
            logger.info(f"  默认语音: {self.default_voice}")
            logger.info(f"  采样率: {self.sample_rate}")
            logger.info(f"  可用语音数: {len(self._available_voices) if self._available_voices else 0}")
            logger.info("=" * 60)

            self.is_loaded = True
            logger.info("Kokoro-82M 模型加载完成")
            return True

        except Exception as e:
            logger.error(f"模型加载失败: {str(e)}", exc_info=True)
            raise

    def _scan_available_voices(self):
        """扫描 voices 目录下可用的语音文件"""
        self._available_voices = []
        if os.path.exists(self.voices_dir):
            for filename in os.listdir(self.voices_dir):
                if filename.endswith('.pt'):
                    voice_name = filename[:-3]
                    self._available_voices.append(voice_name)
            self._available_voices.sort()
            logger.info(f"扫描到 {len(self._available_voices)} 个语音文件: {self._available_voices}")
        else:
            logger.warning(f"语音目录不存在: {self.voices_dir}")

    def voice_exists(self, voice_name: str) -> bool:
        """检查语音文件是否存在"""
        voice_path = os.path.join(self.voices_dir, f"{voice_name}.pt")
        return os.path.exists(voice_path)

    def validate_voice(self, voice_name: str) -> str:
        """验证语音是否可用，不可用则抛出异常"""
        if voice_name is None:
            return self.default_voice
        if not self.voice_exists(voice_name):
            available = self.get_available_voices()
            raise ValueError(f"语音 '{voice_name}' 不存在，可用语音: {available}")
        return voice_name

    def get_voice_path(self, voice_name: str) -> str:
        """获取语音文件完整路径"""
        return os.path.join(self.voices_dir, f"{voice_name}.pt")

    def get_pipeline(self):
        return self.pipeline

    def get_config(self) -> Dict:
        return self.config

    def get_default_voice(self) -> str:
        return self.default_voice

    def get_sample_rate(self) -> int:
        return self.sample_rate

    def get_voices_dir(self) -> str:
        return self.voices_dir

    def get_available_voices(self) -> List[str]:
        """获取可用的语音列表"""
        if self._available_voices is None:
            self._scan_available_voices()
        return self._available_voices.copy() if self._available_voices else []

    def get_inference_params(self) -> Dict:
        """获取推理参数"""
        return {
            'lang_code': self.lang_code,
            'default_voice': self.default_voice,
            'sample_rate': self.sample_rate,
            'speed': getattr(self, "default_speed", 1.0),
            'device': self.device,
            'device_id': self.device_id,
            'voices_dir': self.voices_dir
        }

#!/usr/bin/env bash
# Kokoro-82M：本机模型目录 + 离线启动脚本
# 依赖：api/model_loader.py 中 KOKORO_LOCAL_MODEL 兜底

set -euo pipefail

# ---------- 按你的机器修改 ----------
# 本机模型根目录（含 config.json、kokoro-v1_0.pth、voices/ 等）
export KOKORO_LOCAL_MODEL="${KOKORO_LOCAL_MODEL:-/opt/atomgit/models/Kokoro-82M}"
# 语音目录（一般为模型根下的 voices）
export VOICES_DIR="${VOICES_DIR:-${KOKORO_LOCAL_MODEL}/voices}"
# 本仓库（含 api/、config/）的路径
export KOKORO_ROOT="${KOKORO_ROOT:-/opt/atomgit/Kokoro-82M}"
# ------------------------------------

export PYTHONPATH="${KOKORO_ROOT}:${PYTHONPATH:-}"

export LANG_CODE="${LANG_CODE:-z}"
export DEFAULT_VOICE="${DEFAULT_VOICE:-af_heart}"
export SAMPLE_RATE="${SAMPLE_RATE:-24000}"
export PORT="${PORT:-8000}"
export LOG_LEVEL="${LOG_LEVEL:-INFO}"
export INFERENCE_WORKERS="${INFERENCE_WORKERS:-6}"

# 离线：由 KOKORO_LOCAL_MODEL 提供全部模型文件
export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-1}"

if [[ ! -d "${KOKORO_LOCAL_MODEL}" ]]; then
  echo "错误: KOKORO_LOCAL_MODEL 不是目录: ${KOKORO_LOCAL_MODEL}" >&2
  exit 1
fi
if [[ ! -f "${KOKORO_LOCAL_MODEL}/config.json" ]]; then
  echo "警告: 未找到 ${KOKORO_LOCAL_MODEL}/config.json，KModel 初始化可能失败" >&2
fi
if [[ ! -d "${VOICES_DIR}" ]]; then
  echo "错误: VOICES_DIR 不是目录: ${VOICES_DIR}" >&2
  exit 1
fi

# CANN 环境变量（路径按安装位置调整）
if [[ -f /usr/local/Ascend/ascend-toolkit/set_env.sh ]]; then
  source /usr/local/Ascend/ascend-toolkit/set_env.sh
fi

export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0}"
export CUDA_VISIBLE_DEVICES=

cd "${KOKORO_ROOT}"
exec python3 -m uvicorn api.main:app --host 0.0.0.0 --port "${PORT}" --app-dir "${KOKORO_ROOT}"

# [服务启动] 执行离线启动脚本
!chmod +x /opt/atomgit/offline.sh
!/opt/atomgit/offline.sh

# [健康检查] 验证服务是否正常运行
!curl http://localhost:8000/health

# [TTS 推理] 调用语音合成接口，生成英文语音
!curl -X POST http://localhost:8000/v1/audio/speech \
  -H "Content-Type: application/json" \
  -o output.wav \
  -d '{"input": "The answer to the universe is 42", "voice": "af_heart"}'

export ASCEND_RT_VISIBLE_DEVICES=0  # 昇腾可见设备
export CUDA_VISIBLE_DEVICES=        # 清空，防止 CUDA 回退

GET /health

POST /v1/audio/speech
Content-Type: application/json