|
| 1 | +import queue |
| 2 | +import asyncio |
| 3 | +from typing import Any, Union, Callable, AsyncGenerator, cast |
| 4 | + |
| 5 | +import numpy as np |
| 6 | +import sounddevice as sd |
| 7 | +import numpy.typing as npt |
| 8 | + |
| 9 | +from ... import _legacy_response |
| 10 | +from ..._response import StreamedBinaryAPIResponse, AsyncStreamedBinaryAPIResponse |
| 11 | + |
| 12 | +SAMPLE_RATE = 24000 |
| 13 | + |
| 14 | + |
| 15 | +class LocalAudioPlayer: |
| 16 | + def __init__( |
| 17 | + self, |
| 18 | + should_stop: Union[Callable[[], bool], None] = None, |
| 19 | + ): |
| 20 | + self.channels = 1 |
| 21 | + self.dtype = np.float32 |
| 22 | + self.should_stop = should_stop |
| 23 | + |
| 24 | + async def _tts_response_to_buffer( |
| 25 | + self, |
| 26 | + response: Union[ |
| 27 | + _legacy_response.HttpxBinaryResponseContent, |
| 28 | + AsyncStreamedBinaryAPIResponse, |
| 29 | + StreamedBinaryAPIResponse, |
| 30 | + ], |
| 31 | + ) -> npt.NDArray[np.float32]: |
| 32 | + chunks: list[bytes] = [] |
| 33 | + if isinstance(response, _legacy_response.HttpxBinaryResponseContent) or isinstance( |
| 34 | + response, StreamedBinaryAPIResponse |
| 35 | + ): |
| 36 | + for chunk in response.iter_bytes(chunk_size=1024): |
| 37 | + if chunk: |
| 38 | + chunks.append(chunk) |
| 39 | + else: |
| 40 | + async for chunk in response.iter_bytes(chunk_size=1024): |
| 41 | + if chunk: |
| 42 | + chunks.append(chunk) |
| 43 | + |
| 44 | + audio_bytes = b"".join(chunks) |
| 45 | + audio_np = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32767.0 |
| 46 | + audio_np = audio_np.reshape(-1, 1) |
| 47 | + return audio_np |
| 48 | + |
| 49 | + async def play( |
| 50 | + self, |
| 51 | + input: Union[ |
| 52 | + npt.NDArray[np.int16], |
| 53 | + npt.NDArray[np.float32], |
| 54 | + _legacy_response.HttpxBinaryResponseContent, |
| 55 | + AsyncStreamedBinaryAPIResponse, |
| 56 | + StreamedBinaryAPIResponse, |
| 57 | + ], |
| 58 | + ) -> None: |
| 59 | + audio_content: npt.NDArray[np.float32] |
| 60 | + if isinstance(input, np.ndarray): |
| 61 | + if input.dtype == np.int16 and self.dtype == np.float32: |
| 62 | + audio_content = (input.astype(np.float32) / 32767.0).reshape(-1, self.channels) |
| 63 | + elif input.dtype == np.float32: |
| 64 | + audio_content = cast(npt.NDArray[np.float32], input) |
| 65 | + else: |
| 66 | + raise ValueError(f"Unsupported dtype: {input.dtype}") |
| 67 | + else: |
| 68 | + audio_content = await self._tts_response_to_buffer(input) |
| 69 | + |
| 70 | + loop = asyncio.get_event_loop() |
| 71 | + event = asyncio.Event() |
| 72 | + idx = 0 |
| 73 | + |
| 74 | + def callback( |
| 75 | + outdata: npt.NDArray[np.float32], |
| 76 | + frame_count: int, |
| 77 | + _time_info: Any, |
| 78 | + _status: Any, |
| 79 | + ): |
| 80 | + nonlocal idx |
| 81 | + |
| 82 | + remainder = len(audio_content) - idx |
| 83 | + if remainder == 0 or (callable(self.should_stop) and self.should_stop()): |
| 84 | + loop.call_soon_threadsafe(event.set) |
| 85 | + raise sd.CallbackStop |
| 86 | + valid_frames = frame_count if remainder >= frame_count else remainder |
| 87 | + outdata[:valid_frames] = audio_content[idx : idx + valid_frames] |
| 88 | + outdata[valid_frames:] = 0 |
| 89 | + idx += valid_frames |
| 90 | + |
| 91 | + stream = sd.OutputStream( |
| 92 | + samplerate=SAMPLE_RATE, |
| 93 | + callback=callback, |
| 94 | + dtype=audio_content.dtype, |
| 95 | + channels=audio_content.shape[1], |
| 96 | + ) |
| 97 | + with stream: |
| 98 | + await event.wait() |
| 99 | + |
| 100 | + async def play_stream( |
| 101 | + self, |
| 102 | + buffer_stream: AsyncGenerator[Union[npt.NDArray[np.float32], npt.NDArray[np.int16], None], None], |
| 103 | + ) -> None: |
| 104 | + loop = asyncio.get_event_loop() |
| 105 | + event = asyncio.Event() |
| 106 | + buffer_queue: queue.Queue[Union[npt.NDArray[np.float32], npt.NDArray[np.int16], None]] = queue.Queue(maxsize=50) |
| 107 | + |
| 108 | + async def buffer_producer(): |
| 109 | + async for buffer in buffer_stream: |
| 110 | + if buffer is None: |
| 111 | + break |
| 112 | + await loop.run_in_executor(None, buffer_queue.put, buffer) |
| 113 | + await loop.run_in_executor(None, buffer_queue.put, None) # Signal completion |
| 114 | + |
| 115 | + def callback( |
| 116 | + outdata: npt.NDArray[np.float32], |
| 117 | + frame_count: int, |
| 118 | + _time_info: Any, |
| 119 | + _status: Any, |
| 120 | + ): |
| 121 | + nonlocal current_buffer, buffer_pos |
| 122 | + |
| 123 | + frames_written = 0 |
| 124 | + while frames_written < frame_count: |
| 125 | + if current_buffer is None or buffer_pos >= len(current_buffer): |
| 126 | + try: |
| 127 | + current_buffer = buffer_queue.get(timeout=0.1) |
| 128 | + if current_buffer is None: |
| 129 | + loop.call_soon_threadsafe(event.set) |
| 130 | + raise sd.CallbackStop |
| 131 | + buffer_pos = 0 |
| 132 | + |
| 133 | + if current_buffer.dtype == np.int16 and self.dtype == np.float32: |
| 134 | + current_buffer = (current_buffer.astype(np.float32) / 32767.0).reshape(-1, self.channels) |
| 135 | + |
| 136 | + except queue.Empty: |
| 137 | + outdata[frames_written:] = 0 |
| 138 | + return |
| 139 | + |
| 140 | + remaining_frames = len(current_buffer) - buffer_pos |
| 141 | + frames_to_write = min(frame_count - frames_written, remaining_frames) |
| 142 | + outdata[frames_written : frames_written + frames_to_write] = current_buffer[ |
| 143 | + buffer_pos : buffer_pos + frames_to_write |
| 144 | + ] |
| 145 | + buffer_pos += frames_to_write |
| 146 | + frames_written += frames_to_write |
| 147 | + |
| 148 | + current_buffer = None |
| 149 | + buffer_pos = 0 |
| 150 | + |
| 151 | + producer_task = asyncio.create_task(buffer_producer()) |
| 152 | + |
| 153 | + with sd.OutputStream( |
| 154 | + samplerate=SAMPLE_RATE, |
| 155 | + channels=self.channels, |
| 156 | + dtype=self.dtype, |
| 157 | + callback=callback, |
| 158 | + ): |
| 159 | + await event.wait() |
| 160 | + |
| 161 | + await producer_task |
0 commit comments