File size: 2,784 Bytes
819ec2f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
"""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