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, WHISPER_MODEL_SIZE, WHISPER_MIN_FREE_SPACE_BYTES logger = logging.getLogger(__name__) _model: WhisperModel | None = None _model_lock = asyncio.Lock() BASE_DIR = Path(__file__).resolve().parent.parent.parent MODELS_DIR = BASE_DIR / "models" / "faster-whisper" VOICE_DIR = BASE_DIR / "bot" / "data" / "voice" def _setup_proxy() -> None: if PROXY_ENABLED and PROXY_URL: os.environ["ALL_PROXY"] = PROXY_URL 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 = _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( "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 Whisper model '%s' from %s...", WHISPER_MODEL_SIZE, model_path, ) return WhisperModel( str(model_path), device="cpu", cpu_threads=4, compute_type="int8", ) 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 if _model is None: async with _model_lock: if _model is None: loop = asyncio.get_running_loop() _model = await loop.run_in_executor(None, _load_model) logger.info("Whisper model '%s' loaded on demand", WHISPER_MODEL_SIZE) return _model async def download_voice(bot: Bot, file_id: str) -> str: VOICE_DIR.mkdir(parents=True, exist_ok=True) file = await bot.get_file(file_id) safe_id = re.sub(r"[^A-Za-z0-9_-]", "_", file_id) path = str(VOICE_DIR / f"{safe_id}.ogg") await bot.download_file(file.file_path, destination=path) logger.info("Voice downloaded: %s", path) return path async def convert_to_wav(input_path: str) -> str: output_path = str(Path(input_path).with_suffix(".wav")) proc = await asyncio.create_subprocess_exec( "ffmpeg", "-i", input_path, "-ar", "16000", "-ac", "1", "-c:a", "pcm_s16le", output_path, "-y", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) _, stderr = await proc.communicate() if proc.returncode != 0: logger.error("ffmpeg error: %s", stderr.decode(errors="replace")) raise RuntimeError("ffmpeg conversion failed") logger.info("Converted to WAV: %s", output_path) return output_path async def transcribe_audio(file_path: str) -> str: model = await _get_model() loop = asyncio.get_running_loop() def _transcribe() -> str: segments, _info = model.transcribe(file_path, language="ru", beam_size=5) return " ".join(seg.text for seg in segments).strip() text = await loop.run_in_executor(None, _transcribe) logger.info("Transcription result (%d chars): %s...", len(text), text[:100]) return text def normalize_text(text: str) -> str: text = text.strip() if not text: return text text = text[0].upper() + text[1:] if text[-1] not in ".!?…": text += "." text = re.sub(r"\s+", " ", text) return text async def cleanup_files(*paths: str | None) -> None: for path in paths: if path and os.path.exists(path): try: await asyncio.to_thread(os.unlink, path) logger.debug("Cleaned up: %s", path) except Exception as exc: logger.warning("Cleanup failed for %s: %s", path, exc)