feat: preload Whisper model at startup with disk space check
- Add WHISPER_MODEL_SIZE and WHISPER_MIN_FREE_SPACE_BYTES config options - voice.py now preloads/downloads the Whisper model at bot startup - Detailed logging for model presence, disk space, download progress and readiness - If model is already present locally, preload is skipped - If disk space is insufficient, bot fails fast with a clear error - main.py calls preload_model() after proxy setup and before polling - Document voice recognition model download behavior in README
This commit is contained in:
+80
-7
@@ -2,12 +2,14 @@ import asyncio
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from faster_whisper import WhisperModel
|
||||
from aiogram import Bot
|
||||
|
||||
from config import PROXY_ENABLED, PROXY_URL
|
||||
from config import PROXY_ENABLED, PROXY_URL, WHISPER_MODEL_SIZE, WHISPER_MIN_FREE_SPACE_BYTES
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -25,26 +27,59 @@ def _setup_proxy() -> None:
|
||||
logger.info("Proxy set for model download: %s...", PROXY_URL[:30])
|
||||
|
||||
|
||||
def _get_model_path() -> Path:
|
||||
return MODELS_DIR / WHISPER_MODEL_SIZE
|
||||
|
||||
|
||||
def _is_model_downloaded() -> bool:
|
||||
"""Check whether the Whisper model files are already present locally."""
|
||||
model_path = _get_model_path()
|
||||
model_bin = model_path / "model.bin"
|
||||
return model_path.is_dir() and model_bin.is_file()
|
||||
|
||||
|
||||
def _get_free_disk_space_bytes(path: Path) -> int:
|
||||
"""Return free disk space in bytes for the filesystem containing path."""
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return shutil.disk_usage(path).free
|
||||
|
||||
|
||||
def _load_model() -> WhisperModel:
|
||||
"""Synchronous model load/download. Must run in a thread."""
|
||||
MODELS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
VOICE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
model_path = MODELS_DIR / "base"
|
||||
model_path = _get_model_path()
|
||||
model_bin = model_path / "model.bin"
|
||||
|
||||
if not model_path.is_dir() or not model_bin.is_file():
|
||||
_setup_proxy()
|
||||
logger.info("Model not found locally, downloading to %s...", MODELS_DIR)
|
||||
return WhisperModel(
|
||||
"base",
|
||||
logger.info(
|
||||
"Whisper model '%s' not found locally. Starting download to %s...",
|
||||
WHISPER_MODEL_SIZE,
|
||||
MODELS_DIR,
|
||||
)
|
||||
start = time.monotonic()
|
||||
model = WhisperModel(
|
||||
WHISPER_MODEL_SIZE,
|
||||
device="cpu",
|
||||
cpu_threads=4,
|
||||
compute_type="int8",
|
||||
download_root=str(MODELS_DIR),
|
||||
)
|
||||
elapsed = time.monotonic() - start
|
||||
logger.info(
|
||||
"Whisper model '%s' downloaded and loaded in %.1fs",
|
||||
WHISPER_MODEL_SIZE,
|
||||
elapsed,
|
||||
)
|
||||
return model
|
||||
|
||||
logger.info("Loading model from %s...", model_path)
|
||||
logger.info(
|
||||
"Loading Whisper model '%s' from %s...",
|
||||
WHISPER_MODEL_SIZE,
|
||||
model_path,
|
||||
)
|
||||
return WhisperModel(
|
||||
str(model_path),
|
||||
device="cpu",
|
||||
@@ -53,6 +88,44 @@ def _load_model() -> WhisperModel:
|
||||
)
|
||||
|
||||
|
||||
async def preload_model() -> None:
|
||||
"""Download/load the Whisper model at bot startup if it is not already loaded."""
|
||||
global _model
|
||||
async with _model_lock:
|
||||
if _model is not None:
|
||||
logger.info("Whisper model is already loaded, skipping preload")
|
||||
return
|
||||
|
||||
if _is_model_downloaded():
|
||||
logger.info(
|
||||
"Whisper model '%s' found locally at %s",
|
||||
WHISPER_MODEL_SIZE,
|
||||
_get_model_path(),
|
||||
)
|
||||
else:
|
||||
free_space = _get_free_disk_space_bytes(MODELS_DIR)
|
||||
required_space = WHISPER_MIN_FREE_SPACE_BYTES
|
||||
logger.info(
|
||||
"Free disk space: %.2f GB, required: %.2f GB",
|
||||
free_space / (1024 ** 3),
|
||||
required_space / (1024 ** 3),
|
||||
)
|
||||
if free_space < required_space:
|
||||
raise RuntimeError(
|
||||
f"Not enough disk space to download Whisper model '{WHISPER_MODEL_SIZE}'. "
|
||||
f"Free: {free_space / (1024 ** 3):.2f} GB, "
|
||||
f"required: {required_space / (1024 ** 3):.2f} GB"
|
||||
)
|
||||
logger.info(
|
||||
"Whisper model '%s' will be downloaded at startup",
|
||||
WHISPER_MODEL_SIZE,
|
||||
)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
_model = await loop.run_in_executor(None, _load_model)
|
||||
logger.info("Whisper model '%s' is ready", WHISPER_MODEL_SIZE)
|
||||
|
||||
|
||||
async def _get_model() -> WhisperModel:
|
||||
"""Thread-safe lazy initializer for the Whisper model."""
|
||||
global _model
|
||||
@@ -61,7 +134,7 @@ async def _get_model() -> WhisperModel:
|
||||
if _model is None:
|
||||
loop = asyncio.get_running_loop()
|
||||
_model = await loop.run_in_executor(None, _load_model)
|
||||
logger.info("Whisper model loaded")
|
||||
logger.info("Whisper model '%s' loaded on demand", WHISPER_MODEL_SIZE)
|
||||
return _model
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user