85 lines
3.5 KiB
Python
85 lines
3.5 KiB
Python
"""diora's built-in Whisper: a small OpenAI-compatible transcription server.
|
|
|
|
POST /v1/audio/transcriptions (multipart: file, model, language, prompt, response_format)
|
|
returns {text} or, with response_format=verbose_json, {language, duration, text, segments}.
|
|
The `model` field is ignored; WHISPER_MODEL decides. The model is loaded on the first request
|
|
and freed again after WHISPER_IDLE seconds without work, so it holds no memory while unused.
|
|
One transcription at a time; parallel requests wait (the server has few cores to share).
|
|
See docs/transcription.md.
|
|
"""
|
|
import os
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
|
|
# Windows with a GPU: where the CUDA libraries (cuBLAS, cuDNN) are, separated by ';'.
|
|
for _d in filter(None, os.environ.get('WHISPER_DLL_DIRS', '').split(';')):
|
|
os.add_dll_directory(_d)
|
|
os.environ['PATH'] = _d + os.pathsep + os.environ['PATH']
|
|
|
|
import numpy as np # noqa: E402
|
|
from fastapi import FastAPI, File, Form, HTTPException, UploadFile # noqa: E402
|
|
from faster_whisper import WhisperModel # noqa: E402
|
|
|
|
MODEL = os.environ.get('WHISPER_MODEL', 'large-v3-turbo')
|
|
THREADS = int(os.environ.get('WHISPER_THREADS', '4'))
|
|
DEVICE = os.environ.get('WHISPER_DEVICE', 'cpu') # cpu or cuda
|
|
COMPUTE = os.environ.get('WHISPER_COMPUTE', 'int8') # int8_float16 on a GPU
|
|
IDLE = int(os.environ.get('WHISPER_IDLE', '600'))
|
|
CACHE = os.environ.get('WHISPER_CACHE', '/models')
|
|
|
|
app = FastAPI()
|
|
work = threading.Lock() # one transcription at a time
|
|
state = {'model': None, 'last': 0.0}
|
|
|
|
|
|
def model():
|
|
if state['model'] is None:
|
|
state['model'] = WhisperModel(MODEL, device=DEVICE, compute_type=COMPUTE, cpu_threads=THREADS,
|
|
download_root=CACHE)
|
|
return state['model']
|
|
|
|
|
|
def unloader():
|
|
while True:
|
|
time.sleep(30)
|
|
if state['model'] is not None and time.time() - state['last'] > IDLE and work.acquire(blocking=False):
|
|
try:
|
|
state['model'] = None
|
|
finally:
|
|
work.release()
|
|
|
|
|
|
threading.Thread(target=unloader, daemon=True).start()
|
|
|
|
|
|
def decode(data: bytes) -> np.ndarray:
|
|
p = subprocess.run(['ffmpeg', '-nostdin', '-loglevel', 'error', '-i', 'pipe:0', '-ac', '1', '-ar', '16000',
|
|
'-f', 'f32le', 'pipe:1'], input=data, capture_output=True)
|
|
if p.returncode != 0:
|
|
raise HTTPException(400, 'not audio: ' + p.stderr.decode('utf-8', 'replace')[-200:])
|
|
return np.frombuffer(p.stdout, np.float32)
|
|
|
|
|
|
@app.get('/health')
|
|
def health():
|
|
return {'ok': True, 'model': MODEL, 'device': DEVICE, 'loaded': state['model'] is not None}
|
|
|
|
|
|
@app.post('/v1/audio/transcriptions')
|
|
def transcriptions(file: UploadFile = File(...), language: str = Form(''), prompt: str = Form(''),
|
|
response_format: str = Form('json')):
|
|
audio = decode(file.file.read())
|
|
t0 = time.time()
|
|
with work:
|
|
segs, info = model().transcribe(audio, language=language or None, initial_prompt=prompt or None,
|
|
beam_size=5, vad_filter=True)
|
|
segs = [{'id': i, 'start': round(s.start, 2), 'end': round(s.end, 2), 'text': s.text}
|
|
for i, s in enumerate(segs)]
|
|
state['last'] = time.time()
|
|
dur = len(audio) / 16000
|
|
print(f'{dur:.0f} s audio in {time.time() - t0:.0f} s', flush=True)
|
|
text = ''.join(s['text'] for s in segs).strip()
|
|
if response_format == 'verbose_json':
|
|
return {'task': 'transcribe', 'language': info.language, 'duration': dur, 'text': text, 'segments': segs}
|
|
return {'text': text}
|