|
| 1 | +import asyncio |
| 2 | +import numpy as np |
| 3 | +import contextlib |
| 4 | +from dataclasses import dataclass |
| 5 | +from typing import AsyncIterator, Optional |
| 6 | +from .audio_frame import AudioFrame |
| 7 | +from .log import logger |
| 8 | + |
| 9 | +_Stream = AsyncIterator[AudioFrame] |
| 10 | + |
| 11 | + |
| 12 | +@dataclass |
| 13 | +class _Contribution: |
| 14 | + stream: _Stream |
| 15 | + data: np.ndarray |
| 16 | + buffer: np.ndarray |
| 17 | + had_data: bool |
| 18 | + exhausted: bool |
| 19 | + |
| 20 | + |
| 21 | +class AudioMixer: |
| 22 | + def __init__( |
| 23 | + self, |
| 24 | + sample_rate: int, |
| 25 | + num_channels: int, |
| 26 | + *, |
| 27 | + blocksize: int = 0, |
| 28 | + stream_timeout_ms: int = 100, |
| 29 | + capacity: int = 100, |
| 30 | + ) -> None: |
| 31 | + """ |
| 32 | + Initialize the AudioMixer. |
| 33 | +
|
| 34 | + The mixer accepts multiple async audio streams and mixes them into a single output stream. |
| 35 | + Each output frame is generated with a fixed chunk size determined by the blocksize (in samples). |
| 36 | + If blocksize is not provided (or 0), it defaults to 100ms. |
| 37 | +
|
| 38 | + Each input stream is processed in parallel, accumulating audio data until at least one chunk |
| 39 | + of samples is available. If an input stream does not provide data within the specified timeout, |
| 40 | + a warning is logged. The mixer can be closed immediately |
| 41 | + (dropping unconsumed frames) or allowed to flush remaining data using end_input(). |
| 42 | +
|
| 43 | + Args: |
| 44 | + sample_rate (int): The audio sample rate in Hz. |
| 45 | + num_channels (int): The number of audio channels. |
| 46 | + blocksize (int, optional): The size of the audio block (in samples) for mixing. If not provided, |
| 47 | + defaults to sample_rate // 10. |
| 48 | + stream_timeout_ms (int, optional): The maximum wait time in milliseconds for each stream to provide |
| 49 | + audio data before timing out. Defaults to 100 ms. |
| 50 | + capacity (int, optional): The maximum number of mixed frames to store in the output queue. |
| 51 | + Defaults to 100. |
| 52 | + """ |
| 53 | + self._streams: set[_Stream] = set() |
| 54 | + self._buffers: dict[_Stream, np.ndarray] = {} |
| 55 | + self._sample_rate: int = sample_rate |
| 56 | + self._num_channels: int = num_channels |
| 57 | + self._chunk_size: int = blocksize if blocksize > 0 else int(sample_rate // 10) |
| 58 | + self._stream_timeout_ms: int = stream_timeout_ms |
| 59 | + self._queue: asyncio.Queue[Optional[AudioFrame]] = asyncio.Queue(maxsize=capacity) |
| 60 | + # _ending signals that no new streams will be added, |
| 61 | + # but we continue processing until all streams are exhausted. |
| 62 | + self._ending: bool = False |
| 63 | + self._mixer_task: asyncio.Task = asyncio.create_task(self._mixer()) |
| 64 | + |
| 65 | + def add_stream(self, stream: AsyncIterator[AudioFrame]) -> None: |
| 66 | + """ |
| 67 | + Add an audio stream to the mixer. |
| 68 | +
|
| 69 | + The stream is added to the internal set of streams and an empty buffer is initialized for it, |
| 70 | + if not already present. |
| 71 | +
|
| 72 | + Args: |
| 73 | + stream (AsyncIterator[AudioFrame]): An async iterator that produces AudioFrame objects. |
| 74 | + """ |
| 75 | + if self._ending: |
| 76 | + raise RuntimeError("Cannot add stream after mixer has been closed") |
| 77 | + |
| 78 | + self._streams.add(stream) |
| 79 | + if stream not in self._buffers: |
| 80 | + self._buffers[stream] = np.empty((0, self._num_channels), dtype=np.int16) |
| 81 | + |
| 82 | + def remove_stream(self, stream: AsyncIterator[AudioFrame]) -> None: |
| 83 | + """ |
| 84 | + Remove an audio stream from the mixer. |
| 85 | +
|
| 86 | + This method removes the specified stream and its associated buffer from the mixer. |
| 87 | +
|
| 88 | + Args: |
| 89 | + stream (AsyncIterator[AudioFrame]): The audio stream to remove. |
| 90 | + """ |
| 91 | + self._streams.discard(stream) |
| 92 | + self._buffers.pop(stream, None) |
| 93 | + |
| 94 | + def __aiter__(self) -> "AudioMixer": |
| 95 | + return self |
| 96 | + |
| 97 | + async def __anext__(self) -> AudioFrame: |
| 98 | + item = await self._queue.get() |
| 99 | + if item is None: |
| 100 | + raise StopAsyncIteration |
| 101 | + return item |
| 102 | + |
| 103 | + async def aclose(self) -> None: |
| 104 | + """ |
| 105 | + Immediately stop mixing and close the mixer. |
| 106 | +
|
| 107 | + This cancels the mixing task, and any unconsumed output in the queue may be dropped. |
| 108 | + """ |
| 109 | + self._ending = True |
| 110 | + self._mixer_task.cancel() |
| 111 | + with contextlib.suppress(asyncio.CancelledError): |
| 112 | + await self._mixer_task |
| 113 | + |
| 114 | + def end_input(self) -> None: |
| 115 | + """ |
| 116 | + Signal that no more streams will be added. |
| 117 | +
|
| 118 | + This method marks the mixer as closed so that it flushes any remaining buffered output before ending. |
| 119 | + Note that existing streams will still be processed until exhausted. |
| 120 | + """ |
| 121 | + self._ending = True |
| 122 | + |
| 123 | + async def _mixer(self) -> None: |
| 124 | + while True: |
| 125 | + # If we're in ending mode and there are no more streams, exit. |
| 126 | + if self._ending and not self._streams: |
| 127 | + break |
| 128 | + |
| 129 | + if not self._streams: |
| 130 | + await asyncio.sleep(0.01) |
| 131 | + continue |
| 132 | + |
| 133 | + tasks = [ |
| 134 | + self._get_contribution( |
| 135 | + stream, |
| 136 | + self._buffers.get(stream, np.empty((0, self._num_channels), dtype=np.int16)), |
| 137 | + ) |
| 138 | + for stream in list(self._streams) |
| 139 | + ] |
| 140 | + results = await asyncio.gather(*tasks, return_exceptions=True) |
| 141 | + contributions = [] |
| 142 | + any_data = False |
| 143 | + removals = [] |
| 144 | + for contrib in results: |
| 145 | + if not isinstance(contrib, _Contribution): |
| 146 | + continue |
| 147 | + |
| 148 | + contributions.append(contrib.data.astype(np.float32)) |
| 149 | + self._buffers[contrib.stream] = contrib.buffer |
| 150 | + if contrib.had_data: |
| 151 | + any_data = True |
| 152 | + if contrib.exhausted and contrib.buffer.shape[0] == 0: |
| 153 | + removals.append(contrib.stream) |
| 154 | + |
| 155 | + for stream in removals: |
| 156 | + self.remove_stream(stream) |
| 157 | + |
| 158 | + if not any_data: |
| 159 | + await asyncio.sleep(0.001) |
| 160 | + continue |
| 161 | + |
| 162 | + mixed = np.sum(np.stack(contributions, axis=0), axis=0) |
| 163 | + mixed = np.clip(mixed, -32768, 32767).astype(np.int16) |
| 164 | + frame = AudioFrame( |
| 165 | + mixed.tobytes(), self._sample_rate, self._num_channels, self._chunk_size |
| 166 | + ) |
| 167 | + await self._queue.put(frame) |
| 168 | + |
| 169 | + await self._queue.put(None) |
| 170 | + |
| 171 | + async def _get_contribution( |
| 172 | + self, stream: AsyncIterator[AudioFrame], buf: np.ndarray |
| 173 | + ) -> _Contribution: |
| 174 | + had_data = buf.shape[0] > 0 |
| 175 | + exhausted = False |
| 176 | + while buf.shape[0] < self._chunk_size and not exhausted: |
| 177 | + try: |
| 178 | + frame = await asyncio.wait_for( |
| 179 | + stream.__anext__(), timeout=self._stream_timeout_ms / 1000 |
| 180 | + ) |
| 181 | + except asyncio.TimeoutError: |
| 182 | + logger.warning(f"AudioMixer: stream {stream} timeout, ignoring") |
| 183 | + break |
| 184 | + except StopAsyncIteration: |
| 185 | + exhausted = True |
| 186 | + break |
| 187 | + new_data = np.frombuffer(frame.data.tobytes(), dtype=np.int16).reshape( |
| 188 | + -1, self._num_channels |
| 189 | + ) |
| 190 | + buf = np.concatenate((buf, new_data), axis=0) if buf.size else new_data |
| 191 | + had_data = True |
| 192 | + if buf.shape[0] >= self._chunk_size: |
| 193 | + contrib, buf = buf[: self._chunk_size], buf[self._chunk_size :] |
| 194 | + else: |
| 195 | + pad = np.zeros((self._chunk_size - buf.shape[0], self._num_channels), dtype=np.int16) |
| 196 | + contrib, buf = ( |
| 197 | + np.concatenate((buf, pad), axis=0), |
| 198 | + np.empty((0, self._num_channels), dtype=np.int16), |
| 199 | + ) |
| 200 | + return _Contribution(stream, contrib, buf, had_data, exhausted) |
0 commit comments