Optimize pipeline processing so we don't wait for the completion of one generator to move onto the next.

This commit is contained in:
Moishe Lettvin
2024-03-07 18:59:24 -05:00
parent 3c5f4800d4
commit d0076dd4ee
5 changed files with 31 additions and 19 deletions

View File

@@ -48,6 +48,16 @@ class Pipeline:
raise ValueError("Source queue not set") raise ValueError("Source queue not set")
yield await self.source.get() yield await self.source.get()
async def run_pipeline_recursively(
self, initial_frame: Frame, processors: List[FrameProcessor]
) -> AsyncGenerator[Frame, None]:
if processors:
async for frame in processors[0].process_frame(initial_frame):
async for final_frame in self.run_pipeline_recursively(frame, processors[1:]):
yield final_frame
else:
yield initial_frame
async def run_pipeline(self): async def run_pipeline(self):
""" Run the pipeline. Take each frame from the source queue, pass it to """ Run the pipeline. Take each frame from the source queue, pass it to
the first frame_processor, pass the output of that frame_processor to the the first frame_processor, pass the output of that frame_processor to the
@@ -65,23 +75,12 @@ class Pipeline:
try: try:
while True: while True:
frame_generators = [self.get_next_source_frame()] initial_frame = await self.source.get()
for processor in self.processors: async for frame in self.run_pipeline_recursively(initial_frame, self.processors):
next_frame_generators = [] await self.sink.put(frame)
for frame_generator in frame_generators:
async for frame in frame_generator:
next_frame_generators.append(processor.process_frame(frame))
frame_generators = next_frame_generators
for frame_generator in frame_generators: if isinstance(initial_frame, EndFrame) or isinstance(initial_frame, EndPipeFrame):
async for frame in frame_generator: break
await self.sink.put(frame)
if isinstance(
frame, EndFrame
) or isinstance(
frame, EndPipeFrame
):
return
except asyncio.CancelledError: except asyncio.CancelledError:
# this means there's been an interruption, do any cleanup necessary here. # this means there's been an interruption, do any cleanup necessary here.
for processor in self.processors: for processor in self.processors:

View File

@@ -138,6 +138,7 @@ class TTSService(AIService):
text = frame.text text = frame.text
else: else:
self.current_sentence += frame.text self.current_sentence += frame.text
print(f"current sentence: {self.current_sentence}")
if self.current_sentence.endswith((".", "?", "!")): if self.current_sentence.endswith((".", "?", "!")):
text = self.current_sentence text = self.current_sentence
self.current_sentence = "" self.current_sentence = ""

View File

@@ -1,5 +1,4 @@
import asyncio import asyncio
from doctest import OutputChecker
import unittest import unittest
from dailyai.pipeline.aggregators import SentenceAggregator, StatelessTextTransformer from dailyai.pipeline.aggregators import SentenceAggregator, StatelessTextTransformer
from dailyai.pipeline.frames import EndFrame, TextFrame from dailyai.pipeline.frames import EndFrame, TextFrame

View File

@@ -47,7 +47,20 @@ async def main(room_url):
source_queue = asyncio.Queue() source_queue = asyncio.Queue()
for month in ["January", "February"]: for month in [
"January",
"February",
"March",
"April",
"May",
"June",
"July",
"August",
"September",
"October",
"November",
"December",
]:
messages = [ messages = [
{ {
"role": "system", "role": "system",

View File

@@ -51,8 +51,8 @@ async def main(room_url: str, token):
tma_in, tma_in,
llm, llm,
fl2, fl2,
tts,
tma_out, tma_out,
tts
], ],
) )
await transport.run_uninterruptible_pipeline(pipeline) await transport.run_uninterruptible_pipeline(pipeline)