Shaman
105 lines
2.5 KiB
Python
105 lines
2.5 KiB
Python
import shutil
|
|
import struct
|
|
import subprocess
|
|
from io import BytesIO
|
|
|
|
|
|
def _find_ffmpeg(configured: str = "ffmpeg") -> str:
|
|
if shutil.which(configured):
|
|
return configured
|
|
try:
|
|
from imageio_ffmpeg import get_ffmpeg_exe
|
|
return get_ffmpeg_exe()
|
|
except Exception:
|
|
raise RuntimeError(
|
|
"ffmpeg not found: install it on PATH, set runtime.ffmpeg_path "
|
|
"in config.json, or `pip install imageio-ffmpeg`"
|
|
) from None
|
|
|
|
|
|
_SUPPORTED_FORMATS = {"mp3", "wav", "pcm", "opus", "aac", "flac"}
|
|
|
|
_CONTENT_TYPES = {
|
|
"mp3": "audio/mpeg",
|
|
"wav": "audio/wav",
|
|
"pcm": "audio/pcm",
|
|
"opus": "audio/opus",
|
|
"aac": "audio/aac",
|
|
"flac": "audio/flac",
|
|
}
|
|
|
|
_FFMPEG_CODECS = {
|
|
"mp3": "libmp3lame",
|
|
"opus": "libopus",
|
|
"aac": "aac",
|
|
"flac": "flac",
|
|
}
|
|
|
|
|
|
def pcm_to_wav(pcm: bytes, sample_rate: int = 24000) -> bytes:
|
|
data_size = len(pcm)
|
|
byte_rate = sample_rate * 2
|
|
header = struct.pack(
|
|
"<4sI4s4sI4sI",
|
|
b"RIFF",
|
|
36 + data_size,
|
|
b"WAVE",
|
|
b"fmt ",
|
|
16,
|
|
b"data",
|
|
data_size,
|
|
)
|
|
fmt = struct.pack("<HHIIHH", 1, 1, sample_rate, byte_rate, 2, 16)
|
|
return header[:20] + fmt + header[20:] + pcm
|
|
|
|
|
|
def pcm_to_raw(pcm: bytes) -> bytes:
|
|
return pcm
|
|
|
|
|
|
def pcm_to_ffmpeg(
|
|
pcm: bytes, fmt: str, sample_rate: int = 24000, ffmpeg_path: str = "ffmpeg"
|
|
) -> bytes:
|
|
codec = _FFMPEG_CODECS.get(fmt)
|
|
if codec is None:
|
|
raise ValueError(f"Unsupported ffmpeg format: {fmt}")
|
|
|
|
cmd = [
|
|
ffmpeg_path, "-y",
|
|
"-f", "s16le",
|
|
"-ar", str(sample_rate),
|
|
"-ac", "1",
|
|
"-i", "pipe:0",
|
|
"-f", _ffprobe_fmt(fmt),
|
|
"-codec:a", codec,
|
|
"pipe:1",
|
|
]
|
|
|
|
proc = subprocess.run(
|
|
cmd,
|
|
input=pcm,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
if proc.returncode != 0:
|
|
raise RuntimeError(f"ffmpeg failed: {proc.stderr.decode(errors='replace')}")
|
|
return proc.stdout
|
|
|
|
|
|
def _ffprobe_fmt(fmt: str) -> str:
|
|
return {"mp3": "mp3", "opus": "opus", "aac": "adts", "flac": "flac"}.get(fmt, fmt)
|
|
|
|
|
|
def encode(pcm: bytes, fmt: str, sample_rate: int = 24000, ffmpeg_path: str = "ffmpeg") -> tuple[bytes, str]:
|
|
if fmt not in _SUPPORTED_FORMATS:
|
|
raise ValueError(f"Unsupported response_format: {fmt}")
|
|
|
|
if fmt == "wav":
|
|
data = pcm_to_wav(pcm, sample_rate)
|
|
elif fmt == "pcm":
|
|
data = pcm
|
|
else:
|
|
data = pcm_to_ffmpeg(pcm, fmt, sample_rate, _find_ffmpeg(ffmpeg_path))
|
|
|
|
return data, _CONTENT_TYPES[fmt]
|