Initial commit: Nat20 Notes — TTRPG session transcription & summarization

backend/app/pipeline/transcribe.py

@@ -0,0 +1,104 @@
+"""
+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