load the ASR model once per process

- keep WhisperX, align and GigaAM caches process-wide and thread-safe\n- preload the ASR model at the start of the staged stage 3\n- log that the model is loaded once and then reused\n- never clear caches between files, batches or OOM retries
This commit is contained in:
jadssss
2026-10-28 09:54:06 +03:00
parent e23acf1555
commit 60846e3c6c
2 changed files with 91 additions and 34 deletions

View File

@@ -19,7 +19,9 @@ from .lyrics import LyricsMatch, search_by_identity, search_by_text
from .metadata import TrackMetadata, read_metadata from .metadata import TrackMetadata, read_metadata
from .recognize import ShazamResult, recognize from .recognize import ShazamResult, recognize
from .separate import separate_vocals, separate_vocals_batch from .separate import separate_vocals, separate_vocals_batch
from .transcribe import CudaOOMError, transcribe from .transcribe import (
CudaOOMError, preload_transcription_model, transcribe,
)
from .lrc import is_lrc, read_embedded_lyrics from .lrc import is_lrc, read_embedded_lyrics
@@ -1261,6 +1263,20 @@ def _run_staged_batch(
if config.transcribe.enabled and not cancelled(): if config.transcribe.enabled and not cancelled():
todo = active() todo = active()
log(f"\n══ Этап 3 · Транскрибация ({len(todo)} файлов) ══") log(f"\n══ Этап 3 · Транскрибация ({len(todo)} файлов) ══")
if todo:
# Load once, after separation has finished. Every later file
# and every later batch reuses this instance.
try:
preload_transcription_model(
engine=config.transcribe.engine,
model=config.transcribe.model,
device=config.transcribe.device,
compute_type=config.transcribe.compute_type,
gigaam_provider=config.transcribe.gigaam_provider,
local_files_only=config.transcribe.local_files_only,
)
except Exception as e: # noqa: BLE001
log(f" ⚠ Не удалось загрузить модель заранее: {e}")
progress("transcribe", 0, len(todo)) progress("transcribe", 0, len(todo))
for i, f in enumerate(todo, 1): for i, f in enumerate(todo, 1):
if cancelled(): if cancelled():

View File

@@ -6,11 +6,21 @@ Models are cached; aggressive VRAM cleanup happens after each file.
from __future__ import annotations from __future__ import annotations
import gc import gc
import threading
import time import time
from pathlib import Path from pathlib import Path
from typing import Optional from typing import Optional
# The model caches below are process-wide. They are intentionally never
# cleared between files or batches: reloading WhisperX takes minutes and is
# what made runs look like they reloaded the model every batch. The lock makes
# the "load once, reuse forever" guarantee safe even if ASR is ever called
# from multiple threads.
_MODEL_LOCK = threading.RLock()
_REUSE_LOGGED: set = set()
class CudaOOMError(RuntimeError): class CudaOOMError(RuntimeError):
"""CUDA OOM, переживший все стратегии понижения нагрузки.""" """CUDA OOM, переживший все стратегии понижения нагрузки."""
@@ -110,6 +120,7 @@ class _WhisperXModels:
if compute_type is None: if compute_type is None:
compute_type = "int8_float16" if device == "cuda" else "int8" compute_type = "int8_float16" if device == "cuda" else "int8"
key = (model_name, device, compute_type, local_files_only) key = (model_name, device, compute_type, local_files_only)
with _MODEL_LOCK:
if key not in cls._cache: if key not in cls._cache:
if device == "cuda": if device == "cuda":
_enable_tf32() _enable_tf32()
@@ -125,7 +136,10 @@ class _WhisperXModels:
) )
cls._cache[key] = _retry(load) cls._cache[key] = _retry(load)
print(" [ASR] model loaded.") print(" [ASR] model loaded once — subsequent files will reuse it.")
elif ("whisperx",) + key not in _REUSE_LOGGED:
_REUSE_LOGGED.add(("whisperx",) + key)
print(f" [ASR] reusing loaded WhisperX {model_name} ({device}).")
return cls._cache[key] return cls._cache[key]
@@ -135,6 +149,7 @@ class _AlignModels:
@classmethod @classmethod
def get(cls, language_code: str, device: str, model_name: Optional[str] = None): def get(cls, language_code: str, device: str, model_name: Optional[str] = None):
key = (language_code, device, model_name or "") key = (language_code, device, model_name or "")
with _MODEL_LOCK:
if key not in cls._cache: if key not in cls._cache:
import whisperx import whisperx
print(f" [ALIGN] loading align model for {language_code} " print(f" [ALIGN] loading align model for {language_code} "
@@ -147,7 +162,10 @@ class _AlignModels:
return whisperx.load_align_model(**kwargs) return whisperx.load_align_model(**kwargs)
cls._cache[key] = _retry(load) cls._cache[key] = _retry(load)
print(" [ALIGN] align model loaded.") print(" [ALIGN] align model loaded once — reusing it.")
elif ("align",) + key not in _REUSE_LOGGED:
_REUSE_LOGGED.add(("align",) + key)
print(f" [ALIGN] reusing loaded align model for {language_code}.")
return cls._cache[key] return cls._cache[key]
@@ -162,6 +180,7 @@ class _GigaAMModels:
"auto": ["CUDAExecutionProvider", "CPUExecutionProvider"], "auto": ["CUDAExecutionProvider", "CPUExecutionProvider"],
}.get(provider, ["CPUExecutionProvider"]) }.get(provider, ["CPUExecutionProvider"])
key = tuple(providers) key = tuple(providers)
with _MODEL_LOCK:
if key not in cls._cache: if key not in cls._cache:
import onnx_asr import onnx_asr
@@ -172,6 +191,28 @@ class _GigaAMModels:
return cls._cache[key] return cls._cache[key]
def preload_transcription_model(
engine: str = "whisperx",
model: str = "large-v3",
device: str = "auto",
compute_type: Optional[str] = None,
gigaam_provider: str = "auto",
local_files_only: bool = False,
) -> None:
"""Load the ASR model once so the per-file loop only reuses it.
All model caches are process-wide and are never cleared between files,
batches or pipeline runs. This helper just makes the first load explicit.
"""
device = _resolve_device(device)
if engine == "whisperx":
_WhisperXModels.get(model, device, compute_type, local_files_only)
elif engine == "gigaam":
_GigaAMModels.get(gigaam_provider)
else:
raise ValueError(f"Unknown ASR engine: {engine!r}")
def transcribe( def transcribe(
audio_path: Path, audio_path: Path,
engine: str = "whisperx", engine: str = "whisperx",