From 9829f77052dafec771ca9b8829b0736636d918d7 Mon Sep 17 00:00:00 2001 From: Moishe Lettvin Date: Wed, 28 Feb 2024 13:49:08 -0500 Subject: [PATCH] sample 05 works --- src/dailyai/queue_aggregators.py | 24 +++++-------------- src/dailyai/services/ai_services.py | 1 - .../foundational/05-sync-speech-and-image.py | 15 +++++++----- 3 files changed, 15 insertions(+), 25 deletions(-) diff --git a/src/dailyai/queue_aggregators.py b/src/dailyai/queue_aggregators.py index 3935a056f..ffe756092 100644 --- a/src/dailyai/queue_aggregators.py +++ b/src/dailyai/queue_aggregators.py @@ -14,11 +14,6 @@ from dailyai.services.ai_services import AIService, PipeService from typing import Any, AsyncGenerator, Callable, List, Tuple -@dataclass -class SinkAndQueuePair: - sink: PipeService - queue: asyncio.Queue - class QueueTee(PipeService): def __init__( @@ -26,11 +21,10 @@ class QueueTee(PipeService): ): super().__init__(*args, **kwargs) - self.sinks: List[SinkAndQueuePair] = [] + self.sinks: List[PipeService] = [] for sink in sinks: - pair = SinkAndQueuePair(sink, asyncio.Queue()) - self.sinks.append(pair) - sink.source_queue = pair.queue + sink.source_queue = asyncio.Queue() + self.sinks.append(sink) async def process_queue(self): if not self.source_queue: @@ -38,10 +32,10 @@ class QueueTee(PipeService): while True: frame: QueueFrame = await self.source_queue.get() - print("got frame") for sink in self.sinks: - print("putting frame in sink") - await sink.queue.put(frame) + if sink.source_queue: + await sink.source_queue.put(frame) + if isinstance(frame, EndStreamQueueFrame): break @@ -63,10 +57,6 @@ class QueueFrameAggregator(PipeService): async def process_frame( self, frame: QueueFrame ) -> AsyncGenerator[QueueFrame, None]: - if isinstance(frame, ControlQueueFrame): - yield frame - return - output_frame: QueueFrame | None = None (self.aggregation, output_frame) = self.aggregator( self.aggregation, frame @@ -92,8 +82,6 @@ class QueueMergeGateOnFirst(PipeService): *[source_queue.get() for source_queue in self.source_queues] ) for idx, frame in enumerate(frames): - print("frame", idx, frame) - # if the frame we got from a source is an EndStreamQueueFrame, remove that source if isinstance(frame, EndStreamQueueFrame): self.source_queues.pop(idx) diff --git a/src/dailyai/services/ai_services.py b/src/dailyai/services/ai_services.py index 3bac0ac21..7005a7250 100644 --- a/src/dailyai/services/ai_services.py +++ b/src/dailyai/services/ai_services.py @@ -47,7 +47,6 @@ class PipeService(AbstractPipeService): frame: QueueFrame = await self.source_queue.get() async for output_frame in self.process_frame(frame): if isinstance(frame, EndStreamQueueFrame): - print("end of stream", type(self)) async for final_frame in self.finalize(): await self.sink_queue.put(final_frame) await self.sink_queue.put(output_frame) diff --git a/src/examples/foundational/05-sync-speech-and-image.py b/src/examples/foundational/05-sync-speech-and-image.py index 64699cf52..541fb0781 100644 --- a/src/examples/foundational/05-sync-speech-and-image.py +++ b/src/examples/foundational/05-sync-speech-and-image.py @@ -32,9 +32,9 @@ async def main(room_url): """ - / TTS \ - Month prompt -> LLM -> Fork -> -> Gate -> Transport - \ ImageGen / + / TTS \ + Month prompt -> LLM -> Fork -> -> Gate -> Transport + \ Aggregate -> ImageGen / """ month_description_queue: asyncio.Queue[QueueFrame] = asyncio.Queue() @@ -60,6 +60,9 @@ async def main(room_url): def aggregator( accumulation, frame: QueueFrame ) -> tuple[Any, QueueFrame | None]: + if not accumulation: + accumulation = "" + if isinstance(frame, TextQueueFrame): accumulation += frame.text return (accumulation, None) @@ -70,7 +73,7 @@ async def main(room_url): # This queue service takes chunks from LLM output and merges them into one text frame # that will be used to prompt the image service. - llm_aggregator_for_image = QueueFrameAggregator(source_queue=llm.sink_queue, aggregator=aggregator, finalizer=lambda x: None) + llm_aggregator_for_image = QueueFrameAggregator(aggregator=aggregator, finalizer=lambda x: None) # Set the source queue for the image service to the sink of the aggregator service dalle.source_queue = llm_aggregator_for_image.sink_queue @@ -89,7 +92,7 @@ async def main(room_url): tts_image_gate.sink_queue = transport.send_queue # Queue up all the months in the LLM service source queue - months = ["January", "February"] + months = ["January"] #, "February"] for month in months: messages = [ { @@ -102,7 +105,7 @@ async def main(room_url): await month_description_queue.put(EndStreamQueueFrame()) - await asyncio.gather(transport.run(), *[service.process_queue() for service in [llm, tts, dalle, tee, tts_image_gate]]) + await asyncio.gather(transport.run(), *[service.process_queue() for service in [llm, tts, dalle, tee, tts_image_gate, llm_aggregator_for_image]]) if __name__ == "__main__": (url, token) = configure()