audio_export.py

"""Write generated audio to compressed formats (MP3 / AAC / FLAC / WAV)."""

from __future__ import annotations

import logging
import os
import shutil
import subprocess
import tempfile
from typing import Optional

import numpy as np
import soundfile as sf

log = logging.getLogger("ace-inspire")

AUDIO_FORMATS = (
    "MP3",
    "AAC (M4A)",
    "FLAC",
    "WAV",
)
DEFAULT_AUDIO_FORMAT = "MP3"

_FFMPEG_ARGS = {
    "MP3": (".mp3", ["-c:a", "libmp3lame", "-b:a", "192k", "-ar", "48000"]),
    "AAC (M4A)": (".m4a", ["-c:a", "aac", "-b:a", "192k", "-ar", "48000"]),
}


def _ffmpeg_bin() -> Optional[str]:
    return shutil.which("ffmpeg")


def write_audio(audio: np.ndarray, sample_rate: int, fmt: str) -> str:
    """Encode float audio to a temp file. Returns path. Prefer MP3/AAC/FLAC over WAV."""
    fmt = (fmt or DEFAULT_AUDIO_FORMAT).strip()
    if fmt not in AUDIO_FORMATS:
        fmt = DEFAULT_AUDIO_FORMAT

    audio = np.asarray(audio, dtype=np.float32)
    if audio.ndim > 2:
        audio = audio.reshape(audio.shape[0], -1)

    # Peak-normalize softly so lossy encoders don't clip
    peak = float(np.max(np.abs(audio))) if audio.size else 0.0
    if peak > 1.0:
        audio = audio / peak

    fd, base = tempfile.mkstemp(prefix="ace_inspire_")
    os.close(fd)
    os.unlink(base)

    if fmt == "WAV":
        path = base + ".wav"
        sf.write(path, audio, samplerate=sample_rate, subtype="PCM_16")
        return path

    if fmt == "FLAC":
        path = base + ".flac"
        sf.write(path, audio, samplerate=sample_rate, format="FLAC")
        return path

    # MP3 / AAC via ffmpeg
    ext, ff_args = _FFMPEG_ARGS[fmt]
    out_path = base + ext
    wav_tmp = base + ".__tmp.wav"
    sf.write(wav_tmp, audio, samplerate=sample_rate, subtype="PCM_16")

    ffmpeg = _ffmpeg_bin()
    if not ffmpeg:
        log.warning("ffmpeg not found — falling back to FLAC for %s", fmt)
        os.unlink(wav_tmp)
        path = base + ".flac"
        sf.write(path, audio, samplerate=sample_rate, format="FLAC")
        return path

    cmd = [ffmpeg, "-y", "-hide_banner", "-loglevel", "error", "-i", wav_tmp, *ff_args, out_path]
    try:
        subprocess.run(cmd, check=True, capture_output=True)
    except subprocess.CalledProcessError as e:
        err = (e.stderr or b"").decode("utf-8", errors="replace")[:500]
        log.error("ffmpeg encode failed for %s: %s", fmt, err)
        # Fallback FLAC so the user still gets a file
        try:
            os.unlink(wav_tmp)
        except OSError:
            pass
        path = base + ".flac"
        sf.write(path, audio, samplerate=sample_rate, format="FLAC")
        return path
    finally:
        try:
            os.unlink(wav_tmp)
        except OSError:
            pass

    return out_path