Move some things to AbstractTransport class

This commit is contained in:
Moishe Lettvin
2024-03-27 12:59:08 -04:00
parent e8a6560ac1
commit 4ce140bf84
4 changed files with 37 additions and 22 deletions

View File

@@ -12,9 +12,10 @@ class SequentialMergePipeline(Pipeline):
self.pipelines = pipelines
async def run_pipeline(self):
for pipeline in self.pipelines:
for idx, pipeline in enumerate(self.pipelines):
while True:
frame = await pipeline.sink.get()
print(idx, frame)
if isinstance(
frame, EndFrame) or isinstance(
frame, EndPipeFrame):

View File

@@ -1,12 +1,41 @@
from abc import abstractmethod
import asyncio
import logging
import time
from dailyai.pipeline.frame_processor import FrameProcessor
from dailyai.pipeline.pipeline import Pipeline
class AbstractTransport:
def __init__(self, **kwargs):
pass
self.send_queue = asyncio.Queue()
self.receive_queue = asyncio.Queue()
self.completed_queue = asyncio.Queue()
duration_minutes = kwargs.get("duration_minutes") or 10
self._expiration = time.time() + duration_minutes * 60
self._mic_enabled = kwargs.get("mic_enabled") or False
self._mic_sample_rate = kwargs.get("mic_sample_rate") or 16000
self._camera_enabled = kwargs.get("camera_enabled") or False
self._camera_width = kwargs.get("camera_width") or 1024
self._camera_height = kwargs.get("camera_height") or 768
self._speaker_enabled = kwargs.get("speaker_enabled") or False
self._speaker_sample_rate = kwargs.get("speaker_sample_rate") or 16000
self._fps = kwargs.get("fps") or 8
self._logger: logging.Logger = logging.getLogger("dailyai.transport")
@abstractmethod
async def run(self, pipeline: Pipeline, override_pipeline_source_queue=True):
pass
@abstractmethod
async def run_interruptible_pipeline(
self,
pipeline: Pipeline,
pre_processor: FrameProcessor | None = None,
post_processor: FrameProcessor | None = None,
):
pass

View File

@@ -27,6 +27,7 @@ from dailyai.pipeline.frames import (
)
from dailyai.pipeline.pipeline import Pipeline
from dailyai.services.ai_services import TTSService
from dailyai.transports.abstract_transport import AbstractTransport
torch.set_num_threads(1)
@@ -72,20 +73,13 @@ class VADState(Enum):
STOPPING = 4
class ThreadedTransport:
class ThreadedTransport(AbstractTransport):
def __init__(
self,
**kwargs,
) -> None:
self._mic_enabled = kwargs.get("mic_enabled") or False
self._mic_sample_rate = kwargs.get("mic_sample_rate") or 16000
self._camera_enabled = kwargs.get("camera_enabled") or False
self._camera_width = kwargs.get("camera_width") or 1024
self._camera_height = kwargs.get("camera_height") or 768
self._speaker_enabled = kwargs.get("speaker_enabled") or False
self._speaker_sample_rate = kwargs.get("speaker_sample_rate") or 16000
self._fps = kwargs.get("fps") or 8
super().__init__(**kwargs)
self._vad_start_s = kwargs.get("vad_start_s") or 0.2
self._vad_stop_s = kwargs.get("vad_stop_s") or 0.8
self._context = kwargs.get("context") or []
@@ -105,14 +99,6 @@ class ThreadedTransport:
self._vad_state = VADState.QUIET
self._user_is_speaking = False
duration_minutes = kwargs.get("duration_minutes") or 10
self._expiration = time.time() + duration_minutes * 60
self.send_queue = asyncio.Queue()
self.receive_queue = asyncio.Queue()
self.completed_queue = asyncio.Queue()
self._threadsafe_send_queue = queue.Queue()
self._images = None
@@ -125,8 +111,6 @@ class ThreadedTransport:
self._stop_threads = threading.Event()
self._is_interrupted = threading.Event()
self._logger: logging.Logger = logging.getLogger()
async def run(self, pipeline: Pipeline | None = None, override_pipeline_source_queue=True):
self._prerun()

View File

@@ -7,6 +7,7 @@ from dailyai.pipeline.frame_processor import FrameProcessor
from dailyai.pipeline.frames import AudioFrame, ControlFrame, EndFrame, Frame, TTSEndFrame, TTSStartFrame, TextFrame
from dailyai.pipeline.pipeline import Pipeline
from dailyai.serializers.protobuf_serializer import ProtobufFrameSerializer
from dailyai.transports.abstract_transport import AbstractTransport
from dailyai.transports.threaded_transport import ThreadedTransport
@@ -45,7 +46,7 @@ class WebSocketFrameProcessor(FrameProcessor):
yield frame
class WebsocketTransport(ThreadedTransport):
class WebsocketTransport(AbstractTransport):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self._sample_width = kwargs.get("sample_width", 2)