server.py

import asyncio
import os
import secrets
import threading
import time
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, status
from faster_whisper import WhisperModel

# Picks up .env in the working directory (WHISPER_API_TOKEN and friends) for both bare
# `uvicorn server:app` runs and Docker Compose alike — real env vars still take
# precedence over anything in the file.
load_dotenv()

# Shared-secret token required to open a /ws connection — this endpoint burns real GPU
# time per connection, so it must not be left open to anyone who can reach the port.
# Lives in .env (loaded above) rather than being generated on the fly, so the token is
# stable across restarts and every client doesn't need re-pasting it each time. Checked
# before loading the model below so a missing token fails fast, not after a multi-GB
# model load.
API_TOKEN = os.environ.get("WHISPER_API_TOKEN")
if not API_TOKEN:
    raise RuntimeError(
        "WHISPER_API_TOKEN is not set. Add it to .env (see .env.example) — "
        "generate one with: python3 -c \"import secrets; print(secrets.token_urlsafe(32))\""
    )


def token_is_valid(token: str) -> bool:
    return secrets.compare_digest(token or "", API_TOKEN)


# Local CTranslate2 conversion of Na0s/Medical-Whisper-Large-v3 (large-v3 fine-tuned on
# doctor/patient consultations). Converted via ct2-transformers-converter into
# models/medical-whisper-large-v3-ct2 — see README for the conversion steps.
MODEL_SIZE = os.environ.get("WHISPER_MODEL", "models/medical-whisper-large-v3-ct2")
DEVICE = os.environ.get("WHISPER_DEVICE", "cuda")
COMPUTE_TYPE = os.environ.get("WHISPER_COMPUTE", "float16")
# Comma-separated GPU indices to replicate the model across. Defaults to all 3 GPUs on
# this box so the model runs at full available capacity out of the box — CTranslate2
# does NOT spread across GPUs unless told to via device_index.
DEVICE_INDEX = [int(i) for i in os.environ.get("WHISPER_DEVICE_INDEX", "0,1,2").split(",")]
# CTranslate2 (the backend faster-whisper uses) supports genuinely concurrent
# inference via internal worker threads that share the same loaded weights, so
# several connections' transcribe() calls can run on the GPU at once instead of
# queuing behind each other. This is workers PER entry in DEVICE_INDEX, so total
# concurrent capacity is roughly NUM_WORKERS * len(DEVICE_INDEX). Re-benchmark VRAM
# headroom for this value on current hardware — see README "Concurrency notes".
NUM_WORKERS = int(os.environ.get("WHISPER_NUM_WORKERS", "4"))

SAMPLE_RATE = 16000
WINDOW_SECONDS = 8          # how much trailing audio we re-transcribe each pass
PROCESS_INTERVAL = 1.5      # seconds between partial re-transcriptions
SILENCE_AMPLITUDE = 300     # int16 amplitude below which audio counts as silence
SILENCE_TO_FINALIZE = 1.2   # seconds of silence before we commit a final segment

print(f"Loading faster-whisper model '{MODEL_SIZE}' on {DEVICE} device_index={DEVICE_INDEX} "
      f"({COMPUTE_TYPE}), num_workers={NUM_WORKERS} per device...")
model = WhisperModel(
    MODEL_SIZE,
    device=DEVICE,
    device_index=DEVICE_INDEX,
    compute_type=COMPUTE_TYPE,
    num_workers=NUM_WORKERS,
)
print("Model loaded.")

# Hands transcribe() calls off to Python threads so the asyncio event loop keeps
# serving other connections' audio while GPU work runs. Sized to match total worker
# capacity (NUM_WORKERS per GPU) so that many transcriptions can genuinely be in
# flight across all GPUs at once.
gpu_executor = ThreadPoolExecutor(max_workers=NUM_WORKERS * len(DEVICE_INDEX))

# Total concurrent transcribe() calls the GPU pool can run at once — see README
# "Concurrency notes". Exposed via /load so clients (or a load balancer picking
# between multiple instances of this server) can see how saturated this one is.
TRANSCRIBE_CAPACITY = NUM_WORKERS * len(DEVICE_INDEX)

_load_lock = threading.Lock()
active_connections = 0
active_transcriptions = 0

# Separate port for /load so congestion/status can be polled independently of the
# main API port (e.g. left reachable through a firewall/LB even if WS access to
# port API_PORT is restricted). 0 disables the second listener entirely.
STATUS_PORT = int(os.environ.get("WHISPER_STATUS_PORT", "8001"))
CERT_FILE = os.environ.get("WHISPER_CERT_FILE", "certs/cert.pem")
KEY_FILE = os.environ.get("WHISPER_KEY_FILE", "certs/key.pem")


SUPPORTED_LANGUAGES = {"es", "it", "fr", "de", "en"}
DEFAULT_LANGUAGE = "it"


def run_transcribe(audio: np.ndarray, language: str) -> str:
    global active_transcriptions
    with _load_lock:
        active_transcriptions += 1
    try:
        segments, _ = model.transcribe(audio, beam_size=1, vad_filter=True, language=language)
        return "".join(s.text for s in segments).strip()
    finally:
        with _load_lock:
            active_transcriptions -= 1


def get_load() -> dict:
    with _load_lock:
        connections, transcriptions = active_connections, active_transcriptions
    return {
        "active_connections": connections,
        "active_transcriptions": transcriptions,
        "transcribe_capacity": TRANSCRIBE_CAPACITY,
        "load_ratio": round(transcriptions / TRANSCRIBE_CAPACITY, 3),
    }


app = FastAPI()


@app.get("/load")
async def load():
    return get_load()


# Minimal second app, serving only /load on WHISPER_STATUS_PORT — shares the same
# in-process counters as `app` above (this only works because both run in the same
# process; it's not a separate replica).
status_app = FastAPI()


@status_app.get("/load")
async def status_load():
    return get_load()


@app.on_event("startup")
async def start_status_server():
    if STATUS_PORT == 0:
        return
    config_kwargs = dict(app=status_app, host="0.0.0.0", port=STATUS_PORT, log_level="warning")
    if os.path.exists(CERT_FILE) and os.path.exists(KEY_FILE):
        config_kwargs["ssl_certfile"] = CERT_FILE
        config_kwargs["ssl_keyfile"] = KEY_FILE
    server = uvicorn.Server(uvicorn.Config(**config_kwargs))
    asyncio.create_task(server.serve())
    print(f"Status endpoint (/load) listening on port {STATUS_PORT}"
          f"{' (TLS)' if 'ssl_certfile' in config_kwargs else ''}.")


@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    # Checked before accept() so an invalid/missing token gets the connection refused
    # at the handshake rather than an accepted-then-dropped connection.
    if not token_is_valid(ws.query_params.get("token", "")):
        await ws.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    language = ws.query_params.get("lang", DEFAULT_LANGUAGE)
    if language not in SUPPORTED_LANGUAGES:
        language = DEFAULT_LANGUAGE
    await ws.accept()

    global active_connections
    with _load_lock:
        active_connections += 1

    buffer = np.zeros(0, dtype=np.float32)
    committed_text = ""
    last_process_time = time.time()
    last_voice_time = time.time()
    loop = asyncio.get_running_loop()

    try:
        while True:
            data = await ws.receive_bytes()
            pcm16 = np.frombuffer(data, dtype=np.int16)
            if pcm16.size == 0:
                continue

            audio = pcm16.astype(np.float32) / 32768.0
            buffer = np.concatenate([buffer, audio])

            now = time.time()
            if np.abs(pcm16).mean() > SILENCE_AMPLITUDE:
                last_voice_time = now
            silence_elapsed = now - last_voice_time

            min_samples = int(SAMPLE_RATE * 0.5)

            if silence_elapsed >= SILENCE_TO_FINALIZE and buffer.size > min_samples:
                text = await loop.run_in_executor(gpu_executor, run_transcribe, buffer, language)
                if text:
                    committed_text = (committed_text + " " + text).strip()
                    await ws.send_json({"type": "final", "text": committed_text})
                buffer = np.zeros(0, dtype=np.float32)
                last_process_time = now
                last_voice_time = now
                continue

            if now - last_process_time >= PROCESS_INTERVAL and buffer.size > min_samples:
                last_process_time = now
                window = buffer[-SAMPLE_RATE * WINDOW_SECONDS:]
                text = await loop.run_in_executor(gpu_executor, run_transcribe, window, language)
                await ws.send_json({"type": "partial", "text": (committed_text + " " + text).strip()})

    except WebSocketDisconnect:
        pass
    finally:
        with _load_lock:
            active_connections -= 1