216 lines
9.2 KiB
Python
216 lines
9.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, Form, 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()
|
|
|
|
MODES = {
|
|
"vocals": {"model": MODEL, "stems": ("vocals", "instrumental"), "archive": "athena-vocals-instrumental.zip", "engine": "mdxc"},
|
|
"four_stem": {"model": "htdemucs_ft.yaml", "stems": ("vocals", "drums", "bass", "other"), "archive": "athena-4-stems.zip", "engine": "demucs"},
|
|
"six_stem": {"model": "htdemucs_6s.yaml", "stems": ("vocals", "drums", "bass", "guitar", "piano", "other"), "archive": "athena-6-stems-experimental.zip", "engine": "demucs"},
|
|
"speech": {"model": "MossFormer2_SE_48K", "stems": ("speech", "noise"), "archive": "athena-sprache-und-hintergrund.zip", "engine": "clearvoice"},
|
|
}
|
|
TARGETS = {
|
|
"vocals": {"mode": "vocals", "stem": "vocals", "remainder": "instrumental", "archive": "athena-gesang-und-rest.zip", "rest_file": "instrumental.flac"},
|
|
"drums": {"mode": "four_stem", "stem": "drums", "archive": "athena-schlagzeug-und-rest.zip", "rest_file": "rest-ohne-schlagzeug.flac"},
|
|
"bass": {"mode": "four_stem", "stem": "bass", "archive": "athena-bass-und-rest.zip", "rest_file": "rest-ohne-bass.flac"},
|
|
"guitar": {"mode": "six_stem", "stem": "guitar", "archive": "athena-gitarre-und-rest.zip", "rest_file": "rest-ohne-gitarre.flac"},
|
|
"piano": {"mode": "six_stem", "stem": "piano", "archive": "athena-piano-und-rest.zip", "rest_file": "rest-ohne-piano.flac"},
|
|
"other": {"mode": "six_stem", "stem": "other", "archive": "athena-sonstiges-und-rest.zip", "rest_file": "rest-ohne-sonstiges.flac"},
|
|
"speech": {"mode": "speech", "stem": "speech", "remainder": "noise", "archive": "athena-sprache-und-hintergrund.zip", "rest_file": "hintergrund-ohne-sprache.flac"},
|
|
}
|
|
|
|
app = FastAPI(title="Athena Stem Separator", version="2.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:
|
|
available = {
|
|
name: (
|
|
(MODEL_DIR / "clearvoice" / mode["model"] / "last_best_checkpoint").exists()
|
|
if mode["engine"] == "clearvoice"
|
|
else (MODEL_DIR / mode["model"]).exists()
|
|
)
|
|
for name, mode in MODES.items()
|
|
}
|
|
return {
|
|
"status": "ok" if all(available.values()) else "starting",
|
|
"models": {name: mode["model"] for name, mode in MODES.items()},
|
|
"models_ready": available,
|
|
"targets": list(TARGETS),
|
|
"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, mode: dict) -> None:
|
|
if mode["engine"] == "clearvoice":
|
|
completed = subprocess.run(
|
|
[
|
|
"/opt/clearvoice-venv/bin/python", "/app/speech_enhance.py", str(input_path),
|
|
str(output_dir / "speech.flac"), str(output_dir / "noise.flac"),
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=7200,
|
|
)
|
|
if completed.returncode:
|
|
detail = (completed.stderr or completed.stdout or "unknown ClearVoice error")[-4000:]
|
|
raise RuntimeError(detail)
|
|
return
|
|
args = [
|
|
"audio-separator", str(input_path),
|
|
"--model_filename", mode["model"],
|
|
"--model_file_dir", str(MODEL_DIR),
|
|
"--output_dir", str(output_dir),
|
|
"--output_format", "FLAC",
|
|
"--sample_rate", "44100",
|
|
"--use_autocast",
|
|
]
|
|
if mode["engine"] == "mdxc":
|
|
args.extend(["--mdxc_segment_size", "256", "--mdxc_overlap", "8", "--mdxc_batch_size", "1"])
|
|
else:
|
|
args.extend(["--demucs_segment_size", "40", "--demucs_shifts", "2", "--demucs_overlap", "0.25"])
|
|
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)
|
|
|
|
|
|
def _stem_name(path: Path, expected: tuple[str, ...]) -> str | None:
|
|
lower = path.stem.lower()
|
|
for stem in sorted(expected, key=len, reverse=True):
|
|
if stem in lower:
|
|
return stem
|
|
return None
|
|
|
|
|
|
def _mix_remainder(stems: list[Path], output_path: Path) -> None:
|
|
args = ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y"]
|
|
for stem in stems:
|
|
args.extend(["-i", str(stem)])
|
|
inputs = "".join(f"[{index}:a]" for index in range(len(stems)))
|
|
args.extend([
|
|
"-filter_complex", f"{inputs}amix=inputs={len(stems)}:normalize=0:dropout_transition=0[rest]",
|
|
"-map", "[rest]", "-ar", "44100", "-c:a", "flac", str(output_path),
|
|
])
|
|
completed = subprocess.run(args, capture_output=True, text=True, timeout=1800)
|
|
if completed.returncode:
|
|
detail = (completed.stderr or completed.stdout or "unknown ffmpeg error")[-4000:]
|
|
raise RuntimeError(f"Restspur konnte nicht erzeugt werden: {detail}")
|
|
|
|
|
|
@app.post("/v1/separate")
|
|
async def separate(
|
|
file: UploadFile = File(...),
|
|
target: str | None = Form(None),
|
|
mode: str | None = Form(None),
|
|
) -> FileResponse:
|
|
selected_target = TARGETS.get(target) if target else None
|
|
if target and selected_target is None:
|
|
raise HTTPException(422, f"Unbekannte Zielspur: {target}")
|
|
selected_mode_name = selected_target["mode"] if selected_target else (mode or "vocals")
|
|
selected_mode = MODES.get(selected_mode_name)
|
|
if selected_mode is None:
|
|
raise HTTPException(422, f"Unbekannter Trennmodus: {selected_mode_name}")
|
|
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, selected_mode)
|
|
|
|
stems = sorted(output_dir.glob("*.flac"))
|
|
expected = selected_mode["stems"]
|
|
recognized = {_stem_name(stem, expected): stem for stem in stems}
|
|
recognized.pop(None, None)
|
|
missing = [stem for stem in expected if stem not in recognized]
|
|
if missing:
|
|
found = ", ".join(stem.name for stem in stems) or "keine"
|
|
raise RuntimeError(f"Fehlende Spuren: {', '.join(missing)}; gefunden: {found}")
|
|
archive_name = selected_target["archive"] if selected_target else selected_mode["archive"]
|
|
archive = job / archive_name
|
|
with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_STORED) as bundle:
|
|
if selected_target:
|
|
target_stem = selected_target["stem"]
|
|
bundle.write(recognized[target_stem], f"{target_stem}.flac")
|
|
if "remainder" in selected_target:
|
|
remainder = recognized[selected_target["remainder"]]
|
|
else:
|
|
remainder = job / selected_target["rest_file"]
|
|
await asyncio.to_thread(
|
|
_mix_remainder,
|
|
[recognized[stem] for stem in expected if stem != target_stem],
|
|
remainder,
|
|
)
|
|
bundle.write(remainder, selected_target["rest_file"])
|
|
else:
|
|
# Rückwärtskompatibilität für bestehende API-Clients.
|
|
for stem in expected:
|
|
bundle.write(recognized[stem], f"{stem}.flac")
|
|
return FileResponse(
|
|
archive,
|
|
media_type="application/zip",
|
|
filename=archive_name,
|
|
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)
|