transport(websocket): update audio_frame_size

This commit is contained in:
Aleix Conchillo Flaqué
2024-05-29 23:48:55 -07:00
parent 2957416d90
commit e31e87aabd
2 changed files with 27 additions and 5 deletions

View File

@@ -196,7 +196,7 @@ class ImageGenService(AIService):
super().__init__()
# Renders the image. Returns an Image object.
@ abstractmethod
@abstractmethod
async def run_image_gen(self, prompt: str) -> AsyncGenerator[Frame, None]:
pass
@@ -215,7 +215,7 @@ class VisionService(AIService):
super().__init__()
self._describe_text = None
@ abstractmethod
@abstractmethod
async def run_vision(self, frame: VisionImageRawFrame) -> AsyncGenerator[Frame, None]:
pass

View File

@@ -6,6 +6,8 @@
import asyncio
import io
import wave
import websockets
from typing import Awaitable, Callable
@@ -28,7 +30,8 @@ from loguru import logger
class WebsocketServerParams(TransportParams):
audio_frame_size: int = 16000
add_wav_header: bool = False
audio_frame_size: int = 6400 # 200ms
serializer: FrameSerializer = ProtobufFrameSerializer()
@@ -94,6 +97,8 @@ class WebsocketServerInputTransport(FrameProcessor):
frame = self._params.serializer.deserialize(message)
await self._internal_push_frame(frame)
logger.info(f"Client {websocket.remote_address} disconnected")
async def _start(self):
loop = self.get_event_loop()
self._server_task = loop.create_task(self._server_task_handler())
@@ -159,13 +164,30 @@ class WebsocketServerOutputTransport(FrameProcessor):
while running:
frame = await self._send_queue.get()
if self._websocket and frame:
# We send WAV data so we can easily decoded in the browser.
if self._params.add_wav_header:
content = io.BytesIO()
ww = wave.open(content, "wb")
ww.setsampwidth(2)
ww.setnchannels(frame.num_channels)
ww.setframerate(frame.sample_rate)
ww.writeframes(frame.audio)
ww.close()
content.seek(0)
wav_frame = AudioRawFrame(
content.read(),
sample_rate=frame.sample_rate,
num_channels=frame.num_channels)
frame = wav_frame
proto = self._params.serializer.serialize(frame)
await self._websocket.send(proto)
async def _stop(self):
self._send_queue_task.cancel()
async def process_frame(self, frame: Frame, direction: FrameDirection):
if isinstance(frame, CancelFrame):
# await self.stop()
# We don't queue a CancelFrame since we want to stop ASAP.
await self._stop()
await self.push_frame(frame, direction)
elif isinstance(frame, TTSStartedFrame):
self._in_tts_audio = True