migrated speech to text to an actor

charles_actor.py

@@ -24,8 +24,8 @@ class CharlesActor:
         self._streamlit_av_queue = StreamlitAVQueue()
 
         print("002")
-        from speech_to_text_vosk import SpeechToTextVosk
-        self._speech_to_text_vosk = SpeechToTextVosk()
+        from speech_to_text_vosk_actor import SpeechToTextVoskActor
+        self._speech_to_text_actor = SpeechToTextVoskActor.remote()
 
         from chat_pipeline import ChatPipeline
         self._chat_pipeline = ChatPipeline()
@@ -49,38 +49,49 @@ class CharlesActor:
         total_video_frames = 0
         total_audio_frames = 0
         loops = 0
+        
+        process_speech_to_text_future = []
 
         while True:
             if len(self._debug_queue) > 0:
                 prompt = self._debug_queue.pop(0)
                 await self._chat_pipeline.enqueue(prompt)
-            audio_frames = await self._streamlit_av_queue.get_audio_frames_async()
+            audio_frames = await self._streamlit_av_queue.get_audio_frames_async()    
             if len(audio_frames) > 0:
                 total_audio_frames += len(audio_frames)
                 # Concatenate all audio frames into a single buffer
                 audio_buffer = b"".join([buffer.tobytes() for buffer in audio_frames])
-                self._speech_to_text_vosk.add_speech_bytes(audio_buffer)
-            prompt, speaker_finished = self._speech_to_text_vosk.get_text()
-            if speaker_finished and len(prompt) > 0:
-                print(f"Prompt: {prompt}")
-                system_one_audio_history.append(prompt)
-                if len(system_one_audio_history) > 10:
-                    system_one_audio_history = system_one_audio_history[-10:]
-                table_content = "| System 1 Audio History |\n| --- |\n"
-                table_content += "\n".join([f"| {item} |" for item in reversed(system_one_audio_history)])
-                self._system_one_audio_history_output = table_content
-                await self._chat_pipeline.enqueue(prompt)
-            video_frames = await self._streamlit_av_queue.get_video_frames_async()
-            if len(video_frames) > 0:
-                total_video_frames += len(video_frames)
-            #     for video_frame in video_frames:
-            #         system_one_video_output.image(video_frame.to_ndarray())
-            #         pass
+                future = self._speech_to_text_actor.process_speech.remote(audio_buffer)
+                process_speech_to_text_future.append(future)
+            # audio_frames_task = None
+
+            if len(process_speech_to_text_future) > 0:
+                ready, _ = ray.wait([process_speech_to_text_future[0]], timeout=0)
+                if ready:
+                    prompt, speaker_finished = ray.get(process_speech_to_text_future[0])
+                    del process_speech_to_text_future[0]
+
+                    if speaker_finished and len(prompt) > 0:
+                        print(f"Prompt: {prompt}")
+                        system_one_audio_history.append(prompt)
+                        if len(system_one_audio_history) > 10:
+                            system_one_audio_history = system_one_audio_history[-10:]
+                        table_content = "| System 1 Audio History |\n| --- |\n"
+                        table_content += "\n".join([f"| {item} |" for item in reversed(system_one_audio_history)])
+                        self._system_one_audio_history_output = table_content
+                        await self._chat_pipeline.enqueue(prompt)
+                        
+            # video_frames = await self._streamlit_av_queue.get_video_frames_async()
+            # if len(video_frames) > 0:
+            #     total_video_frames += len(video_frames)
+            # #     for video_frame in video_frames:
+            # #         system_one_video_output.image(video_frame.to_ndarray())
+            # #         pass
             
             # update debug output
             if (total_video_frames >0 or total_audio_frames > 0):
                 self._state = f"Processed {total_video_frames} video frames and {total_audio_frames} audio frames"
-            await asyncio.sleep(0.1)
+            await asyncio.sleep(0.01)
             loops+=1
             self._state = f"Processed {total_video_frames} video frames and {total_audio_frames} audio frames, loops: {loops}"
 
@@ -115,6 +126,6 @@ if __name__ == "__main__":
                 # The start method is still running. You can poll for debug information here.
                 time.sleep(1)
                 state = charles_actor.get_state.remote()
-                # print(f"Charles is in state: {ray.get(state)}")
+                print(f"Charles is in state: {ray.get(state)}")
     except KeyboardInterrupt:
         print("Script was manually terminated")

speech_to_text_vosk.py → speech_to_text_vosk_actor.py

@@ -1,62 +1,57 @@
 import json
 import os
-import asyncio
 from vosk import SetLogLevel, Model, KaldiRecognizer
-from multiprocessing import Process, Queue
-from queue import Empty
-SetLogLevel(-1) # mutes vosk verbosity
+import ray
+SetLogLevel(-1)  # mutes vosk verbosity
 
-class SpeechToTextVosk:
+@ray.remote
+class SpeechToTextVoskActor:
     def __init__(self, model='small', audio_bit_rate=16000) -> None:
         self.model = model
         self.audio_bit_rate = audio_bit_rate
 
-        # Create a Queue for inter-process communication
-        self.queue = Queue()
-        self.result_queue = Queue()
-
-        # Create and start a new Process with the worker function
-        self.process = Process(target=self.worker)
-        self.process.start()
-
-    def worker(self):
         # load vosk model
         # get path of current file
         current_file_path = os.path.abspath(__file__)
         current_directory = os.path.dirname(current_file_path)
         _path = os.path.join(current_directory, 'models', 'vosk', self.model)
-        model_voice = Model(_path)
-        vosk = KaldiRecognizer(model_voice, self.audio_bit_rate)
-
-        while True:
-            try:
-                # Get the next item from the queue. Blocks for 1s if necessary.
-                data = self.queue.get(timeout=1)
+        self.model_voice = Model(_path)
+        self.vosk = KaldiRecognizer(self.model_voice, self.audio_bit_rate)
 
-                # Stop the worker if the sentinel None is received
-                if data is None:
-                    break
-                
-                text, speaker_finished = self._process_speech(vosk, data)
-
-                # put the result into result_queue
-                self.result_queue.put((text, speaker_finished))
-            except Empty:
-                pass
+        self.text_queue = []
+        self.finished_queue = []
+        
+    def process_speech(self, data: bytearray) -> tuple[str, bool]:
+        text = ''
+        speaker_finished = False
+        if self.vosk.AcceptWaveform(data):
+            result = self.vosk.Result()
+            result_json = json.loads(result)
+            text = result_json['text']
+            speaker_finished = True
+        else:
+            result = self.vosk.PartialResult()
+            result_json = json.loads(result)
+            text = result_json['partial']
+        return text, speaker_finished
+        
 
     def add_speech_bytes(self, data: bytearray):
-        self.queue.put(data)
+        text, speaker_finished = self._process_speech(data)
+        self.text_queue.append(text)
+        if speaker_finished:
+            self.finished_queue.append(speaker_finished)
 
-    def _process_speech(self, vosk: KaldiRecognizer, data: bytearray) -> tuple[str, bool]:
+    def _process_speech(self, data: bytearray) -> tuple[str, bool]:
         text = ''
         speaker_finished = False
-        if vosk.AcceptWaveform(data):
-            result = vosk.Result()
+        if self.vosk.AcceptWaveform(data):
+            result = self.vosk.Result()
             result_json = json.loads(result)
             text = result_json['text']
             speaker_finished = True
         else:
-            result = vosk.PartialResult()
+            result = self.vosk.PartialResult()
             result_json = json.loads(result)
             text = result_json['partial']
         return text, speaker_finished
@@ -64,27 +59,13 @@ class SpeechToTextVosk:
     def get_text(self):
         text = ''
         speaker_finished = False
-        while not self.result_queue.empty():
-            result, speaker_finished = self.result_queue.get()
+        while self.text_queue:
+            result = self.text_queue.pop(0)
             text += result
-            if speaker_finished:
+            if self.finished_queue:
+                speaker_finished = self.finished_queue.pop(0)
                 break
-        return (text, speaker_finished)
-    
+        return text, speaker_finished
+
     def get_audio_bit_rate(self):
         return self.audio_bit_rate
-
-
-    def shutdown(self):
-        # Send sentinel value to stop the worker
-        self.queue.put(None)
-        # Wait for the worker process to finish
-        self.process.join()
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self.shutdown()
-
-    def __del__(self):
-        self.shutdown()