pause_control25.py

# -*- coding: utf-8 -*-
"""IndexTTS 2.5 波形级停顿控制模块

链路（12 段实测验证）：
    [pause:Nms] → 解析标记（替换为逗号生成）→ 生成 wav（调用方做）
    → ① 检测器预检（可选）：停没停 + 段中心码位置
    → ② whisper 双锚定位：前锚=标记前最近定位单元（zh 字 / en 词）token end、
         后锚=标记后最近定位单元 token start（zh 同音匹配 fallback，en 关闭）
    → ③ 区间能量谷 = 插入点
    → ④ 停没停判定（whisper 间隙 / 检测器预检交叉）
    → ⑤ 执行：停了 → 增量插入（原停顿段中心补差 N-X，边界不动）；
         没停 → 直接插入 N ms

坐标系约定：
- parse_pause_marks 返回的标记位置为 clean 文本（标记已替换为逗号）的字符下标。
- locate/apply_pauses 内部统一为秒（float）与 wav 采样点，wav 采样率按文件实际值。

依赖：whisper / pypinyin / librosa / soundfile / numpy（检测器预检另需 joblib / scikit-learn / torch）。
"""
import os
import re
import shutil

import numpy as np

# ==================== 可调参数区 ====================
VALLEY_WINDOW_MS = 150      # 能量谷搜索窗口：锚点区间中心 ±150ms（区间更窄时用整个区间）
GAP_THRESHOLD_MS = 50       # 停没停判定阈值：whisper 间隙 ≥ 50ms 视为"停了"
SIL_RMS_THRESHOLD = 0.01    # 静音段能量阈值（RMS，16bit 满幅归一化后）
SIL_FRAME_MS = 10           # 能量分析帧长（ms）
SIL_HOP_MS = 5              # 能量分析帧移（ms）
ANCHOR_FALLBACK_S = 0.25    # 后锚缺失时，搜索区间 = 前锚 + 0.25s
DETECTOR_NEAR_MS = 400      # 检测器预检：预期码位置 ±400ms 内的停顿段视为"本标记的停顿"
WHISPER_MODEL = 'base'      # whisper 模型档位
WHISPER_PROMPT = '以下是简体中文普通话的转录。'  # 简体引导（防繁体转录导致锚点匹配失败）
WHISPER_PROMPT_EN = None    # 英文 initial_prompt（None = 不传，英文转写本身够准）
# 推理设备：'cpu'（默认，零显存，每句额外 2~5 秒）或 'cuda'（几乎不减速，
# whisper base 额外显存约 1GB）。可用环境变量 PAUSE_DEVICE 覆盖，如 set PAUSE_DEVICE=cuda。
DEVICE = os.environ.get('PAUSE_DEVICE', 'cpu')
# ===================================================

_PAUSE_RE = re.compile(r'\[pause:(\d+)ms\]')
_PUNCT = '，。、；：！？,.;:!? \t'
_EN_WORD_RE = re.compile(r"[A-Za-z0-9']")


def parse_pause_marks(text):
    """解析 [pause:Nms] 标记。

    标记在 clean 文本中替换为逗号（模型最熟悉的标点，定案），
    返回 (clean_text, marks)；marks = [(字符位置, Nms)]，
    字符位置 = 替换后逗号在 clean 文本中的下标（即停顿应出现的语义位置）。
    """
    marks = []
    out = []
    pos = 0
    for m in _PAUSE_RE.finditer(text):
        out.append(text[pos:m.start()])
        out.append('，')
        marks.append((len(''.join(out)) - 1, int(m.group(1))))
        pos = m.end()
    out.append(text[pos:])
    return ''.join(out), marks


def _py(ch):
    """无声调拼音（中文同音匹配用）。"""
    from pypinyin import lazy_pinyin
    return lazy_pinyin(ch)[0] if ch.strip() else ''


def _norm_en(s):
    """英文 token/词归一化（小写 + 去非词字符），用于词级精确匹配。"""
    return re.sub(r"[^a-z0-9']", '', s.lower())


def _units_around(text, pos, lang='zh'):
    """取标记前后的定位单元：不依赖"字"概念。

    zh：单元 = 标记前/后最近的非标点非空白字符（对应 whisper 中文字 token）；
    en：单元 = 标记前/后最近的完整单词（对应 whisper 英文词 token）。

    Returns:
        (pre_unit, post_unit)，取不到则为 ''。
    """
    i = pos - 1
    while i >= 0 and text[i] in _PUNCT:
        i -= 1
    if i < 0:
        pre = ''
    elif lang == 'en':
        j = i
        while j - 1 >= 0 and _EN_WORD_RE.match(text[j - 1]):
            j -= 1
        pre = text[j:i + 1]
    else:
        pre = text[i]
    i = pos + 1
    while i < len(text) and text[i] in _PUNCT:
        i += 1
    if i >= len(text):
        post = ''
    elif lang == 'en':
        j = i
        while j + 1 < len(text) and _EN_WORD_RE.match(text[j + 1]):
            j += 1
        post = text[i:j + 1]
    else:
        post = text[i]
    return pre, post


def _find_anchor(tokens, unit, prefer, lang='zh'):
    """在 whisper token 中找 unit 对应 token（精确优先；zh 启用同音 fallback）。

    中文：unit 为单字，匹配"token 含该字"；英文：unit 为单词，
    匹配"归一化后词形相等"。zh 精确失败时按无声调拼音同音匹配兜底；
    en 不做同音 fallback（whisper 英文转写够准，同音匹配反而引入误锚）。

    Args:
        tokens: [(word, start, end)]（秒）
        unit: 定位单元（中文字 / 英文词）
        prefer: 'end'（前锚，取 token 结束）或 'start'（后锚，取 token 开始）
        lang: 'zh' / 'en'
    Returns:
        (边界秒, 命中的 token 文本)；未命中返回 (None, None)
    """
    if not unit:
        return None, None
    if lang == 'en':
        nu = _norm_en(unit)
        cands = [(t, s, e) for t, s, e in tokens if _norm_en(t) == nu]
    else:
        cands = [(t, s, e) for t, s, e in tokens if unit in t]
        if not cands:
            cands = [(t, s, e) for t, s, e in tokens
                     if any(_py(c) == _py(unit) for c in t if c.strip())]
    if not cands:
        return None, None
    t, s, e = cands[0]
    return (e if prefer == 'end' else s), t


def _frame_rms(wav, sr, frame_ms=SIL_FRAME_MS, hop_ms=SIL_HOP_MS):
    """分帧 RMS，返回 (rms 数组, 每帧起始采样点数组)。"""
    frame = max(1, int(frame_ms / 1000 * sr))
    hop = max(1, int(hop_ms / 1000 * sr))
    starts = np.arange(0, max(1, len(wav) - frame + 1), hop)
    rms = np.array([np.sqrt(np.mean(wav[s:s + frame] ** 2)) if len(wav[s:s + frame]) else 0.0
                    for s in starts])
    return rms, starts


def _valley_in_window(wav, sr, lo_s, hi_s):
    """在 [lo_s, hi_s]（秒）内找能量谷中心（秒）。

    窗口过宽时收缩为区间中心 ±VALLEY_WINDOW_MS。
    """
    if hi_s <= lo_s:
        hi_s = lo_s + 0.01
    center = (lo_s + hi_s) / 2
    half = VALLEY_WINDOW_MS / 1000
    if hi_s - lo_s > 2 * half:
        lo_s, hi_s = max(lo_s, center - half), min(hi_s, center + half)
    a0 = max(0, int(lo_s * sr))
    a1 = min(len(wav), int(hi_s * sr))
    if a1 - a0 < int(SIL_FRAME_MS / 1000 * sr):
        return center
    hop = max(1, int(SIL_HOP_MS / 1000 * sr))
    frame = max(1, int(SIL_FRAME_MS / 1000 * sr))
    idx = np.arange(a0, a1 - frame + 1, hop)
    rms = np.array([np.sqrt(np.mean(wav[i:i + frame] ** 2)) for i in idx])
    return (idx[int(np.argmin(rms))] + frame / 2) / sr


def _silence_run_at(wav, sr, point_s, threshold=SIL_RMS_THRESHOLD):
    """从 point_s（秒，应位于静音内）向两侧扩展，返回完整静音段 (start_s, end_s)。

    用于测量原停顿时长 X。若 point 处并非静音，返回以 point 为中心的 0 长度段。
    """
    rms, starts = _frame_rms(wav, sr)
    if len(starts) == 0:
        return point_s, point_s
    hop_s = (starts[1] - starts[0]) / sr if len(starts) > 1 else SIL_HOP_MS / 1000
    i0 = int(np.searchsorted(starts / sr, point_s))
    i0 = min(max(i0, 0), len(rms) - 1)
    if rms[i0] >= threshold:
        return point_s, point_s
    lo = i0
    while lo > 0 and rms[lo - 1] < threshold:
        lo -= 1
    hi = i0
    while hi < len(rms) - 1 and rms[hi + 1] < threshold:
        hi += 1
    return starts[lo] / sr, starts[hi] / sr + SIL_FRAME_MS / 1000


class PauseLocator:
    """停顿定位器：whisper 双锚 + 能量谷 + （可选）检测器预检交叉。"""

    def __init__(self, whisper_model=WHISPER_MODEL, device=DEVICE):
        self._wm_name = whisper_model
        self._device = device
        self._wm = None
        self._detector = None

    def _load_whisper(self):
        if self._wm is None:
            import whisper
            self._wm = whisper.load_model(self._wm_name, device=self._device)
        return self._wm

    def _load_detector(self):
        if self._detector is None:
            from detector_pause25 import PauseDetector
            self._detector = PauseDetector(device=self._device).load()
        return self._detector

    def _transcribe(self, wav_path, lang='zh'):
        """whisper 转录（16k 重采样 + 逐词时间戳）。

        zh 用简体引导 initial_prompt（防繁体转录导致锚点匹配失败）；
        en 不传 initial_prompt。
        """
        import librosa
        wav16, _ = librosa.load(wav_path, sr=16000)
        wm = self._load_whisper()
        kwargs = dict(word_timestamps=True, language=lang,
                      condition_on_previous_text=False)
        prompt = WHISPER_PROMPT if lang == 'zh' else WHISPER_PROMPT_EN
        if prompt:
            kwargs['initial_prompt'] = prompt
        res = wm.transcribe(wav16, **kwargs)
        return [(w['word'].strip(), w['start'], w['end'])
                for s in res['segments'] for w in s.get('words', [])]

    def _detector_precheck(self, wav_path, marks, clean_text, dur_s,
                           codes=None, codes_npz=None):
        """检测器预检：每个标记附近的停顿段（码域 → 秒域）。

        输入 codes（序列）或 codes_npz（训练数据 npz 路径，含组装 124D
        特征所需的全部字段）。两者都没有则跳过（返回 None）。

        Returns:
            None 或 [{'seg': (start_s, end_s), 'center_s': …} 或 None, ...]
            与 marks 等长；元素为 None 表示该标记附近未检出停顿段。
        """
        if codes is None and codes_npz is None:
            return None
        det = self._load_detector()
        if codes_npz is not None:
            feat, codes = det.feat124_from_npz(codes_npz)
        else:
            raise ValueError('检测器预检需要 codes_npz（含注意力/logits 特征）；'
                             '裸 codes 不足以组装 124D 特征')
        segs = det.detect(codes, feat=feat)  # [(start_code, end_code, conf)]
        ms_per_code = dur_s * 1000 / len(codes) if len(codes) else 40.0
        out = []
        n_chars = max(len(clean_text), 1)
        for pos, _n_ms in marks:
            expect_s = pos / n_chars * dur_s
            best = None
            for a, b, _conf in segs:
                s_s, e_s = a * ms_per_code / 1000, b * ms_per_code / 1000
                c_s = (s_s + e_s) / 2
                if abs(c_s - expect_s) <= DETECTOR_NEAR_MS / 1000:
                    if best is None or abs(c_s - expect_s) < abs(best['center_s'] - expect_s):
                        best = {'seg': (s_s, e_s), 'center_s': c_s}
            out.append(best)
        return out

    def locate(self, wav_path, text, marks, use_detector=False,
               codes=None, codes_npz=None, lang='zh'):
        """定位每个标记的插入点。

        Args:
            wav_path: 生成的 wav（clean 文本合成）
            text: clean 文本（parse_pause_marks 的第一个返回值）
            marks: parse_pause_marks 的第二个返回值 [(字符位置, Nms)]
            use_detector: 是否启用码级检测器预检交叉
            codes / codes_npz: 检测器输入（npz 优先）
            lang: 'zh'（按字 token 双锚 + 同音 fallback）或
                  'en'（按词 token 双锚，无同音 fallback）

        Returns:
            list[dict]，与 marks 等长：
            {char_pos, target_ms, insert_point（秒）, paused（bool）,
             orig_pause_ms（X）, anchor_prev, anchor_next, valley,
             anchor_prev_tok, anchor_next_tok, detector（预检结果或 None）}
            定位失败的标记 dict 中 insert_point 为 None（apply_pauses 跳过）。
        """
        import librosa
        import soundfile as sf

        wav, sr = sf.read(wav_path)
        if wav.ndim > 1:
            wav = wav.mean(axis=1)
        wav = wav.astype(np.float64)
        if np.abs(wav).max() > 1.5:  # int16 读数
            wav /= 32768.0
        dur_s = len(wav) / sr

        tokens = self._transcribe(wav_path, lang=lang)
        prechecks = (self._detector_precheck(wav_path, marks, text, dur_s,
                                             codes=codes, codes_npz=codes_npz)
                     if use_detector else None)

        locs = []
        for mi, (pos, n_ms) in enumerate(marks):
            pre_unit, post_unit = _units_around(text, pos, lang=lang)
            a_prev, tok_prev = _find_anchor(tokens, pre_unit, 'end', lang=lang)
            a_next, tok_next = (_find_anchor(tokens, post_unit, 'start', lang=lang)
                                if post_unit else (None, None))
            rec = {'char_pos': pos, 'target_ms': n_ms,
                   'insert_point': None, 'paused': False, 'orig_pause_ms': 0.0,
                   'anchor_prev': a_prev, 'anchor_next': a_next, 'valley': None,
                   'anchor_prev_tok': tok_prev, 'anchor_next_tok': tok_next,
                   'detector': prechecks[mi] if prechecks else None}
            if a_prev is None:
                # whisper 完全定位失败：检测器段中心兜底
                if rec['detector'] is not None:
                    s_s, e_s = rec['detector']['seg']
                    rec.update(insert_point=rec['detector']['center_s'], paused=True,
                               orig_pause_ms=(e_s - s_s) * 1000,
                               valley=rec['detector']['center_s'])
                locs.append(rec)
                continue
            lo = a_prev
            hi = a_next if (a_next is not None and a_next > lo) else lo + ANCHOR_FALLBACK_S
            valley = _valley_in_window(wav, sr, lo, hi)
            # 停没停判定：whisper 间隙优先，检测器交叉
            gap_ms = (hi - lo) * 1000 if a_next is not None else 0.0
            paused = gap_ms >= GAP_THRESHOLD_MS
            det = rec['detector']
            if det is not None and not paused:
                # whisper 判没停但检测器在附近检出停顿段 → 采信检测器
                paused = True
            # 原停顿时长 X：能量静音段实测；检测器段长兜底
            if paused:
                s_s, e_s = _silence_run_at(wav, sr, valley)
                x_ms = (e_s - s_s) * 1000
                if x_ms <= 0 and det is not None:
                    ds, de = det['seg']
                    x_ms = (de - ds) * 1000
            else:
                x_ms = 0.0
            rec.update(insert_point=valley, paused=paused,
                       orig_pause_ms=x_ms, valley=valley)
            locs.append(rec)
        return locs


def apply_pauses(wav_path, locs, out_path):
    """执行停顿控制：停了 → 原停顿段中心增量补差 N-X；没停 → 谷点直接插入 N ms。

    操作按插入点从后往前执行，防坐标漂移。边界不动（只插入，不删改语音）。
    只延长不缩短：X ≥ N 时不操作。

    Returns:
        操作记录 list[dict]：{char_pos, action('extend'/'insert'/'skip'/'failed'),
        insert_point, inserted_ms, target_ms, orig_pause_ms}
    """
    import soundfile as sf

    wav, sr = sf.read(wav_path)
    ch_shape = wav.shape[1:]  # () 单声道 / (C,) 多声道
    x = wav
    ops = []
    for rec in sorted(locs, key=lambda r: (r['insert_point'] or 0), reverse=True):
        op = {'char_pos': rec['char_pos'], 'target_ms': rec['target_ms'],
              'orig_pause_ms': rec['orig_pause_ms'],
              'insert_point': rec['insert_point'], 'inserted_ms': 0}
        if rec['insert_point'] is None:
            op['action'] = 'failed'
            ops.append(op)
            continue
        if rec['paused']:
            delta = rec['target_ms'] - rec['orig_pause_ms']
            if delta <= GAP_THRESHOLD_MS:
                op['action'] = 'skip'  # 已达到目标（或只差一点，不值得动）
                ops.append(op)
                continue
            ins_ms = delta
            op['action'] = 'extend'
        else:
            ins_ms = rec['target_ms']
            op['action'] = 'insert'
        pos = int(round(rec['insert_point'] * sr))
        sil = np.zeros((int(ins_ms / 1000 * sr),) + ch_shape, dtype=x.dtype)
        x = np.concatenate([x[:pos], sil, x[pos:]], axis=0)
        op['inserted_ms'] = ins_ms
        ops.append(op)
    ops.reverse()
    os.makedirs(os.path.dirname(os.path.abspath(out_path)), exist_ok=True)
    sf.write(out_path, x, sr)
    return ops


def process(text, wav_path, out_path, use_detector=False, codes=None,
            codes_npz=None, locator=None, lang='zh'):
    """一站式入口：解析标记 → 定位 → 执行。

    Args:
        text: 含 [pause:Nms] 标记的原始文本
        wav_path: 由 clean 文本（标记已替换为逗号）生成的 wav
        out_path: 输出 wav
        use_detector: 启用码级检测器预检（需 codes_npz）
        codes / codes_npz: 检测器输入
        locator: 复用的 PauseLocator 实例（批量调用时避免重复加载 whisper）
        lang: 'zh' / 'en'（透传至 PauseLocator.locate / whisper）

    Returns:
        (clean_text, marks, locs, ops)；无标记时直接复制 wav 并返回空记录。
    """
    clean, marks = parse_pause_marks(text)
    if not marks:
        if os.path.abspath(wav_path) != os.path.abspath(out_path):
            shutil.copyfile(wav_path, out_path)
        return clean, marks, [], []
    if locator is None:
        locator = PauseLocator()
    locs = locator.locate(wav_path, clean, marks, use_detector=use_detector,
                          codes=codes, codes_npz=codes_npz, lang=lang)
    ops = apply_pauses(wav_path, locs, out_path)
    return clean, marks, locs, ops