Update app.py

app.py

@@ -1,12 +1,10 @@
-import os
-import io
-import math
-import tempfile
+import os, io, math, tempfile
 from typing import List, Tuple
 
 import numpy as np
 import gradio as gr
 import librosa
+import torch
 
 try:
     from scipy.ndimage import median_filter
@@ -14,27 +12,31 @@ try:
 except Exception:
     _HAS_SCIPY = False
 
-import torch
 from transformers import pipeline
+import spaces  # 关键：用于 ZeroGPU
 
-# ================== 默认参数（UI 初始值 & 限制） ==================
+# ================== 默认参数 ==================
 MODEL_NAME = "openai/whisper-large-v3"
 BATCH_SIZE = 8
 FILE_LIMIT_MB = 1000
 
-DEF_SILENCE_MIN_LEN = 0.45   # 停顿（静音段）最短持续秒数
-DEF_DB_DROP         = 25.0   # 相对峰值下落阈值（max_db - DB_DROP）
-DEF_PCTL_FLOOR      = 20.0   # 能量分位（dB）下限（越大越保守）
-DEF_MIN_SEG_DUR     = 1.00   # 每段最短显示时长
-DEF_FRAME_LEN_MS    = 25     # 能量分析帧长
-DEF_HOP_LEN_MS      = 10     # 帧移
-DEF_CUT_OFFSET_SEC  = 0.00   # 切分偏移（整体校准）
-DEF_CHUNK_LEN_S     = 20     # ASR 分块长度，越小时间漂移越小
-DEF_STRIDE_LEN_S    = 2      # ASR 重叠长度，帮助跨块稳定时间戳
-SR_TARGET           = 16000  # 统一采样率（Whisper 期望 16k）
-
-# ================== ASR Pipeline ==================
-def _get_device_and_dtype():
+DEF_SILENCE_MIN_LEN = 0.45
+DEF_DB_DROP         = 25.0
+DEF_PCTL_FLOOR      = 20.0
+DEF_MIN_SEG_DUR     = 1.00
+DEF_FRAME_LEN_MS    = 25
+DEF_HOP_LEN_MS      = 10
+DEF_CUT_OFFSET_SEC  = 0.00
+DEF_CHUNK_LEN_S     = 20
+DEF_STRIDE_LEN_S    = 2
+SR_TARGET           = 16000
+
+# ================== 全局懒加载 ==================
+_ASR = None
+_ASR_DEVICE = None
+_ASR_DTYPE = None
+
+def _pick_device_dtype():
     if torch.cuda.is_available():
         return 0, torch.float16
     elif torch.backends.mps.is_available():
@@ -42,32 +44,38 @@ def _get_device_and_dtype():
     else:
         return -1, torch.float32
 
-DEVICE, DTYPE = _get_device_and_dtype()
-
-asr = pipeline(
-    task="automatic-speech-recognition",
-    model=MODEL_NAME,
-    device=DEVICE,
-    torch_dtype=DTYPE,
-    # 关键：词级时间戳
-    return_timestamps="word",
-    chunk_length_s=DEF_CHUNK_LEN_S,
-    stride_length_s=DEF_STRIDE_LEN_S,
-)
-
-# ================== 工具函数 ==================
+def _get_asr():
+    """
+    在 ZeroGPU 下必须在 @spaces.GPU 修饰的函数内首次调用，才能拿到 cuda。
+    CPU/常规 GPU 也兼容。
+    """
+    global _ASR, _ASR_DEVICE, _ASR_DTYPE
+    dev, dt = _pick_device_dtype()
+    if _ASR is None or _ASR_DEVICE != dev:
+        _ASR = pipeline(
+            task="automatic-speech-recognition",
+            model=MODEL_NAME,
+            device=dev,
+            torch_dtype=dt,
+            return_timestamps="word",
+            chunk_length_s=DEF_CHUNK_LEN_S,
+            stride_length_s=DEF_STRIDE_LEN_S,
+            ignore_warning=True,
+        )
+        _ASR_DEVICE, _ASR_DTYPE = dev, dt
+        print(f"[ASR] Initialized on device={dev} dtype={dt}")
+    return _ASR
+
+# ================== 音频 & 工具 ==================
 def _load_audio(path: str, sr: int = SR_TARGET):
     y, sr = librosa.load(path, sr=sr, mono=True)
     return y, sr
 
 def _to_db(rms: np.ndarray):
-    # librosa 的 amplitude_to_db 以幅值为输入
-    # 这里确保不会除以 0
     ref = np.maximum(np.max(rms), 1e-10)
-    db = 20.0 * np.log10(np.maximum(rms, 1e-10) / ref)
-    return db
+    return 20.0 * np.log10(np.maximum(rms, 1e-10) / ref)
 
-def _format_ts(sec: float) -> str:
+def _fmt_ts(sec: float) -> str:
     if sec < 0: sec = 0.0
     h = int(sec // 3600)
     m = int((sec % 3600) // 60)
@@ -76,10 +84,6 @@ def _format_ts(sec: float) -> str:
     return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"
 
 def _extract_word_stream(chunks) -> List[Tuple[str, float, float]]:
-    """
-    将 pipeline 返回的 chunks 统一抽取为 [(text, start, end), ...]
-    兼容字段名：timestamp/timestamps/start/end/time_start/time_end
-    """
     out = []
     if not chunks:
         return out
@@ -87,9 +91,8 @@ def _extract_word_stream(chunks) -> List[Tuple[str, float, float]]:
         txt = (ch.get("text") or "").strip()
         ts = ch.get("timestamp", ch.get("timestamps", None))
         if ts is None:
-            # 有些实现用 start/end 或 time_start/time_end
             s = ch.get("start", ch.get("time_start", None))
-            e = ch.get("end", ch.get("time_end", None))
+            e = ch.get("end",   ch.get("time_end",   None))
             if s is not None and e is not None and txt:
                 s = float(s); e = float(e)
                 if e < s: e = s
@@ -104,75 +107,53 @@ def _extract_word_stream(chunks) -> List[Tuple[str, float, float]]:
 def _detect_silence_cuts(
     y: np.ndarray,
     sr: int,
-    silence_min_len: float = DEF_SILENCE_MIN_LEN,
-    db_drop: float = DEF_DB_DROP,
-    pctl_floor: float = DEF_PCTL_FLOOR,
-    frame_len_ms: int = DEF_FRAME_LEN_MS,
-    hop_len_ms: int = DEF_HOP_LEN_MS,
-) -> Tuple[List[float], float]:
-    """
-    用 RMS(dB) + 最低点策略找切分点；返回 [cut_times], total_dur
-    """
-    frame_len = int(sr * frame_len_ms / 1000)
-    hop_len   = int(sr * hop_len_ms / 1000)
-    frame_len = max(256, frame_len)
-    hop_len   = max(64, hop_len)
+    silence_min_len: float,
+    db_drop: float,
+    pctl_floor: float,
+    frame_len_ms: int,
+    hop_len_ms: int,
+):
+    frame_len = max(256, int(sr * frame_len_ms / 1000))
+    hop_len   = max( 64, int(sr * hop_len_ms  / 1000))
 
     rms = librosa.feature.rms(y=y, frame_length=frame_len, hop_length=hop_len, center=True)[0]
     rms_db = _to_db(rms)
-
     if _HAS_SCIPY:
-        rms_db = median_filter(rms_db, size=5)  # 轻微平滑
+        rms_db = median_filter(rms_db, size=5)
 
     max_db = float(np.max(rms_db))
     floor_db = float(np.percentile(rms_db, pctl_floor))
     thr = max(max_db - db_drop, floor_db)
 
-    # 低于阈值视为“静音/低能”
-    low_mask = rms_db <= thr
-
-    # 找连续的低能区间
+    low = rms_db <= thr
     cut_times = []
+    n = len(low)
     i = 0
-    n = len(low_mask)
-    min_frames = int(silence_min_len * sr / hop_len)
+    min_frames = max(1, int(silence_min_len * sr / hop_len))
     while i < n:
-        if not low_mask[i]:
-            i += 1
-            continue
+        if not low[i]:
+            i += 1; continue
         j = i + 1
-        while j < n and low_mask[j]:
+        while j < n and low[j]:
             j += 1
-        span = j - i
-        if span >= max(1, min_frames):
+        if (j - i) >= min_frames:
             local = rms_db[i:j]
             k = int(np.argmin(local))
             best = i + k
-            cut_t = best * hop_len / sr
-            cut_times.append(float(cut_t))
+            cut_times.append(best * hop_len / sr)
         i = j
 
-    total_dur = float(len(y) / sr)
-    # 去重+排序+边界裁剪
-    cut_times = sorted(set(t for t in cut_times if 0.05 <= t <= total_dur - 0.05))
-    return cut_times, total_dur
-
-def _snap_to_word_boundaries(
-    cut_times: List[float],
-    word_stream: List[Tuple[str, float, float]],
-    max_dist: float = 0.25
-) -> List[float]:
-    if not cut_times or not word_stream:
-        return cut_times
-    bounds = []
-    for _, s, e in word_stream:
-        bounds.append(s); bounds.append(e)
-    bounds = sorted(set(bounds))
+    total = float(len(y) / sr)
+    cut_times = sorted(set(t for t in cut_times if 0.05 <= t <= total - 0.05))
+    return cut_times, total
+
+def _snap_to_word_bounds(cuts: List[float], words: List[Tuple[str, float, float]], max_dist=0.25):
+    if not cuts or not words: return cuts
+    bounds = sorted({b for _, s, e in words for b in (s, e)})
     snapped = []
-    for t in cut_times:
+    for t in cuts:
         idx = min(range(len(bounds)), key=lambda i: abs(bounds[i]-t))
         snapped.append(bounds[idx] if abs(bounds[idx]-t) <= max_dist else t)
-    # 去重并保证最小间隔
     snapped = sorted(set(snapped))
     out = []
     for t in snapped:
@@ -180,101 +161,60 @@ def _snap_to_word_boundaries(
             out.append(t)
     return out
 
-def _segment_by_energy(
-    word_stream: List[Tuple[str, float, float]],
-    cut_times: List[float],
-    total_dur: float,
-    min_seg_dur: float = DEF_MIN_SEG_DUR,
-) -> List[Tuple[float, float, str]]:
-    """
-    根据 cut_times 把词流切段；不足 min_seg_dur 的与邻段合并（优先并右）
-    返回 [(st, en, text), ...]
-    """
-    if not word_stream:
-        # 没有词流时，返回整段空文本占位
-        return [(0.0, total_dur, "").copy()]
-
-    bnds = [0.0] + [t for t in cut_times if 0.0 < t < total_dur] + [total_dur]
+def _segment(words: List[Tuple[str,float,float]], cuts: List[float], total: float, min_seg: float):
+    if not words:
+        return [(0.0, total, "")]
+    bnds = [0.0] + [t for t in cuts if 0.0 < t < total] + [total]
     segs = []
-    wi = 0
-    W = len(word_stream)
-
-    for i in range(len(bnds) - 1):
+    wi, W = 0, len(words)
+    for i in range(len(bnds)-1):
         L, R = bnds[i], bnds[i+1]
         texts, starts, ends = [], [], []
-        # 收集与 [L,R] 有交集的词
-        while wi < W and word_stream[wi][2] <= L:
+        while wi < W and words[wi][2] <= L:
             wi += 1
         wj = wi
-        while wj < W and word_stream[wj][1] < R:
-            txt, s, e = word_stream[wj]
+        while wj < W and words[wj][1] < R:
+            txt, s, e = words[wj]
             if e > L and s < R:
-                texts.append(txt)
-                starts.append(s)
-                ends.append(e)
+                texts.append(txt); starts.append(s); ends.append(e)
             wj += 1
-
         if texts:
-            st = float(min(starts)); en = float(max(ends))
-            # 防止越界
-            st = max(st, L); en = min(en, R)
-            tx = " ".join(texts).strip()
-            segs.append([st, en, tx])
-        else:
-            # 没有词但留下时间窗，避免生成过短空段
-            if (R - L) >= max(0.25, min_seg_dur * 0.5):
-                segs.append([L, R, ""])
-
-    # 合并过短段（优先向右）
-    def has_punc(t: str) -> bool:
-        return any(p in t for p in "，。！？,.!?;；：:")
+            st, en = max(min(starts), L), min(max(ends), R)
+            segs.append([float(st), float(en), " ".join(texts).strip()])
+        elif (R - L) >= max(0.25, min_seg * 0.5):
+            segs.append([L, R, ""])
 
+    def has_punc(t): return any(p in t for p in "，。！？,.!?;；：:")
     i = 0
     while i < len(segs):
         st, en, tx = segs[i]
-        if (en - st) < min_seg_dur and len(segs) > 1:
-            # 选择合并目标：先右后左；若右/左都存在，则优先含标点的
-            target = None
+        if (en - st) < min_seg and len(segs) > 1:
             cand = []
             if i + 1 < len(segs): cand.append(i + 1)
             if i - 1 >= 0:        cand.append(i - 1)
-            if not cand:
-                i += 1
-                continue
-            # 优先含标点
             cand.sort(key=lambda j: (not has_punc(segs[j][2]), abs(j - i)))
-            target = cand[0]
-            # 合并
-            nst = min(segs[target][0], st)
-            nen = max(segs[target][1], en)
-            ntx = " ".join([segs[target][2], tx]).strip() if target < i else " ".join([tx, segs[target][2]]).strip()
-            # 放在较小的索引位，删除另一段
-            keep, drop = (target, i) if target < i else (i, target)
+            t = cand[0]
+            nst, nen = min(segs[t][0], st), max(segs[t][1], en)
+            ntx = (" ".join([segs[t][2], tx]) if t < i else " ".join([tx, segs[t][2]])).strip()
+            keep, drop = (t, i) if t < i else (i, t)
             segs[keep] = [nst, nen, ntx]
             del segs[drop]
-            # 合并后从较早索引重新审视
-            i = max(0, keep - 1)
-            continue
+            i = max(0, keep - 1); continue
         i += 1
 
-    # 去除空文字但极短的段
-    out = []
-    for st, en, tx in segs:
-        if (en - st) < 0.12:
-            continue
-        out.append((float(st), float(en), tx.strip()))
-    return out
+    return [(st, en, tx.strip()) for st, en, tx in segs if (en - st) >= 0.12]
 
-def _build_srt(segs: List[Tuple[float, float, str]]) -> str:
+def _build_srt(segs: List[Tuple[float,float,str]]) -> str:
     lines = []
     for idx, (st, en, tx) in enumerate(segs, start=1):
         lines.append(str(idx))
-        lines.append(f"{_format_ts(st)} --> {_format_ts(en)}")
-        lines.append(tx if tx else "")
+        lines.append(f"{_fmt_ts(st)} --> {_fmt_ts(en)}")
+        lines.append(tx)
         lines.append("")
     return "\n".join(lines).strip() + "\n"
 
-# ================== 主流程 ==================
+# ================== 推理核心（放在 GPU 上执行） ==================
+@spaces.GPU  # 关键：ZeroGPU 运行入口（按钮点击会调用它）
 def transcribe_and_split(
     audio_path: str,
     silence_min_len: float = DEF_SILENCE_MIN_LEN,
@@ -288,30 +228,26 @@ def transcribe_and_split(
     if not audio_path:
         raise gr.Error("请先上传或录制音频。")
 
-    # 体积限制
     try:
-        fsize_mb = os.path.getsize(audio_path) / (1024 * 1024)
-        if fsize_mb > FILE_LIMIT_MB:
-            raise gr.Error(f"文件过大：{fsize_mb:.1f} MB，超过上限 {FILE_LIMIT_MB} MB。")
+        if os.path.getsize(audio_path) / (1024*1024) > FILE_LIMIT_MB:
+            raise gr.Error(f"文件过大，超过 {FILE_LIMIT_MB} MB。")
     except Exception:
         pass
 
-    # ASR
+    asr = _get_asr()  # 在 GPU 上首次创建
+
     result = asr(
         audio_path,
-        # 以下参数也可传入调用（已在构造时设置了默认）
         return_timestamps="word",
         chunk_length_s=DEF_CHUNK_LEN_S,
         stride_length_s=DEF_STRIDE_LEN_S,
         batch_size=BATCH_SIZE,
     )
     text = (result.get("text") or "").strip()
-    chunks = result.get("chunks") or []
-    words = _extract_word_stream(chunks)
+    words = _extract_word_stream(result.get("chunks") or [])
 
-    # 能量切分（在 16k 单声道下计算）
     y, sr = _load_audio(audio_path, sr=SR_TARGET)
-    cut_times, total_dur = _detect_silence_cuts(
+    cuts, total = _detect_silence_cuts(
         y, sr,
         silence_min_len=silence_min_len,
         db_drop=db_drop,
@@ -320,55 +256,49 @@ def transcribe_and_split(
         hop_len_ms=hop_len_ms,
     )
 
-    # 切点整体偏移
     if abs(cut_offset_sec) > 1e-6:
-        cut_times = [max(0.0, min(total_dur, t + cut_offset_sec)) for t in cut_times]
+        cuts = [max(0.0, min(total, t + cut_offset_sec)) for t in cuts]
 
-    # 切点吸附到最近词边界
-    cut_times = _snap_to_word_boundaries(cut_times, words, max_dist=0.25)
-
-    # 用切点+词流生成分段
-    segs = _segment_by_energy(words, cut_times, total_dur, min_seg_dur=min_seg_dur)
-
-    # 兜底：如果没有词（极端情况），给整段
+    cuts = _snap_to_word_bounds(cuts, words, max_dist=0.25)
+    segs = _segment(words, cuts, total, min_seg_dur)
     if not segs:
-        segs = [(0.0, total_dur, text)]
+        segs = [(0.0, total, text)]
+    srt = _build_srt(segs)
 
-    srt_text = _build_srt(segs)
-
-    # 保存到临时 .srt 文件供下载
     tmpf = tempfile.NamedTemporaryFile(delete=False, suffix=".srt")
-    tmpf.write(srt_text.encode("utf-8"))
-    tmpf.flush(); tmpf.close()
+    tmpf.write(srt.encode("utf-8")); tmpf.flush(); tmpf.close()
+    return srt, tmpf.name
 
-    return srt_text, tmpf.name
+# 让启动检查看到 GPU 入口（可选，不调用也行）
+@spaces.GPU
+def gpu_warmup():
+    return "ok"
 
-# ================== Gradio UI ==================
+# ================== UI ==================
 with gr.Blocks(title="Whisper Large V3 · 智能切分 SRT", theme=gr.themes.Soft()) as demo:
     gr.Markdown("### 🎧 Whisper Large V3 · 更稳的 SRT 切分\n"
                 "- 词级时间戳 + 能量最低点切分 + 词边界吸附\n"
-                "- 片段时长不足将自动与邻段合并（优先右侧）\n")
+                "- 片段过短自动合并，SRT 含序号行\n")
+
+    audio = gr.Audio(sources=["upload", "microphone"], type="filepath", label="音频（上传或录制）")
 
-    with gr.Row():
-        audio = gr.Audio(sources=["upload", "microphone"], type="filepath", label="音频（上传或录制）")
     with gr.Accordion("高级参数", open=False):
         with gr.Row():
             silence_min_len = gr.Slider(0.1, 1.0, value=DEF_SILENCE_MIN_LEN, step=0.05, label="静音最短时长 (s)")
-            db_drop         = gr.Slider(10, 40, value=DEF_DB_DROP, step=1.0, label="相对峰值下落 (dB)")
-            pctl_floor      = gr.Slider(0, 50, value=DEF_PCTL_FLOOR, step=1.0, label="能量分位下限 (dB)")
+            db_drop         = gr.Slider(10, 40, value=DEF_DB_DROP, step=1.0,  label="相对峰值下落 (dB)")
+            pctl_floor      = gr.Slider(0,  50, value=DEF_PCTL_FLOOR, step=1.0,  label="能量分位下限 (dB)")
         with gr.Row():
             min_seg_dur     = gr.Slider(0.3, 3.0, value=DEF_MIN_SEG_DUR, step=0.05, label="最短片段时长 (s)")
-            frame_len_ms    = gr.Slider(10, 50, value=DEF_FRAME_LEN_MS, step=1, label="帧长 (ms)")
-            hop_len_ms      = gr.Slider(5, 25, value=DEF_HOP_LEN_MS, step=1, label="帧移 (ms)")
+            frame_len_ms    = gr.Slider(10, 50, value=DEF_FRAME_LEN_MS, step=1,   label="帧长 (ms)")
+            hop_len_ms      = gr.Slider(5,  25, value=DEF_HOP_LEN_MS,  step=1,   label="帧移 (ms)")
             cut_offset_sec  = gr.Slider(-0.20, 0.20, value=DEF_CUT_OFFSET_SEC, step=0.01, label="切分整体偏移 (s)")
 
     btn = gr.Button("开始识别并生成 SRT", variant="primary")
-    with gr.Row():
-        srt_preview = gr.Textbox(lines=16, label="SRT 预览", show_copy_button=True)
+    srt_preview = gr.Textbox(lines=16, label="SRT 预览", show_copy_button=True)
     srt_file = gr.File(label="下载 SRT 文件", file_count="single")
 
     btn.click(
-        fn=transcribe_and_split,
+        fn=transcribe_and_split,  # 注意：绑定的是 @spaces.GPU 函数
         inputs=[audio, silence_min_len, db_drop, pctl_floor, min_seg_dur, frame_len_ms, hop_len_ms, cut_offset_sec],
         outputs=[srt_preview, srt_file],
     )