wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

transports: fully use asyncio in all read/write operations

Browse Source
This commit is contained in:
Aleix Conchillo Flaqué
2024-06-17 17:44:19 -07:00
parent c544b2474b
commit d3e410b2ac
12 changed files with 335 additions and 346 deletions
Download Patch File Download Diff File Expand all files Collapse all files

172
src/pipecat/transports/base_output.py
View File

@@ -7,11 +7,6 @@
import asyncio
import itertools
import queue
import time
import threading
from concurrent.futures import ThreadPoolExecutor
from PIL import Image
from typing import List
@@ -42,67 +37,51 @@ class BaseOutputTransport(FrameProcessor):
self._params = params
self._running = False
self._executor = ThreadPoolExecutor(max_workers=5)
# These are the images that we should send to the camera at our desired
# framerate.
self._camera_images = None
# Create media threads queues.
if self._params.camera_out_enabled:
self._camera_out_queue = queue.Queue()
self._sink_queue = queue.Queue()
self._sink_thread = None
self._stopped_event = asyncio.Event()
self._is_interrupted = threading.Event()
# We will write 20ms audio at a time. If we receive long audio frames we
# will chunk them. This will help with interruption handling.
audio_bytes_10ms = int(self._params.audio_out_sample_rate / 100) * \
self._params.audio_out_channels * 2
self._audio_chunk_size = audio_bytes_10ms * 2
self._stopped_event = asyncio.Event()
# Create sink frame task. This is the task that will actually write
# audio or video frames. We write audio/video in a task so we can keep
# generating frames upstream while, for example, the audio is playing.
self._create_sink_task()
# 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):
if self._running:
return
self._running = True
loop = self.get_event_loop()
# Create queues and threads.
# Create media threads queues.
if self._params.camera_out_enabled:
self._camera_out_thread = loop.run_in_executor(
self._executor, self._camera_out_thread_handler)
self._sink_thread = loop.run_in_executor(self._executor, self._sink_thread_handler)
self._camera_out_queue = asyncio.Queue()
self._camera_out_task = self.get_event_loop().create_task(self._camera_out_task_handler())
async def stop(self):
if not self._running:
return
# This will exit all threads.
self._running = False
# Wait on the threads to finish.
if self._params.camera_out_enabled:
self._camera_out_task.cancel()
await self._camera_out_task
self._stopped_event.set()
def send_message(self, frame: TransportMessageFrame):
async def send_message(self, frame: TransportMessageFrame):
pass
def send_metrics(self, frame: MetricsFrame):
async def send_metrics(self, frame: MetricsFrame):
pass
def write_frame_to_camera(self, frame: ImageRawFrame):
async def write_frame_to_camera(self, frame: ImageRawFrame):
pass
def write_raw_audio_frames(self, frames: bytes):
async def write_raw_audio_frames(self, frames: bytes):
pass
#
@@ -110,12 +89,12 @@ class BaseOutputTransport(FrameProcessor):
#
async def cleanup(self):
# Wait on the threads to finish.
if self._params.camera_out_enabled:
await self._camera_out_thread
if self._sink_task:
self._sink_task.cancel()
await self._sink_task
if self._sink_thread:
await self._sink_thread
self._push_frame_task.cancel()
await self._push_frame_task
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)
@@ -128,7 +107,7 @@ class BaseOutputTransport(FrameProcessor):
if isinstance(frame, StartFrame):
await self.start(frame)
await self.push_frame(frame, direction)
# EndFrame is managed in the queue handler.
# EndFrame is managed in the sink queue handler.
elif isinstance(frame, CancelFrame):
await self.stop()
await self.push_frame(frame, direction)
@@ -136,14 +115,14 @@ class BaseOutputTransport(FrameProcessor):
await self._handle_interruptions(frame)
await self.push_frame(frame, direction)
elif isinstance(frame, MetricsFrame):
self.send_metrics(frame)
await self.send_metrics(frame)
await self.push_frame(frame, direction)
elif isinstance(frame, SystemFrame):
await self.push_frame(frame, direction)
elif isinstance(frame, AudioRawFrame):
await self._handle_audio(frame)
else:
self._sink_queue.put_nowait(frame)
await self._sink_queue.put(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
@@ -156,50 +135,51 @@ class BaseOutputTransport(FrameProcessor):
return
if isinstance(frame, StartInterruptionFrame):
self._is_interrupted.set()
# Stop sink task.
self._sink_task.cancel()
await self._sink_task
self._create_sink_task()
# Stop push task.
self._push_frame_task.cancel()
await self._push_frame_task
self._create_push_task()
elif isinstance(frame, StopInterruptionFrame):
self._is_interrupted.clear()
async def _handle_audio(self, frame: AudioRawFrame):
audio = frame.audio
for i in range(0, len(audio), self._audio_chunk_size):
chunk = AudioRawFrame(audio[i: i + self._audio_chunk_size],
sample_rate=frame.sample_rate, num_channels=frame.num_channels)
self._sink_queue.put_nowait(chunk)
await self._sink_queue.put(chunk)
def _sink_thread_handler(self):
def _create_sink_task(self):
loop = self.get_event_loop()
self._sink_queue = asyncio.Queue()
self._sink_task = loop.create_task(self._sink_task_handler())
async def _sink_task_handler(self):
# Audio accumlation buffer
buffer = bytearray()
while self._running:
while True:
try:
frame = self._sink_queue.get(timeout=1)
if not self._is_interrupted.is_set():
if isinstance(frame, AudioRawFrame) and self._params.audio_out_enabled:
buffer.extend(frame.audio)
buffer = self._maybe_send_audio(buffer)
elif isinstance(frame, ImageRawFrame) and self._params.camera_out_enabled:
self._set_camera_image(frame)
elif isinstance(frame, SpriteFrame) and self._params.camera_out_enabled:
self._set_camera_images(frame.images)
elif isinstance(frame, TransportMessageFrame):
self.send_message(frame)
else:
future = asyncio.run_coroutine_threadsafe(
self._internal_push_frame(frame), self.get_event_loop())
future.result()
frame = await self._sink_queue.get()
if isinstance(frame, AudioRawFrame) and self._params.audio_out_enabled:
buffer.extend(frame.audio)
buffer = await self._maybe_send_audio(buffer)
elif isinstance(frame, ImageRawFrame) and self._params.camera_out_enabled:
await self._set_camera_image(frame)
elif isinstance(frame, SpriteFrame) and self._params.camera_out_enabled:
await self._set_camera_images(frame.images)
elif isinstance(frame, TransportMessageFrame):
await self.send_message(frame)
else:
# If we get interrupted just clear the output buffer.
buffer = bytearray()
await self._internal_push_frame(frame)
if isinstance(frame, EndFrame):
future = asyncio.run_coroutine_threadsafe(self.stop(), self.get_event_loop())
future.result()
await self.stop()
self._sink_queue.task_done()
except queue.Empty:
pass
except asyncio.CancelledError:
break
except BaseException as e:
logger.error(f"{self} error processing sink queue: {e}")
@@ -209,8 +189,8 @@ class BaseOutputTransport(FrameProcessor):
def _create_push_task(self):
loop = self.get_event_loop()
self._push_frame_task = loop.create_task(self._push_frame_task_handler())
self._push_queue = asyncio.Queue()
self._push_frame_task = loop.create_task(self._push_frame_task_handler())
async def _internal_push_frame(
self,
@@ -233,7 +213,7 @@ class BaseOutputTransport(FrameProcessor):
async def send_image(self, frame: ImageRawFrame | SpriteFrame):
await self.process_frame(frame, FrameDirection.DOWNSTREAM)
def _draw_image(self, frame: ImageRawFrame):
async def _draw_image(self, frame: ImageRawFrame):
desired_size = (self._params.camera_out_width, self._params.camera_out_height)
if frame.size != desired_size:
@@ -243,32 +223,32 @@ class BaseOutputTransport(FrameProcessor):
f"{frame} does not have the expected size {desired_size}, resizing")
frame = ImageRawFrame(resized_image.tobytes(), resized_image.size, resized_image.format)
self.write_frame_to_camera(frame)
await self.write_frame_to_camera(frame)
def _set_camera_image(self, image: ImageRawFrame):
async def _set_camera_image(self, image: ImageRawFrame):
if self._params.camera_out_is_live:
self._camera_out_queue.put_nowait(image)
await self._camera_out_queue.put(image)
else:
self._camera_images = itertools.cycle([image])
def _set_camera_images(self, images: List[ImageRawFrame]):
async def _set_camera_images(self, images: List[ImageRawFrame]):
self._camera_images = itertools.cycle(images)
def _camera_out_thread_handler(self):
while self._running:
async def _camera_out_task_handler(self):
while True:
try:
if self._params.camera_out_is_live:
image = self._camera_out_queue.get(timeout=1)
self._draw_image(image)
image = await self._camera_out_queue.get()
await self._draw_image(image)
self._camera_out_queue.task_done()
elif self._camera_images:
image = next(self._camera_images)
self._draw_image(image)
time.sleep(1.0 / self._params.camera_out_framerate)
await self._draw_image(image)
await asyncio.sleep(1.0 / self._params.camera_out_framerate)
else:
time.sleep(1.0 / self._params.camera_out_framerate)
except queue.Empty:
pass
await asyncio.sleep(1.0 / self._params.camera_out_framerate)
except asyncio.CancelledError:
break
except Exception as e:
logger.error(f"{self} error writing to camera: {e}")
@@ -279,12 +259,8 @@ class BaseOutputTransport(FrameProcessor):
async def send_audio(self, frame: AudioRawFrame):
await self.process_frame(frame, FrameDirection.DOWNSTREAM)
def _maybe_send_audio(self, buffer: bytearray) -> bytearray:
try:
if len(buffer) >= self._audio_chunk_size:
self.write_raw_audio_frames(bytes(buffer[:self._audio_chunk_size]))
buffer = buffer[self._audio_chunk_size:]
return buffer
except BaseException as e:
logger.error(f"{self} error writing audio frames: {e}")
return buffer
async def _maybe_send_audio(self, buffer: bytearray) -> bytearray:
if len(buffer) >= self._audio_chunk_size:
await self.write_raw_audio_frames(bytes(buffer[:self._audio_chunk_size]))
buffer = buffer[self._audio_chunk_size:]
return buffer