71 lines
1.9 KiB
Python
71 lines
1.9 KiB
Python
_V = "/runpod-volume"
|
|
import os
|
|
if os.path.isdir(_V):
|
|
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
|
|
|
|
import base64
|
|
import subprocess
|
|
import tempfile
|
|
|
|
import runpod
|
|
|
|
# Path where the Dockerfile places the RNNoise model file.
|
|
RNNN_MODEL = "/app/models/rnnoise.rnnn"
|
|
|
|
|
|
def _decode_b64(s):
|
|
if isinstance(s, str) and "," in s and s.strip().lower().startswith("data:"):
|
|
s = s.split(",", 1)[1]
|
|
return base64.b64decode(s)
|
|
|
|
|
|
def _denoise(in_path, out_path):
|
|
"""Run ffmpeg. Prefer RNNoise (arnndn) when a model file is present,
|
|
otherwise fall back to the built-in FFT denoiser (afftdn)."""
|
|
if os.path.isfile(RNNN_MODEL):
|
|
af = "arnndn=m=" + RNNN_MODEL
|
|
else:
|
|
af = "afftdn=nf=-25"
|
|
|
|
cmd = [
|
|
"ffmpeg", "-y", "-hide_banner", "-nostdin",
|
|
"-i", in_path,
|
|
"-af", af,
|
|
"-ar", "48000", "-ac", "1",
|
|
"-f", "wav",
|
|
out_path,
|
|
]
|
|
proc = subprocess.run(cmd, capture_output=True)
|
|
if proc.returncode != 0:
|
|
raise RuntimeError(
|
|
"ffmpeg failed: " + proc.stderr.decode("utf-8", "replace")[-800:]
|
|
)
|
|
|
|
|
|
def handler(event):
|
|
try:
|
|
inp = (event or {}).get("input") or {}
|
|
audio_b64 = inp.get("audio")
|
|
if not audio_b64:
|
|
return {"error": "Missing 'audio' in input"}
|
|
|
|
raw = _decode_b64(audio_b64)
|
|
|
|
with tempfile.TemporaryDirectory() as td:
|
|
in_path = os.path.join(td, "in")
|
|
out_path = os.path.join(td, "out.wav")
|
|
with open(in_path, "wb") as f:
|
|
f.write(raw)
|
|
|
|
_denoise(in_path, out_path)
|
|
|
|
with open(out_path, "rb") as f:
|
|
out_bytes = f.read()
|
|
|
|
return {"audio": base64.b64encode(out_bytes).decode("ascii")}
|
|
except Exception as exc:
|
|
return {"error": f"{type(exc).__name__}: {exc}"}
|
|
|
|
|
|
runpod.serverless.start({"handler": handler})
|