125 lines
4.2 KiB
Python
125 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
from fastapi import FastAPI, File, HTTPException, UploadFile
|
|
from fastapi.responses import FileResponse, HTMLResponse
|
|
from starlette.background import BackgroundTask
|
|
|
|
|
|
MODEL = os.getenv("MODEL_FILENAME", "model_bs_roformer_ep_317_sdr_12.9755.ckpt")
|
|
MODEL_DIR = Path(os.getenv("MODEL_DIR", "/models"))
|
|
JOB_DIR = Path(os.getenv("JOB_DIR", "/data/jobs"))
|
|
MAX_UPLOAD = int(os.getenv("MAX_UPLOAD_BYTES", str(1024 ** 3)))
|
|
ALLOWED = {".wav", ".flac", ".mp3", ".m4a", ".aac", ".ogg", ".opus", ".wma"}
|
|
SEPARATION_LOCK = asyncio.Lock()
|
|
STARTED = time.time()
|
|
|
|
app = FastAPI(title="Athena Vocal Separator", version="1.0")
|
|
|
|
|
|
@app.get("/", response_class=HTMLResponse)
|
|
def index() -> str:
|
|
return Path("/app/index.html").read_text(encoding="utf-8")
|
|
|
|
|
|
@app.get("/health")
|
|
def health() -> dict:
|
|
checkpoint = MODEL_DIR / MODEL
|
|
return {
|
|
"status": "ok" if checkpoint.exists() else "starting",
|
|
"model": MODEL,
|
|
"model_ready": checkpoint.exists(),
|
|
"busy": SEPARATION_LOCK.locked(),
|
|
"uptime_seconds": round(time.time() - STARTED, 1),
|
|
}
|
|
|
|
|
|
def _cleanup(path: Path) -> None:
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
|
|
|
|
def _run_separator(input_path: Path, output_dir: Path) -> None:
|
|
args = [
|
|
"audio-separator", str(input_path),
|
|
"--model_filename", MODEL,
|
|
"--model_file_dir", str(MODEL_DIR),
|
|
"--output_dir", str(output_dir),
|
|
"--output_format", "FLAC",
|
|
"--sample_rate", "44100",
|
|
"--use_soundfile",
|
|
"--use_autocast",
|
|
"--mdxc_segment_size", "256",
|
|
"--mdxc_overlap", "8",
|
|
"--mdxc_batch_size", "1",
|
|
]
|
|
completed = subprocess.run(args, capture_output=True, text=True, timeout=7200)
|
|
if completed.returncode:
|
|
detail = (completed.stderr or completed.stdout or "unknown error")[-4000:]
|
|
raise RuntimeError(detail)
|
|
|
|
|
|
@app.post("/v1/separate")
|
|
async def separate(file: UploadFile = File(...)) -> FileResponse:
|
|
suffix = Path(file.filename or "upload.wav").suffix.lower()
|
|
if suffix not in ALLOWED:
|
|
raise HTTPException(415, "Dieses Audioformat wird nicht unterstützt.")
|
|
if SEPARATION_LOCK.locked():
|
|
raise HTTPException(409, "Eine Trennung läuft bereits.")
|
|
|
|
job = Path(tempfile.mkdtemp(prefix="separate-", dir=JOB_DIR))
|
|
input_path = job / f"input{suffix}"
|
|
output_dir = job / "output"
|
|
output_dir.mkdir()
|
|
size = 0
|
|
try:
|
|
with input_path.open("wb") as handle:
|
|
while chunk := await file.read(1024 * 1024):
|
|
size += len(chunk)
|
|
if size > MAX_UPLOAD:
|
|
raise HTTPException(413, "Datei ist größer als 1 GiB.")
|
|
handle.write(chunk)
|
|
async with SEPARATION_LOCK:
|
|
await asyncio.to_thread(_run_separator, input_path, output_dir)
|
|
|
|
stems = sorted(output_dir.glob("*.flac"))
|
|
if len(stems) != 2:
|
|
raise RuntimeError(f"Erwartet wurden zwei FLAC-Dateien, gefunden: {len(stems)}")
|
|
archive = job / "athena-vocals-instrumental.zip"
|
|
with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_STORED) as bundle:
|
|
for stem in stems:
|
|
lower = stem.name.lower()
|
|
target = "vocals.flac" if "vocal" in lower else "instrumental.flac"
|
|
bundle.write(stem, target)
|
|
return FileResponse(
|
|
archive,
|
|
media_type="application/zip",
|
|
filename="athena-vocals-instrumental.zip",
|
|
background=BackgroundTask(_cleanup, job),
|
|
)
|
|
except HTTPException:
|
|
_cleanup(job)
|
|
raise
|
|
except subprocess.TimeoutExpired:
|
|
_cleanup(job)
|
|
raise HTTPException(504, "Die Trennung hat das Zeitlimit überschritten.")
|
|
except Exception as exc:
|
|
_cleanup(job)
|
|
raise HTTPException(500, f"Trennung fehlgeschlagen: {exc}")
|
|
|
|
|
|
@app.on_event("startup")
|
|
def prepare() -> None:
|
|
JOB_DIR.mkdir(parents=True, exist_ok=True)
|
|
MODEL_DIR.mkdir(parents=True, exist_ok=True)
|
|
for old in JOB_DIR.glob("separate-*"):
|
|
if old.is_dir() and time.time() - old.stat().st_mtime > 86400:
|
|
_cleanup(old)
|