// execution for diarization and transcription

whisper_fastapi_online_server.py

@@ -70,6 +70,78 @@ BYTES_PER_SEC = SAMPLES_PER_SEC * BYTES_PER_SAMPLE
 MAX_BYTES_PER_SEC = 32000 * 5  # 5 seconds of audio at 32 kHz
 
 
+class SharedState:
+    def __init__(self):
+        self.tokens = []
+        self.buffer_transcription = ""
+        self.buffer_diarization = ""
+        self.full_transcription = ""
+        self.end_buffer = 0
+        self.end_attributed_speaker = 0
+        self.lock = asyncio.Lock()
+        self.beg_loop = time()
+        self.sep = " "  # Default separator
+        
+    async def update_transcription(self, new_tokens, buffer, end_buffer, full_transcription, sep):
+        async with self.lock:
+            self.tokens.extend(new_tokens)
+            self.buffer_transcription = buffer
+            self.end_buffer = end_buffer
+            self.full_transcription = full_transcription
+            self.sep = sep
+            
+    async def update_diarization(self, end_attributed_speaker, buffer_diarization=""):
+        async with self.lock:
+            self.end_attributed_speaker = end_attributed_speaker
+            if buffer_diarization:
+                self.buffer_diarization = buffer_diarization
+            
+    async def add_dummy_token(self):
+        async with self.lock:
+            current_time = time() - self.beg_loop
+            dummy_token = ASRToken(
+                start=current_time,
+                end=current_time + 0.5,
+                text="",
+                speaker=-1
+            )
+            self.tokens.append(dummy_token)
+            
+    async def get_current_state(self):
+        async with self.lock:
+            current_time = time()
+            remaining_time_transcription = 0
+            remaining_time_diarization = 0
+            
+            # Calculate remaining time for transcription buffer
+            if self.end_buffer > 0:
+                remaining_time_transcription = max(0, round(current_time - self.beg_loop - self.end_buffer, 2))
+                
+            # Calculate remaining time for diarization
+            if self.end_attributed_speaker > 0:
+                remaining_time_diarization = max(0, round(current_time - self.beg_loop - self.end_attributed_speaker, 2))
+                
+            return {
+                "tokens": self.tokens.copy(),
+                "buffer_transcription": self.buffer_transcription,
+                "buffer_diarization": self.buffer_diarization,
+                "end_buffer": self.end_buffer,
+                "end_attributed_speaker": self.end_attributed_speaker,
+                "sep": self.sep,
+                "remaining_time_transcription": remaining_time_transcription,
+                "remaining_time_diarization": remaining_time_diarization
+            }
+            
+    async def reset(self):
+        """Reset the state."""
+        async with self.lock:
+            self.tokens = []
+            self.buffer_transcription = ""
+            self.buffer_diarization = ""
+            self.end_buffer = 0
+            self.end_attributed_speaker = 0
+            self.full_transcription = ""
+            self.beg_loop = time()
 
 ##### LOAD APP #####
 
@@ -120,112 +192,89 @@ async def start_ffmpeg_decoder():
     )
     return process
 
-
+async def transcription_processor(shared_state, pcm_queue, online):
-##### ENDPOINTS #####
-
-@app.get("/")
-async def get():
-    return HTMLResponse(html)
-
-@app.websocket("/asr")
-async def websocket_endpoint(websocket: WebSocket):
-    await websocket.accept()
-    logger.info("WebSocket connection opened.")
-
-    ffmpeg_process = None
-    pcm_buffer = bytearray()
-    online = online_factory(args, asr, tokenizer) if args.transcription else None
-    
-
-    async def restart_ffmpeg():
-        nonlocal ffmpeg_process, online, pcm_buffer
-        if ffmpeg_process:
-            try:
-                ffmpeg_process.kill()
-                await asyncio.get_event_loop().run_in_executor(None, ffmpeg_process.wait)
-            except Exception as e:
-                logger.warning(f"Error killing FFmpeg process: {e}")
-        ffmpeg_process = await start_ffmpeg_decoder()
-        pcm_buffer = bytearray()
-        online = online_factory(args, asr, tokenizer) if args.transcription else None
-        logger.info("FFmpeg process started.")
-
-    await restart_ffmpeg()
-
-    async def ffmpeg_stdout_reader():
-        nonlocal ffmpeg_process, online, pcm_buffer
-        loop = asyncio.get_event_loop()
     full_transcription = ""
-        beg = time()
-        beg_loop = time()
-        tokens = []
-        end_attributed_speaker = 0
     sep = online.asr.sep
     
     while True:
         try:
-                elapsed_time = math.floor((time() - beg) * 10) / 10 # Round to 0.1 sec
+            pcm_array = await pcm_queue.get()
-                ffmpeg_buffer_from_duration = max(int(32000 * elapsed_time), 4096)
-                beg = time()
             
-                # Read chunk with timeout
-                try:
-                    chunk = await asyncio.wait_for(
-                        loop.run_in_executor(
-                            None, ffmpeg_process.stdout.read, ffmpeg_buffer_from_duration
-                        ),
-                        timeout=5.0
-                    )
-                except asyncio.TimeoutError:
-                    logger.warning("FFmpeg read timeout. Restarting...")
-                    await restart_ffmpeg()
-                    full_transcription = ""
-                    beg = time()
-                    continue  # Skip processing and read from new process
-
-                if not chunk:
-                    logger.info("FFmpeg stdout closed.")
-                    break
-                pcm_buffer.extend(chunk)
-                if len(pcm_buffer) >= BYTES_PER_SEC:
-                    if len(pcm_buffer) > MAX_BYTES_PER_SEC:
-                        logger.warning(
-                            f"""Audio buffer is too large: {len(pcm_buffer) / BYTES_PER_SEC:.2f} seconds.
-                            The model probably struggles to keep up. Consider using a smaller model.
-                            """)
-                    # Convert int16 -> float32
-                    pcm_array = (
-                        np.frombuffer(pcm_buffer[:MAX_BYTES_PER_SEC], dtype=np.int16).astype(np.float32)
-                        / 32768.0
-                    )
-                    pcm_buffer = pcm_buffer[MAX_BYTES_PER_SEC:]
-                    
-                    if args.transcription:
             logger.info(f"{len(online.audio_buffer) / online.SAMPLING_RATE} seconds of audio will be processed by the model.")
+            
+            # Process transcription
             online.insert_audio_chunk(pcm_array)
             new_tokens = online.process_iter()
-                        tokens.extend(new_tokens)
+            
+            if new_tokens:
                 full_transcription += sep.join([t.text for t in new_tokens])
+                
             _buffer = online.get_buffer()
             buffer = _buffer.text
-                        end_buffer = _buffer.end if _buffer.end else tokens[-1].end if tokens else 0
+            end_buffer = _buffer.end if _buffer.end else (new_tokens[-1].end if new_tokens else 0)
-                        if buffer in full_transcription: # With VAC, the buffer is not updated until the next chunk is processed
+            
+            if buffer in full_transcription:
                 buffer = ""
-                    else:
-                        tokens.append(
-                            ASRToken(
-                                start = time() - beg_loop,
-                                end = time() - beg_loop + 0.5))
-                        sleep(0.5)
-                        buffer = ''
                 
-                    if args.diarization:
+            await shared_state.update_transcription(
-                        await diarization.diarize(pcm_array)
+                new_tokens, buffer, end_buffer, full_transcription, sep)
-                        end_attributed_speaker = diarization.assign_speakers_to_tokens(end_attributed_speaker, tokens)
             
+        except Exception as e:
+            logger.warning(f"Exception in transcription_processor: {e}")
+        finally:
+            pcm_queue.task_done()
+
+async def diarization_processor(shared_state, pcm_queue, diarization_obj):
+    buffer_diarization = ""
+    
+    while True:
+        try:
+            pcm_array = await pcm_queue.get()
+            
+            # Process diarization
+            await diarization_obj.diarize(pcm_array)
+            
+            # Get current state
+            state = await shared_state.get_current_state()
+            tokens = state["tokens"]
+            end_attributed_speaker = state["end_attributed_speaker"]
+            
+            # Update speaker information
+            new_end_attributed_speaker = diarization_obj.assign_speakers_to_tokens(
+                end_attributed_speaker, tokens)
+            
+            await shared_state.update_diarization(new_end_attributed_speaker, buffer_diarization)
+            
+        except Exception as e:
+            logger.warning(f"Exception in diarization_processor: {e}")
+        finally:
+            pcm_queue.task_done()
+
+async def results_formatter(shared_state, websocket):
+    while True:
+        try:
+            # Get the current state
+            state = await shared_state.get_current_state()
+            tokens = state["tokens"]
+            buffer_transcription = state["buffer_transcription"]
+            buffer_diarization = state["buffer_diarization"]
+            end_attributed_speaker = state["end_attributed_speaker"]
+            remaining_time_transcription = state["remaining_time_transcription"]
+            remaining_time_diarization = state["remaining_time_diarization"]
+            sep = state["sep"]
+            
+            # If diarization is enabled but no transcription, add dummy tokens periodically
+            if not tokens and not args.transcription and args.diarization:
+                await shared_state.add_dummy_token()
+                # Re-fetch tokens after adding dummy
+                state = await shared_state.get_current_state()
+                tokens = state["tokens"]
+            
+            # Process tokens to create response
             previous_speaker = -10
             lines = []
             last_end_diarized = 0
+            
             for token in tokens:
                 speaker = token.speaker
                 if args.diarization:
@@ -248,15 +297,129 @@ async def websocket_endpoint(websocket: WebSocket):
                         }
                     )
                     previous_speaker = speaker
-                        else:
+                elif token.text:  # Only append if text isn't empty
                     lines[-1]["text"] += sep + token.text
                     lines[-1]["end"] = format_time(token.end)
                     lines[-1]["diff"] = round(token.end - last_end_diarized, 2)
             
-                    response = {"lines": lines, "buffer": buffer}
+            # Prepare response object
-                    # response = {"lines": lines, "buffer": buffer, "time_buffer_transcription": time() + beg_loop - end_buffer, "time_buffer_diarization": time() + beg_loop - end_attributed_speaker}
+            response = {
+                "lines": lines, 
+                "buffer_transcription": buffer_transcription,
+                "buffer_diarization": buffer_diarization,
+                "remaining_time_transcription": remaining_time_transcription,
+                "remaining_time_diarization": remaining_time_diarization
+            }
+            
             await websocket.send_json(response)
             
+            # Add a small delay to avoid overwhelming the client
+            await asyncio.sleep(0.1)
+            
+        except Exception as e:
+            logger.warning(f"Exception in results_formatter: {e}")
+            await asyncio.sleep(0.5)  # Back off on error
+
+##### ENDPOINTS #####
+
+@app.get("/")
+async def get():
+    return HTMLResponse(html)
+
+@app.websocket("/asr")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    logger.info("WebSocket connection opened.")
+
+    ffmpeg_process = None
+    pcm_buffer = bytearray()
+    shared_state = SharedState()
+    
+    transcription_queue = asyncio.Queue() if args.transcription else None
+    diarization_queue = asyncio.Queue() if args.diarization else None
+    
+    online = None
+
+    async def restart_ffmpeg():
+        nonlocal ffmpeg_process, online, pcm_buffer
+        if ffmpeg_process:
+            try:
+                ffmpeg_process.kill()
+                await asyncio.get_event_loop().run_in_executor(None, ffmpeg_process.wait)
+            except Exception as e:
+                logger.warning(f"Error killing FFmpeg process: {e}")
+        ffmpeg_process = await start_ffmpeg_decoder()
+        pcm_buffer = bytearray()
+        
+        if args.transcription:
+            online = online_factory(args, asr, tokenizer)
+        
+        await shared_state.reset()
+        logger.info("FFmpeg process started.")
+
+    await restart_ffmpeg()
+
+    tasks = []    
+    if args.transcription and online:
+        tasks.append(asyncio.create_task(
+            transcription_processor(shared_state, transcription_queue, online)))    
+    if args.diarization and diarization:
+        tasks.append(asyncio.create_task(
+            diarization_processor(shared_state, diarization_queue, diarization)))
+    formatter_task = asyncio.create_task(results_formatter(shared_state, websocket))
+    tasks.append(formatter_task)
+
+    async def ffmpeg_stdout_reader():
+        nonlocal ffmpeg_process, pcm_buffer
+        loop = asyncio.get_event_loop()
+        beg = time()
+        
+        while True:
+            try:
+                elapsed_time = math.floor((time() - beg) * 10) / 10 # Round to 0.1 sec
+                ffmpeg_buffer_from_duration = max(int(32000 * elapsed_time), 4096)
+                beg = time()
+
+                # Read chunk with timeout
+                try:
+                    chunk = await asyncio.wait_for(
+                        loop.run_in_executor(
+                            None, ffmpeg_process.stdout.read, ffmpeg_buffer_from_duration
+                        ),
+                        timeout=5.0
+                    )
+                except asyncio.TimeoutError:
+                    logger.warning("FFmpeg read timeout. Restarting...")
+                    await restart_ffmpeg()
+                    beg = time()
+                    continue  # Skip processing and read from new process
+
+                if not chunk:
+                    logger.info("FFmpeg stdout closed.")
+                    break
+                pcm_buffer.extend(chunk)
+                if len(pcm_buffer) >= BYTES_PER_SEC:
+                    if len(pcm_buffer) > MAX_BYTES_PER_SEC:
+                        logger.warning(
+                            f"""Audio buffer is too large: {len(pcm_buffer) / BYTES_PER_SEC:.2f} seconds.
+                            The model probably struggles to keep up. Consider using a smaller model.
+                            """)
+                    # Convert int16 -> float32
+                    pcm_array = (
+                        np.frombuffer(pcm_buffer[:MAX_BYTES_PER_SEC], dtype=np.int16).astype(np.float32)
+                        / 32768.0
+                    )
+                    pcm_buffer = pcm_buffer[MAX_BYTES_PER_SEC:]
+                    
+                    if args.transcription and transcription_queue:
+                        await transcription_queue.put(pcm_array.copy())
+                    
+                    if args.diarization and diarization_queue:
+                        await diarization_queue.put(pcm_array.copy())
+                    
+                    if not args.transcription and not args.diarization:
+                        await asyncio.sleep(0.1)
+                    
             except Exception as e:
                 logger.warning(f"Exception in ffmpeg_stdout_reader: {e}")
                 break
@@ -264,7 +427,7 @@ async def websocket_endpoint(websocket: WebSocket):
         logger.info("Exiting ffmpeg_stdout_reader...")
 
     stdout_reader_task = asyncio.create_task(ffmpeg_stdout_reader())
-
+    tasks.append(stdout_reader_task)    
     try:
         while True:
             # Receive incoming WebM audio chunks from the client
@@ -280,16 +443,20 @@ async def websocket_endpoint(websocket: WebSocket):
     except WebSocketDisconnect:
         logger.warning("WebSocket disconnected.")
     finally:
-        stdout_reader_task.cancel()
+        for task in tasks:
+            task.cancel()
+            
         try:
+            await asyncio.gather(*tasks, return_exceptions=True)
             ffmpeg_process.stdin.close()
             ffmpeg_process.wait()
-        except:
+        except Exception as e:
-            pass
+            logger.warning(f"Error during cleanup: {e}")
-        if args.diarization:
+        
+        if args.diarization and diarization:
             diarization.close()
         
-
+        logger.info("WebSocket endpoint cleaned up.")
 
 if __name__ == "__main__":
     import uvicorn