feat: initial commit

2026-04-14 19:19:45 -07:00
commit 349ad168d4
14 changed files with 2788 additions and 0 deletions
@@ -0,0 +1,4 @@
+"""ACE-Step service constants and configuration."""
+
+ACESTEP_PORT = 8001
+ACESTEP_BASE_URL = f"http://127.0.0.1:{ACESTEP_PORT}"
@@ -0,0 +1,4 @@
+"""Foundation-1 service constants and configuration."""
+
+FOUNDATION1_PORT = 8002
+FOUNDATION1_BASE_URL = f"http://127.0.0.1:{FOUNDATION1_PORT}"
@@ -0,0 +1,126 @@
+"""GPU model manager — ensures only one model subprocess runs at a time."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import subprocess
+import time
+from pathlib import Path
+from typing import Optional
+
+import httpx
+
+ROOT_DIR = Path(__file__).parent.parent
+
+ACESTEP_PORT = 8001
+FOUNDATION1_PORT = 8002
+HEALTH_POLL_INTERVAL = 2.0
+HEALTH_TIMEOUT_SECONDS = 180  # model loading can be slow on first run
+
+logger = logging.getLogger(__name__)
+
+
+def _python_exe(venv_dir: Path) -> Path:
+    """Return the Python executable for a given venv directory."""
+    win = venv_dir / "Scripts" / "python.exe"
+    if win.exists():
+        return win
+    return venv_dir / "bin" / "python"
+
+
+class ModelNotActiveError(Exception):
+    def __init__(self, requested: str, active: Optional[str]) -> None:
+        self.requested = requested
+        self.active = active
+        super().__init__(f"Model '{requested}' is not loaded (active: {active!r})")
+
+
+class ModelManager:
+    def __init__(self) -> None:
+        self.active_model: Optional[str] = None
+        self._lock = asyncio.Lock()
+        self._process: Optional[subprocess.Popen] = None
+
+    async def load(self, model: str) -> None:
+        """Load a model, unloading the currently active one first if needed."""
+        async with self._lock:
+            if self.active_model == model:
+                logger.info("Model '%s' is already active.", model)
+                return
+            if self.active_model is not None:
+                logger.info("Unloading model '%s'...", self.active_model)
+                await asyncio.get_event_loop().run_in_executor(None, self._stop_process)
+            logger.info("Starting model '%s'...", model)
+            await asyncio.get_event_loop().run_in_executor(None, self._start_process, model)
+            self.active_model = model
+            logger.info("Model '%s' is ready.", model)
+
+    async def unload(self) -> None:
+        """Unload the active model."""
+        async with self._lock:
+            if self.active_model is None:
+                return
+            logger.info("Unloading model '%s'...", self.active_model)
+            await asyncio.get_event_loop().run_in_executor(None, self._stop_process)
+
+    def _stop_process(self) -> None:
+        if self._process is None:
+            return
+        try:
+            self._process.terminate()
+            self._process.wait(timeout=30)
+        except subprocess.TimeoutExpired:
+            logger.warning("Process did not terminate gracefully, killing.")
+            self._process.kill()
+            self._process.wait()
+        except Exception as exc:
+            logger.error("Error stopping process: %s", exc)
+        finally:
+            self._process = None
+        self.active_model = None
+
+    def _start_process(self, model: str) -> None:
+        if model == "fullsong":
+            self._process = self._launch_acestep()
+            self._wait_healthy(f"http://127.0.0.1:{ACESTEP_PORT}/health")
+        elif model == "clip":
+            self._process = self._launch_foundation1()
+            self._wait_healthy(f"http://127.0.0.1:{FOUNDATION1_PORT}/health")
+        elif model == "separator":
+            pass  # no persistent process; previous model already stopped by load(), VRAM is free
+        else:
+            raise ValueError(f"Unknown model: '{model}'. Must be 'fullsong', 'clip', or 'separator'.")
+
+    def _launch_acestep(self) -> subprocess.Popen:
+        venv_python = _python_exe(ROOT_DIR / "ace-step" / ".venv")
+        cmd = [str(venv_python), "-c", "from acestep.api_server import main; main()"]
+        logger.info("Launching ACE-Step: %s", " ".join(cmd))
+        return subprocess.Popen(cmd, cwd=str(ROOT_DIR / "ace-step"))
+
+    def _launch_foundation1(self) -> subprocess.Popen:
+        venv_python = _python_exe(ROOT_DIR / "foundation1" / ".venv")
+        server_script = str(ROOT_DIR / "foundation1_server" / "server.py")
+        cmd = [str(venv_python), server_script]
+        logger.info("Launching Foundation-1 server: %s", " ".join(cmd))
+        return subprocess.Popen(cmd, cwd=str(ROOT_DIR))
+
+    def _wait_healthy(self, url: str) -> None:
+        deadline = time.monotonic() + HEALTH_TIMEOUT_SECONDS
+        logger.info("Waiting for service at %s ...", url)
+        while time.monotonic() < deadline:
+            try:
+                with httpx.Client(timeout=5.0) as client:
+                    resp = client.get(url)
+                    if resp.status_code == 200:
+                        logger.info("Service healthy: %s", url)
+                        return
+            except Exception:
+                pass
+            time.sleep(HEALTH_POLL_INTERVAL)
+        raise TimeoutError(
+            f"Service at {url} did not become healthy within {HEALTH_TIMEOUT_SECONDS}s"
+        )
+
+
+model_manager = ModelManager()
@@ -0,0 +1,120 @@
+"""Stem separator runner — invokes audio-separator CLI per-request.
+
+Uses a single-worker ThreadPoolExecutor to serialize GPU use.
+Coordinates with model_manager: callers should check model_manager.active_model
+before submitting, and model load routes should check separator_runner.active.
+"""
+
+from __future__ import annotations
+
+import logging
+import subprocess
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Optional
+
+ROOT_DIR = Path(__file__).parent.parent
+SEPARATOR_DIR = ROOT_DIR / "separator"
+OUTPUT_DIR = ROOT_DIR / "outputs" / "separator"
+UPLOAD_DIR = ROOT_DIR / "uploads" / "separator"
+
+ALLOWED_MODELS = frozenset([
+    "UVR-MDX-NET-Inst_HQ_3.onnx",
+    "MDX23C-8KFFT-InstVoc_HQ.ckpt",
+    "htdemucs_6s.yaml",
+])
+
+logger = logging.getLogger(__name__)
+
+
+class SeparatorRunner:
+    def __init__(self) -> None:
+        self._tasks: dict[str, dict] = {}
+        self._executor = ThreadPoolExecutor(max_workers=1)
+        self._running_count = 0
+        self._lock = threading.Lock()
+
+    @property
+    def active(self) -> bool:
+        with self._lock:
+            return self._running_count > 0
+
+    def submit(self, task_id: str, file_path: Path, model_filename: str) -> None:
+        """Schedule a separation task. Returns immediately; poll get_task() for status."""
+        self._tasks[task_id] = {"status": "pending", "created_at": time.time()}
+        with self._lock:
+            self._running_count += 1
+        self._executor.submit(self._run, task_id, file_path, model_filename)
+
+    def get_task(self, task_id: str) -> Optional[dict]:
+        return self._tasks.get(task_id)
+
+    def _run(self, task_id: str, file_path: Path, model_filename: str) -> None:
+        self._tasks[task_id]["status"] = "running"
+        try:
+            cmd = [
+                "uv", "run", "audio-separator",
+                str(file_path),
+                "--model_filename", model_filename,
+                "--output_dir", str(OUTPUT_DIR),
+                "--output_format", "MP3",
+            ]
+            logger.info("Running separator: %s", " ".join(cmd))
+            proc = subprocess.Popen(
+                cmd,
+                cwd=str(SEPARATOR_DIR),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,  # merge stderr into stdout
+                text=True,
+                bufsize=1,
+            )
+            output_lines: list[str] = []
+
+            def _read_output() -> None:
+                for line in proc.stdout:
+                    stripped = line.rstrip()
+                    logger.info("[separator] %s", stripped)
+                    output_lines.append(stripped)
+
+            reader = threading.Thread(target=_read_output, daemon=True)
+            reader.start()
+            try:
+                proc.wait(timeout=600)
+            except subprocess.TimeoutExpired:
+                proc.kill()
+                proc.wait()
+                raise RuntimeError("audio-separator timed out after 600s")
+            finally:
+                reader.join(timeout=5)
+
+            if proc.returncode != 0:
+                detail = "\n".join(output_lines[-20:]) or "unknown error"
+                raise RuntimeError(f"audio-separator exited {proc.returncode}: {detail}")
+
+            # Collect output files — named {input_stem}_{stem_type}_{model_base}.mp3
+            stem = file_path.stem
+            files = sorted(p.name for p in OUTPUT_DIR.glob(f"{stem}_*"))
+            if not files:
+                raise RuntimeError("audio-separator completed but produced no output files")
+
+            self._tasks[task_id].update({"status": "complete", "files": files})
+            logger.info("Separation complete for task %s: %s", task_id, files)
+        except Exception:
+            logger.exception("Separation failed for task %s", task_id)
+            self._tasks[task_id].update({
+                "status": "error",
+                "error": "Separation failed — check server logs.",
+            })
+        finally:
+            with self._lock:
+                self._running_count -= 1
+            # Clean up upload file after processing
+            try:
+                file_path.unlink(missing_ok=True)
+            except Exception:
+                pass
+
+
+separator_runner = SeparatorRunner()