Nat20-Notes/backend/app/pipeline/transcribe.py

"""
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__)

# Module-level model cache — kept alive across jobs so back-to-back
# transcriptions don't reload from disk each time.
_whisper_model = None
_whisper_config = {}  # {"model_size": str, "compute_type": str, "device": str}


def _get_model(model_size: str, compute_type: str, device: str):
    global _whisper_model, _whisper_config
    if (
        _whisper_model is not None
        and _whisper_config.get("model_size") == model_size
        and _whisper_config.get("compute_type") == compute_type
        and _whisper_config.get("device") == device
    ):
        log.info("Reusing cached Whisper model (%s, %s, %s)", model_size, compute_type, device)
        return _whisper_model
    log.info("Loading Whisper model (%s, %s, %s)", model_size, compute_type, device)
    _unload_cached()
    _whisper_model = whisperx.load_model(model_size, device, compute_type=compute_type)
    _whisper_config = {"model_size": model_size, "compute_type": compute_type, "device": device}
    return _whisper_model


def _unload_cached():
    global _whisper_model, _whisper_config
    if _whisper_model is not None:
        log.info("Unloading cached Whisper model")
        del _whisper_model
        _whisper_model = None
        _whisper_config = {}
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.synchronize()
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats()


def unload_models():
    """Release all GPU memory held by transcription models.
    Called by the job runner when no more GPU jobs are queued."""
    _unload_cached()
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.synchronize()
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()


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 = _get_model(model_size, compute_type, device)
    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)

    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)
    except Exception as e:
        raise TranscriptionError(f"Alignment step failed: {e}", cause=e)
    finally:
        try:
            del align_model
        except NameError:
            pass
        gc.collect()
        if device == "cuda":
            torch.cuda.empty_cache()

    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)
        finally:
            try:
                del diarize_model
            except NameError:
                pass
            gc.collect()
            if device == "cuda":
                torch.cuda.empty_cache()
    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