Files
Retrieval-based-Voice-Conve…/tools/pymss/audio_io.py

281 lines
10 KiB
Python

import json
import subprocess
import av
import numpy as np
def _frame_to_audio(frame, mono):
"""Implement the frame to audio helper.
Args:
frame (Any): Frame value.
mono (bool): Mono value.
Returns:
Any: Computed result."""
audio = frame.to_ndarray()
audio = audio[None, :] if audio.ndim == 1 else audio
return (audio.mean(axis=0, keepdims=True) if mono and audio.shape[0] > 1 else audio).astype(np.float32, copy=False)
def _ffmpeg_audio_stream_info(path):
"""Return basic audio stream information from ffprobe."""
command = [
"ffprobe",
"-v",
"error",
"-select_streams",
"a:0",
"-show_entries",
"stream=sample_rate,channels",
"-of",
"json",
str(path),
]
result = subprocess.run(command, check=True, capture_output=True, text=True)
streams = json.loads(result.stdout or "{}").get("streams") or []
if not streams:
raise ValueError(f"No audio stream found in {path!s}.")
stream = streams[0]
return int(stream["sample_rate"]), int(stream["channels"])
def _load_audio_ffmpeg(path, sr=None, mono=False, offset=0.0, duration=None):
"""Load audio through the ffmpeg CLI as a fallback for damaged streams."""
source_rate, source_channels = _ffmpeg_audio_stream_info(path)
out_rate = int(sr or source_rate)
channels = 1 if mono else source_channels
command = ["ffmpeg", "-nostdin", "-v", "error"]
if offset:
command += ["-ss", str(float(offset))]
command += ["-i", str(path), "-map", "0:a:0", "-vn"]
if duration is not None:
command += ["-t", str(float(duration))]
command += ["-f", "f32le", "-acodec", "pcm_f32le", "-ar", str(out_rate), "-ac", str(channels), "-"]
result = subprocess.run(command, check=True, capture_output=True)
audio = np.frombuffer(result.stdout, dtype="<f4")
complete_samples = audio.size // channels
audio = audio[: complete_samples * channels]
audio = audio.reshape(complete_samples, channels).T
audio = np.ascontiguousarray(audio.astype(np.float32, copy=False))
return (audio[0] if mono or channels == 1 else audio), out_rate
def _load_audio_librosa(path, sr=None, mono=False, offset=0.0, duration=None):
"""Load audio through librosa as a fallback before the ffmpeg CLI."""
import librosa
audio, out_rate = librosa.load(
path,
sr=sr,
mono=mono,
offset=float(offset or 0.0),
duration=None if duration is None else float(duration),
dtype=np.float32,
)
audio = np.ascontiguousarray(np.asarray(audio, dtype=np.float32))
return (audio[0] if audio.ndim > 1 and (mono or audio.shape[0] == 1) else audio), int(out_rate)
def _load_audio_av(path, sr=None, mono=False, offset=0.0, duration=None):
"""Load audio through PyAV."""
chunks = []
out_rate = None
with av.open(path) as container:
stream = container.streams.audio[0]
out_rate = int(sr or stream.rate)
resampler = None
stop_samples = None if duration is None else int(round((offset + duration) * out_rate))
decoded = 0
for frame in container.decode(stream):
if resampler is None:
resampler = av.AudioResampler(format="fltp", layout=frame.layout.name, rate=out_rate)
for out in resampler.resample(frame):
chunks.append(audio := _frame_to_audio(out, mono))
decoded += audio.shape[-1]
if stop_samples is not None and decoded >= stop_samples:
break
if resampler is not None:
for out in resampler.resample(None):
chunks.append(_frame_to_audio(out, mono))
start = int(round(offset * out_rate))
stop = None if duration is None else start + int(round(duration * out_rate))
channels = 1 if mono else 0
audio = np.ascontiguousarray(
(np.concatenate(chunks, axis=-1) if chunks else np.empty((channels, 0), dtype=np.float32))[..., start:stop]
)
return (audio[0] if mono or audio.shape[0] == 1 else audio), out_rate
def load_audio(path, sr=None, mono=False, offset=0.0, duration=None):
"""Load an audio file as float32 NumPy samples.
Audio decoding is attempted in this order: PyAV, librosa, then the
ffmpeg CLI fallback. Stereo or multi-channel output is returned
channel-first as ``(channels, samples)``. Mono output is returned as a
one-dimensional array.
Args:
path (str | os.PathLike): Input audio file path. Any format supported
by the local FFmpeg/PyAV build can be decoded.
sr (int | None, optional): Target sample rate. ``None`` keeps the
source stream sample rate. Defaults to None.
mono (bool, optional): Whether to downmix multi-channel audio to mono.
Defaults to False.
offset (float, optional): Start offset in seconds. Defaults to 0.0.
duration (float | None, optional): Maximum duration to return in
seconds after ``offset``. ``None`` reads to the end. Defaults to
None.
Returns:
tuple[np.ndarray, int]: Audio samples and sample rate. The array is
channel-first for multi-channel audio and one-dimensional for mono
output.
Example:
>>> from pymss import load_audio
>>> audio, sample_rate = load_audio("song.wav", sr=44100)
>>> sample_rate
44100
Example:
>>> clip, sample_rate = load_audio(
... "song.wav",
... mono=True,
... offset=30.0,
... duration=10.0,
... )
>>> clip.ndim
1"""
loaders = [
("PyAV", _load_audio_av),
("librosa", _load_audio_librosa),
("ffmpeg CLI", _load_audio_ffmpeg),
]
errors = []
for name, loader in loaders:
try:
return loader(path, sr=sr, mono=mono, offset=offset, duration=duration)
except Exception as e:
errors.append(f"{name}: {e}")
continue
raise RuntimeError(f"All audio loading methods failed: {'; '.join(errors)}")
def _bitrate_to_int(value):
"""Implement the bitrate to int helper.
Args:
value (Any): Value value.
Returns:
Any: Computed result."""
if value is None:
return None
if isinstance(value, int):
return value
value = str(value).strip().lower()
return int(float(value[:-1]) * 1000) if value.endswith("k") else int(value)
def _format_audio(audio):
"""Format audio.
Args:
audio (np.ndarray): Audio samples.
Returns:
Any: Computed result."""
audio = np.asarray(audio)
audio = np.ascontiguousarray(audio[:, None] if audio.ndim == 1 else audio)
# We can use "fltp" container for all output formats, while the final result is determined by the codec.
# Using the fltp sample format can also help avoid some clipping distortion that occurs with integer formats.
return np.ascontiguousarray(audio.astype(np.float32).T)
def save_audio(path, audio, sr, output_format, audio_params):
"""Save a NumPy audio array to wav, flac, mp3, or m4a.
Audio is expected as sample-major data, either ``(samples,)`` for mono or
``(samples, channels)`` for multi-channel audio. The output codec is chosen
from ``output_format`` and ``audio_params``.
Args:
path (str | os.PathLike): Output file path.
audio (np.ndarray): Audio samples. Mono arrays may be one-dimensional;
stereo arrays should be shaped as ``(samples, 2)``.
sr (int): Sample rate in Hz.
output_format (str): Output format. Supported values are ``wav``,
``flac``, ``mp3``, and ``m4a``.
audio_params (dict): Encoding options. Supported keys include
``wav_bit_depth`` (``FLOAT``, ``PCM_16``, ``PCM_24``),
``flac_bit_depth`` (currently ``PCM_24`` uses soundfile),
``mp3_bit_rate`` (for example ``"320k"``), ``m4a_bit_rate``,
``m4a_codec``, and ``m4a_aac_at_quality``.
Returns:
None: The file is written to ``path``.
Example:
>>> from pymss import save_audio
>>> save_audio(
... "vocals.wav",
... vocals,
... 44100,
... "wav",
... {"wav_bit_depth": "FLOAT"},
... )
Example:
>>> save_audio(
... "instrumental.flac",
... instrumental,
... 44100,
... "flac",
... {"flac_bit_depth": "PCM_24"},
... )"""
output_format = output_format.lower()
audio_array = np.asarray(audio)
layout = "stereo" if audio_array.ndim > 1 and audio_array.shape[1] == 2 else "mono"
if output_format == "mp3":
codec = "libmp3lame"
elif output_format == "m4a":
codec = audio_params.get("m4a_codec", "aac")
elif output_format == "flac":
# PyAV's FLAC encoder only exposes a single "flac" codec in the current version.
# In the current version, without access to bits_per_raw_sample in PyAV, PCM_24 may still be encoded as 16-bit.
# Use soundfile to save 24-bit FLAC
codec = "flac"
if audio_params.get("flac_bit_depth", "PCM_24") == "PCM_24":
import soundfile as sf
return sf.write(path, audio_array, int(sr), format="FLAC", subtype="PCM_24")
else:
wav_codecs = {"PCM_16": "pcm_s16le", "PCM_24": "pcm_s24le", "FLOAT": "pcm_f32le"}
codec = wav_codecs.get(audio_params.get("wav_bit_depth", "FLOAT"), wav_codecs["FLOAT"])
with av.open(path, "w") as container:
stream = container.add_stream(codec, rate=int(sr))
stream.layout = layout
if output_format == "mp3":
stream.bit_rate = _bitrate_to_int(audio_params.get("mp3_bit_rate", "320k"))
elif output_format == "m4a":
stream.bit_rate = _bitrate_to_int(audio_params.get("m4a_bit_rate", "512k"))
if codec == "aac_at":
stream.codec_context.options = {"aac_at_quality": str(audio_params.get("m4a_aac_at_quality", 2))}
frame = av.AudioFrame.from_ndarray(_format_audio(audio_array), format="fltp", layout=layout)
frame.sample_rate = int(sr)
for packet in stream.encode(frame):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)