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,23 +120,27 @@ 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)
if key not in cls._cache: with _MODEL_LOCK:
if device == "cuda": if key not in cls._cache:
_enable_tf32() if device == "cuda":
import whisperx _enable_tf32()
print(f" [ASR] loading WhisperX {model_name} on {device} " import whisperx
f"({compute_type}, local_files_only={local_files_only})...") print(f" [ASR] loading WhisperX {model_name} on {device} "
f"({compute_type}, local_files_only={local_files_only})...")
def load(): def load():
return whisperx.load_model( return whisperx.load_model(
model_name, device, model_name, device,
compute_type=compute_type, compute_type=compute_type,
local_files_only=local_files_only, local_files_only=local_files_only,
) )
cls._cache[key] = _retry(load) cls._cache[key] = _retry(load)
print(" [ASR] model loaded.") print(" [ASR] model loaded once — subsequent files will reuse it.")
return cls._cache[key] 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]
class _AlignModels: class _AlignModels:
@@ -135,20 +149,24 @@ 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 "")
if key not in cls._cache: with _MODEL_LOCK:
import whisperx if key not in cls._cache:
print(f" [ALIGN] loading align model for {language_code} " import whisperx
f"({model_name or 'default'}) on {device}...") print(f" [ALIGN] loading align model for {language_code} "
f"({model_name or 'default'}) on {device}...")
def load(): def load():
kwargs = {"language_code": language_code, "device": device} kwargs = {"language_code": language_code, "device": device}
if model_name: if model_name:
kwargs["model_name"] = model_name kwargs["model_name"] = model_name
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.")
return cls._cache[key] 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]
class _GigaAMModels: class _GigaAMModels:
@@ -162,14 +180,37 @@ class _GigaAMModels:
"auto": ["CUDAExecutionProvider", "CPUExecutionProvider"], "auto": ["CUDAExecutionProvider", "CPUExecutionProvider"],
}.get(provider, ["CPUExecutionProvider"]) }.get(provider, ["CPUExecutionProvider"])
key = tuple(providers) key = tuple(providers)
if key not in cls._cache: with _MODEL_LOCK:
import onnx_asr if key not in cls._cache:
import onnx_asr
def load(): def load():
return onnx_asr.load_model("gigaam-v3-ctc", providers=list(providers)) return onnx_asr.load_model("gigaam-v3-ctc", providers=list(providers))
cls._cache[key] = _retry(load) cls._cache[key] = _retry(load)
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(