visualnovel/stt.py

"""The Ear — speech to text.

Player utterances are short, so Whisper `small`/`base` transcribes near-instantly even on CPU.
On Apple Silicon, `mlx-whisper` or `whisper.cpp` use the GPU; on AMD, `whisper.cpp`
(Vulkan/ROCm) or CPU. `faster-whisper` is the simplest default and is fine on CPU for short clips.
"""

from __future__ import annotations

from typing import Protocol

from . import config


class STTBackend(Protocol):
    def transcribe(self, audio_path: str) -> str: ...


class MockSTT:
    def transcribe(self, audio_path: str) -> str:
        return ""  # the UI falls back to whatever was typed


class WhisperSTT:
    def __init__(self) -> None:
        from faster_whisper import WhisperModel  # noqa: PLC0415

        device = config.detect_device()
        # CTranslate2 runs on CUDA or CPU (no Metal/ROCm) -> use CPU off NVIDIA.
        ct2_device = "cuda" if device == "cuda" else "cpu"
        self.model = WhisperModel(config.WHISPER_SIZE, device=ct2_device, compute_type="auto")

    @staticmethod
    def _ensure_extension(audio_path: str) -> tuple[str, bool]:
        """Return (path, should_delete). Copies extensionless blobs to .webm so PyAV can probe format."""
        import pathlib
        import shutil
        import tempfile

        if pathlib.Path(audio_path).suffix.lower():
            return audio_path, False
        # Gradio saves browser recordings as bare 'blob' files — give it the right extension.
        tmp = tempfile.mktemp(suffix=".webm")
        shutil.copy2(audio_path, tmp)
        return tmp, True

    def transcribe(self, audio_path: str) -> str:
        import os

        path, temp = self._ensure_extension(audio_path)
        try:
            segments, _ = self.model.transcribe(
                path,
                beam_size=1,
                no_speech_threshold=0.6,
                condition_on_previous_text=False,
            )
            return " ".join(s.text for s in segments if s.no_speech_prob <= 0.6).strip()
        finally:
            if temp:
                os.unlink(path)


def get_stt() -> STTBackend:
    if config.USE_MOCK and not config.REAL_STT:
        return MockSTT()
    return WhisperSTT()


# Debug: record from microphone and transcribe
if __name__ == "__main__":
    import sys
    import tempfile
    import wave

    import numpy as np

    RATE = 16_000
    SECONDS = int(sys.argv[1]) if len(sys.argv) > 1 else 5

    try:
        import sounddevice as sd

        print(f"Recording {SECONDS}s … speak now!")
        audio = sd.rec(RATE * SECONDS, samplerate=RATE, channels=1, dtype="float32")
        sd.wait()
        audio = audio[:, 0]
        print("Done recording.")
    except ImportError:
        print("sounddevice not installed — using 5 s of silence as fallback.")
        print("Install with:  uv add sounddevice")
        audio = np.zeros(RATE * SECONDS, dtype=np.float32)

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        tmp_path = f.name
    pcm = (audio * 32767).astype(np.int16)
    with wave.open(tmp_path, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(RATE)
        wf.writeframes(pcm.tobytes())

    stt = WhisperSTT()
    result = stt.transcribe(tmp_path)
    print(f"Transcription: {result!r}")