Use lifespan to load the model just one

This commit is contained in:
Quentin Fuxa 2025-02-12 05:53:55 +01:00
parent adaca751ce
commit 149d2ee44c

View file

@ -4,6 +4,7 @@ import asyncio
import numpy as np import numpy as np
import ffmpeg import ffmpeg
from time import time from time import time
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse from fastapi.responses import HTMLResponse
@ -11,15 +12,8 @@ from fastapi.middleware.cors import CORSMiddleware
from src.whisper_streaming.whisper_online import backend_factory, online_factory, add_shared_args from src.whisper_streaming.whisper_online import backend_factory, online_factory, add_shared_args
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
##### LOAD ARGS #####
parser = argparse.ArgumentParser(description="Whisper FastAPI Online Server") parser = argparse.ArgumentParser(description="Whisper FastAPI Online Server")
parser.add_argument( parser.add_argument(
@ -49,28 +43,37 @@ parser.add_argument(
add_shared_args(parser) add_shared_args(parser)
args = parser.parse_args() args = parser.parse_args()
asr, tokenizer = backend_factory(args)
if args.diarization:
from src.diarization.diarization_online import DiartDiarization
# Load demo HTML for the root endpoint
with open("src/web/live_transcription.html", "r", encoding="utf-8") as f:
html = f.read()
@app.get("/")
async def get():
return HTMLResponse(html)
SAMPLE_RATE = 16000 SAMPLE_RATE = 16000
CHANNELS = 1 CHANNELS = 1
SAMPLES_PER_SEC = SAMPLE_RATE * int(args.min_chunk_size) SAMPLES_PER_SEC = SAMPLE_RATE * int(args.min_chunk_size)
BYTES_PER_SAMPLE = 2 # s16le = 2 bytes per sample BYTES_PER_SAMPLE = 2 # s16le = 2 bytes per sample
BYTES_PER_SEC = SAMPLES_PER_SEC * BYTES_PER_SAMPLE BYTES_PER_SEC = SAMPLES_PER_SEC * BYTES_PER_SAMPLE
if args.diarization:
from src.diarization.diarization_online import DiartDiarization
##### LOAD APP #####
@asynccontextmanager
async def lifespan(app: FastAPI):
global asr, tokenizer
asr, tokenizer = backend_factory(args)
yield
app = FastAPI(lifespan=lifespan)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Load demo HTML for the root endpoint
with open("src/web/live_transcription.html", "r", encoding="utf-8") as f:
html = f.read()
async def start_ffmpeg_decoder(): async def start_ffmpeg_decoder():
""" """
@ -91,6 +94,11 @@ async def start_ffmpeg_decoder():
return process return process
##### ENDPOINTS #####
@app.get("/")
async def get():
return HTMLResponse(html)
@app.websocket("/asr") @app.websocket("/asr")
async def websocket_endpoint(websocket: WebSocket): async def websocket_endpoint(websocket: WebSocket):