a6b5c7231c
faster-whisper 的 GPU 利用率呈尖刺波（峰=批量解码满载，谷=CPU 提取 Mel 特征 + 处理结果时 GPU 空闲），平均利用率低。瓶颈不在算力而在 CPU/GPU 未重叠。 - batch_size 16→32：拉长单次 GPU 解码时间，相对掩盖 CPU 特征提取间隙， 尖刺变宽变平，平均利用率上升。turbo FP16 仅 ~1.6GB，3090 24G 充裕。 - beam_size 5→2：turbo 模型鲁棒，候选数 5→2 大幅减少解码步数，让 GPU 峰更密、间隙更短。保留 1 个候选做歧义发音保险，质量损失小。 - beam_size 从硬编码提到 config 可调，CPU/CPU 模板/GPU/示例 四份配置对齐 - /health 增加 asr_beam_size，模型加载日志同步输出 batch+beam word_timestamps 保留 True：segmenter 强依赖词级时间戳做精确断句， 关闭会触发匀速估算退化路径，得不偿失。
143 lines
5.0 KiB
Python
143 lines
5.0 KiB
Python
"""模型管理器：ASR 与翻译模型不共驻，任一时刻 GPU 上只有一个模型。
|
||
|
||
设计：
|
||
- 单例 ModelManager 跟踪当前加载的模型类型（none / asr / translator）
|
||
- get_asr(): 若翻译器在内存 → 先卸载（del + gc + empty_cache）→ 加载 faster-whisper
|
||
- get_translator(): 若 ASR 在内存 → 先卸载 → 加载 NLLB
|
||
- 翻译阶段独占显存，可用大 batch_size；ASR 阶段同理
|
||
|
||
这样在 24G 3090 上无需担心显存叠加，CPU dev 时也省内存。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import gc
|
||
import logging
|
||
import threading
|
||
from typing import Any
|
||
|
||
from ..config import get_settings
|
||
|
||
logger = logging.getLogger("audio2text.models")
|
||
|
||
# 全局单例 + 锁：模型加载/卸载必须串行
|
||
_lock = threading.Lock()
|
||
|
||
|
||
def _free_memory() -> None:
|
||
"""释放 Python 对象与 GPU 缓存。"""
|
||
gc.collect()
|
||
try:
|
||
import torch
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
torch.cuda.synchronize()
|
||
except Exception: # pragma: no cover
|
||
pass
|
||
|
||
|
||
class ModelManager:
|
||
"""ASR / 翻译模型的不共驻管理器。"""
|
||
|
||
def __init__(self) -> None:
|
||
self._asr: Any = None # faster_whisper.WhisperModel
|
||
self._translator: Any = None # transformers pipeline
|
||
self._current: str = "none" # none | asr | translator
|
||
|
||
# ---------------- ASR ----------------
|
||
|
||
def get_asr(self) -> Any:
|
||
"""返回已加载的 faster-whisper 模型；必要时先卸载翻译器。"""
|
||
with _lock:
|
||
if self._asr is not None:
|
||
return self._asr
|
||
if self._translator is not None:
|
||
self._unload_translator_locked()
|
||
s = get_settings().asr
|
||
logger.info("加载 ASR 模型 model=%s device=%s compute_type=%s",
|
||
s.model, s.device, s.compute_type)
|
||
from faster_whisper import WhisperModel, BatchedInferencePipeline
|
||
# device/compute_type 组合：cpu+int8 / cuda+float16
|
||
# BatchedInferencePipeline 包装 WhisperModel，使 transcribe() 支持 batch_size，
|
||
# 多个音频块（chunk_length=30s）一次性送 GPU 解码，配合内部 prefill 提高利用率。
|
||
whisper = WhisperModel(
|
||
s.model, device=s.device, compute_type=s.compute_type,
|
||
)
|
||
self._asr = BatchedInferencePipeline(model=whisper)
|
||
self._current = "asr"
|
||
logger.info("ASR 模型已就绪（batched, batch_size=%d, beam_size=%d）。",
|
||
s.batch_size, s.beam_size)
|
||
return self._asr
|
||
|
||
def unload_asr(self) -> None:
|
||
with _lock:
|
||
self._unload_asr_locked()
|
||
|
||
def _unload_asr_locked(self) -> None:
|
||
if self._asr is None:
|
||
return
|
||
logger.info("卸载 ASR 模型（释放显存供翻译器独占）。")
|
||
# faster-whisper 模型无显式 close，del 即可
|
||
del self._asr
|
||
self._asr = None
|
||
self._current = "none" if self._translator is None else "translator"
|
||
_free_memory()
|
||
|
||
# ---------------- 翻译器 ----------------
|
||
|
||
def get_translator(self) -> Any:
|
||
"""返回已加载的 NLLB 翻译 pipeline；必要时先卸载 ASR。"""
|
||
with _lock:
|
||
if self._translator is not None:
|
||
return self._translator
|
||
if self._asr is not None:
|
||
self._unload_asr_locked()
|
||
s = get_settings().translation
|
||
logger.info("加载翻译模型 model=%s device=%s", s.model, s.device)
|
||
from transformers import pipeline
|
||
self._translator = pipeline(
|
||
"translation",
|
||
model=s.model,
|
||
device=s.device,
|
||
src_lang=s.src_lang,
|
||
tgt_lang=s.tgt_lang,
|
||
batch_size=s.batch_size, # pipeline 内部批大小，与 translate_service 分块对齐
|
||
)
|
||
self._current = "translator"
|
||
logger.info("翻译模型已就绪（显存独占，batch_size=%d，sort_by_length=%s）。",
|
||
s.batch_size, s.sort_by_length)
|
||
return self._translator
|
||
|
||
def unload_translator(self) -> None:
|
||
with _lock:
|
||
self._unload_translator_locked()
|
||
|
||
def _unload_translator_locked(self) -> None:
|
||
if self._translator is None:
|
||
return
|
||
logger.info("卸载翻译模型。")
|
||
# 释放 pipeline 持有的 model + tokenizer
|
||
mdl = getattr(self._translator, "model", None)
|
||
tok = getattr(self._translator, "tokenizer", None)
|
||
del self._translator, mdl, tok
|
||
self._translator = None
|
||
self._current = "none" if self._asr is None else "asr"
|
||
_free_memory()
|
||
|
||
# ---------------- 状态 ----------------
|
||
|
||
@property
|
||
def current(self) -> str:
|
||
return self._current
|
||
|
||
|
||
# 进程级单例
|
||
_manager: ModelManager | None = None
|
||
|
||
|
||
def get_model_manager() -> ModelManager:
|
||
global _manager
|
||
if _manager is None:
|
||
_manager = ModelManager()
|
||
return _manager
|