From a674b43243f8d70a57c2e2103623e958992c7394 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Mon, 3 Jun 2024 20:54:55 -0700 Subject: [PATCH] transport: remove redundant camera thread and switch audio pull for push --- src/pipecat/transports/base_input.py | 41 ++++++++-------- src/pipecat/transports/base_output.py | 21 +++++---- src/pipecat/transports/local/audio.py | 26 ++++++---- src/pipecat/transports/local/tk.py | 27 ++++++----- src/pipecat/transports/services/daily.py | 60 ++++++++++++------------ 5 files changed, 92 insertions(+), 83 deletions(-) diff --git a/src/pipecat/transports/base_input.py b/src/pipecat/transports/base_input.py index ba4f8ae4e..6fb1e2ee8 100644 --- a/src/pipecat/transports/base_input.py +++ b/src/pipecat/transports/base_input.py @@ -36,11 +36,7 @@ class BaseInputTransport(FrameProcessor): self._running = False self._allow_interruptions = False - self._in_executor = ThreadPoolExecutor(max_workers=5) - - # Create audio input queue if needed. - if self._params.audio_in_enabled or self._params.vad_enabled: - self._audio_in_queue = queue.Queue() + self._executor = ThreadPoolExecutor(max_workers=5) # Create push frame task. This is the task that will push frames in # order. We also guarantee that all frames are pushed in the same task. @@ -57,9 +53,11 @@ class BaseInputTransport(FrameProcessor): self._running = True + # Create audio input queue and thread if needed. if self._params.audio_in_enabled or self._params.vad_enabled: - loop = self.get_event_loop() - self._audio_thread = loop.run_in_executor(self._in_executor, self._audio_thread_handler) + self._audio_in_queue = queue.Queue() + self._audio_thread = self._loop.run_in_executor( + self._executor, self._audio_thread_handler) async def stop(self): if not self._running: @@ -77,8 +75,8 @@ class BaseInputTransport(FrameProcessor): def vad_analyzer(self) -> VADAnalyzer | None: return self._params.vad_analyzer - def read_next_audio_frame(self) -> AudioRawFrame | None: - pass + def push_audio_frame(self, frame: AudioRawFrame): + self._audio_in_queue.put_nowait(frame) # # Frame processor @@ -174,21 +172,22 @@ class BaseInputTransport(FrameProcessor): vad_state: VADState = VADState.QUIET while self._running: try: - frame = self.read_next_audio_frame() + frame: AudioRawFrame = self._audio_in_queue.get(timeout=1) - if frame: - audio_passthrough = True + audio_passthrough = True - # Check VAD and push event if necessary. We just care about - # changes from QUIET to SPEAKING and vice versa. - if self._params.vad_enabled: - vad_state = self._handle_vad(frame.audio, vad_state) - audio_passthrough = self._params.vad_audio_passthrough + # Check VAD and push event if necessary. We just care about + # changes from QUIET to SPEAKING and vice versa. + if self._params.vad_enabled: + vad_state = self._handle_vad(frame.audio, vad_state) + audio_passthrough = self._params.vad_audio_passthrough # Push audio downstream if passthrough. - if audio_passthrough: - future = asyncio.run_coroutine_threadsafe( - self._internal_push_frame(frame), self.get_event_loop()) - future.result() + if audio_passthrough: + future = asyncio.run_coroutine_threadsafe( + self._internal_push_frame(frame), self._loop) + future.result() + except queue.Empty: + pass except BaseException as e: logger.error(f"Error reading audio frames: {e}") diff --git a/src/pipecat/transports/base_output.py b/src/pipecat/transports/base_output.py index 146c58496..9e99ce736 100644 --- a/src/pipecat/transports/base_output.py +++ b/src/pipecat/transports/base_output.py @@ -43,7 +43,7 @@ class BaseOutputTransport(FrameProcessor): self._running = False self._allow_interruptions = False - self._out_executor = ThreadPoolExecutor(max_workers=5) + self._executor = ThreadPoolExecutor(max_workers=5) # These are the images that we should send to the camera at our desired # framerate. @@ -57,6 +57,10 @@ class BaseOutputTransport(FrameProcessor): self._stopped_event = asyncio.Event() self._is_interrupted = threading.Event() + # Create push frame task. This is the task that will push frames in + # order. We also guarantee that all frames are pushed in the same task. + self._create_push_task() + async def start(self, frame: StartFrame): # Make sure we have the latest params. Note that this transport might # have been started on another task that might not need interruptions, @@ -70,15 +74,12 @@ class BaseOutputTransport(FrameProcessor): loop = self.get_event_loop() + # Create queues and threads. if self._params.camera_out_enabled: self._camera_out_thread = loop.run_in_executor( - self._out_executor, self._camera_out_thread_handler) + self._executor, self._camera_out_thread_handler) - self._sink_thread = loop.run_in_executor(self._out_executor, self._sink_thread_handler) - - # Create push frame task. This is the task that will push frames in - # order. We also guarantee that all frames are pushed in the same task. - self._create_push_task() + self._sink_thread = loop.run_in_executor(self._executor, self._sink_thread_handler) async def stop(self): if not self._running: @@ -117,7 +118,7 @@ class BaseOutputTransport(FrameProcessor): # if isinstance(frame, StartFrame): await self.start(frame) - self._sink_queue.put(frame) + self._sink_queue.put_nowait(frame) # EndFrame is managed in the queue handler. elif isinstance(frame, CancelFrame): await self.stop() @@ -126,7 +127,7 @@ class BaseOutputTransport(FrameProcessor): await self._handle_interruptions(frame) await self.push_frame(frame, direction) else: - self._sink_queue.put(frame) + self._sink_queue.put_nowait(frame) # If we are finishing, wait here until we have stopped, otherwise we might # close things too early upstream. We need this event because we don't @@ -233,7 +234,7 @@ class BaseOutputTransport(FrameProcessor): def _set_camera_image(self, image: ImageRawFrame): if self._params.camera_out_is_live: - self._camera_out_queue.put(image) + self._camera_out_queue.put_nowait(image) else: self._camera_images = itertools.cycle([image]) diff --git a/src/pipecat/transports/local/audio.py b/src/pipecat/transports/local/audio.py index 14b8bd5d3..6ac461b3b 100644 --- a/src/pipecat/transports/local/audio.py +++ b/src/pipecat/transports/local/audio.py @@ -28,22 +28,17 @@ class LocalAudioInputTransport(BaseInputTransport): def __init__(self, py_audio: pyaudio.PyAudio, params: TransportParams): super().__init__(params) + sample_rate = self._params.audio_in_sample_rate + num_frames = int(sample_rate / 100) # 10ms of audio + self._in_stream = py_audio.open( format=py_audio.get_format_from_width(2), channels=params.audio_in_channels, rate=params.audio_in_sample_rate, - frames_per_buffer=params.audio_in_sample_rate, + frames_per_buffer=num_frames, + stream_callback=self._audio_in_callback, input=True) - def read_next_audio_frame(self) -> AudioRawFrame | None: - sample_rate = self._params.audio_in_sample_rate - num_channels = self._params.audio_in_channels - num_frames = int(sample_rate / 100) # 10ms of audio - - audio = self._in_stream.read(num_frames, exception_on_overflow=False) - - return AudioRawFrame(audio=audio, sample_rate=sample_rate, num_channels=num_channels) - async def start(self, frame: StartFrame): await super().start(frame) self._in_stream.start_stream() @@ -60,6 +55,17 @@ class LocalAudioInputTransport(BaseInputTransport): await super().cleanup() + def _audio_in_callback(self, in_data, frame_count, time_info, status): + if not self._running: + return (None, pyaudio.paAbort) + + frame = AudioRawFrame(audio=in_data, + sample_rate=self._params.audio_in_sample_rate, + num_channels=self._params.audio_in_channels) + self.push_audio_frame(frame) + + return (None, pyaudio.paContinue) + class LocalAudioOutputTransport(BaseOutputTransport): diff --git a/src/pipecat/transports/local/tk.py b/src/pipecat/transports/local/tk.py index 808837998..b8de9ce71 100644 --- a/src/pipecat/transports/local/tk.py +++ b/src/pipecat/transports/local/tk.py @@ -38,22 +38,17 @@ class TkInputTransport(BaseInputTransport): def __init__(self, py_audio: pyaudio.PyAudio, params: TransportParams): super().__init__(params) + sample_rate = self._params.audio_in_sample_rate + num_frames = int(sample_rate / 100) # 10ms of audio + self._in_stream = py_audio.open( format=py_audio.get_format_from_width(2), channels=params.audio_in_channels, rate=params.audio_in_sample_rate, - frames_per_buffer=params.audio_in_sample_rate, + frames_per_buffer=num_frames, + stream_callback=self._audio_in_callback, input=True) - def read_next_audio_frame(self) -> AudioRawFrame | None: - sample_rate = self._params.audio_in_sample_rate - num_channels = self._params.audio_in_channels - num_frames = int(sample_rate / 100) # 10ms of audio - - audio = self._in_stream.read(num_frames, exception_on_overflow=False) - - return AudioRawFrame(audio=audio, sample_rate=sample_rate, num_channels=num_channels) - async def start(self, frame: StartFrame): await super().start(frame) self._in_stream.start_stream() @@ -63,12 +58,22 @@ class TkInputTransport(BaseInputTransport): self._in_stream.stop_stream() async def cleanup(self): + await super().cleanup() # This is not very pretty (taken from PyAudio docs). while self._in_stream.is_active(): await asyncio.sleep(0.1) self._in_stream.close() - await super().cleanup() + def _audio_in_callback(self, in_data, frame_count, time_info, status): + if not self._running: + return (None, pyaudio.paAbort) + + frame = AudioRawFrame(audio=in_data, + sample_rate=self._params.audio_in_sample_rate, + num_channels=self._params.audio_in_channels) + self.push_audio_frame(frame) + + return (None, pyaudio.paContinue) class TkOutputTransport(BaseOutputTransport): diff --git a/src/pipecat/transports/services/daily.py b/src/pipecat/transports/services/daily.py index f20eaa4dc..5ce4a6c84 100644 --- a/src/pipecat/transports/services/daily.py +++ b/src/pipecat/transports/services/daily.py @@ -37,7 +37,7 @@ from pipecat.processors.frame_processor import FrameDirection, FrameProcessor from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport, TransportParams -from pipecat.vad.vad_analyzer import VADAnalyzer, VADParams, VADState +from pipecat.vad.vad_analyzer import VADAnalyzer, VADParams from loguru import logger @@ -193,7 +193,7 @@ class DailyTransportClient(EventHandler): num_channels = self._params.audio_in_channels if self._other_participant_has_joined: - num_frames = int(sample_rate / 100) # 10ms of audio + num_frames = int(sample_rate / 100) * 2 # 20ms of audio audio = self._speaker.read_frames(num_frames) @@ -472,7 +472,6 @@ class DailyInputTransport(BaseInputTransport): self._client = client self._video_renderers = {} - self._camera_in_queue = queue.Queue() self._vad_analyzer: VADAnalyzer | None = params.vad_analyzer if params.vad_enabled and not params.vad_analyzer: @@ -483,23 +482,26 @@ class DailyInputTransport(BaseInputTransport): async def start(self, frame: StartFrame): if self._running: return + # Parent start. + await super().start(frame) # Join the room. await self._client.join() - # This will set _running=True - await super().start(frame) - # Create camera in thread (runs if _running is true). - self._camera_in_thread = self._loop.run_in_executor( - self._in_executor, self._camera_in_thread_handler) + # Create audio task. It reads audio frames from Daily and push them + # internally for VAD processing. + if self._params.audio_in_enabled or self._params.vad_enabled: + self._audio_in_thread = self._loop.run_in_executor( + self._executor, self._audio_in_thread_handler) async def stop(self): if not self._running: return + # Parent stop. This will set _running to False. + await super().stop() # Leave the room. await self._client.leave() - # This will set _running=False - await super().stop() - # The thread will stop. - await self._camera_in_thread + # Stop audio thread. + if self._params.audio_in_enabled or self._params.vad_enabled: + await self._audio_in_thread async def cleanup(self): await super().cleanup() @@ -508,9 +510,6 @@ class DailyInputTransport(BaseInputTransport): def vad_analyzer(self) -> VADAnalyzer | None: return self._vad_analyzer - def read_next_audio_frame(self) -> AudioRawFrame | None: - return self._client.read_next_audio_frame() - # # FrameProcessor # @@ -536,6 +535,16 @@ class DailyInputTransport(BaseInputTransport): self._internal_push_frame(frame), self.get_event_loop()) future.result() + # + # Audio in + # + + def _audio_in_thread_handler(self): + while self._running: + frame = self._client.read_next_audio_frame() + if frame: + self.push_audio_frame(frame) + # # Camera in # @@ -584,23 +593,12 @@ class DailyInputTransport(BaseInputTransport): image=buffer, size=size, format=format) - self._camera_in_queue.put(frame) + future = asyncio.run_coroutine_threadsafe( + self._internal_push_frame(frame), self.get_event_loop()) + future.result() self._video_renderers[participant_id]["timestamp"] = curr_time - def _camera_in_thread_handler(self): - while self._running: - try: - frame = self._camera_in_queue.get(timeout=1) - future = asyncio.run_coroutine_threadsafe( - self._internal_push_frame(frame), self.get_event_loop()) - future.result() - self._camera_in_queue.task_done() - except queue.Empty: - pass - except BaseException as e: - logger.error(f"Error capturing video: {e}") - class DailyOutputTransport(BaseOutputTransport): @@ -612,7 +610,7 @@ class DailyOutputTransport(BaseOutputTransport): async def start(self, frame: StartFrame): if self._running: return - # This will set _running=True + # Parent start. await super().start(frame) # Join the room. await self._client.join() @@ -620,7 +618,7 @@ class DailyOutputTransport(BaseOutputTransport): async def stop(self): if not self._running: return - # This will set _running=False + # Parent stop. This will set _running to False. await super().stop() # Leave the room. await self._client.leave()