make long-file CUDA OOM recoverable per file

- replace permanent _CUDA_BROKEN flag with CudaOOMError and retry helper
- retry WhisperX with configured settings, then batch=1 without VAD, then once more
- keep the loaded model cached between retries instead of reloading it
- actually pass vad_filter through transcribe() and use it in aligned ASR
- skip only the OOM file in sequential, staged and LRC pipelines
- mark OOM files with oom_skipped / lrc+oom for the summary
- fix staged vocal separation writing input_dir/vocals.wav when convert is off
- route original inputs and resumed files through safe per-file separation
- validate ffmpeg before starting the LRC task
This commit is contained in:
jadssss
2026-10-27 23:03:29 +03:00
parent 999934f174
commit f88ca913ce
4 changed files with 86 additions and 47 deletions

View File

@@ -14,7 +14,7 @@ from .metadata import read_metadata
from .pipeline import (
FileResult, _artifact_belongs_to, _save_unrecognized, collect_audio_files,
)
from .transcribe import transcribe_aligned
from .transcribe import CudaOOMError, transcribe_aligned
LogFn = Callable[[str], None]
@@ -155,6 +155,10 @@ def run_lrc_pipeline(
metadata=read_metadata(f),
)
results.append(FileResult(f, None, track_out, None, "lrc"))
except CudaOOMError as e:
error = f"CudaOOMError: {e}"
log(f" ⊘ {f.name}: OOM — пропущен, идём дальше")
results.append(FileResult(f, None, None, error, "lrc+oom"))
except Exception as e: # noqa: BLE001
error = f"{type(e).__name__}: {e}"
log(f" ✗ Ошибка: {error}")

View File

@@ -18,7 +18,7 @@ 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 transcribe
from .transcribe import CudaOOMError, transcribe
from .lrc import is_lrc, read_embedded_lyrics
@@ -729,6 +729,10 @@ def _run_sequential(
log("⏹ Отменено")
was_cancelled = True
break
except CudaOOMError as e:
error = f"CudaOOMError: {e}"
source = "oom_skipped"
log(f" ⊘ {f.name}: OOM после 3 попыток — пропущен, идём дальше")
except Exception as e: # noqa: BLE001
error = f"{type(e).__name__}: {e}"
log(f" ✗ Ошибка: {error}")
@@ -1010,7 +1014,22 @@ def _run_staged_batch(
def separate_single(f: Path) -> None:
try:
src = state[f]["audio_path"]
target = src.parent / "vocals.wav"
# Never write next to the original input: with
# convert.enabled=False src == f, so every track would
# overwrite input_dir/vocals.wav. Only reuse the track
# output folder when the source already lives there.
track_dir = output_root / f.stem
if src.parent == track_dir:
# Already in the persistent track folder (resume).
target = src.parent / "vocals.wav"
elif src.parent == f.parent:
# src is the original input file (convert disabled).
# Never write next to it: all tracks share input_dir
# and would overwrite input_dir/vocals.wav.
target = _mkdtemp("at_sep_") / "vocals.wav"
else:
# Freshly converted file in its own unique temp dir.
target = src.parent / "vocals.wav"
with tempfile.TemporaryDirectory() as tmp:
raw = separate_vocals(
src, Path(tmp), model=config.separate.model,
@@ -1027,11 +1046,15 @@ def _run_staged_batch(
# Resumed convert.wav files already live in output_root/<stem>,
# while freshly converted files live in unique temp dirs.
# Original inputs (convert disabled) must also go through the
# one-by-one path: batching them shares input_dir and can
# collide on vocals.wav.
# unblend's `{parent}` batch layout would mix those roots, so
# separate resumed files one by one and batch the rest.
# separate resumed/original files one by one and batch the rest.
resumed = [
f for f in todo
if state[f]["audio_path"].parent == output_root / f.stem
or state[f]["audio_path"].parent == f.parent
]
resumed_set = set(resumed)
batchable = [f for f in todo if f not in resumed_set]
@@ -1111,6 +1134,10 @@ def _run_staged_batch(
state[f]["text"] = text
state[f]["source"] = config.transcribe.engine
log(f" ✓ {f.name} ({len(text)} символов)")
except CudaOOMError as e:
state[f]["error"] = f"CudaOOMError: {e}"
state[f]["source"] = "oom_skipped"
log(f" ⊘ {f.name}: OOM — пропущен")
except Exception as e: # noqa: BLE001
state[f]["error"] = f"{type(e).__name__}: {e}"
log(f" ✗ {f.name}: {e}")

View File

@@ -11,10 +11,8 @@ from pathlib import Path
from typing import Optional
# Set permanently once the CUDA context has died (e.g. after an unrecoverable
# OOM). Every later CUDA transcription fails instantly instead of re-loading
# the model forever and hanging the pipeline.
_CUDA_BROKEN = False
class CudaOOMError(RuntimeError):
"""CUDA OOM, переживший все стратегии понижения нагрузки."""
def _is_cuda_error(exc: BaseException) -> bool:
@@ -28,14 +26,6 @@ def _is_cuda_error(exc: BaseException) -> bool:
)
def _check_cuda_usable(device: str) -> None:
if device == "cuda" and _CUDA_BROKEN:
raise RuntimeError(
"CUDA context is broken (previous OOM was unrecoverable). "
"Restart the app or switch device to CPU."
)
def _resolve_device(device: str) -> str:
if device in ("cuda", "cpu"):
return device
@@ -68,6 +58,36 @@ def _free_vram() -> None:
pass
def _model_transcribe_with_retry(
model_obj, audio, language, batch_size, vad_filter, name="<audio>",
):
"""Try progressively lighter settings so one long file can still finish.
The loaded model is intentionally kept in the cache between attempts:
reloading it costs minutes and does not help with an OOM caused by a long
audio file. If even batch=1 without VAD fails, raise CudaOOMError so the
caller can skip this file and continue with the rest of the batch.
"""
attempts = [
(batch_size, vad_filter), # as configured
(1, False), # minimum memory: no Pyannote VAD, batch=1
(1, False), # one more time after empty_cache()
]
last: Optional[Exception] = None
for i, (bs, vad) in enumerate(attempts, 1):
try:
return model_obj.transcribe(
audio, batch_size=bs, language=language, vad_filter=vad,
)
except Exception as e: # noqa: BLE001
if not _is_cuda_error(e):
raise
last = e
print(f" [ASR] OOM on {name} (attempt {i}/3, bs={bs}, vad={vad}): {e}")
_free_vram()
raise CudaOOMError(f"OOM after 3 attempts on {name}: {last}") from last
class _WhisperXModels:
_cache: dict = {}
@@ -149,37 +169,13 @@ def transcribe(
local_files_only: bool = False,
vad_filter: bool = True,
) -> str:
global _CUDA_BROKEN
device = _resolve_device(device)
_check_cuda_usable(device)
if engine == "whisperx":
try:
return _transcribe_whisperx(
audio_path, model, language, device, compute_type,
batch_size, local_files_only, vad_filter,
)
except Exception as e:
if not _is_cuda_error(e):
raise
print(f" [ASR] CUDA error: {e}. Clearing model cache and retrying...")
_WhisperXModels._cache.clear()
_free_vram()
try:
return _transcribe_whisperx(
audio_path, model, language, device, compute_type,
batch_size, local_files_only, vad_filter,
)
except Exception as retry_error:
if not _is_cuda_error(retry_error):
raise
_CUDA_BROKEN = True
print(" [ASR] CUDA context is unrecoverable — "
"GPU disabled for the rest of this run.")
raise RuntimeError(
"CUDA context is broken after OOM recovery failed. "
"Restart the app or switch device to CPU."
) from retry_error
return _transcribe_whisperx(
audio_path, model, language, device, compute_type,
batch_size, local_files_only, vad_filter,
)
if engine == "gigaam":
return _transcribe_gigaam(audio_path, gigaam_provider)
@@ -199,7 +195,9 @@ def _transcribe_whisperx(
import whisperx
model = _WhisperXModels.get(model_name, device, compute_type, local_files_only)
audio = whisperx.load_audio(str(audio_path))
result = model.transcribe(audio, batch_size=batch_size, language=language)
result = _model_transcribe_with_retry(
model, audio, language, batch_size, vad_filter, name=audio_path.name,
)
text = " ".join(seg["text"].strip() for seg in result["segments"]).strip()
del audio, result
_free_vram()
@@ -220,6 +218,7 @@ def transcribe_aligned(
batch_size: int = 8,
align_model: Optional[str] = None,
local_files_only: bool = False,
vad_filter: bool = True,
) -> tuple[list[dict], str]:
"""Transcribe + word-level align. Returns (aligned_segments, detected_language).
@@ -228,10 +227,11 @@ def transcribe_aligned(
"""
import whisperx
device = _resolve_device(device)
_check_cuda_usable(device)
model_obj = _WhisperXModels.get(model, device, compute_type, local_files_only)
audio = whisperx.load_audio(str(audio_path))
result = model_obj.transcribe(audio, batch_size=batch_size, language=language)
result = _model_transcribe_with_retry(
model_obj, audio, language, batch_size, vad_filter, name=audio_path.name,
)
detected = result.get("language") or language or "en"