admin-valentin/WhisperLiveKit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

DiartDiarization loaded in asynccontextmanager lifespan

Browse source
This commit is contained in:
Quentin Fuxa 2025-02-19 16:47:34 +01:00
parent 09d40a7de8
commit fa39eda923
1 changed files with 10 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
whisper_fastapi_online_server.py
View file

@ -68,19 +68,23 @@ 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
MAX_BYTES_PER_SEC = 32000 * 5 # 5 seconds of audio at 32 kHz MAX_BYTES_PER_SEC = 32000 * 5 # 5 seconds of audio at 32 kHz
if args.diarization:
from src.diarization.diarization_online import DiartDiarization
##### LOAD APP ##### ##### LOAD APP #####
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
global asr, tokenizer global asr, tokenizer, diarization
if args.transcription: if args.transcription:
asr, tokenizer = backend_factory(args) asr, tokenizer = backend_factory(args)
else: else:
asr, tokenizer = None, None asr, tokenizer = None, None
if args.diarization:
from src.diarization.diarization_online import DiartDiarization
diarization = DiartDiarization(SAMPLE_RATE)
else :
diarization = None
yield yield
app = FastAPI(lifespan=lifespan) app = FastAPI(lifespan=lifespan)
@ -130,10 +134,10 @@ async def websocket_endpoint(websocket: WebSocket):
ffmpeg_process = None ffmpeg_process = None
pcm_buffer = bytearray() pcm_buffer = bytearray()
online = online_factory(args, asr, tokenizer) if args.transcription else None online = online_factory(args, asr, tokenizer) if args.transcription else None
diarization = DiartDiarization(SAMPLE_RATE) if args.diarization else None
async def restart_ffmpeg(): async def restart_ffmpeg():
nonlocal ffmpeg_process, online, diarization, pcm_buffer nonlocal ffmpeg_process, online, pcm_buffer
if ffmpeg_process: if ffmpeg_process:
try: try:
ffmpeg_process.kill() ffmpeg_process.kill()
@ -143,14 +147,12 @@ async def websocket_endpoint(websocket: WebSocket):
ffmpeg_process = await start_ffmpeg_decoder() ffmpeg_process = await start_ffmpeg_decoder()
pcm_buffer = bytearray() pcm_buffer = bytearray()
online = online_factory(args, asr, tokenizer) if args.transcription else None online = online_factory(args, asr, tokenizer) if args.transcription else None
if args.diarization:
diarization = DiartDiarization(SAMPLE_RATE)
logger.info("FFmpeg process started.") logger.info("FFmpeg process started.")
await restart_ffmpeg() await restart_ffmpeg()
async def ffmpeg_stdout_reader(): async def ffmpeg_stdout_reader():
nonlocal ffmpeg_process, online, diarization, pcm_buffer nonlocal ffmpeg_process, online, pcm_buffer
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
full_transcription = "" full_transcription = ""
beg = time() beg = time()