Update AudioStreamTrack to use stateful av.AudioResampler for better quality

getstream/video/rtc/audio_track.py

@@ -1,257 +1,181 @@
-import time
 import asyncio
+import fractions
 import logging
+import time
+from collections import deque
 
 import aiortc
-from av import AudioFrame
-from av.frame import Frame
-import fractions
+import av
+import numpy as np
 
-from getstream.video.rtc.track_util import PcmData
+from .track_util import AudioFormat, AudioFormatType, FrameResampler, PcmData
 
 logger = logging.getLogger(__name__)
 
 
 class AudioStreamTrack(aiortc.mediastreams.MediaStreamTrack):
-    """
-    Audio stream track that accepts PcmData objects and buffers them as bytes.
-
-    Works with PcmData objects instead of raw bytes, avoiding format conversion issues.
-    Internally buffers as bytes for efficient memory usage.
+    """aiortc audio track that streams buffered PCM into a WebRTC call.
 
-    Usage:
-        track = AudioStreamTrack(sample_rate=48000, channels=2)
+    `write()` accepts PcmData at any sample rate, channel layout, or format, and converts it to fixed 20ms packed-s16 frames at the track's output
+    rate/layout and queues them (dropping the oldest once the queue exceeds
+    audio_buffer_size_ms).
 
-        # Write PcmData objects (any format, any sample rate, any channels)
-        await track.write(pcm_data)
+    `recv()` paces the queue out in real time via _FramePacer,
+    handing back one frame per call with its pts stamped, and synthesizes silence
+    when the queue is empty so the RTP timeline never stalls.
 
-        # The track will automatically resample/convert to the configured format
+    `flush()` clears the queue to support interruption/barge-in.
     """
 
     kind = "audio"
 
     def __init__(
         self,
-        sample_rate: int = 48000,
-        channels: int = 1,
-        format: str = "s16",
-        audio_buffer_size_ms: int = 30000,  # 30 seconds default
+        sample_rate: int = 48000,  # rate of emitted frames; 48kHz matches Opus (avoids a re-resample)
+        channels: int = 1,  # output channel count (1=mono, 2=stereo)
+        format: AudioFormatType = AudioFormat.S16,  # output sample format; must be s16 (aiortc's Opus encoder requirement)
+        audio_buffer_size_ms: int = 30000,  # max audio to hold before dropping oldest
     ):
-        """
-        Initialize an AudioStreamTrack that accepts PcmData objects.
-
-        Args:
-            sample_rate: Target sample rate in Hz (default: 48000)
-            channels: Number of channels - 1=mono, 2=stereo (default: 1)
-            format: Audio format - "s16" or "f32" (default: "s16")
-            audio_buffer_size_ms: Maximum buffer size in milliseconds (default: 30000ms = 30s)
-        """
         super().__init__()
+        if format != AudioFormat.S16:
+            raise ValueError(
+                f"AudioStreamTrack output format must be 's16', got {format!r}; "
+                "aiortc's Opus encoder only accepts s16 frames."
+            )
         self.sample_rate = sample_rate
         self.channels = channels
         self.format = format
         self.audio_buffer_size_ms = audio_buffer_size_ms
 
-        logger.debug(
-            "Initialized AudioStreamTrack",
-            extra={
-                "sample_rate": sample_rate,
-                "channels": channels,
-                "format": format,
-                "audio_buffer_size_ms": audio_buffer_size_ms,
-            },
+        self._frame_buffer: deque[av.AudioFrame] = deque()
+        # Running per-channel sample total, to enforce the size cap cheaply.
+        self._buffered_samples = 0
+        # Serializes buffer mutation between write() and recv().
+        self._frame_lock = asyncio.Lock()
+
+        self._layout = "stereo" if channels == 2 else "mono"
+        # Samples-per-channel in one 20ms frame (e.g. 0.02 * 48000 = 960).
+        self._samples_per_frame = int(aiortc.mediastreams.AUDIO_PTIME * sample_rate)
+        # Cap in per-channel samples; beyond this, the oldest frames are dropped.
+        self._max_samples = int((audio_buffer_size_ms / 1000) * sample_rate)
+        # Pre-built zeros reused to synthesize a silence frame on starvation.
+        self._silence = np.zeros(
+            (1, self._samples_per_frame * channels), dtype=np.int16
         )
 
-        # Internal bytearray buffer for audio data
-        self._buffer = bytearray()
-        self._buffer_lock = asyncio.Lock()
-
-        # Timing for frame pacing
-        self._start = None
-        self._timestamp = None
-        self._last_frame_time = None
-
-        # Calculate bytes per sample based on format
-        self._bytes_per_sample = 2 if format == "s16" else 4  # s16=2 bytes, f32=4 bytes
-        self._bytes_per_frame = int(
-            aiortc.mediastreams.AUDIO_PTIME
-            * self.sample_rate
-            * self.channels
-            * self._bytes_per_sample
+        # Resampling and pacing state live in their own helpers; the track only
+        # calls their methods.
+        self._resampler = FrameResampler(
+            rate=sample_rate,
+            layout=self._layout,
+            format=format,
+            frame_size=self._samples_per_frame,
         )
+        self._pacer = _FramePacer(sample_rate, self._samples_per_frame)
 
-    async def write(self, pcm: PcmData):
-        """
-        Add PcmData to the buffer.
+    async def write(self, pcm: PcmData, final: bool = False) -> None:
+        """Write PCM data to the track.
 
-        The PcmData will be automatically resampled/converted to match
-        the track's configured sample_rate, channels, and format,
-        then converted to bytes and stored in the buffer.
+        Under the hood, it resamples to fixed 20ms frames and buffers them.
+        When final is True, drain the resampler's tail so the utterance plays out.
 
         Args:
-            pcm: PcmData object with audio data
+            pcm: PCM data to write
+            final: if True, drain the resampler's tail and add it to the buffer
+                (e.g. on end-of-utterance event).
         """
-        # Normalize the PCM data to target format immediately
-        pcm_normalized = self._normalize_pcm(pcm)
-
-        # Convert to bytes
-        audio_bytes = pcm_normalized.to_bytes()
-
-        async with self._buffer_lock:
-            # Check buffer size before adding
-            max_buffer_bytes = int(
-                (self.audio_buffer_size_ms / 1000)
-                * self.sample_rate
-                * self.channels
-                * self._bytes_per_sample
-            )
-
-            # Add new data to buffer first
-            self._buffer.extend(audio_bytes)
-
-            # Check if we exceeded the limit
-            if len(self._buffer) > max_buffer_bytes:
-                # Calculate how many bytes to drop from the beginning
-                bytes_to_drop = len(self._buffer) - max_buffer_bytes
-                dropped_ms = (
-                    bytes_to_drop
-                    / (self.sample_rate * self.channels * self._bytes_per_sample)
-                ) * 1000
-
-                logger.debug(
-                    "Audio buffer overflow, dropping %.1fms of audio. Buffer max is %dms",
-                    dropped_ms,
-                    self.audio_buffer_size_ms,
-                    extra={
-                        "buffer_size_bytes": len(self._buffer),
-                        "incoming_bytes": len(audio_bytes),
-                        "dropped_bytes": bytes_to_drop,
-                    },
-                )
-
-                # Drop from the beginning of the buffer to keep latest data
-                del self._buffer[:bytes_to_drop]
-
-            buffer_duration_ms = (
-                len(self._buffer)
-                / (self.sample_rate * self.channels * self._bytes_per_sample)
-            ) * 1000
-
-        logger.debug(
-            "Added audio to buffer",
-            extra={
-                "pcm_duration_ms": pcm.duration_ms,
-                "buffer_duration_ms": buffer_duration_ms,
-                "buffer_size_bytes": len(self._buffer),
-            },
-        )
+        # Resample under the lock so an interleaved flush() (barge-in) can't reset
+        # the resampler between resampling and enqueuing, leaking pre-flush audio.
+        async with self._frame_lock:
+            # Zero or more finished frames (empty while the resampler is buffering).
+            frames = self._resampler.resample(pcm, flush=final)
+            if not frames:
+                return
+
+            for frame in frames:
+                self._frame_buffer.append(frame)
+                self._buffered_samples += frame.samples
+
+            # Bound latency/memory: if the producer outran the consumer, drop the
+            # oldest frames until back under the cap.
+            while self._buffered_samples > self._max_samples and self._frame_buffer:
+                self._buffered_samples -= self._frame_buffer.popleft().samples
 
     async def flush(self) -> None:
-        """
-        Clear any pending audio from the buffer.
-        Playback stops immediately.
-        """
-        async with self._buffer_lock:
-            bytes_cleared = len(self._buffer)
-            self._buffer.clear()
-
-        logger.debug("Flushed audio buffer", extra={"cleared_bytes": bytes_cleared})
-
-    async def recv(self) -> Frame:
-        """
-        Receive the next 20ms audio frame.
-
-        Returns:
-            AudioFrame with the configured sample_rate, channels, and format
-        """
+        """Drop pending frames and reset the resampler (interruption)."""
+        async with self._frame_lock:
+            # Cut off queued playback (barge-in) and reset the resampler so no
+            # samples bleed into the next utterance.
+            self._frame_buffer.clear()
+            self._buffered_samples = 0
+            self._resampler.flush()
+
+    async def recv(self) -> av.AudioFrame:
+        """Return the next 20ms frame, synthesizing silence when starved."""
+        # aiortc calls recv() in a loop; once the track is stopped, signal EOF.
         if self.readyState != "live":
             raise aiortc.mediastreams.MediaStreamError
 
-        # Calculate samples needed for 20ms frame
-        samples_per_frame = int(aiortc.mediastreams.AUDIO_PTIME * self.sample_rate)
+        # Block until this frame is due, and get the pts to stamp on it.
+        pts = await self._pacer.next_pts()
 
-        # Initialize timestamp if not already done
-        if self._timestamp is None:
-            self._start = time.time()
-            self._timestamp = 0
-            self._last_frame_time = time.time()
-        else:
-            # Use timestamp-based pacing to avoid drift over time
-            # This ensures we stay synchronized with the expected audio rate
-            # even if individual frames have slight timing variations
-            self._timestamp += samples_per_frame
-            start_ts = self._start or time.time()
-            wait = start_ts + (self._timestamp / self.sample_rate) - time.time()
-            if wait > 0:
-                await asyncio.sleep(wait)
-
-            self._last_frame_time = time.time()
-
-        # Get 20ms of audio data from buffer
-        async with self._buffer_lock:
-            if len(self._buffer) >= self._bytes_per_frame:
-                # We have enough data
-                audio_bytes = bytes(self._buffer[: self._bytes_per_frame])
-                del self._buffer[: self._bytes_per_frame]
-            elif len(self._buffer) > 0:
-                # We have some data but not enough - pad with silence
-                audio_bytes = bytes(self._buffer)
-                padding_needed = self._bytes_per_frame - len(audio_bytes)
-                audio_bytes += bytes(padding_needed)  # Pad with zeros (silence)
-                self._buffer.clear()
-
-                logger.debug(
-                    "Padded audio frame with silence",
-                    extra={
-                        "available_bytes": len(audio_bytes) - padding_needed,
-                        "required_bytes": self._bytes_per_frame,
-                        "padding_bytes": padding_needed,
-                    },
-                )
+        async with self._frame_lock:
+            if self._frame_buffer:
+                # Normal path: hand out the next ready frame.
+                frame = self._frame_buffer.popleft()
+                self._buffered_samples -= frame.samples
             else:
-                # No data at all - emit silence
-                audio_bytes = bytes(self._bytes_per_frame)
-
-        # Create AudioFrame
-        layout = "stereo" if self.channels == 2 else "mono"
-
-        # Convert format name: "s16" -> "s16", "f32" -> "flt"
-        if self.format == "s16":
-            av_format = "s16"  # Packed int16
-        elif self.format == "f32":
-            av_format = "flt"  # Packed float32
-        else:
-            av_format = "s16"  # Default to s16
-
-        frame = AudioFrame(format=av_format, layout=layout, samples=samples_per_frame)
+                # Starved (gap between utterances): emit a silence frame so the RTP
+                # timeline stays continuous instead of stalling.
+                frame = av.AudioFrame.from_ndarray(
+                    self._silence, format="s16", layout=self._layout
+                )
 
-        # Fill frame with data
-        frame.planes[0].update(audio_bytes)
+        # A flushed tail can be shorter than a full frame; pad it with trailing
+        # silence so every emitted frame fills its fixed-rate pts slot.
+        if frame.samples < self._samples_per_frame:
+            frame = self._pad_to_full_frame(frame)
 
-        # Set frame properties
-        frame.pts = self._timestamp
+        # Stamp the frame's timing so the encoder/RTP sender place it correctly.
+        frame.pts = pts
         frame.sample_rate = self.sample_rate
         frame.time_base = fractions.Fraction(1, self.sample_rate)
-
         return frame
 
-    def _normalize_pcm(self, pcm: PcmData) -> PcmData:
-        """
-        Normalize PcmData to match the track's target format.
+    def _pad_to_full_frame(self, frame: av.AudioFrame) -> av.AudioFrame:
+        """Pad a short (flushed-tail) frame up to a full frame with trailing silence."""
+        data = frame.to_ndarray()
+        pad = self._samples_per_frame * self.channels - data.shape[1]
+        padded = np.pad(data, ((0, 0), (0, pad)))
+        return av.AudioFrame.from_ndarray(padded, format="s16", layout=self._layout)
 
-        Args:
-            pcm: Input PcmData
 
-        Returns:
-            PcmData resampled/converted to target sample_rate, channels, and format
-        """
-
-        pcm = pcm.resample(self.sample_rate, target_channels=self.channels)
+class _FramePacer:
+    """Real-time clock for fixed-size frames.
 
-        # Convert format if needed
-        if self.format == "s16" and pcm.format != "s16":
-            pcm = pcm.to_int16()
-        elif self.format == "f32" and pcm.format != "f32":
-            pcm = pcm.to_float32()
+    aiortc sends whatever recv() returns immediately, so recv() must block until each
+    frame is due. next_pts() sleeps the right amount and returns the pts to stamp.
+    """
 
-        return pcm
+    def __init__(self, sample_rate: int, samples_per_frame: int):
+        self._sample_rate = sample_rate
+        self._samples_per_frame = samples_per_frame
+        # Monotonic clock anchor and sample cursor (= next pts); None until start.
+        self._start: float | None = None
+        self._ts: int | None = None
+
+    async def next_pts(self) -> int:
+        if self._ts is None:
+            # First frame: anchor the clock and emit without waiting.
+            self._start = time.monotonic()
+            self._ts = 0
+        else:
+            # Advance one frame and sleep until that time. Anchoring to _start
+            # (not the previous frame) keeps drift from accumulating. monotonic()
+            # is used so a wall-clock step (NTP, VM migration) can't stall pacing.
+            self._ts += self._samples_per_frame
+            start = self._start or time.monotonic()
+            wait = start + self._ts / self._sample_rate - time.monotonic()
+            if wait > 0:
+                await asyncio.sleep(wait)
+        return self._ts