228 lines
8.7 KiB
Python
228 lines
8.7 KiB
Python
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()
|