首次提交
This commit is contained in:
227
tool/whisper_manager.py
Normal file
227
tool/whisper_manager.py
Normal file
@@ -0,0 +1,227 @@
|
||||
import os
|
||||
import hashlib
|
||||
import subprocess
|
||||
import tempfile
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
from typing import Optional
|
||||
|
||||
from pydub import AudioSegment
|
||||
|
||||
from config.config_utils import get_ffmpeg, get_ffprobe
|
||||
|
||||
# 让 pydub 也用项目自带的 ffmpeg/ffprobe(显式指定，不依赖 PATH)
|
||||
AudioSegment.converter = get_ffmpeg()
|
||||
AudioSegment.ffmpeg = get_ffmpeg()
|
||||
AudioSegment.ffprobe = get_ffprobe()
|
||||
|
||||
_MAX_CACHE_SIZE = 512
|
||||
|
||||
|
||||
class WhisperManager:
|
||||
_instance = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._initialized = False
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if not self._initialized:
|
||||
self.model = None
|
||||
self.model_name = None
|
||||
self._load_lock = threading.Lock() # 仅保护模型加载，推理不加锁
|
||||
self._cache = OrderedDict()
|
||||
self._cache_lock = threading.Lock()
|
||||
self._initialized = True
|
||||
|
||||
@staticmethod
|
||||
def _cuda_available() -> bool:
|
||||
"""nvidia-smi 可用即认为有可用 GPU。"""
|
||||
try:
|
||||
r = subprocess.run(["nvidia-smi"], capture_output=True, timeout=10)
|
||||
return r.returncode == 0
|
||||
except (FileNotFoundError, OSError, subprocess.TimeoutExpired):
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _gpu_count() -> int:
|
||||
"""数 GPU 张数（nvidia-smi -L 行数）。"""
|
||||
try:
|
||||
r = subprocess.run(["nvidia-smi", "-L"], capture_output=True, text=True, timeout=10)
|
||||
if r.returncode == 0:
|
||||
return sum(1 for ln in r.stdout.splitlines() if ln.strip().startswith("GPU "))
|
||||
except (FileNotFoundError, OSError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
return 0
|
||||
|
||||
def _resolve_device(self):
|
||||
"""决定 whisper 跑在 cpu 还是 gpu。读 config.audio.device：auto/cuda/gpu/cpu。
|
||||
device_index=auto/last 时用【最后一张卡】，与抽帧(用前 N-1 张)错开。
|
||||
返回 (device, compute_type, device_index)。"""
|
||||
try:
|
||||
from config.config_utils import config
|
||||
acfg = config.get_audio_config() or {}
|
||||
except Exception:
|
||||
acfg = {}
|
||||
dev = str(acfg.get("device", "auto")).lower()
|
||||
idx_cfg = acfg.get("device_index", "auto")
|
||||
ctype = str(acfg.get("compute_type", "float16"))
|
||||
if dev in ("cuda", "gpu") or (dev == "auto" and self._cuda_available()):
|
||||
if str(idx_cfg).lower() in ("auto", "last"):
|
||||
idx = max(0, self._gpu_count() - 1) # 用最后一张
|
||||
else:
|
||||
try:
|
||||
idx = int(idx_cfg)
|
||||
except (TypeError, ValueError):
|
||||
idx = 0
|
||||
return "cuda", ctype, idx
|
||||
return "cpu", "int8", 0
|
||||
|
||||
def load_model(self, model_name: str = "base"):
|
||||
"""加载 faster-whisper 模型：有 GPU 用 cuda(float16)，否则 CPU(int8)；GPU 失败自动回退 CPU。"""
|
||||
with self._load_lock:
|
||||
if self.model is None or self.model_name != model_name:
|
||||
from faster_whisper import WhisperModel
|
||||
device, ctype, idx = self._resolve_device()
|
||||
# CPU 推理限制内部线程数：防止进程池下每个 worker 的 whisper 各抓一堆核 → 线程爆炸。
|
||||
# whisper 只识别首尾各5秒,2 线程足够;可用 config.audio.cpu_threads 调整。
|
||||
try:
|
||||
from config.config_utils import config as _cfg
|
||||
cpu_threads = int((_cfg.get_audio_config() or {}).get("cpu_threads", 2) or 2)
|
||||
except Exception:
|
||||
cpu_threads = 2
|
||||
cpu_threads = max(1, cpu_threads)
|
||||
try:
|
||||
print(f"[WhisperManager] 加载 faster-whisper '{model_name}' (device={device}, {ctype})...")
|
||||
if device == "cuda":
|
||||
self.model = WhisperModel(model_name, device="cuda",
|
||||
device_index=idx, compute_type=ctype)
|
||||
else:
|
||||
self.model = WhisperModel(model_name, device="cpu",
|
||||
compute_type="int8", cpu_threads=cpu_threads)
|
||||
except Exception as e:
|
||||
print(f"[WhisperManager] GPU 加载失败({type(e).__name__}: {e})，回退 CPU(int8, cpu_threads={cpu_threads})")
|
||||
self.model = WhisperModel(model_name, device="cpu",
|
||||
compute_type="int8", cpu_threads=cpu_threads)
|
||||
self.model_name = model_name
|
||||
print(f"[WhisperManager] 模型 '{model_name}' 加载完成。")
|
||||
return self.model
|
||||
|
||||
def extract_audio_from_video(
|
||||
self,
|
||||
video_path: str,
|
||||
output_audio_path: str,
|
||||
logger: Optional[object] = None
|
||||
) -> Optional[float]:
|
||||
if not os.path.exists(video_path):
|
||||
if logger:
|
||||
logger.log(f"视频文件不存在: {video_path}", level="error")
|
||||
return None
|
||||
|
||||
cmd = [
|
||||
get_ffmpeg(), "-loglevel", "error",
|
||||
"-i", video_path,
|
||||
"-vn", "-acodec", "pcm_s16le",
|
||||
"-ar", "16000", "-ac", "1",
|
||||
"-y", output_audio_path
|
||||
]
|
||||
try:
|
||||
subprocess.run(cmd, check=True, capture_output=True, text=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
if logger:
|
||||
logger.log(f"ffmpeg 提取音频失败: {e.stderr}", level="error")
|
||||
return None
|
||||
|
||||
try:
|
||||
audio = AudioSegment.from_file(output_audio_path)
|
||||
duration = len(audio) / 1000.0
|
||||
if logger:
|
||||
logger.log(f"音频提取成功，时长: {duration:.2f} 秒")
|
||||
return duration
|
||||
except Exception as e:
|
||||
if logger:
|
||||
logger.log(f"获取音频时长失败: {e}", level="error")
|
||||
return None
|
||||
|
||||
def extract_audio_text_with_whisper(
|
||||
self,
|
||||
audio_path: str,
|
||||
start_time: float,
|
||||
end_time: float,
|
||||
logger: Optional[object] = None,
|
||||
use_cache: bool = True
|
||||
) -> str:
|
||||
if self.model is None:
|
||||
raise RuntimeError("Whisper 模型未加载，请先调用 load_model()")
|
||||
|
||||
if not os.path.exists(audio_path):
|
||||
if logger:
|
||||
logger.log(f"音频文件不存在: {audio_path}", level="error")
|
||||
return ""
|
||||
|
||||
cache_key = None
|
||||
if use_cache:
|
||||
mtime = os.path.getmtime(audio_path)
|
||||
cache_key = hashlib.md5(
|
||||
f"{audio_path}_{mtime}_{start_time}_{end_time}".encode()
|
||||
).hexdigest()
|
||||
with self._cache_lock:
|
||||
if cache_key in self._cache:
|
||||
self._cache.move_to_end(cache_key)
|
||||
if logger:
|
||||
logger.log("使用缓存识别结果")
|
||||
return self._cache[cache_key]
|
||||
|
||||
try:
|
||||
full_audio = AudioSegment.from_file(audio_path)
|
||||
start_ms = int(start_time * 1000)
|
||||
end_ms = int(end_time * 1000) + 500
|
||||
segment = full_audio[start_ms:end_ms]
|
||||
|
||||
if len(segment) == 0:
|
||||
if logger:
|
||||
logger.log("音频片段为空，跳过识别")
|
||||
return ""
|
||||
|
||||
segment = segment.set_frame_rate(16000).set_channels(1).set_sample_width(2)
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
|
||||
temp_path = tmp.name
|
||||
segment.export(temp_path, format="wav")
|
||||
|
||||
try:
|
||||
# faster-whisper CPU 模式线程安全，无需加锁，多线程并发推理
|
||||
segments, _ = self.model.transcribe(
|
||||
temp_path, language="zh", beam_size=5
|
||||
)
|
||||
text = "".join(seg.text for seg in segments).strip()
|
||||
finally:
|
||||
try:
|
||||
os.unlink(temp_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
if use_cache and cache_key:
|
||||
with self._cache_lock:
|
||||
self._cache[cache_key] = text
|
||||
self._cache.move_to_end(cache_key)
|
||||
while len(self._cache) > _MAX_CACHE_SIZE:
|
||||
self._cache.popitem(last=False)
|
||||
|
||||
if logger:
|
||||
logger.log(f"音频片段识别结果: {text}")
|
||||
|
||||
return text
|
||||
|
||||
except Exception as e:
|
||||
if logger:
|
||||
logger.log(f"音频识别失败: {e}", level="error")
|
||||
return ""
|
||||
|
||||
|
||||
whisper_manager = WhisperManager()
|
||||
Reference in New Issue
Block a user