diff --git a/src/pipecat/audio/utils.py b/src/pipecat/audio/utils.py index f2260c57b..057942e04 100644 --- a/src/pipecat/audio/utils.py +++ b/src/pipecat/audio/utils.py @@ -18,6 +18,37 @@ def resample_audio(audio: bytes, original_rate: int, target_rate: int) -> bytes: return resampled_audio.astype(np.int16).tobytes() +def mix_audio(audio1: bytes, audio2: bytes) -> bytes: + data1 = np.frombuffer(audio1, dtype=np.int16) + data2 = np.frombuffer(audio2, dtype=np.int16) + + # Max length + max_length = max(len(data1), len(data2)) + + # Zero-pad the arrays to the same length + padded1 = np.pad(data1, (0, max_length - len(data1)), mode="constant") + padded2 = np.pad(data2, (0, max_length - len(data2)), mode="constant") + + # Mix the arrays + mixed_audio = padded1.astype(np.int32) + padded2.astype(np.int32) + mixed_audio = np.clip(mixed_audio, -32768, 32767).astype(np.int16) + + return mixed_audio.astype(np.int16).tobytes() + + +def interleave_stereo_audio(left_audio: bytes, right_audio: bytes) -> bytes: + left = np.frombuffer(left_audio, dtype=np.int16) + right = np.frombuffer(right_audio, dtype=np.int16) + + min_length = min(len(left), len(right)) + left = left[:min_length] + right = right[:min_length] + + stereo = np.column_stack((left, right)) + + return stereo.astype(np.int16).tobytes() + + def normalize_value(value, min_value, max_value): normalized = (value - min_value) / (max_value - min_value) normalized_clamped = max(0, min(1, normalized)) diff --git a/src/pipecat/processors/audio/audio_buffer_processor.py b/src/pipecat/processors/audio/audio_buffer_processor.py index 25a1f0237..d7c07c611 100644 --- a/src/pipecat/processors/audio/audio_buffer_processor.py +++ b/src/pipecat/processors/audio/audio_buffer_processor.py @@ -7,8 +7,8 @@ import wave from io import BytesIO +from pipecat.audio.utils import interleave_stereo_audio, mix_audio, resample_audio from pipecat.frames.frames import ( - AudioRawFrame, Frame, InputAudioRawFrame, OutputAudioRawFrame, @@ -17,84 +17,78 @@ from pipecat.processors.frame_processor import FrameDirection, FrameProcessor class AudioBufferProcessor(FrameProcessor): - def __init__(self, **kwargs): - """ - Initialize the AudioBufferProcessor. + """This processor buffers audio raw frames (input and output) that can later + be obtained as an in-memory WAV. You can provide the desired output + `sample_rate` and incoming audio frames will resampled to match it. Also, + you can provide the number of channels, 1 for mono and 2 for stereo. With + mono audio user and bot audio will be mixed, in the case of stereo the left + channel will be used for the user's audio and the right channel for the bot. - This constructor sets up the initial state for audio processing: - - audio_buffer: A bytearray to store incoming audio data. - - num_channels: The number of audio channels (initialized as None). - - sample_rate: The sample rate of the audio (initialized as None). + """ - The num_channels and sample_rate are set to None initially and will be - populated when the first audio frame is processed. - """ + def __init__(self, *, sample_rate: int = 24000, num_channels: int = 1, **kwargs): super().__init__(**kwargs) - self._user_audio_buffer = bytearray() - self._assistant_audio_buffer = bytearray() - self._num_channels = None - self._sample_rate = None + self._sample_rate = sample_rate + self._num_channels = num_channels - def _buffer_has_audio(self, buffer: bytearray): + self._user_audio_buffer = bytearray() + self._bot_audio_buffer = bytearray() + + def _buffer_has_audio(self, buffer: bytearray) -> bool: return buffer is not None and len(buffer) > 0 - def has_audio(self): - return ( - self._buffer_has_audio(self._user_audio_buffer) - and self._buffer_has_audio(self._assistant_audio_buffer) - and self._sample_rate is not None + def has_audio(self) -> bool: + return self._buffer_has_audio(self._user_audio_buffer) and self._buffer_has_audio( + self._bot_audio_buffer ) def reset_audio_buffer(self): self._user_audio_buffer = bytearray() - self._assistant_audio_buffer = bytearray() + self._bot_audio_buffer = bytearray() - def merge_audio_buffers(self): + def merge_audio_buffers(self) -> bytes: + if self._num_channels == 1: + return self._merge_mono() + elif self._num_channels == 2: + return self._merge_stereo() + else: + return b"" + + async def process_frame(self, frame: Frame, direction: FrameDirection): + await super().process_frame(frame, direction) + + # Include all audio from the user. + if isinstance(frame, InputAudioRawFrame): + resampled = resample_audio(frame.audio, frame.sample_rate, self._sample_rate) + self._user_audio_buffer.extend(resampled) + # Sync the bot's buffer to the user's buffer by adding silence if needed + if len(self._user_audio_buffer) > len(self._bot_audio_buffer): + silence = b"\x00" * len(resampled) + self._bot_audio_buffer.extend(silence) + + # If the bot is speaking, include all audio from the bot. + if isinstance(frame, OutputAudioRawFrame): + resampled = resample_audio(frame.audio, frame.sample_rate, self._sample_rate) + self._bot_audio_buffer.extend(resampled) + + def _merge_mono(self) -> bytes: + with BytesIO() as buffer: + with wave.open(buffer, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(self._sample_rate) + mixed = mix_audio(bytes(self._user_audio_buffer), bytes(self._bot_audio_buffer)) + wf.writeframes(mixed) + return buffer.getvalue() + + def _merge_stereo(self) -> bytes: with BytesIO() as buffer: with wave.open(buffer, "wb") as wf: wf.setnchannels(2) wf.setsampwidth(2) wf.setframerate(self._sample_rate) - # Interleave the two audio streams - max_length = max(len(self._user_audio_buffer), len(self._assistant_audio_buffer)) - interleaved = bytearray(max_length * 2) - - for i in range(0, max_length, 2): - if i < len(self._user_audio_buffer): - interleaved[i * 2] = self._user_audio_buffer[i] - interleaved[i * 2 + 1] = self._user_audio_buffer[i + 1] - else: - interleaved[i * 2] = 0 - interleaved[i * 2 + 1] = 0 - - if i < len(self._assistant_audio_buffer): - interleaved[i * 2 + 2] = self._assistant_audio_buffer[i] - interleaved[i * 2 + 3] = self._assistant_audio_buffer[i + 1] - else: - interleaved[i * 2 + 2] = 0 - interleaved[i * 2 + 3] = 0 - - wf.writeframes(interleaved) + stereo = interleave_stereo_audio( + bytes(self._user_audio_buffer), bytes(self._bot_audio_buffer) + ) + wf.writeframes(stereo) return buffer.getvalue() - - async def process_frame(self, frame: Frame, direction: FrameDirection): - await super().process_frame(frame, direction) - if isinstance(frame, AudioRawFrame) and self._sample_rate is None: - self._sample_rate = frame.sample_rate - - # include all audio from the user - if isinstance(frame, InputAudioRawFrame): - self._user_audio_buffer.extend(frame.audio) - # Sync the assistant's buffer to the user's buffer by adding silence if needed - if len(self._user_audio_buffer) > len(self._assistant_audio_buffer): - silence_length = len(self._user_audio_buffer) - len(self._assistant_audio_buffer) - silence = b"\x00" * silence_length - self._assistant_audio_buffer.extend(silence) - - # if the assistant is speaking, include all audio from the assistant, - if isinstance(frame, OutputAudioRawFrame): - self._assistant_audio_buffer.extend(frame.audio) - - # do not push the user's audio frame, doing so will result in echo - if not isinstance(frame, InputAudioRawFrame): - await self.push_frame(frame, direction)