pyr0ball
a6f60c9e07
Full FastAPI backend for the AI music continuation editor:
Services
- chain.py: chain + node CRUD, commit/discard, recursive CTE spine query
- musicgen.py: MusicGenClient with cf-orch allocation + mock mode (CF_MUSICGEN_MOCK=1)
- stems.py: Demucs 4-stem separation subprocess wrapper + mock mode
- export.py: ffmpeg concat demuxer to stitch committed spine into WAV/MP3
API endpoints
- chains: CRUD, multipart audio upload (WAV/MP3/FLAC/OGG/M4A/AIFF)
- nodes: branch creation (202 + BackgroundTasks), commit, discard, audio stream
- gpu: cf-orch capacity status; session allocation stubbed pending cf-orch#43
- stems: Paid-tier stem separation (Demucs, gated via tiers.py)
- export: POST /{chain_id}/export → FileResponse download
- events: SSE stream (node-status events) per chain via asyncio Queue pub/sub
Infrastructure
- lifespan: reads SPARROW_DB_PATH/DATA_DIR at startup (not import time)
- events_store: subscribe/unsubscribe/broadcast pattern for SSE
- CORS: open in dev, SPARROW_CORS_ORIGINS in production
- Background generation opens its own DB connection (WAL-safe)
Tests: 30/30 passing across service units and API integration
75 lines
2.7 KiB
Python
75 lines
2.7 KiB
Python
# app/services/stems.py — Demucs 4-stem separation subprocess wrapper
|
|
#
|
|
# Mock mode (CF_STEMS_MOCK=1): copies source file to all 4 stem paths.
|
|
# Real mode: runs demucs htdemucs model via asyncio.create_subprocess_exec
|
|
# (list-form, no shell=True — not vulnerable to injection).
|
|
#
|
|
# Demucs output layout: {output_dir}/htdemucs/{track.stem}/{stem}.wav
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_MOCK = os.environ.get("CF_STEMS_MOCK", "") == "1"
|
|
_STEMS = ("vocals", "drums", "bass", "other")
|
|
|
|
|
|
async def separate(source_audio_path: str, output_dir: str) -> dict[str, str]:
|
|
"""
|
|
Split source_audio_path into 4 stems.
|
|
|
|
Returns a dict of stem name to absolute file path.
|
|
Raises RuntimeError on demucs failure.
|
|
"""
|
|
if _MOCK:
|
|
return await _mock_separate(source_audio_path, output_dir)
|
|
return await _real_separate(source_audio_path, output_dir)
|
|
|
|
|
|
async def _real_separate(source_audio_path: str, output_dir: str) -> dict[str, str]:
|
|
track = Path(source_audio_path)
|
|
logger.info("Running demucs on %s -> %s", track.name, output_dir)
|
|
# create_subprocess_exec takes a list of args — no shell interpolation
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"python", "-m", "demucs",
|
|
"--model", "htdemucs",
|
|
"--out", output_dir,
|
|
str(track),
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
_, stderr = await proc.communicate()
|
|
if proc.returncode != 0:
|
|
raise RuntimeError(f"demucs failed (exit {proc.returncode}): {stderr.decode()}")
|
|
return _stem_paths(output_dir, track.stem)
|
|
|
|
|
|
async def _mock_separate(source_audio_path: str, output_dir: str) -> dict[str, str]:
|
|
"""Mock: copy source to all 4 stems with a short simulated delay."""
|
|
await asyncio.sleep(1.5)
|
|
track = Path(source_audio_path)
|
|
stem_dir = Path(output_dir) / "htdemucs" / track.stem
|
|
stem_dir.mkdir(parents=True, exist_ok=True)
|
|
paths: dict[str, str] = {}
|
|
for stem in _STEMS:
|
|
dest = str(stem_dir / f"{stem}.wav")
|
|
shutil.copy2(source_audio_path, dest)
|
|
paths[stem] = dest
|
|
logger.info("Mock stems: copied %s -> %s", track.name, stem_dir)
|
|
return paths
|
|
|
|
|
|
def _stem_paths(output_dir: str, track_stem: str) -> dict[str, str]:
|
|
"""Build expected output paths from demucs's standard directory layout."""
|
|
stem_dir = Path(output_dir) / "htdemucs" / track_stem
|
|
return {s: str(stem_dir / f"{s}.wav") for s in _STEMS}
|
|
|
|
|
|
def make_stems_dir(data_dir: str, chain_id: str, node_id: str) -> str:
|
|
"""Standard stems output directory for a node."""
|
|
return str(Path(data_dir) / "chains" / chain_id / "stems" / node_id)
|