f023653c71
- add align_text_to_audio(): forced alignment without ASR, fast - add _StableTSModels cache; detect language via stable-ts when language=None - LRC stage: use existing text.txt, copy it if already LRC, fall back to ASR only - filter junk lyrics from tags and ASR hallucinations (promo phrases, tics) - mark empty ASR results with metadata.empty_text and skip them on resume - warn in GUI when LRC is run with transcription disabled - add mark_no_speech.py one-shot script for existing tracks
484 lines
16 KiB
Python
484 lines
16 KiB
Python
"""Speech recognition via WhisperX or GigaAM (onnx-asr).
|
|
|
|
Also exposes `transcribe_aligned()` for word-level alignment (used by LRC).
|
|
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()
|
|
|
|
# CUDA context failures that cannot be recovered by retrying or empty_cache().
|
|
# They require restarting the app (or switching to CPU).
|
|
_DEAD_CUDA_MARKERS = (
|
|
"invalid device ordinal",
|
|
"cudaerrorinvaliddevice",
|
|
"device-side assert",
|
|
"unspecified launch failure",
|
|
)
|
|
|
|
|
|
def _is_dead_cuda(exc: BaseException) -> bool:
|
|
err = str(exc).lower()
|
|
return any(m in err for m in _DEAD_CUDA_MARKERS)
|
|
|
|
|
|
_WhisperXModels._cache.clear()
|
|
_AlignModels._cache.clear()
|
|
_GigaAMModels._cache.clear()
|
|
_StableTSModels._cache.clear()
|
|
_REUSE_LOGGED.clear()
|
|
|
|
|
|
# Phrases that are strong signals of ASR hallucination (end credits)
|
|
_HALLUCINATION_PHRASES = (
|
|
"продолжение следует",
|
|
"спасибо за внимание",
|
|
"спасибо за просмотр",
|
|
"спасибо за подписку",
|
|
"подписывайтесь на канал",
|
|
"ставьте лайки",
|
|
"не забудьте подписаться",
|
|
"субтитры создавал",
|
|
"субтитры сделал",
|
|
"субтитры подготовил",
|
|
"редактор субтитров",
|
|
"корректор",
|
|
"thank you for watching",
|
|
"thanks for watching",
|
|
"please subscribe",
|
|
"subscribe to my channel",
|
|
"like and subscribe",
|
|
)
|
|
|
|
|
|
def _is_hallucination(text: str) -> bool:
|
|
"""Return True if the text looks like an ASR hallucination.
|
|
|
|
Checks:
|
|
- known call-to-action / promo phrases;
|
|
- repeated-word tic ("Музыка Музыка").
|
|
"""
|
|
if not text:
|
|
return False
|
|
t = " ".join(text.lower().split()).strip(" .!?…,;:\"'«»")
|
|
for phrase in _HALLUCINATION_PHRASES:
|
|
if phrase in t:
|
|
return True
|
|
words = t.split()
|
|
if len(words) == 2 and words[0] == words[1]:
|
|
return True
|
|
return False
|
|
|
|
|
|
def _detect_language(model_obj, audio_or_path, fallback: str = "ru") -> str:
|
|
"""Detect the spoken language from audio using the ASR model itself.
|
|
|
|
Returns a lowercase ISO code or `fallback` on failure.
|
|
"""
|
|
try:
|
|
lang, _prob = model_obj.detect_language(audio_or_path)
|
|
if lang:
|
|
return lang.lower()
|
|
except Exception:
|
|
pass
|
|
return fallback
|
|
|
|
|
|
class _StableTSModels:
|
|
_cache: dict = {}
|
|
|
|
@classmethod
|
|
def get(cls, model_name: str, device: str, compute_type: Optional[str],
|
|
local_files_only: bool = False) -> Any:
|
|
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:
|
|
import stable_whisper
|
|
|
|
print(f" [ALIGN] loading stable-ts {model_name} on {device} "
|
|
f"({compute_type}, local_files_only={local_files_only})...")
|
|
|
|
def load() -> 'faster_whisper.WhisperModel':
|
|
return stable_whisper.load_faster_whisper(
|
|
model_name, device=device, compute_type=compute_type,
|
|
local_files_only=local_files_only,
|
|
)
|
|
|
|
cls._cache[key] = _retry(load)
|
|
print(" [ALIGN] stable-ts model loaded once — reuse.")
|
|
return cls._cache[key]
|
|
|
|
|
|
def _clear_model_caches() -> None:
|
|
_WhisperXModels._cache.clear()
|
|
_AlignModels._cache.clear()
|
|
_GigaAMModels._cache.clear()
|
|
_StableTSModels._cache.clear()
|
|
_REUSE_LOGGED.clear()
|
|
|
|
|
|
class CudaOOMError(RuntimeError):
|
|
"""CUDA OOM, переживший все стратегии понижения нагрузки."""
|
|
|
|
|
|
def _is_cuda_error(exc: BaseException) -> bool:
|
|
err = str(exc).lower()
|
|
return (
|
|
"out of memory" in err
|
|
or "cuda" in err
|
|
or "device ordinal" in err
|
|
or "cublas" in err
|
|
or "cudnn" in err
|
|
)
|
|
|
|
|
|
def _resolve_device(device: str) -> str:
|
|
if device in ("cuda", "cpu"):
|
|
return device
|
|
try:
|
|
import torch
|
|
return "cuda" if torch.cuda.is_available() else "cpu"
|
|
except Exception:
|
|
return "cpu"
|
|
|
|
|
|
def _retry(fn, attempts: int = 3, base_delay: float = 2.0):
|
|
last = None
|
|
for i in range(attempts):
|
|
try:
|
|
return fn()
|
|
except Exception as e:
|
|
last = e
|
|
if i < attempts - 1:
|
|
time.sleep(base_delay * (i + 1))
|
|
raise last
|
|
|
|
|
|
def _free_vram() -> None:
|
|
gc.collect()
|
|
try:
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
except Exception: # noqa: BLE001 — cleanup must never mask the original error
|
|
pass
|
|
|
|
|
|
def _enable_tf32() -> None:
|
|
"""Enable TF32 matmul on Ampere+ GPUs (RTX 30/40).
|
|
|
|
Called lazily right before a WhisperX model is loaded, so runs that do not
|
|
use ASR never import Torch or create a CUDA context at startup.
|
|
"""
|
|
try:
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
torch.backends.cudnn.allow_tf32 = True
|
|
except Exception: # noqa: BLE001 — purely an optimization
|
|
pass
|
|
|
|
|
|
def _model_transcribe_with_retry(
|
|
model_obj, audio, language, batch_size, 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. VAD is not touched here — WhisperX expects it at model-load
|
|
time (``vad_model``/``vad_options``), not per transcription call. A dead
|
|
CUDA context is never retried: caches are dropped and CudaOOMError is
|
|
raised immediately. If even batch=1 fails with a normal OOM, raise
|
|
CudaOOMError so the caller can skip this file and continue with the rest
|
|
of the batch.
|
|
"""
|
|
attempts = [batch_size, 1]
|
|
last: Optional[Exception] = None
|
|
for i, bs in enumerate(attempts, 1):
|
|
try:
|
|
return model_obj.transcribe(
|
|
audio, batch_size=bs, language=language,
|
|
)
|
|
except Exception as e: # noqa: BLE001
|
|
if _is_dead_cuda(e):
|
|
_clear_model_caches()
|
|
raise CudaOOMError(
|
|
f"CUDA context died on {name}: {e}. "
|
|
"GPU state is unrecoverable — restart the app or switch to CPU."
|
|
) from e
|
|
if not _is_cuda_error(e):
|
|
raise
|
|
last = e
|
|
print(f" [ASR] OOM on {name} (attempt {i}/{len(attempts)}, bs={bs}): {e}")
|
|
_free_vram()
|
|
raise CudaOOMError(
|
|
f"OOM after {len(attempts)} attempts on {name}: {last}"
|
|
) from last
|
|
|
|
|
|
class _WhisperXModels:
|
|
_cache: dict = {}
|
|
|
|
@classmethod
|
|
def get(cls, model_name: str, device: str, compute_type: Optional[str],
|
|
local_files_only: bool = False):
|
|
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()
|
|
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,
|
|
)
|
|
|
|
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:
|
|
_cache: dict = {}
|
|
|
|
@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} "
|
|
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)
|
|
|
|
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:
|
|
_cache: dict = {}
|
|
|
|
@classmethod
|
|
def get(cls, provider: str):
|
|
providers = {
|
|
"cuda": ["CUDAExecutionProvider"],
|
|
"cpu": ["CPUExecutionProvider"],
|
|
"auto": ["CUDAExecutionProvider", "CPUExecutionProvider"],
|
|
}.get(provider, ["CPUExecutionProvider"])
|
|
key = tuple(providers)
|
|
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))
|
|
|
|
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(
|
|
audio_path: Path,
|
|
engine: str = "whisperx",
|
|
model: str = "large-v3",
|
|
language: Optional[str] = "ru",
|
|
device: str = "auto",
|
|
compute_type: Optional[str] = None,
|
|
batch_size: int = 8,
|
|
gigaam_provider: str = "auto",
|
|
local_files_only: bool = False,
|
|
vad_filter: bool = True,
|
|
) -> str:
|
|
device = _resolve_device(device)
|
|
|
|
if engine == "whisperx":
|
|
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)
|
|
raise ValueError(f"Unknown ASR engine: {engine!r}")
|
|
|
|
|
|
def _transcribe_whisperx(
|
|
audio_path: Path,
|
|
model_name: str,
|
|
language: Optional[str],
|
|
device: str,
|
|
compute_type: Optional[str],
|
|
batch_size: int,
|
|
local_files_only: bool,
|
|
vad_filter: bool,
|
|
) -> str:
|
|
import whisperx
|
|
model = _WhisperXModels.get(model_name, device, compute_type, local_files_only)
|
|
audio = whisperx.load_audio(str(audio_path))
|
|
result = _model_transcribe_with_retry(
|
|
model, audio, language, batch_size, name=audio_path.name,
|
|
)
|
|
|
|
# Filter out hallucinated / spammy segments.
|
|
segments = result.get("segments") or []
|
|
kept = []
|
|
for seg in segments:
|
|
t = (seg.get("text") or "").strip()
|
|
if not t or _is_hallucination(t):
|
|
continue
|
|
if seg.get("avg_logprob", 0.0) < -1.2:
|
|
continue
|
|
kept.append(seg)
|
|
text = " ".join(s["text"].strip() for s in kept).strip()
|
|
if _is_hallucination(text):
|
|
text = ""
|
|
|
|
del audio, result
|
|
_free_vram()
|
|
return text
|
|
|
|
|
|
def _transcribe_gigaam(audio_path: Path, provider: str) -> str:
|
|
model = _GigaAMModels.get(provider)
|
|
return model.recognize(str(audio_path), channel="mean").strip()
|
|
|
|
|
|
def transcribe_aligned(
|
|
audio_path: Path,
|
|
model: str = "large-v3",
|
|
language: Optional[str] = "ru",
|
|
device: str = "auto",
|
|
compute_type: Optional[str] = None,
|
|
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).
|
|
|
|
Each segment has 'start', 'end', 'text', and optionally 'words'.
|
|
Requires whisperx engine.
|
|
"""
|
|
import whisperx
|
|
device = _resolve_device(device)
|
|
model_obj = _WhisperXModels.get(model, device, compute_type, local_files_only)
|
|
audio = whisperx.load_audio(str(audio_path))
|
|
result = _model_transcribe_with_retry(
|
|
model_obj, audio, language, batch_size, name=audio_path.name,
|
|
)
|
|
|
|
detected = result.get("language") or language or "en"
|
|
|
|
try:
|
|
align_model_obj, align_metadata = _AlignModels.get(detected, device, align_model)
|
|
aligned = whisperx.align(
|
|
result["segments"], align_model_obj, align_metadata,
|
|
audio, device, return_char_alignments=False,
|
|
)
|
|
segments = aligned["segments"]
|
|
except Exception as e:
|
|
print(f" [ALIGN] failed: {e}. Falling back to raw WhisperX segments.")
|
|
segments = result["segments"]
|
|
|
|
# Filter out hallucinated segments.
|
|
segments = [s for s in segments if not _is_hallucination(s.get("text", ""))]
|
|
|
|
del audio, result
|
|
_free_vram()
|
|
return segments, detected
|
|
|
|
|
|
def align_text_to_audio(
|
|
audio_path: Path,
|
|
text: str,
|
|
language: Optional[str] = "ru",
|
|
model: str = "large-v3",
|
|
device: str = "auto",
|
|
compute_type: Optional[str] = None,
|
|
local_files_only: bool = False,
|
|
) -> tuple[list[dict], str]:
|
|
"""Forced alignment of EXISTING text to audio. No ASR.
|
|
|
|
When `language` is None, the language is auto-detected from the audio
|
|
using the alignment model itself (which is a faster-whisper model and
|
|
exposes `detect_language`).
|
|
"""
|
|
import stable_whisper as _sw
|
|
device = _resolve_device(device)
|
|
model_obj = _StableTSModels.get(model, device, compute_type, local_files_only)
|
|
|
|
if language:
|
|
lang = language
|
|
else:
|
|
lang = _detect_language(model_obj, str(audio_path), fallback="ru")
|
|
|
|
try:
|
|
result = model_obj.align(str(audio_path), text, language=lang)
|
|
except Exception as e: # noqa: BLE001
|
|
raise RuntimeError(f"align failed on {audio_path.name}: {e}") from e
|
|
|
|
segments = [
|
|
{"start": float(seg.start), "end": float(seg.end),
|
|
"text": (seg.text or "").strip()}
|
|
for seg in result.segments
|
|
if (seg.text or "").strip()
|
|
]
|
|
return segments, lang
|