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:
@@ -19,7 +19,9 @@ from .lyrics import LyricsMatch, search_by_identity, search_by_text
|
||||
from .metadata import TrackMetadata, read_metadata
|
||||
from .recognize import ShazamResult, recognize
|
||||
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
|
||||
|
||||
|
||||
@@ -1261,6 +1263,20 @@ def _run_staged_batch(
|
||||
if config.transcribe.enabled and not cancelled():
|
||||
todo = active()
|
||||
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))
|
||||
for i, f in enumerate(todo, 1):
|
||||
if cancelled():
|
||||
|
||||
@@ -6,11 +6,21 @@ Models are cached; aggressive VRAM cleanup happens after each file.
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
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):
|
||||
"""CUDA OOM, переживший все стратегии понижения нагрузки."""
|
||||
|
||||
@@ -110,6 +120,7 @@ class _WhisperXModels:
|
||||
if compute_type is None:
|
||||
compute_type = "int8_float16" if device == "cuda" else "int8"
|
||||
key = (model_name, device, compute_type, local_files_only)
|
||||
with _MODEL_LOCK:
|
||||
if key not in cls._cache:
|
||||
if device == "cuda":
|
||||
_enable_tf32()
|
||||
@@ -125,7 +136,10 @@ class _WhisperXModels:
|
||||
)
|
||||
|
||||
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]
|
||||
|
||||
|
||||
@@ -135,6 +149,7 @@ class _AlignModels:
|
||||
@classmethod
|
||||
def get(cls, language_code: str, device: str, model_name: Optional[str] = None):
|
||||
key = (language_code, device, model_name or "")
|
||||
with _MODEL_LOCK:
|
||||
if key not in cls._cache:
|
||||
import whisperx
|
||||
print(f" [ALIGN] loading align model for {language_code} "
|
||||
@@ -147,7 +162,10 @@ class _AlignModels:
|
||||
return whisperx.load_align_model(**kwargs)
|
||||
|
||||
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]
|
||||
|
||||
|
||||
@@ -162,6 +180,7 @@ class _GigaAMModels:
|
||||
"auto": ["CUDAExecutionProvider", "CPUExecutionProvider"],
|
||||
}.get(provider, ["CPUExecutionProvider"])
|
||||
key = tuple(providers)
|
||||
with _MODEL_LOCK:
|
||||
if key not in cls._cache:
|
||||
import onnx_asr
|
||||
|
||||
@@ -172,6 +191,28 @@ class _GigaAMModels:
|
||||
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(
|
||||
audio_path: Path,
|
||||
engine: str = "whisperx",
|
||||
|
||||
Reference in New Issue
Block a user