muscriptor 0.1.0__py3-none-any.whl

muscriptor/server.py

@@ -0,0 +1,267 @@
+"""FastAPI server exposing transcription as an SSE event stream.
+
+POST /transcribe with an audio file (multipart/form-data field `file`; WAV,
+or any format soundfile/libsndfile can read — mp3, flac, ogg, m4a, …) returns
+`text/event-stream`. Each event's data is a JSON dict tagged by `type`:
+`start` / `end` note events (same shape as `muscriptor.main._event_to_dict`),
+`progress` chunk anchors (`{completed, total}`), and a final `midi` event
+carrying the base64-encoded .mid file.
+"""
+
+import asyncio
+import base64
+import dataclasses
+import io
+import json
+import os
+import tempfile
+import threading
+import time
+import wave
+from pathlib import Path
+from typing import Annotated
+
+from fastapi import FastAPI, File, Form, HTTPException, UploadFile
+from fastapi.responses import Response, StreamingResponse
+from fastapi.staticfiles import StaticFiles
+from starlette.background import BackgroundTask
+
+from muscriptor.events import NoteEndEvent, NoteStartEvent, ProgressEvent
+from muscriptor.tokenizer.mt3 import MT3_FULL_PLUS_GROUP_NAMES
+from muscriptor.transcription_model import TranscriptionModel
+from muscriptor.utils.audio import _read_non_wav_file, _read_wav_file
+
+
+def _make_release_once(lock: threading.Lock):
+    """Return a callable that releases `lock` at most once.
+
+    Safe to call from multiple cleanup paths (generator finally + response
+    background task), possibly from different threads, without risking a
+    double-release RuntimeError.
+    """
+    guard = threading.Lock()
+    released = False
+
+    def release():
+        nonlocal released
+        with guard:
+            if released:
+                return
+            released = True
+        lock.release()
+
+    return release
+
+
+def event_to_dict(ev: NoteStartEvent | NoteEndEvent) -> dict:
+    if isinstance(ev, NoteStartEvent):
+        return {"type": "start", **dataclasses.asdict(ev)}
+    return {
+        "type": "end",
+        "end_time": ev.end_time,
+        "start_event_index": ev.start_event_index,
+    }
+
+
+def create_app(model: TranscriptionModel, web_dir: str | Path | None = None) -> FastAPI:
+    app = FastAPI(title="muscriptor")
+
+    transcribe_lock = threading.Lock()
+    lock_timeout_s = 60.0
+    # Cancel event of the run currently holding the lock (or the last one to
+    # have held it). A new /transcribe request sets it so an in-flight run
+    # stops at its next event boundary instead of transcribing to completion
+    # for a client that has moved on. This must not rely on TCP disconnects:
+    # aborts don't always reach us (e.g. port forwards / proxies that keep the
+    # upstream connection open after the browser aborts).
+    cancel_guard = threading.Lock()
+    current_cancel: threading.Event | None = None
+
+    @app.get("/health")
+    async def health():
+        return {"status": "ok"}
+
+    @app.get("/instruments")
+    async def list_instruments():
+        return {"instruments": list(MT3_FULL_PLUS_GROUP_NAMES.keys())}
+
+    @app.post("/transcribe")
+    async def transcribe(
+        file: Annotated[UploadFile, File()],
+        instruments: Annotated[list[str], Form(default_factory=list)],
+    ) -> StreamingResponse:
+        data = await file.read()
+        # PCM WAV goes through the stdlib reader (keeps WAV decoding byte-for-byte
+        # identical to the CLI); anything that isn't a readable WAV (mp3, flac,
+        # ogg, m4a, …) falls back to soundfile/libsndfile. A genuinely
+        # undecodable upload (corrupt/truncated file, or a format libsndfile
+        # can't read) is the client's fault, so report it as a 400 rather than
+        # letting it surface as a 500.
+        try:
+            wav, sr = _read_wav_file(io.BytesIO(data))
+        except (wave.Error, EOFError):
+            try:
+                wav, sr = _read_non_wav_file(io.BytesIO(data))
+            except Exception as e:
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"could not decode audio file '{file.filename}': {e}",
+                ) from e
+
+        unknown = [n for n in instruments if n not in MT3_FULL_PLUS_GROUP_NAMES]
+        if unknown:
+            raise HTTPException(
+                status_code=400,
+                detail=f"unknown instrument name(s): {', '.join(unknown)}",
+            )
+
+        # Preempt whoever holds the lock, then wait for it. The short acquire
+        # timeout re-signals each second so that even a run that started while
+        # we were already waiting (another preempting request that beat us to
+        # the lock) gets cancelled too — the newest request always wins.
+        nonlocal current_cancel
+        deadline = time.monotonic() + lock_timeout_s
+        while True:
+            with cancel_guard:
+                if current_cancel is not None:
+                    current_cancel.set()
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise HTTPException(
+                    status_code=503,
+                    detail="server busy: another transcription is in progress",
+                )
+            acquired = await asyncio.to_thread(
+                transcribe_lock.acquire, True, min(1.0, remaining)
+            )
+            if acquired:
+                break
+        cancel = threading.Event()
+        with cancel_guard:
+            current_cancel = cancel
+
+        # Release exactly once, from whichever path runs first. The generator's
+        # finally covers normal completion, errors and mid-stream disconnects;
+        # the StreamingResponse background task covers the case where the client
+        # disconnects *before* the generator is ever iterated (so its finally
+        # would never run) — which otherwise leaks the lock forever.
+        release_lock = _make_release_once(transcribe_lock)
+
+        def gen():
+            try:
+                events: list[NoteStartEvent | NoteEndEvent] = []
+                # batch_size=1 so each chunk's notes stream out as soon as it is
+                # generated, instead of waiting for a whole batch of chunks.
+                # no_eos_is_ok=True so one runaway chunk that never emits EOS only
+                # warns (and keeps its notes) instead of aborting the whole stream.
+                for ev in model.transcribe(
+                    (wav, sr),
+                    instruments=instruments or None,
+                    batch_size=1,
+                    no_eos_is_ok=True,
+                ):
+                    # A newer request preempted this run — stop generating
+                    # (closing the model.transcribe generator) and release the
+                    # lock via the finally, at most one chunk after the signal.
+                    if cancel.is_set():
+                        return
+                    if isinstance(ev, ProgressEvent):
+                        # Coarse chunk-completion anchor — forward it but keep it
+                        # out of the note list the MIDI file is built from.
+                        payload = json.dumps(
+                            {
+                                "type": "progress",
+                                "completed": ev.completed,
+                                "total": ev.total,
+                            }
+                        )
+                        yield f"data: {payload}\n\n"
+                        continue
+                    events.append(ev)
+                    payload = json.dumps(event_to_dict(ev))
+                    yield f"data: {payload}\n\n"
+                # All notes streamed — build the MIDI file in memory (reusing the
+                # exact `muscriptor transcribe` logic) and send it as a final event
+                # with the bytes base64-encoded.
+                if cancel.is_set():
+                    return
+                midi_bytes = model.events_to_midi_bytes(iter(events))
+                midi_b64 = base64.b64encode(midi_bytes).decode("ascii")
+                payload = json.dumps({"type": "midi", "data": midi_b64})
+                yield f"data: {payload}\n\n"
+            finally:
+                release_lock()
+
+        return StreamingResponse(
+            gen(),
+            media_type="text/event-stream",
+            background=BackgroundTask(release_lock),
+            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+        )
+
+    @app.post("/auralize")
+    async def auralize(
+        midi: Annotated[UploadFile, File()],
+        audio: Annotated[UploadFile | None, File()] = None,
+        mode: Annotated[str, Form()] = "mix",
+    ):
+        """Render a transcription as WAV.
+
+        mode="mix": stereo, original audio (L) + FluidSynth synthesis (R);
+        requires the `audio` upload. mode="synth": mono, just the synthesis.
+        """
+        from muscriptor.utils.auralization import auralize as do_auralize
+        from muscriptor.utils.auralization import synthesize
+
+        if mode not in ("mix", "synth"):
+            raise HTTPException(status_code=400, detail=f"unknown mode: {mode!r}")
+        if mode == "mix" and audio is None:
+            raise HTTPException(
+                status_code=400, detail="mode='mix' requires an audio file"
+            )
+
+        midi_data = await midi.read()
+        tmp_paths: list[str] = []
+
+        with tempfile.NamedTemporaryFile(suffix=".mid", delete=False) as tmp_midi:
+            tmp_midi.write(midi_data)
+            midi_tmp = tmp_midi.name
+            tmp_paths.append(midi_tmp)
+
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_out:
+            out_tmp = tmp_out.name
+            tmp_paths.append(out_tmp)
+
+        try:
+            if mode == "synth":
+                synthesize(midi_path=midi_tmp, output_path=out_tmp)
+            else:
+                audio_data = await audio.read()
+                suffix = Path(audio.filename or "audio.wav").suffix.lower() or ".wav"
+                with tempfile.NamedTemporaryFile(
+                    suffix=suffix, delete=False
+                ) as tmp_audio:
+                    tmp_audio.write(audio_data)
+                    tmp_paths.append(tmp_audio.name)
+                do_auralize(
+                    midi_path=midi_tmp,
+                    original_audio_path=tmp_audio.name,
+                    output_path=out_tmp,
+                )
+            with open(out_tmp, "rb") as f:
+                wav_bytes = f.read()
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
+        finally:
+            for p in tmp_paths:
+                if os.path.exists(p):
+                    os.unlink(p)
+
+        return Response(content=wav_bytes, media_type="audio/wav")
+
+    if web_dir is not None:
+        web_path = Path(web_dir)
+        if web_path.is_dir():
+            app.mount("/", StaticFiles(directory=web_path, html=True), name="web")
+
+    return app

muscriptor/tokenizer/__init__.py

@@ -0,0 +1,3 @@
+from .notes import note_event2midi, note2note_event
+
+__all__ = ["note_event2midi", "note2note_event"]

muscriptor/tokenizer/mt3.py

@@ -0,0 +1,231 @@
+"""MT3 MIDI tokenizer.
+
+Adapted from YourMT3+ (https://github.com/mimbres/YourMT3) and audiocraft_trans.
+"""
+
+import difflib
+import logging
+from collections.abc import Iterable
+
+from muscriptor.tokenizer.notes import (
+    DRUM_PROGRAM,
+    SPECIAL_TOKENS,
+    build_event_vocab,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def get_group_program_map(
+    instrument_vocabulary: str,
+    misc_programs: str,
+    is_mt3: bool = False,
+    include_drums: bool = False,
+) -> dict[int, list[int]]:
+    if instrument_vocabulary == "ONLY_PIANO":
+        ret = {0: list(range(128))}
+    elif instrument_vocabulary == "FULL":
+        ret = {i: [i] for i in range(128)}
+    elif instrument_vocabulary == "MT3_MIDI_PLUS":
+        ret = {
+            0: list(range(8)),
+            1: list(range(8, 16)),
+            2: list(range(16, 24)),
+            3: list(range(24, 32)),
+            4: list(range(32, 40)),
+            5: list(range(40, 56)),
+            6: list(range(56, 64)),
+            7: list(range(64, 72)),
+            8: list(range(72, 80)),
+            9: list(range(80, 88)),
+            10: list(range(88, 96)),
+            11: list(range(100, 102)),
+        }
+    elif instrument_vocabulary == "MT3_FULL_PLUS":
+        ret = {
+            0: [0, 1, 3, 6, 7],
+            1: [2, 4, 5],
+            2: list(range(8, 16)),
+            3: list(range(16, 24)),
+            4: [24, 25],
+            5: [26, 27, 28],
+            6: [29, 30, 31],
+            7: [32, 35],
+            8: [33, 34, 36, 37, 38, 39],
+            9: [40],
+            10: [41],
+            11: [42],
+            12: [43],
+            13: [46],
+            14: [47],
+            15: [48, 49, 44, 45],
+            16: [50, 51],
+            17: [52, 53, 54],
+            18: [55],
+            19: [56, 59],
+            20: [57],
+            21: [58],
+            22: [60],
+            23: [61, 62, 63],
+            24: [64, 65],
+            25: [66],
+            26: [67],
+            27: [68],
+            28: [69],
+            29: [70],
+            30: [71],
+            31: list(range(72, 80)),
+            32: list(range(80, 88)),
+            33: list(range(88, 96)),
+            34: [100],
+            35: [101],
+        }
+    elif instrument_vocabulary == "OURS_INSTRUMENT_GROUPS":
+        ret = {
+            0: list(range(8)),
+            1: list(range(24, 32)),
+            2: list(range(32, 40)),
+            3: list(range(40, 56)),
+            4: list(range(56, 64)),
+            5: list(range(16, 24)) + list(range(64, 80)),
+            6: list(range(80, 96)),
+            7: list(range(8, 16)) + list(range(112, 119)),
+        }
+    else:
+        assert False, instrument_vocabulary
+
+    if instrument_vocabulary == "MT3_FULL_PLUS" and not is_mt3:
+        not_assigned = set(range(130)) - set([v for vs in ret.values() for v in vs])
+    else:
+        not_assigned = set(range(128)) - set([v for vs in ret.values() for v in vs])
+    if include_drums:
+        not_assigned = not_assigned.union({DRUM_PROGRAM})
+    if misc_programs == "ONE_GROUP":
+        ret[len(ret)] = list(not_assigned)
+    elif misc_programs == "SINGLETON_GROUPS":
+        for p in not_assigned:
+            ret[len(ret)] = [p]
+    else:
+        assert misc_programs == "OMIT", misc_programs
+    return ret
+
+
+# Human-readable names for the MT3_FULL_PLUS instrument groups (see
+# get_group_program_map). Used by the CLI's --instruments option and the
+# web app's /instruments endpoint. The group IDs index the model's learned
+# program groups and must not change; only the user-facing names do.
+# Notes the model still decodes into an omitted group surface as "program_<n>".
+MT3_FULL_PLUS_GROUP_NAMES: dict[str, int] = {
+    "acoustic_piano": 0,
+    "electric_piano": 1,
+    "chromatic_percussion": 2,
+    "organ": 3,
+    "acoustic_guitar": 4,
+    "clean_electric_guitar": 5,
+    "distorted_electric_guitar": 6,
+    "acoustic_bass": 7,
+    "electric_bass": 8,
+    "violin": 9,
+    "viola": 10,
+    "cello": 11,
+    "contrabass": 12,
+    "orchestral_harp": 13,
+    "timpani": 14,
+    "string_ensemble": 15,
+    "synth_strings": 16,
+    "voice": 17,
+    "orchestra_hit": 18,
+    "trumpet": 19,
+    "trombone": 20,
+    "tuba": 21,
+    "french_horn": 22,
+    "brass_section": 23,
+    "soprano_and_alto_sax": 24,
+    "tenor_sax": 25,
+    "baritone_sax": 26,
+    "oboe": 27,
+    "english_horn": 28,
+    "bassoon": 29,
+    "clarinet": 30,
+    "flutes": 31,
+    "synth_lead": 32,
+    "synth_pad": 33,
+    "drums": 36,
+}
+
+
+def instrument_group_from_names(names: Iterable[str]) -> str:
+    """Map exact instrument group names to the model's conditioning string.
+
+    The strict counterpart of :func:`resolve_instrument_names`: every name
+    must appear verbatim in ``MT3_FULL_PLUS_GROUP_NAMES``. Raises ValueError
+    listing the unknown names otherwise.
+    """
+    names = list(names)
+    unknown = [n for n in names if n not in MT3_FULL_PLUS_GROUP_NAMES]
+    if unknown:
+        raise ValueError(
+            f"unknown instrument name(s): {', '.join(map(repr, unknown))}; "
+            f"valid names: {', '.join(MT3_FULL_PLUS_GROUP_NAMES)}"
+        )
+    return " ".join(str(MT3_FULL_PLUS_GROUP_NAMES[n]) for n in names)
+
+
+def resolve_instrument_names(tokens: Iterable[str]) -> list[str]:
+    """Resolve loosely-typed instrument tokens to canonical group names.
+
+    Matching is case-insensitive; a token that is not an exact name may be
+    any substring that matches exactly one group name (``"timp"`` →
+    ``"timpani"``). Raises ValueError when a token is ambiguous (listing the
+    candidates) or matches nothing (suggesting close spellings).
+    """
+    resolved = []
+    for token in tokens:
+        t = token.strip().lower()
+        if t in MT3_FULL_PLUS_GROUP_NAMES:
+            resolved.append(t)
+            continue
+        hits = [n for n in MT3_FULL_PLUS_GROUP_NAMES if t in n]
+        if len(hits) == 1:
+            resolved.append(hits[0])
+        elif hits:
+            raise ValueError(
+                f"ambiguous instrument name {token!r}: "
+                f"matches {', '.join(hits)}"
+            )
+        else:
+            # Compare against each name AND its underscore-separated words,
+            # so a typo like "pinao" still surfaces "acoustic_piano".
+            def closeness(name: str) -> float:
+                return max(
+                    difflib.SequenceMatcher(None, t, part).ratio()
+                    for part in (name, *name.split("_"))
+                )
+
+            ranked = sorted(MT3_FULL_PLUS_GROUP_NAMES, key=closeness, reverse=True)
+            suggestions = [n for n in ranked[:3] if closeness(n) >= 0.6]
+            hint = (
+                f" — did you mean {', '.join(suggestions)}?"
+                if suggestions
+                else ""
+            )
+            raise ValueError(f"unknown instrument name {token!r}{hint}")
+    return resolved
+
+
+class MT3Tokenizer:
+    def __init__(
+        self,
+        instrument_vocabulary: str = "FULL",
+        max_shift_steps: int = 1001,
+        frame_rate: int = 100,
+    ):
+        self.group_program_map = get_group_program_map(
+            instrument_vocabulary, misc_programs="SINGLETON_GROUPS", is_mt3=True
+        )
+        self.frame_rate = frame_rate
+        self._vocab = build_event_vocab(max_shift_steps)
+        self.num_tokens = len(self._vocab)
+        self.eos_id = SPECIAL_TOKENS.index("EOS")
+
+        logger.info(f"MT3Tokenizer: {self.num_tokens} tokens")