stt_server.py

"""STT API server using faster-whisper."""

import tempfile
from pathlib import Path

import av
import numpy as np
import soundfile as sf
from fastapi import FastAPI, File, Form, UploadFile
from faster_whisper import WhisperModel

from config import STT_BEAM_SIZE, STT_COMPUTE_TYPE, STT_DEVICE, STT_MODEL

app = FastAPI(title="GPU STT Server", version="1.0.0")

model: WhisperModel | None = None


def get_model() -> WhisperModel:
    global model
    if model is not None:
        return model
    model = WhisperModel(STT_MODEL, device=STT_DEVICE, compute_type=STT_COMPUTE_TYPE)
    return model


@app.get("/health")
def health():
    return {"status": "ok", "model": STT_MODEL, "device": STT_DEVICE}


@app.post("/stt")
async def transcribe(
    audio: UploadFile = File(...),
    language: str | None = Form(None),
    beam_size: int = Form(STT_BEAM_SIZE),
):
    """Transcribe audio to text."""
    whisper = get_model()

    content = await audio.read()
    filename = audio.filename or ""

    # Browser mic recordings arrive as WebM/Opus — decode with PyAV and save as WAV
    if not filename.endswith(".wav"):
        raw = tempfile.NamedTemporaryFile(suffix=".webm", delete=False)
        raw.write(content)
        raw.flush()
        raw.close()

        container = av.open(raw.name)
        frames = []
        src_rate = 48000
        for frame in container.decode(audio=0):
            frames.append(frame.to_ndarray())
            src_rate = frame.sample_rate
        container.close()
        Path(raw.name).unlink(missing_ok=True)

        audio_data = np.concatenate(frames, axis=1)
        if audio_data.shape[0] > 1:
            audio_data = audio_data.mean(axis=0, keepdims=True)

        tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
        tmp.close()
        sf.write(tmp.name, audio_data.squeeze(), src_rate)
    else:
        tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
        tmp.write(content)
        tmp.flush()
        tmp.close()

    segments, info = whisper.transcribe(
        tmp.name,
        beam_size=beam_size,
        language=language,
    )

    Path(tmp.name).unlink(missing_ok=True)

    results = []
    full_text_parts: list[str] = []
    for seg in segments:
        results.append({
            "start": round(seg.start, 2),
            "end": round(seg.end, 2),
            "text": seg.text.strip(),
        })
        full_text_parts.append(seg.text.strip())

    return {
        "language": info.language,
        "language_probability": round(info.language_probability, 3),
        "text": " ".join(full_text_parts),
        "segments": results,
    }


if __name__ == "__main__":
    import uvicorn
    from config import HOST, PORT

    get_model()
    uvicorn.run(app, host=HOST, port=int(PORT))