async queue pipeline

asr.py

@@ -58,18 +58,18 @@ class ASRPipeline:
 
         self.last_text = ""
         self.is_active = True
-        self.is_processing = False
         self.key_index = 0
+        self.key_lock = threading.Lock()
 
         print("ASR Pipeline siap ✓")
 
     def get_groq_client(self):
-        client = Groq(api_key=GROQ_KEYS[self.key_index])
-        self.key_index = (self.key_index + 1) % len(GROQ_KEYS)
+        with self.key_lock:
+            client = Groq(api_key=GROQ_KEYS[self.key_index])
+            self.key_index = (self.key_index + 1) % len(GROQ_KEYS)
         return client
 
     def stop(self):
-        self.is_processing = False
         self.is_active = True
         print("Pipeline stopped ✓")
 
@@ -141,22 +141,16 @@ class ASRPipeline:
         if not self.is_active:
             return None
 
-        if self.is_processing:
-            return None
-
         vad_conf = self.get_vad_confidence(audio_chunk)
         if vad_conf <= VAD_THRESHOLD:
             return None
 
-        self.is_processing = True
-        try:
-            text = self.groq_transcribe(audio_chunk)
-            if not text:
-                return None
-            if not self.is_valid_text(text):
-                print(f"Filtered: {text}")
-                return None
-            self.last_text = text
-            return text.strip()
-        finally:
-            self.is_processing = False
+        text = self.groq_transcribe(audio_chunk)
+        if not text:
+            return None
+        if not self.is_valid_text(text):
+            print(f"Filtered: {text}")
+            return None
+
+        self.last_text = text
+        return text.strip()

main.py

@@ -3,6 +3,7 @@ import numpy as np
 import logging
 import warnings
 import os
+import threading
 
 logging.getLogger("nemo_logger").setLevel(logging.ERROR)
 logging.getLogger("nemo").setLevel(logging.ERROR)
@@ -22,7 +23,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-MODEL_DIR = "/app/models"
+MODEL_DIR = "C:/Users/YEREMIA/Documents/captionai_server/models"
 pipeline = ASRPipeline(MODEL_DIR)
 
 @app.get("/")
@@ -36,26 +37,45 @@ async def websocket_endpoint(websocket: WebSocket):
     print("Client connected")
     print("=" * 50)
 
+    loop = asyncio.get_event_loop()
+    caption_queue = asyncio.Queue()
+
+    async def send_results():
+        while True:
+            text = await caption_queue.get()
+            if text is None:
+                break
+            try:
+                await websocket.send_text(text)
+            except:
+                break
+
+    sender = asyncio.create_task(send_results())
+
     try:
         while True:
             data = await websocket.receive_bytes()
             audio_chunk = np.frombuffer(data, dtype=np.float32)
 
-            result = await asyncio.get_event_loop().run_in_executor(
-                None, pipeline.transcribe_chunk, audio_chunk
-            )
+            def process(chunk=audio_chunk):
+                result = pipeline.transcribe_chunk(chunk)
+                if result:
+                    asyncio.run_coroutine_threadsafe(
+                        caption_queue.put(result),
+                        loop
+                    )
 
-            if result:
-                await websocket.send_text(result)
-            else:
-                await websocket.send_text("")
+            threading.Thread(target=process, daemon=True).start()
 
     except WebSocketDisconnect:
         pipeline.stop()
+        await caption_queue.put(None)
+        await sender
         print("=" * 50)
         print("Client disconnected")
         print("=" * 50)
     except Exception as e:
         print(f"Error: {e}")
         pipeline.stop()
+        await caption_queue.put(None)
         await websocket.close()