install_gpu.py

# -*- coding: utf-8 -*-
"""IndexTTS 2.5 停顿控制补丁 — 一键安装脚本

用法：
    python install.py [官方仓库目录] [--python 解释器] [--no-detector] [--yes] [--mode cpu|gpu]

流程：
    1. 定位官方 IndexTTS 2.5 仓库（webui.py 所在目录）
    2. 检查 Python 环境（默认用当前解释器，可 --python 指定官方 venv）
    3. 安装依赖（whisper/pypinyin/librosa/soundfile；检测器组件另装 joblib/sklearn/torch）
    4. 复制 pause_control25.py（可选：detector_pause25.py + models/）
    5. 选择后处理模式：CPU（默认，零显存，稍慢）或 GPU（加速，约 +1GB 显存），
       并生成对应启动脚本（启动脚本自带 PAUSE_DEVICE 环境变量）
    6. 应用补丁（git apply 优先，patch -p1 兜底，均失败则给出手动方案）
    7. 冒烟自检 + 提示重启 webui

退出码：0 = 成功；1 = 失败（有提示）。
"""
import argparse
import os
import shutil
import subprocess
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
PATCH_HOOK_MARK = 'pause-control hook'  # 已打补丁的判定标记（webui.py 内）
PY = sys.executable


def log(msg: str, ok: bool = True):
    tag = '[OK]' if ok else '[!!]'
    print(f'{tag} {msg}')


def find_repo(explicit: str) -> Path:
    """定位官方仓库根目录（含 webui.py 且含 indextts 包）。"""
    if explicit:
        p = Path(explicit).resolve()
        if (p / 'webui.py').is_file() and (p / 'indextts').is_dir():
            return p
        raise SystemExit(f'指定的目录不是官方仓库根目录（缺 webui.py 或 indextts/）：{p}')
    # 自动探测：当前目录 / 当前目录的上两级
    for c in (Path.cwd(), Path.cwd().parent, Path.cwd().parent.parent):
        if (c / 'webui.py').is_file() and (c / 'indextts').is_dir():
            return c
    raise SystemExit(
        '未找到官方 IndexTTS 2.5 仓库（需要同时有 webui.py 与 indextts/ 目录）。\n'
        '请 cd 到官方仓库根目录后重跑，或用：python install.py <官方仓库目录>')


def pip_install(py: str, pkgs: list, label: str):
    print(f'\n>>> 安装 {label} 依赖：{" ".join(pkgs)}')
    r = subprocess.run([py, '-m', 'pip', 'install', *pkgs])
    if r.returncode == 0:
        log(f'{label} 依赖安装完成')
        return
    # 兜底：uv 创建的虚拟环境不带 pip，改用 uv pip
    if shutil.which('uv'):
        print(f'  当前环境无 pip，改用 uv pip（--python {py}）')
        r = subprocess.run(['uv', 'pip', 'install', '--python', py, *pkgs])
        if r.returncode == 0:
            log(f'{label} 依赖安装完成（uv pip）')
            return
    raise SystemExit(f'pip 安装失败（{label}），请检查网络后重试')


def has_pkgs(py: str, mods: list) -> bool:
    code = ';'.join(f'import {m}' for m in mods)
    r = subprocess.run([py, '-c', code], capture_output=True)
    return r.returncode == 0


def apply_patch(repo: Path):
    """应用补丁：git apply 优先，patch -p1 兜底，最后提示手动。"""
    diff = HERE / 'patch_webui_pause.diff'
    webui = repo / 'webui.py'
    text = webui.read_bytes()
    if PATCH_HOOK_MARK.encode() in text:
        log('补丁已应用过（webui.py 含 pause-control hook），跳过')
        return True
    print('\n>>> 应用补丁')
    # 1) git apply
    r = subprocess.run(['git', 'apply', '--check', str(diff)], cwd=repo,
                       capture_output=True)
    if r.returncode == 0:
        r = subprocess.run(['git', 'apply', str(diff)], cwd=repo)
        if r.returncode == 0:
            log('补丁已应用（git apply）')
            return True
    # 2) patch（LF diff，对 LF/CRLF 目标都适用；勿加 --binary——LF diff 配 --binary 会失败）
    if shutil.which('patch'):
        r = subprocess.run(['patch', '-p1', '-i', str(diff)], cwd=repo)
        if r.returncode == 0:
            log('补丁已应用（patch -p1）')
            return True
    # 3) 手动方案
    print('[!!] 自动打补丁失败（需要 git 或 patch 命令，且仓库不能是只读）。')
    print('     手动方案：打开 webui.py，找到 gen_single() 里的')
    print('         output = tts.infer(**infer_kwargs)')
    print('     在它后面插入以下 7 行（4 空格缩进，行尾与文件其他行保持一致）：')
    print('''\
        # --- pause-control hook: [pause:Nms] waveform post-processing (pause_control25.py) ---
        if output and "[pause:" in (text or ""):
            try:
                import pause_control25
                pause_control25.process(text, output, output)  # in-place: locate + insert silence
            except Exception as e:
                print(f"[pause_control25] failed, keep original audio: {e}")''')
    return False


def smoke_test(py: str, repo: Path):
    print('\n>>> 冒烟自检')
    code = (
        "import sys; sys.path.insert(0, r'%s'); "
        "import pause_control25 as pc; "
        "assert pc.parse_pause_marks('你好[pause:500ms]世界')[1] == [(2, 500)], 'parse 异常'; "
        "print('pause_control25 导入与解析 OK')"
    ) % repo
    r = subprocess.run([py, '-c', code])
    if r.returncode == 0:
        log('自检通过')
        return True
    print('[!!] 自检失败，请把上面的报错发给 AI 助手排查')
    return False


def make_launcher(repo: Path, py: str, mode: str):
    """生成带 PAUSE_DEVICE 的启动脚本（Windows .bat / 其他 .sh）。

    注意（Windows 坑）：
    - 用 write_bytes 写，避免文本模式把 \\n 再转成 \\r\\n（CRCRLF）
    - .bat 必须 chcp 65001 切 UTF-8 代码页，否则 cmd 按 GBK 解析中文路径会乱码
    - 注释用英文（保守），变量与路径行在 chcp 65001 后按 UTF-8 解析
    """
    env = 'cuda' if mode == 'gpu' else 'cpu'
    desc = 'GPU (faster, ~+1GB VRAM)' if mode == 'gpu' else 'CPU (zero VRAM, slower)'
    is_win = sys.platform.startswith('win') or str(py).lower().endswith('.exe')
    if is_win:
        name = repo / 'start_webui.bat'
        lines = [
            '@echo off',
            'chcp 65001 >nul',
            'REM IndexTTS 2.5 PauseControl - %s launcher (generated by install.py)' % desc,
            'cd /d "%~dp0"',
            'set PAUSE_DEVICE=%s' % env,
            '"%s" webui.py %%*' % py,
            'pause',
        ]
        name.write_bytes(('\r\n'.join(lines) + '\r\n').encode('utf-8'))
    else:
        name = repo / 'start_webui.sh'
        lines = [
            '#!/usr/bin/env bash',
            '# IndexTTS 2.5 PauseControl - %s launcher (generated by install.py)' % desc,
            'cd "$(dirname "$0")"',
            'export PAUSE_DEVICE=%s' % env,
            '"%s" webui.py "$@"' % py,
        ]
        name.write_bytes(('\n'.join(lines) + '\n').encode('utf-8'))
    log(f'已生成启动脚本：{name}（{desc}）')
    return name


def main():
    ap = argparse.ArgumentParser(description='IndexTTS 2.5 停顿控制一键安装')
    ap.add_argument('repo', nargs='?', default='',
                    help='官方仓库目录（默认自动探测）')
    ap.add_argument('--python', default=PY,
                    help='Python 解释器路径（默认当前解释器；建议官方 venv）')
    ap.add_argument('--no-detector', action='store_true',
                    help='不安装码级检测器预检组件（detector_pause25.py + models/ + 依赖）')
    ap.add_argument('--mode', choices=['cpu', 'gpu'], default=None,
                    help='后处理模式：cpu（默认，零显存，稍慢）/ gpu（加速，约 +1GB 显存）；不传则询问')
    ap.add_argument('--yes', action='store_true', help='全部默认（CPU 模式），不询问')
    args = ap.parse_args()

    print('=' * 56)
    print(' IndexTTS 2.5 停顿控制补丁 — 一键安装')
    print('=' * 56)

    repo = find_repo(args.repo)
    log(f'官方仓库：{repo}')
    py = args.python
    log(f'Python 解释器：{py}')

    # ---- 依赖 ----
    need_detector = not args.no_detector
    if need_detector and not args.yes:
        ans = input('\n是否安装可选的"码级检测器预检"组件？（更稳的停顿判定，'
                    '需要 torch/joblib/scikit-learn）[Y/n] ').strip().lower()
        if ans in ('n', 'no'):
            need_detector = False
    base = ['openai-whisper', 'pypinyin', 'librosa', 'soundfile']
    pip_install(py, base, '基础')
    if need_detector:
        extra = [m for m in ('torch', 'joblib', 'scikit-learn')
                 if not has_pkgs(py, [m])]
        if extra:
            pip_install(py, extra, '检测器')
        else:
            log('检测器依赖已具备（torch/joblib/scikit-learn）')

    # ---- 后处理模式（CPU / GPU）----
    # 文件名决定默认模式：install_cpu.py → cpu，install_gpu.py → gpu；
    # install.py 本体 → 询问（--yes 默认 cpu）；--mode 参数始终优先
    mode = args.mode or {
        'install_cpu.py': 'cpu',
        'install_gpu.py': 'gpu',
    }.get(Path(sys.argv[0]).name, None)
    if mode is None and not args.yes:
        ans = input('\n停顿后处理用哪种模式？\n'
                    '  1) CPU（默认，零显存，每句稍慢 2~5 秒）\n'
                    '  2) GPU（几乎不减速，额外显存约 1GB）\n'
                    '选择 [1/2，回车=1] ').strip()
        mode = 'gpu' if ans == '2' else 'cpu'
    mode = mode or 'cpu'
    launcher = make_launcher(repo, py, mode)
    if mode == 'gpu':
        print('\n[提示] GPU 模式需本机有 NVIDIA 显卡且 torch 可用 CUDA；'
              '若启动时报 CUDA 错误，请改用 install_cpu.py 重新安装。')

    # ---- 复制文件 ----
    print('\n>>> 复制模块文件')
    shutil.copy2(HERE / 'pause_control25.py', repo / 'pause_control25.py')
    log('pause_control25.py → webui.py 同目录')
    if need_detector:
        shutil.copy2(HERE / 'detector_pause25.py', repo / 'detector_pause25.py')
        models_dst = repo / 'models'
        if (models_dst / 'pause_detector_lr.pkl').is_file():
            log('models/ 已存在，跳过复制')
        else:
            models_dst.mkdir(exist_ok=True)
            for w in ('pause_detector_lr.pkl', 'pause_detector_lstm.pt'):
                shutil.copy2(HERE / 'models' / w, models_dst / w)
            log('detector_pause25.py + models/ 已复制')
        if not os.environ.get('INDEXTTS_CODEC_PATH') and \
                not (repo / 'checkpoints' / 'codec.pth').is_file() and \
                not (models_dst / 'codec.pth').is_file():
            print('\n[提示] 检测器预检需要官方 codec.pth（约 607MB）。未检测到，'
                  '可任选其一：')
            print('   1) 官方仓库 checkpoints/ 下已有则忽略本提示')
            print('   2) 设环境变量 INDEXTTS_CODEC_PATH 指向 codec.pth 路径')
            print('   3) 把 codec.pth 复制到 models/ 目录')

    # ---- 打补丁 ----
    patched = apply_patch(repo)

    # ---- 自检 ----
    tested = smoke_test(py, repo)

    print('\n' + '=' * 56)
    if patched and tested:
        print(' 安装完成！')
        print(f' 启动方式：双击 {launcher.name}（自动带 PAUSE_DEVICE={mode}）')
        print(' 若你习惯用官方原来的启动器（如整合包 GUI），'
              '请先手动设环境变量 set PAUSE_DEVICE=' + ('cuda' if mode == 'gpu' else 'cpu'))
        print(' 重启后文本里写 [pause:800ms] 即可生成精确停顿。')
    else:
        print(' 安装未完全成功，请按上方提示处理（可把输出发给 AI 助手）。')
    print('=' * 56)
    sys.exit(0 if (patched and tested) else 1)


if __name__ == '__main__':
    main()