Spaces:
Running
Running
| # SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: BSD 2-Clause License | |
| """Audio utilities.""" | |
| import os | |
| import wave | |
| from pathlib import Path | |
| from loguru import logger | |
| from pipecat.frames.frames import AudioRawFrame, Frame | |
| from pipecat.processors.frame_processor import FrameDirection, FrameProcessor | |
| from pipecat.transports.base_transport import TransportParams | |
| # ruff: noqa: SIM115 | |
| class AudioRecorder(FrameProcessor): | |
| """Records audio frames to a file. | |
| Args: | |
| output_file (str): The path to the output file. | |
| params (TransportParams): The transport parameters. | |
| bit_per_sample (int): The number of bits per sample. Defaults to 2. | |
| **kwargs: Additional keyword arguments passed to the parent FrameProcessor. | |
| """ | |
| def __init__( | |
| self, | |
| output_file: str, | |
| params: TransportParams, | |
| bit_per_sample: int = 2, | |
| frame_type: type[Frame] = AudioRawFrame, | |
| override_sample_rate: int | None = None, | |
| **kwargs, | |
| ): | |
| """Initialize the AudioRecorder. | |
| Args: | |
| output_file (str): Path to the output WAV file. | |
| params (TransportParams): Transport parameters. | |
| bit_per_sample (int, optional): Bits per sample. Defaults to 2. | |
| frame_type (Frame): Type of frame to record. Defaults to AudioRawFrame. | |
| **kwargs: Additional keyword arguments. | |
| """ | |
| super().__init__(**kwargs) | |
| self._output_file = output_file | |
| Path(os.path.dirname(self._output_file)).mkdir(parents=True, exist_ok=True) | |
| self._writer = wave.open(self._output_file, "wb") | |
| self._writer.setnchannels(params.audio_out_channels) | |
| self._writer.setsampwidth(bit_per_sample) # 2 bytes - 16 bits PCM | |
| # Coerce sample rate to a valid positive int for WAV header | |
| desired_rate = override_sample_rate if override_sample_rate is not None else params.audio_out_sample_rate | |
| try: | |
| sr_int = int(desired_rate) | |
| if sr_int <= 0: | |
| raise ValueError("non-positive sample rate") | |
| except Exception: | |
| logger.warning( | |
| f"AudioRecorder: invalid sample rate '{desired_rate}', falling back to transport rate {params.audio_out_sample_rate}" | |
| ) | |
| sr_int = int(params.audio_out_sample_rate) | |
| self._writer.setframerate(sr_int) | |
| self._frame_type = frame_type | |
| async def process_frame(self, frame: Frame, direction: FrameDirection): | |
| """Process a frame. | |
| Args: | |
| frame (Frame): The frame to process. | |
| direction (FrameDirection): The direction of frame processing. | |
| """ | |
| await super().process_frame(frame, direction) | |
| logger.trace(f"AudioFileSaver::process_frame - {frame}") | |
| if isinstance(frame, self._frame_type): | |
| logger.trace(f"writing audio frame (length: {len(frame.audio)})") | |
| self._writer.writeframes(frame.audio) | |
| await super().push_frame(frame, direction) | |
| async def cleanup(self): | |
| """Clean up the audio recorder. | |
| Closes the audio file writer and performs necessary cleanup operations. | |
| """ | |
| await super().cleanup() | |
| logger.info("Finalizing audio file.") | |
| if self._writer: | |
| self._writer.close() | |
| self._writer = None | |