Initial commit: audio2text 双语字幕生成服务
- 音频/视频转双语(英/中)SRT 字幕,Docker 容器化,CPU 开发/GPU 生产同一份代码 - faster-whisper ASR(词级时间戳) + 断句时间戳重算 + NLLB 翻译(模型不共驻) - 分片上传(断点续传) + SQLite 持久化 + 主页/历史/日志页面 - 历史页文件名搜索;缓存定时清理(默认保留7天,可配置) - 双 Dockerfile(cpu/gpu) + setup/start/stop 脚本
This commit is contained in:
136
app/services/model_manager.py
Normal file
136
app/services/model_manager.py
Normal file
@@ -0,0 +1,136 @@
|
||||
"""模型管理器: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.debug("加载 ASR 模型 model=%s device=%s compute_type=%s",
|
||||
s.model, s.device, s.compute_type)
|
||||
from faster_whisper import WhisperModel
|
||||
# device/compute_type 组合:cpu+int8 / cuda+float16
|
||||
self._asr = WhisperModel(
|
||||
s.model, device=s.device, compute_type=s.compute_type,
|
||||
)
|
||||
self._current = "asr"
|
||||
logger.debug("ASR 模型已就绪。")
|
||||
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.debug("卸载 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.debug("加载翻译模型 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,
|
||||
)
|
||||
self._current = "translator"
|
||||
logger.debug("翻译模型已就绪(独占显存,可用大 batch)。")
|
||||
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.debug("卸载翻译模型。")
|
||||
# 释放 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
|
||||
Reference in New Issue
Block a user