sourav-das/stem-separator

1import asyncio2from dataclasses import dataclass3 4from backend.separator import StemSeparatorService5from backend import file_manager6 7 8@dataclass9class JobProgress:10    state: str = "queued"11    progress: float = 0.012    message: str = "Waiting in queue..."13    stems: dict[str, str] | None = None14    error: str | None = None15 16 17# Shared state18jobs: dict[str, JobProgress] = {}19_queue: asyncio.Queue | None = None20 21 22def get_queue() -> asyncio.Queue:23    global _queue24    if _queue is None:25        _queue = asyncio.Queue(maxsize=5)26    return _queue27 28 29def get_job_progress(job_id: str) -> JobProgress | None:30    return jobs.get(job_id)31 32 33async def enqueue_job(job_id: str, stems: list[str], output_format: str) -> bool:34    """Enqueue a separation job. Returns False if queue is full."""35    q = get_queue()36    if q.full():37        return False38    jobs[job_id] = JobProgress()39    await q.put((job_id, stems, output_format))40    return True41 42 43async def worker_loop():44    """Single worker that processes separation jobs sequentially."""45    separator = StemSeparatorService()46    q = get_queue()47 48    while True:49        job_id, stems, output_format = await q.get()50        try:51            progress = jobs.get(job_id)52            if progress is None:53                progress = JobProgress()54                jobs[job_id] = progress55 56            def update_progress(state: str, pct: float):57                progress.state = state58                progress.progress = pct59                messages = {60                    "loading_model": "Loading BS-RoFormer model...",61                    "separating": "Separating stems...",62                    "finalizing": "Finalizing output files...",63                    "done": "Separation complete!",64                }65                progress.message = messages.get(state, f"{state}...")66 67            input_file = file_manager.get_input_file(job_id)68            if input_file is None:69                progress.state = "error"70                progress.error = "Input file not found"71                progress.message = "Error: input file not found"72                continue73 74            output_dir = str(file_manager.get_output_dir(job_id))75 76            # Run separation in a thread to avoid blocking the event loop77            loop = asyncio.get_event_loop()78            result = await loop.run_in_executor(79                None,80                separator.separate,81                str(input_file),82                output_dir,83                stems,84                output_format,85                update_progress,86            )87 88            progress.state = "done"89            progress.progress = 1.090            progress.message = "Separation complete!"91            progress.stems = result92 93        except Exception as e:94            progress = jobs.get(job_id, JobProgress())95            progress.state = "error"96            progress.error = str(e)97            progress.message = f"Error: {e}"98            jobs[job_id] = progress99        finally:100            q.task_done()101