muscriptor_api.py

"""Persistent FastAPI queue for MuScriptor transcription."""

import argparse
import asyncio
import hashlib
import json
import os
import re
import subprocess
import threading
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from uuid import uuid4

import uvicorn
from fastapi import FastAPI, File, HTTPException, Query, UploadFile
from fastapi.responses import FileResponse

from muscriptor_cli import MuScriptorLlama


def _convert_to_wav(input_path: Path, output_path: Path) -> None:
    cmd = [
        "ffmpeg", "-y", "-i", str(input_path),
        "-ac", "1", "-ar", "16000",
        "-c:a", "pcm_s16le",
        str(output_path),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg failed: {result.stderr}")


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _md5(path: Path) -> str:
    digest = hashlib.md5()
    with path.open("rb") as file:
        for chunk in iter(lambda: file.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _safe_name(name: str | None) -> str:
    name = Path(name or "audio.bin").name
    name = re.sub(r"[^A-Za-z0-9._-]", "_", name).strip("._")
    return name[:160] or "audio.bin"


def _task_response(task: dict[str, Any]) -> dict[str, Any]:
    response = {
        "sequence": task["sequence"],
        "status": task["status"],
        "completed": task["status"] == "completed",
        "filename": task.get("display_name", ""),
    }
    if task["status"] == "completed":
        response["md5"] = task["midi_md5"]
    elif task["status"] == "error":
        response["error"] = task["error"]
    return response


class TaskStore:
    def __init__(self, data_dir: Path):
        self.data_dir = data_dir.resolve()
        self.audio_dir = self.data_dir / "audio"
        self.midi_dir = self.data_dir / "midi"
        self.audio_dir.mkdir(parents=True, exist_ok=True)
        self.midi_dir.mkdir(parents=True, exist_ok=True)
        self.path = self.data_dir / "muscriptor_tasks.json"
        self._lock = threading.Lock()
        self._data: dict[str, Any] = {"next_sequence": 1, "tasks": {}}
        self._load()

    def _load(self) -> None:
        if not self.path.exists():
            return
        with self.path.open("r", encoding="utf-8") as file:
            data = json.load(file)
        if not isinstance(data.get("tasks"), dict):
            raise ValueError(f"invalid task database: {self.path}")
        self._data = data

    def _save_locked(self) -> None:
        temp = self.path.with_name(f".{self.path.name}.{uuid4().hex}.tmp")
        with temp.open("w", encoding="utf-8") as file:
            json.dump(self._data, file, ensure_ascii=False, indent=2)
            file.flush()
            os.fsync(file.fileno())
        temp.replace(self.path)

    def add(self, original_name: str) -> dict[str, Any]:
        with self._lock:
            sequence = int(self._data["next_sequence"])
            self._data["next_sequence"] = sequence + 1
            suffix = Path(original_name or "audio.bin").suffix or ".bin"
            safe_stem = _safe_name(original_name).rsplit(".", 1)[0] if "." in _safe_name(original_name) else _safe_name(original_name)
            audio_name = f"audio_{sequence:08d}_{safe_stem}{suffix}"
            midi_name = f"midi_{sequence:08d}.mid"
            display_name = Path(original_name or "audio").stem
            task = {
                "sequence": sequence,
                "status": "uploading",
                "audio_path": audio_name,
                "midi_path": midi_name,
                "display_name": display_name,
                "midi_md5": None,
                "error": None,
                "created_at": _now(),
                "updated_at": _now(),
            }
            self._data["tasks"][str(sequence)] = task
            self._save_locked()
            return dict(task)

    def get(self, sequence: int) -> dict[str, Any] | None:
        with self._lock:
            task = self._data["tasks"].get(str(sequence))
            return dict(task) if task else None

    def page(self, page: int, page_size: int = 30) -> tuple[list[dict[str, Any]], int]:
        with self._lock:
            tasks = sorted(
                self._data["tasks"].values(),
                key=lambda task: int(task["sequence"]),
                reverse=True,
            )
            start = (page - 1) * page_size
            return [dict(task) for task in tasks[start : start + page_size]], len(tasks)

    def update(self, sequence: int, **values: Any) -> dict[str, Any]:
        with self._lock:
            task = self._data["tasks"][str(sequence)]
            task.update(values, updated_at=_now())
            self._save_locked()
            return dict(task)

    def recoverable(self) -> list[int]:
        with self._lock:
            result = []
            changed = False
            for task in self._data["tasks"].values():
                if task["status"] == "uploading":
                    task.update(
                        status="error",
                        error="upload was interrupted before completion",
                        updated_at=_now(),
                    )
                    changed = True
                    continue
                if task["status"] in {"queued", "processing"}:
                    task.update(status="queued", error=None, updated_at=_now())
                    result.append(int(task["sequence"]))
                    changed = True
            if changed:
                self._save_locked()
            return sorted(result)

    def find_by_md5(self, midi_md5: str) -> dict[str, Any] | None:
        with self._lock:
            for task in self._data["tasks"].values():
                if task["status"] == "completed" and task["midi_md5"] == midi_md5:
                    return dict(task)
        return None


class TranscriptionQueue:
    def __init__(self, store: TaskStore, weights: str, server_url: str, parallel: int):
        self.store = store
        self.weights = weights
        self.server_url = server_url
        self.parallel = parallel
        self.queue: asyncio.Queue[int] = asyncio.Queue()
        self.model: MuScriptorLlama | None = None

    async def start(self) -> None:
        for sequence in self.store.recoverable():
            self.queue.put_nowait(sequence)

    def _transcribe(self, sequence: int) -> None:
        task = self.store.get(sequence)
        if task is None:
            return
        if self.model is None:
            self.model = MuScriptorLlama(
                self.weights,
                server_url=self.server_url,
                parallel=self.parallel,
            )
        audio_path = self.store.audio_dir / task["audio_path"]
        wav_path = audio_path.with_suffix(".wav")
        _convert_to_wav(audio_path, wav_path)
        try:
            midi_path = self.store.midi_dir / task["midi_path"]
            midi = self.model.transcribe_to_midi(str(wav_path))
        finally:
            wav_path.unlink(missing_ok=True)
        midi_path.write_bytes(midi)
        self.store.update(sequence, status="completed", midi_md5=_md5(midi_path))

    async def worker(self) -> None:
        while True:
            sequence = await self.queue.get()
            try:
                self.store.update(sequence, status="processing", error=None)
                await asyncio.to_thread(self._transcribe, sequence)
            except asyncio.CancelledError:
                raise
            except Exception as exc:
                self.store.update(sequence, status="error", error=str(exc))
            finally:
                self.queue.task_done()


def create_app(
    data_dir: Path = Path.cwd(),
    weights: str = "model.safetensors",
    server_url: str = "http://127.0.0.1:8081",
    parallel: int = 8,
) -> FastAPI:
    store = TaskStore(data_dir)
    manager = TranscriptionQueue(store, weights, server_url, parallel)

    @asynccontextmanager
    async def lifespan(_: FastAPI):
        await manager.start()
        worker = asyncio.create_task(manager.worker())
        try:
            yield
        finally:
            worker.cancel()
            await asyncio.gather(worker, return_exceptions=True)

    app = FastAPI(title="MuScriptor API", lifespan=lifespan)
    app.state.task_store = store
    app.state.transcription_queue = manager

    @app.post("/tasks", status_code=202)
    async def submit_audio(file: UploadFile = File(...)):
        task = store.add(file.filename or "audio.bin")
        audio_path = store.audio_dir / task["audio_path"]
        try:
            with audio_path.open("wb") as output:
                while chunk := await file.read(1024 * 1024):
                    output.write(chunk)
            if audio_path.stat().st_size == 0:
                raise ValueError("uploaded audio file is empty")
        except Exception as exc:
            audio_path.unlink(missing_ok=True)
            store.update(task["sequence"], status="error", error=str(exc))
            return {"success": False, "sequence": task["sequence"], "error": str(exc)}
        finally:
            await file.close()

        store.update(task["sequence"], status="queued")
        manager.queue.put_nowait(task["sequence"])
        return {"success": True, "sequence": task["sequence"]}

    @app.get("/tasks")
    async def list_tasks(page: int = Query(default=1, ge=1)):
        page_size = 30
        tasks, total = store.page(page, page_size)
        return {
            "page": page,
            "page_size": page_size,
            "total": total,
            "total_pages": (total + page_size - 1) // page_size,
            "tasks": [_task_response(task) for task in tasks],
        }

    @app.get("/tasks/{sequence}")
    async def task_status(sequence: int):
        task = store.get(sequence)
        if task is None:
            raise HTTPException(status_code=404, detail="task not found")
        return _task_response(task)

    @app.get("/midi/{midi_md5}")
    async def download_midi(midi_md5: str):
        if not re.fullmatch(r"[0-9a-f]{32}", midi_md5):
            raise HTTPException(status_code=404, detail="MIDI not found")
        task = store.find_by_md5(midi_md5)
        if task is None:
            raise HTTPException(status_code=404, detail="MIDI not found")
        path = store.midi_dir / task["midi_path"]
        if not path.is_file() or _md5(path) != midi_md5:
            raise HTTPException(status_code=404, detail="MIDI not found")
        return FileResponse(path, media_type="audio/midi", filename=path.name)

    return app


app = create_app()


def main() -> None:
    parser = argparse.ArgumentParser(description="MuScriptor FastAPI service")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=1240)
    parser.add_argument("--weights", default="conditioning_weights.safetensors")
    parser.add_argument("--server", default="http://127.0.0.1:8081")
    parser.add_argument("--parallel", type=int, default=8)
    args = parser.parse_args()
    uvicorn.run(
        create_app(weights=args.weights, server_url=args.server, parallel=args.parallel),
        host=args.host,
        port=args.port,
    )


if __name__ == "__main__":
    main()