Files
stemify/stemroller/worker/audio_io.py
T
ClericalErrorandClaude Opus 5.5 1845f03c41 StemRoller: desktop stem separator, Android stems player, and server daemons
- stemroller/: PySide6/QML desktop app and player for separating music into stems
- android/: Plex music player that plays per-stem mixes, with Android Auto support,
  a graphic + parametric (Equalizer APO-style) EQ, and offline downloads
- server/: stemserve (serves stems to the app) and stemd (separates the library on
  the server's GPU, with "Stem Now" priority requests)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-14 21:17:17 -07:00

80 lines
3.4 KiB
Python

"""Audio decode/encode and waveform peak generation for the worker."""
import subprocess
import numpy as np
import soundfile as sf
PEAK_BASE = 256 # samples per bucket at mip level 0
def _source_channels(path: str) -> int:
try:
out = subprocess.run(["ffprobe", "-v", "error", "-select_streams", "a:0", "-show_entries",
"stream=channels", "-of", "csv=p=0", path],
capture_output=True, text=True, timeout=10).stdout
return int(out.strip().splitlines()[0])
except (OSError, ValueError, IndexError, subprocess.SubprocessError):
return 0
def load_audio(path: str, samplerate: int, channels: int = 2) -> np.ndarray:
"""Decode any ffmpeg-readable file to float32 (channels, samples)."""
# ffmpeg's -ac upmix treats mono as centre and spreads it at -3 dB per side;
# decode mono as mono and duplicate it so stems keep the source level.
decode_ch = 1 if _source_channels(path) == 1 else channels
cmd = ["ffmpeg", "-nostdin", "-v", "error", "-i", path, "-map", "0:a:0",
"-f", "f32le", "-acodec", "pcm_f32le", "-ac", str(decode_ch), "-ar", str(samplerate), "-"]
proc = subprocess.run(cmd, capture_output=True)
if proc.returncode != 0:
lines = [l.strip() for l in proc.stderr.decode(errors="replace").splitlines() if l.strip()]
reason = lines[-1].split(": ", 1)[-1] if lines else f"ffmpeg exited with {proc.returncode}"
raise RuntimeError(f"Could not decode the file ({reason[:160]})")
data = np.frombuffer(proc.stdout, dtype=np.float32)
if data.size == 0:
raise RuntimeError("file contains no audio")
data = data.reshape(-1, decode_ch).T
if decode_ch != channels:
data = np.repeat(data, channels, axis=0)
return data.copy()
def write_stem(path: str, audio: np.ndarray, samplerate: int):
"""Write (channels, samples) audio. 24-bit unless it would clip, then float."""
peak = float(np.abs(audio).max()) if audio.size else 0.0
subtype = "FLOAT" if peak > 0.999 else "PCM_24"
sf.write(path, audio.T, samplerate, subtype=subtype, format="WAV")
def compute_peaks(audio: np.ndarray) -> list[np.ndarray]:
"""Min/max mip pyramid of a (channels, samples) signal, mono-folded.
Level k has PEAK_BASE * 2**k samples per bucket; each level is (n, 2) float16.
"""
lo = audio.min(axis=0)
hi = audio.max(axis=0)
n = -(-lo.size // PEAK_BASE)
pad = n * PEAK_BASE - lo.size
if pad:
lo = np.concatenate([lo, np.zeros(pad, lo.dtype)])
hi = np.concatenate([hi, np.zeros(pad, hi.dtype)])
lvl_lo = lo.reshape(n, PEAK_BASE).min(axis=1)
lvl_hi = hi.reshape(n, PEAK_BASE).max(axis=1)
levels = [np.stack([lvl_lo, lvl_hi], axis=1).astype(np.float16)]
while lvl_lo.size > 64:
if lvl_lo.size % 2:
lvl_lo = np.append(lvl_lo, 0)
lvl_hi = np.append(lvl_hi, 0)
lvl_lo = lvl_lo.reshape(-1, 2).min(axis=1)
lvl_hi = lvl_hi.reshape(-1, 2).max(axis=1)
levels.append(np.stack([lvl_lo, lvl_hi], axis=1).astype(np.float16))
return levels
def save_peaks(path: str, stems: dict[str, np.ndarray], samplerate: int):
arrays = {"samplerate": np.array(samplerate), "base": np.array(PEAK_BASE)}
for name, audio in stems.items():
for k, level in enumerate(compute_peaks(audio)):
arrays[f"{name}/{k}"] = level
np.savez(path, **arrays)