105 lines
3.8 KiB
Python
105 lines
3.8 KiB
Python
"""
|
|
Transcription + diarization using WhisperX as a Python library.
|
|
"""
|
|
import gc
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import whisperx
|
|
|
|
from ..errors import TranscriptionError, DiarizationError
|
|
from ..logging_config import get_logger
|
|
|
|
log = get_logger(__name__)
|
|
|
|
|
|
def transcribe_and_diarize(
|
|
audio_path: Path,
|
|
output_json: Path,
|
|
model_size: str,
|
|
compute_type: str,
|
|
hf_token: str,
|
|
progress_cb=None,
|
|
) -> dict:
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
if device == "cpu":
|
|
log.warning("No CUDA GPU available - running on CPU, this will be much slower")
|
|
|
|
def report(msg):
|
|
log.info(msg)
|
|
if progress_cb:
|
|
progress_cb(msg)
|
|
|
|
report("Loading transcription model...")
|
|
try:
|
|
model = whisperx.load_model(model_size, device, compute_type=compute_type)
|
|
except RuntimeError as e:
|
|
if "out of memory" in str(e).lower():
|
|
raise TranscriptionError(
|
|
f"Ran out of GPU memory loading the '{model_size}' model. "
|
|
f"Try a smaller model size in Settings (e.g. 'small' or 'medium'), "
|
|
f"or check nothing else is using the GPU right now.", cause=e
|
|
)
|
|
raise TranscriptionError(f"Failed to load Whisper model '{model_size}': {e}", cause=e)
|
|
except Exception as e:
|
|
raise TranscriptionError(f"Failed to load Whisper model '{model_size}': {e}", cause=e)
|
|
|
|
try:
|
|
report("Loading audio...")
|
|
audio = whisperx.load_audio(str(audio_path))
|
|
|
|
report("Transcribing...")
|
|
result = model.transcribe(audio, batch_size=16)
|
|
except RuntimeError as e:
|
|
if "out of memory" in str(e).lower():
|
|
raise TranscriptionError(
|
|
f"Ran out of GPU memory during transcription. Try a smaller Whisper model size in Settings.", cause=e
|
|
)
|
|
raise TranscriptionError(f"Transcription failed: {e}", cause=e)
|
|
except Exception as e:
|
|
raise TranscriptionError(f"Transcription failed: {e}", cause=e)
|
|
finally:
|
|
del model
|
|
gc.collect()
|
|
if device == "cuda":
|
|
torch.cuda.empty_cache()
|
|
|
|
try:
|
|
report("Aligning...")
|
|
align_model, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
|
|
result = whisperx.align(result["segments"], align_model, metadata, audio, device, return_char_alignments=False)
|
|
del align_model
|
|
gc.collect()
|
|
if device == "cuda":
|
|
torch.cuda.empty_cache()
|
|
except Exception as e:
|
|
raise TranscriptionError(f"Alignment step failed: {e}", cause=e)
|
|
|
|
if hf_token:
|
|
try:
|
|
report("Diarizing...")
|
|
diarize_model = whisperx.diarize.DiarizationPipeline(use_auth_token=hf_token, device=device)
|
|
diarize_segments = diarize_model(audio)
|
|
result = whisperx.assign_word_speakers(diarize_segments, result)
|
|
except Exception as e:
|
|
msg = str(e)
|
|
if "gated" in msg.lower() or "403" in msg or "access" in msg.lower():
|
|
raise DiarizationError(
|
|
"HuggingFace rejected access to the diarization model. Make sure you've accepted the "
|
|
"terms on the gated model page (linked in the error below) using the SAME account "
|
|
f"your HF token belongs to.\n{msg}", cause=e
|
|
)
|
|
raise DiarizationError(f"Diarization failed: {msg}", cause=e)
|
|
else:
|
|
report("No HF token set - skipping speaker diarization (all speech will be unattributed)")
|
|
|
|
try:
|
|
output_json.parent.mkdir(parents=True, exist_ok=True)
|
|
output_json.write_text(json.dumps(result, ensure_ascii=False), encoding='utf-8')
|
|
except OSError as e:
|
|
raise TranscriptionError(f"Transcribed successfully, but failed to write output to {output_json}: {e}", cause=e)
|
|
|
|
report("Done")
|
|
return result
|