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()