start of moving to ray

d_app.py

@@ -24,6 +24,12 @@ load_dotenv()
 
 webrtc_ctx = None
 
+# Initialize Ray
+import ray
+if not ray.is_initialized():
+    ray.init()
+
+
 
 async def main():
 
@@ -98,21 +104,11 @@ async def main():
             if len(st.session_state.debug_queue) > 0:
                 prompt = st.session_state.debug_queue.pop(0)
                 await st.session_state.chat_pipeline.enqueue(prompt)
-            sound_chunk = pydub.AudioSegment.empty()
             audio_frames = st.session_state.streamlit_av_queue.get_audio_frames()
             if len(audio_frames) > 0:
-                for audio_frame in audio_frames:
-                    sound = pydub.AudioSegment(
-                        data=audio_frame.to_ndarray().tobytes(),
-                        sample_width=audio_frame.format.bytes,
-                        frame_rate=audio_frame.sample_rate,
-                        channels=len(audio_frame.layout.channels),
-                    )
-                    sound = sound.set_channels(1)
-                    sound = sound.set_frame_rate(st.session_state.speech_to_text_vosk.get_audio_bit_rate())
-                    sound_chunk += sound
-                buffer = np.array(sound_chunk.get_array_of_samples())
-                st.session_state.speech_to_text_vosk.add_speech_bytes(buffer.tobytes())
+                # Concatenate all audio frames into a single buffer
+                audio_buffer = b"".join([buffer.tobytes() for buffer in audio_frames])
+                st.session_state.speech_to_text_vosk.add_speech_bytes(audio_buffer)
             prompt, speaker_finished = st.session_state.speech_to_text_vosk.get_text()
             if speaker_finished and len(prompt) > 0:
                 print(f"Prompt: {prompt}")

input_av_queue_actor.py

@@ -0,0 +1,37 @@
+import ray
+from ray.util.queue import Queue
+from ray.actor import ActorHandle
+import torch
+import numpy as np
+
+
+@ray.remote
+class InputAVQueueActor:
+    def __init__(self):
+        self.audio_queue = Queue(maxsize=100)  # Adjust the size as needed
+        self.video_queue = Queue(maxsize=100)  # Adjust the size as needed
+
+    def enqueue_video_frame(self, shared_tensor_ref):
+        self.video_queue.put(shared_tensor_ref)
+
+    def enqueue_audio_frame(self, shared_buffer_ref):
+        self.audio_queue.put(shared_buffer_ref)
+
+
+    def get_audio_frames(self):
+        audio_frames = []
+        if self.audio_queue.empty():
+            return audio_frames
+        while not self.audio_queue.empty():
+            shared_tensor_ref = self.audio_queue.get()
+            audio_frames.append(shared_tensor_ref)
+        return audio_frames
+
+    def get_video_frames(self):
+        video_frames = []
+        if self.video_queue.empty():
+            return video_frames
+        while not self.video_queue.empty():
+            shared_tensor_ref = self.video_queue.get()
+            video_frames.append(shared_tensor_ref)
+        return video_frames

streamlit_av_queue.py

@@ -5,29 +5,55 @@ from collections import deque
 import threading
 
 import numpy as np
+import ray
+from input_av_queue_actor import InputAVQueueActor
+import pydub
+import torch
 
 class StreamlitAVQueue:
-    def __init__(self):
-        self.audio_frames_deque_lock = threading.Lock()
-        self.audio_frames_deque: deque = deque([])
-
-        self.video_frames_deque_lock = threading.Lock()
-        self.video_frames_deque: deque = deque([])
+    def __init__(self, audio_bit_rate=16000):
+        self._audio_bit_rate = audio_bit_rate
+        try:
+            self.queue_actor = ray.get_actor("InputAVQueueActor")
+        except ValueError as e:
+            self.queue_actor = InputAVQueueActor.options(name="InputAVQueueActor").remote() 
 
     async def queued_video_frames_callback(
                 self,
                 frames: List[av.AudioFrame],
             ) -> av.AudioFrame:
-        with self.video_frames_deque_lock:
-            self.video_frames_deque.extend(frames)
+        try:
+            for frame in frames:
+                shared_tensor = torch.from_numpy(frame.to_ndarray())
+                shared_tensor_ref = ray.put(shared_tensor)
+                self.queue_actor.enqueue_video_frame.remote(shared_tensor_ref)
+        except Exception as e:
+            print (e)
         return frames
 
     async def queued_audio_frames_callback(
                 self,
                 frames: List[av.AudioFrame],
             ) -> av.AudioFrame:
-        with self.audio_frames_deque_lock:
-            self.audio_frames_deque.extend(frames)
+        try:
+            sound_chunk = pydub.AudioSegment.empty()
+            if len(frames) > 0:
+                for frame in frames:
+                    sound = pydub.AudioSegment(
+                        data=frame.to_ndarray().tobytes(),
+                        sample_width=frame.format.bytes,
+                        frame_rate=frame.sample_rate,
+                        channels=len(frame.layout.channels),
+                    )
+                    sound = sound.set_channels(1)
+                    sound = sound.set_frame_rate(self._audio_bit_rate)
+                    sound_chunk += sound
+                shared_buffer = np.array(sound_chunk.get_array_of_samples())
+                shared_buffer_ref = ray.put(shared_buffer)
+                self.queue_actor.enqueue_audio_frame.remote(shared_buffer_ref)
+        except Exception as e:
+            print (e)
+            
         # return empty frames to avoid echo
         new_frames = []
         for frame in frames:
@@ -41,15 +67,9 @@ class StreamlitAVQueue:
         return new_frames
     
     def get_audio_frames(self) -> List[av.AudioFrame]:
-        audio_frames = []
-        with self.audio_frames_deque_lock:
-            audio_frames = list(self.audio_frames_deque)
-            self.audio_frames_deque.clear()
-        return audio_frames
+        shared_buffers = ray.get(self.queue_actor.get_audio_frames.remote())
+        return shared_buffers
     
     def get_video_frames(self) -> List[av.AudioFrame]:
-        video_frames = []
-        with self.video_frames_deque_lock:
-            video_frames = list(self.video_frames_deque)
-            self.video_frames_deque.clear()
-        return video_frames
+        shared_tensors = ray.get(self.queue_actor.get_video_frames.remote())
+        return shared_tensors