mirror of
https://github.com/JefferyHcool/BiliNote.git
synced 2026-09-05 15:46:48 +08:00
Merge pull request #425 from pumpkinperson996/agent/fix-whisper-model-selection
Fix Whisper model selection caching
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
import os
|
||||
import platform
|
||||
import threading
|
||||
from enum import Enum
|
||||
|
||||
from app.transcriber.groq import GroqTranscriber
|
||||
@@ -38,17 +39,29 @@ _transcribers = {
|
||||
TranscriberType.GROQ: None,
|
||||
}
|
||||
|
||||
# Cache instances together with their constructor configuration. The
|
||||
# transcriber choice and Whisper model size can be changed from the frontend,
|
||||
# so caching by transcriber type alone would keep using the first loaded model.
|
||||
_transcriber_configs = {key: None for key in _transcribers}
|
||||
_transcriber_init_lock = threading.Lock()
|
||||
|
||||
# 公共实例初始化函数
|
||||
def _init_transcriber(key: TranscriberType, cls, *args, **kwargs):
|
||||
if _transcribers[key] is None:
|
||||
logger.info(f'创建 {cls.__name__} 实例: {key}')
|
||||
try:
|
||||
_transcribers[key] = cls(*args, **kwargs)
|
||||
init_config = (args, tuple(sorted(kwargs.items())))
|
||||
with _transcriber_init_lock:
|
||||
instance = _transcribers[key]
|
||||
if instance is None or _transcriber_configs[key] != init_config:
|
||||
action = "创建" if instance is None else "按新配置重新创建"
|
||||
logger.info(f'{action} {cls.__name__} 实例: {key}')
|
||||
try:
|
||||
new_instance = cls(*args, **kwargs)
|
||||
except Exception as e:
|
||||
logger.error(f"{cls.__name__} 创建失败: {e}")
|
||||
raise
|
||||
_transcribers[key] = new_instance
|
||||
_transcriber_configs[key] = init_config
|
||||
logger.info(f'{cls.__name__} 创建成功')
|
||||
except Exception as e:
|
||||
logger.error(f"{cls.__name__} 创建失败: {e}")
|
||||
raise
|
||||
return _transcribers[key]
|
||||
return _transcribers[key]
|
||||
|
||||
# 各类型获取方法
|
||||
def get_groq_transcriber():
|
||||
@@ -70,13 +83,13 @@ def get_mlx_whisper_transcriber(model_size="base"):
|
||||
return _init_transcriber(TranscriberType.MLX_WHISPER, MLXWhisperTranscriber, model_size=model_size)
|
||||
|
||||
# 通用入口
|
||||
def get_transcriber(transcriber_type="fast-whisper", model_size="base", device="cuda"):
|
||||
def get_transcriber(transcriber_type="fast-whisper", model_size=None, device="cuda"):
|
||||
"""
|
||||
获取指定类型的转录器实例
|
||||
|
||||
参数:
|
||||
transcriber_type: 支持 "fast-whisper", "mlx-whisper", "bcut", "kuaishou", "groq"
|
||||
model_size: 模型大小,适用于 whisper 类
|
||||
model_size: 模型大小,适用于 whisper 类;未提供时才读取环境变量默认值
|
||||
device: 设备类型(如 cuda / cpu),仅 whisper 使用
|
||||
|
||||
返回:
|
||||
@@ -90,7 +103,9 @@ def get_transcriber(transcriber_type="fast-whisper", model_size="base", device="
|
||||
logger.warning(f'未知转录器类型 "{transcriber_type}",默认使用 fast-whisper')
|
||||
transcriber_enum = TranscriberType.FAST_WHISPER
|
||||
|
||||
whisper_model_size = os.environ.get("WHISPER_MODEL_SIZE", model_size)
|
||||
# The explicit value normally comes from the persisted frontend setting and
|
||||
# must take precedence over Docker's startup default.
|
||||
whisper_model_size = model_size or os.environ.get("WHISPER_MODEL_SIZE", "base")
|
||||
|
||||
if transcriber_enum == TranscriberType.FAST_WHISPER:
|
||||
return get_whisper_transcriber(whisper_model_size, device=device)
|
||||
|
||||
Reference in New Issue
Block a user