forked from open-webui/open-webui
Introduce Whisper model auto-update control.
* Introduce WHISPER_MODEL_AUTO_UPDATE env var * Pass local_files_only to WhisperModel() * Handle cases where auto-update is disabled but model is non-existent
This commit is contained in:
parent
9e726a32e6
commit
429242b4d3
2 changed files with 20 additions and 6 deletions
|
@ -28,6 +28,7 @@ from config import (
|
|||
UPLOAD_DIR,
|
||||
WHISPER_MODEL,
|
||||
WHISPER_MODEL_DIR,
|
||||
WHISPER_MODEL_AUTO_UPDATE,
|
||||
DEVICE_TYPE,
|
||||
)
|
||||
|
||||
|
@ -69,12 +70,22 @@ def transcribe(
|
|||
f.write(contents)
|
||||
f.close()
|
||||
|
||||
model = WhisperModel(
|
||||
WHISPER_MODEL,
|
||||
device=whisper_device_type,
|
||||
compute_type="int8",
|
||||
download_root=WHISPER_MODEL_DIR,
|
||||
)
|
||||
whisper_kwargs = {
|
||||
"model_size_or_path": WHISPER_MODEL,
|
||||
"device": whisper_device_type,
|
||||
"compute_type": "int8",
|
||||
"download_root": WHISPER_MODEL_DIR,
|
||||
"local_files_only": not WHISPER_MODEL_AUTO_UPDATE,
|
||||
}
|
||||
|
||||
log.debug(f"whisper_kwargs: {whisper_kwargs}")
|
||||
|
||||
try:
|
||||
model = WhisperModel(**whisper_kwargs)
|
||||
except:
|
||||
log.debug("WhisperModel initialization failed, attempting download with local_files_only=False")
|
||||
whisper_kwargs["local_files_only"] = False
|
||||
model = WhisperModel(**whisper_kwargs)
|
||||
|
||||
segments, info = model.transcribe(file_path, beam_size=5)
|
||||
log.info(
|
||||
|
|
|
@ -446,6 +446,9 @@ Query: [query]"""
|
|||
|
||||
WHISPER_MODEL = os.getenv("WHISPER_MODEL", "base")
|
||||
WHISPER_MODEL_DIR = os.getenv("WHISPER_MODEL_DIR", f"{CACHE_DIR}/whisper/models")
|
||||
WHISPER_MODEL_AUTO_UPDATE = (
|
||||
os.environ.get("WHISPER_MODEL_AUTO_UPDATE", "").lower() == "true"
|
||||
)
|
||||
|
||||
|
||||
####################################
|
||||
|
|
Loading…
Reference in a new issue