processors(audio): fix AudioBufferProcessor interruptions

This commit is contained in:
Aleix Conchillo Flaqué
2024-12-02 13:06:30 -08:00
parent a5c7b02a73
commit d6c944cdc1
2 changed files with 90 additions and 65 deletions

View File

@@ -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))

View File

@@ -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)