serve.py

"""Forced-alignment microservice for the Ankui app (job-based).

Wraps align.py's MMS Kazakh CTC aligner (with HDEMUCS vocal separation)
behind an HTTP API. Designed to run anywhere — a local Mac, a VPS, or a
free Hugging Face Space — so the iOS client can align a user's own audio
+ lyrics directly, with no DB / auth / prior upload.

Because a full-song alignment can take a couple of minutes on CPU (longer
than many reverse proxies allow for a single request), the API is
*asynchronous*:

  GET  /                       -> 200 {"status":"ok","ready":bool}   (health)
  GET  /health                 -> same
  POST /align      (JSON body) -> 202 {"jobId": "..."}
       body: {"lyrics": "<text>", "audio_b64": "<base64>", "ext": "m4a"}
  POST /transcribe (JSON body) -> 202 {"jobId": "..."}
       body: {"audio_b64": "<base64>", "ext": "m4a"}   (no lyrics — STT)
  POST /instrumental (JSON body) -> 202 {"jobId": "..."}
       body: {"audio_b64": "<base64>", "ext": "m4a"}   (no lyrics — karaoke минус)
  GET  /jobs/{id}              -> {"status":"running|done|error", "progress":0..1,
                                   "lines":[...] | "text":"<draft>" |
                                   "audio_b64":"<m4a>", "error":?}

`/transcribe` extracts a *draft* lyric transcript from the audio (Kazakh
kk-turbo Whisper over the isolated vocal stem — the inverse of /align). It's
meant to pre-fill the editor; the user corrects it, then /align times it. It
also returns `lines` with the draft's own rough word timings.

An /align job additionally reports `meanConfidence` (0…1), and every line and
word carries a `confidence`, so the client can point the user at the timings
worth checking instead of implying they are all equally certain.

`/instrumental` returns the song's backing track with the vocals removed
(HDEMUCS, keeping drums+bass+other — the complement of the stem /align
isolates) as base64 m4a, for the app's real karaoke ("минус") mode.

When ANKUI_ALIGN_KEY is set, POST /align and GET /jobs require a matching
`X-Ankui-Key` header (basic abuse guard for a public endpoint).

Audio is processed in a temp file and deleted immediately (PRD §6.5).

Env: ANKUI_ALIGN_HOST (0.0.0.0), ANKUI_ALIGN_PORT or PORT (8765),
     ANKUI_ALIGN_DEVICE (cpu), ANKUI_ALIGN_KEY (shared secret, optional).
"""
from __future__ import annotations

import base64
import json
import os
import re
import sys
import tempfile
import threading
import time
import traceback
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import align  # noqa: E402

HOST = os.environ.get("ANKUI_ALIGN_HOST", "0.0.0.0")
PORT = int(os.environ.get("PORT") or os.environ.get("ANKUI_ALIGN_PORT") or "8765")
DEVICE = os.environ.get("ANKUI_ALIGN_DEVICE", "cpu")
SECRET = os.environ.get("ANKUI_ALIGN_KEY", "")

_READY = False

# jobId -> {"status","progress","lines","duration","alignMode","error","ts"}
_JOBS: dict = {}
_JOBS_LOCK = threading.Lock()
_JOB_TTL = 3600  # forget finished jobs after an hour

# Serialize the heavy model work (HDEMUCS separation, MMS alignment, Whisper
# transcription) across ALL job types. The Space is CPU-only with ~2 vCPUs;
# running two full-song passes at once doesn't parallelize — it thrashes the
# cores (each already ~3-4x realtime) and stacks hundreds of MB of tensors,
# risking an OOM kill. One worker at a time bounds both. Jobs queue here;
# their status stays "running" (progress 0.05) until they acquire it.
_WORK = threading.Semaphore(1)


def _split_lines(raw: str) -> list[str]:
    return [ln.strip() for ln in raw.replace("\r", "").split("\n") if ln.strip()]


def _set(job_id: str, **kw):
    with _JOBS_LOCK:
        _JOBS.setdefault(job_id, {}).update(kw)


def _get(job_id: str):
    with _JOBS_LOCK:
        return dict(_JOBS.get(job_id, {})) or None


def _gc():
    now = time.time()
    with _JOBS_LOCK:
        for k in [k for k, v in _JOBS.items()
                  if v.get("status") in ("done", "error") and now - v.get("ts", now) > _JOB_TTL]:
            _JOBS.pop(k, None)


def _run_job(job_id: str, lines: list[str], audio: bytes, ext: str,
             map_structure: bool = False):
    path = None
    plan = None
    try:
        fd, path = tempfile.mkstemp(suffix="." + ext)
        try:
            os.write(fd, audio)
        finally:
            os.close(fd)
        with _WORK:   # serialize heavy model work across all job types
            _set(job_id, status="running", progress=0.1, ts=time.time())
            if map_structure:
                # Opt-in: a lyric sheet writes the chorus once, the recording
                # sings it three times, and forced alignment cannot invent the
                # repeat. Transcribing first tells us the *performance* order —
                # the transcript's words are poor but its structure is right.
                # Costs a full extra pass, which is why the caller has to ask.
                try:
                    import structure
                    draft = align.transcribe_audio(path, device=DEVICE)
                    expanded, plan = structure.expand_reference(
                        lines, draft.get("lines") or [])
                    if len(expanded) >= len(lines):
                        lines = expanded          # additive only, never shorter
                except Exception as exc:  # noqa: BLE001 — fall back to the sheet
                    traceback.print_exc()
                    sys.stderr.write(f"[serve] structure mapping skipped: {exc}\n")
                    plan = None
                _set(job_id, status="running", progress=0.5, ts=time.time())
            result = align.run_alignment(path, lines, device=DEVICE)
        out_lines = result.get("lines", [])
        if not out_lines:
            _set(job_id, status="error", error="alignment produced no lines", ts=time.time())
            return
        # `meanConfidence` (and the per-line `confidence` inside each line) let
        # the client tell the user which timings to check rather than presenting
        # every alignment as equally certain.
        _set(job_id, status="done", progress=1.0, lines=out_lines,
             duration=result.get("duration", 0.0),
             meanConfidence=result.get("meanConfidence", 0.0),
             structurePlan=plan,
             alignMode=result.get("alignMode", "ctc-mms-fl102-kaz"), ts=time.time())
    except Exception as exc:  # noqa: BLE001
        traceback.print_exc()
        _set(job_id, status="error", error=f"{type(exc).__name__}: {exc}", ts=time.time())
    finally:
        if path:
            try:
                os.unlink(path)
            except OSError:
                pass


def _run_transcribe_job(job_id: str, audio: bytes, ext: str):
    """STT draft: audio → text (no reference lyrics). Mirrors _run_job."""
    path = None
    try:
        fd, path = tempfile.mkstemp(suffix="." + ext)
        try:
            os.write(fd, audio)
        finally:
            os.close(fd)
        with _WORK:   # serialize heavy model work across all job types
            _set(job_id, status="running", progress=0.1, ts=time.time())
            result = align.transcribe_audio(path, device=DEVICE)
        text = (result.get("text") or "").strip()
        if not text:
            _set(job_id, status="error", error="no speech detected in audio", ts=time.time())
            return
        # `lines` carries the draft's own rough per-line/per-word timings from
        # Whisper. `text` is unchanged, so an older client keeps working.
        _set(job_id, status="done", progress=1.0, text=text,
             lines=result.get("lines") or [],
             duration=result.get("duration", 0.0),
             model=result.get("model", "kk-turbo-ksc2"), ts=time.time())
    except Exception as exc:  # noqa: BLE001
        traceback.print_exc()
        _set(job_id, status="error", error=f"{type(exc).__name__}: {exc}", ts=time.time())
    finally:
        if path:
            try:
                os.unlink(path)
            except OSError:
                pass


def _run_instrumental_job(job_id: str, audio: bytes, ext: str):
    """Karaoke минус: audio → instrumental stem (vocals removed). Mirrors
    _run_job; the result carries the separated audio as base64 m4a, which
    the client decodes and caches as a local file for karaoke playback."""
    path = None
    try:
        fd, path = tempfile.mkstemp(suffix="." + ext)
        try:
            os.write(fd, audio)
        finally:
            os.close(fd)
        with _WORK:   # serialize heavy model work across all job types
            _set(job_id, status="running", progress=0.1, ts=time.time())
            # Real per-chunk progress over the multi-minute separation so
            # the client shows an honest bar instead of a wall-clock guess.
            def _prog(frac):
                _set(job_id, status="running",
                     progress=0.1 + 0.85 * float(frac), ts=time.time())
            result = align.render_instrumental(path, device=DEVICE, progress=_prog)
        data = result.get("audio") or b""
        if not data:
            _set(job_id, status="error", error="instrumental separation produced no audio", ts=time.time())
            return
        _set(job_id, status="done", progress=1.0,
             audio_b64=base64.b64encode(data).decode("ascii"),
             ext=result.get("ext", "m4a"),
             duration=result.get("duration", 0.0), ts=time.time())
    except Exception as exc:  # noqa: BLE001
        traceback.print_exc()
        _set(job_id, status="error", error=f"{type(exc).__name__}: {exc}", ts=time.time())
    finally:
        if path:
            try:
                os.unlink(path)
            except OSError:
                pass


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, fmt, *args):
        sys.stderr.write("[serve] " + (fmt % args) + "\n")

    def _send(self, code: int, obj: dict):
        body = json.dumps(obj, ensure_ascii=False).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _authorized(self) -> bool:
        if not SECRET:
            return True
        return self.headers.get("X-Ankui-Key", "") == SECRET

    def do_GET(self):
        if self.path in ("/", "/health"):
            self._send(200, {"status": "ok", "ready": _READY})
            return
        if self.path.startswith("/jobs/"):
            if not self._authorized():
                self._send(401, {"error": "unauthorized"})
                return
            job = _get(self.path[len("/jobs/"):])
            if not job:
                self._send(404, {"error": "unknown job"})
                return
            self._send(200, job)
            return
        if self.path == "/stats":
            if not self._authorized():
                self._send(401, {"error": "unauthorized"})
                return
            with _JOBS_LOCK:
                recent = sorted(_JOBS.values(), key=lambda v: v.get("ts", 0), reverse=True)[:10]
                summary = [{"status": j.get("status"), "progress": j.get("progress"),
                            "mode": j.get("alignMode"), "lines": len(j.get("lines", [])),
                            "error": j.get("error")} for j in recent]
                total = len(_JOBS)
            self._send(200, {"ready": _READY, "totalJobs": total, "recent": summary})
            return
        self._send(404, {"error": "not found"})

    def do_POST(self):
        # Always drain the request body FIRST. Returning early (401/503)
        # without consuming Content-Length leaves the body in the socket;
        # on a kept-alive connection (HF's proxy reuses upstream conns)
        # those bytes corrupt the next request's start line.
        length = int(self.headers.get("Content-Length", "0") or 0)
        raw = self.rfile.read(length) if length > 0 else b""
        if self.path not in ("/align", "/transcribe", "/instrumental"):
            self._send(404, {"error": "not found"})
            return
        if not self._authorized():
            self._send(401, {"error": "unauthorized"})
            return
        if not _READY:
            self._send(503, {"error": "model still loading, retry shortly"})
            return
        try:
            if not raw:
                self._send(400, {"error": "empty body"})
                return
            payload = json.loads(raw)
            audio_b64 = payload.get("audio_b64", "")
            # Whitelist, not lstrip(".") — this string becomes a temp-file
            # suffix, and "../x", a NUL, or a 300-char value makes mkstemp raise
            # inside a daemon thread, stranding the job as "running" forever
            # (the client then polls it for 20 minutes before giving up).
            ext = re.sub(r"[^A-Za-z0-9]", "", (payload.get("ext") or "")).lower()[:8] or "m4a"
            if not audio_b64:
                self._send(400, {"error": "audio missing"})
                return
            audio = base64.b64decode(audio_b64)
            _gc()
            job_id = uuid.uuid4().hex
            _set(job_id, status="running", progress=0.05, ts=time.time())

            if self.path == "/transcribe":
                # STT draft — no lyrics needed.
                threading.Thread(target=_run_transcribe_job,
                                 args=(job_id, audio, ext), daemon=True).start()
            elif self.path == "/instrumental":
                # Karaoke минус — separate the backing track (no lyrics).
                threading.Thread(target=_run_instrumental_job,
                                 args=(job_id, audio, ext), daemon=True).start()
            else:
                lines = _split_lines(payload.get("lyrics", ""))
                if not lines:
                    self._send(400, {"error": "lyrics are empty"})
                    return
                threading.Thread(
                    target=_run_job,
                    args=(job_id, lines, audio, ext,
                          bool(payload.get("mapStructure"))),
                    daemon=True).start()
            self._send(202, {"jobId": job_id})
        except Exception as exc:  # noqa: BLE001
            traceback.print_exc()
            self._send(500, {"error": f"{type(exc).__name__}: {exc}"})


def _warmup():
    """Load models in the background so the HTTP port comes up immediately
    (HF Spaces' health check needs the port live within its startup grace
    period; models are baked into the image so this is a fast disk load)."""
    global _READY
    t0 = time.time()
    align.load_model(DEVICE)
    try:
        align.load_separator(DEVICE)
    except Exception as exc:  # noqa: BLE001
        sys.stderr.write(f"[serve] separator preload failed (will run raw): {exc}\n")
    _READY = True
    print(f"[serve] models ready in {time.time()-t0:.1f}s", flush=True)


def main():
    threading.Thread(target=_warmup, daemon=True).start()
    server = ThreadingHTTPServer((HOST, PORT), Handler)
    print(f"[serve] listening on http://{HOST}:{PORT}  (POST /align, GET /jobs/<id>, GET /health)", flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()