From 60846e3c6cae867453482a4ac22a644d2fe8c808 Mon Sep 17 00:00:00 2001 From: jadssss Date: Wed, 28 Oct 2026 09:54:06 +0300 Subject: [PATCH] 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 --- core/pipeline.py | 18 +++++++- core/transcribe.py | 107 +++++++++++++++++++++++++++++++-------------- 2 files changed, 91 insertions(+), 34 deletions(-) diff --git a/core/pipeline.py b/core/pipeline.py index 823dc1a..13de14e 100644 --- a/core/pipeline.py +++ b/core/pipeline.py @@ -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(): diff --git a/core/transcribe.py b/core/transcribe.py index d8a1cd9..ab0572f 100644 --- a/core/transcribe.py +++ b/core/transcribe.py @@ -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,23 +120,27 @@ 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) - if key not in cls._cache: - if device == "cuda": - _enable_tf32() - import whisperx - print(f" [ASR] loading WhisperX {model_name} on {device} " - f"({compute_type}, local_files_only={local_files_only})...") + with _MODEL_LOCK: + if key not in cls._cache: + if device == "cuda": + _enable_tf32() + import whisperx + print(f" [ASR] loading WhisperX {model_name} on {device} " + f"({compute_type}, local_files_only={local_files_only})...") - def load(): - return whisperx.load_model( - model_name, device, - compute_type=compute_type, - local_files_only=local_files_only, - ) + def load(): + return whisperx.load_model( + model_name, device, + compute_type=compute_type, + local_files_only=local_files_only, + ) - cls._cache[key] = _retry(load) - print(" [ASR] model loaded.") - return cls._cache[key] + cls._cache[key] = _retry(load) + 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] class _AlignModels: @@ -135,20 +149,24 @@ class _AlignModels: @classmethod def get(cls, language_code: str, device: str, model_name: Optional[str] = None): key = (language_code, device, model_name or "") - if key not in cls._cache: - import whisperx - print(f" [ALIGN] loading align model for {language_code} " - f"({model_name or 'default'}) on {device}...") + with _MODEL_LOCK: + if key not in cls._cache: + import whisperx + print(f" [ALIGN] loading align model for {language_code} " + f"({model_name or 'default'}) on {device}...") - def load(): - kwargs = {"language_code": language_code, "device": device} - if model_name: - kwargs["model_name"] = model_name - return whisperx.load_align_model(**kwargs) + def load(): + kwargs = {"language_code": language_code, "device": device} + if model_name: + kwargs["model_name"] = model_name + return whisperx.load_align_model(**kwargs) - cls._cache[key] = _retry(load) - print(" [ALIGN] align model loaded.") - return cls._cache[key] + cls._cache[key] = _retry(load) + 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] class _GigaAMModels: @@ -162,14 +180,37 @@ class _GigaAMModels: "auto": ["CUDAExecutionProvider", "CPUExecutionProvider"], }.get(provider, ["CPUExecutionProvider"]) key = tuple(providers) - if key not in cls._cache: - import onnx_asr + with _MODEL_LOCK: + if key not in cls._cache: + import onnx_asr - def load(): - return onnx_asr.load_model("gigaam-v3-ctc", providers=list(providers)) + def load(): + return onnx_asr.load_model("gigaam-v3-ctc", providers=list(providers)) - cls._cache[key] = _retry(load) - return cls._cache[key] + cls._cache[key] = _retry(load) + 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(