ShanitG/stem-splitter

1from fastapi import FastAPI, UploadFile, HTTPException
2from fastapi.middleware.cors import CORSMiddleware
3from fastapi.responses import FileResponse, JSONResponse
4import uvicorn
5import os
6import uuid
7import torch
8import soundfile as sf
9import numpy as np
10from demucs.pretrained import get_pretrained
11from demucs.apply import apply_model
12import asyncio
13import aiofiles
14from pathlib import Path
15import logging
16
17# Configure logging
18logging.basicConfig(level=logging.INFO)
19logger = logging.getLogger(__name__)
20
21app = FastAPI()
22
23# CORS configuration
24app.add_middleware(
25    CORSMiddleware,
26    allow_origins=["*"],
27    allow_credentials=False,
28    allow_methods=["*"],
29    allow_headers=["*"],
30)
31
32# Create directories
33UPLOAD_DIR = Path("uploads")
34PROCESSED_DIR = Path("processed")
35UPLOAD_DIR.mkdir(exist_ok=True)
36PROCESSED_DIR.mkdir(exist_ok=True)
37
38# Global variables
39processing_status = {}
40model = None
41
42@app.on_event("startup")
43async def startup_event():
44    global model
45    try:
46        logger.info("Loading Demucs model...")
47        model = get_pretrained("htdemucs")
48        model.cuda() if torch.cuda.is_available() else model.cpu()
49        logger.info("Model loaded successfully")
50    except Exception as e:
51        logger.error(f"Error loading model: {e}")
52        raise
53
54async def process_audio_file(input_path: Path, job_id: str):
55    try:
56        logger.info(f"Processing file: {input_path}")
57        
58        # Create output directory
59        output_dir = PROCESSED_DIR / job_id
60        output_dir.mkdir(exist_ok=True)
61        
62        # Load and process audio
63        audio, sr = sf.read(str(input_path))
64        if len(audio.shape) == 1:
65            audio = np.expand_dims(audio, axis=1)
66        
67        # Convert to tensor
68        audio_tensor = torch.tensor(audio.T, dtype=torch.float32)
69        audio_tensor = audio_tensor.unsqueeze(0)
70
71        # Process with model
72        with torch.no_grad():
73            sources = apply_model(model, audio_tensor, device="cuda" if torch.cuda.is_available() else "cpu")[0]
74            sources = sources.cpu().numpy()
75
76        # Save stems
77        stem_names = ["vocals", "drums", "bass", "other"]
78        for source, name in zip(sources, stem_names):
79            output_path = output_dir / f"{name}.wav"
80            sf.write(str(output_path), source.T, sr)
81            processing_status[job_id]["progress"] += 25
82
83        processing_status[job_id]["status"] = "completed"
84        logger.info(f"Processing completed for job {job_id}")
85        
86    except Exception as e:
87        logger.error(f"Error processing file: {e}")
88        processing_status[job_id]["status"] = "failed"
89        processing_status[job_id]["error"] = str(e)
90        raise
91
92@app.get("/health")
93async def health_check():
94    return {"status": "ok", "model_loaded": model is not None}
95
96@app.post("/process")
97async def process_file(audio: UploadFile):
98    try:
99        if not audio.filename.lower().endswith(('.mp3', '.wav', '.ogg', '.m4a')):
100            raise HTTPException(400, detail="Unsupported file format")
101
102        # Generate job ID and save file
103        job_id = str(uuid.uuid4())
104        input_path = UPLOAD_DIR / f"{job_id}_{audio.filename}"
105        
106        async with aiofiles.open(input_path, 'wb') as f:
107            content = await audio.read()
108            await f.write(content)
109
110        # Initialize processing status
111        processing_status[job_id] = {
112            "status": "processing",
113            "progress": 0,
114            "error": None
115        }
116
117        # Start processing in background
118        asyncio.create_task(process_audio_file(input_path, job_id))
119        
120        return JSONResponse(content={"id": job_id})
121        
122    except Exception as e:
123        logger.error(f"Error handling upload: {e}")
124        raise HTTPException(status_code=500, detail=str(e))
125
126@app.get("/status/{job_id}")
127async def get_status(job_id: str):
128    if job_id not in processing_status:
129        raise HTTPException(404, detail="Job not found")
130    return processing_status[job_id]
131
132@app.get("/download/{job_id}/{stem_name}")
133async def download_stem(job_id: str, stem_name: str):
134    try:
135        if stem_name not in ["vocals", "drums", "bass", "other"]:
136            raise HTTPException(400, detail="Invalid stem name")
137
138        stem_path = PROCESSED_DIR / job_id / f"{stem_name}.wav"
139        if not stem_path.exists():
140            raise HTTPException(404, detail="Stem not found")
141
142        return FileResponse(
143            path=stem_path,
144            filename=f"{stem_name}.wav",
145            media_type="audio/wav"
146        )
147    except Exception as e:
148        logger.error(f"Error downloading stem: {e}")
149        raise