muscriptor 0.1.0__py3-none-any.whl

muscriptor/__init__.py

@@ -0,0 +1,7 @@
+"""muscriptor — audio-to-MIDI transcription."""
+
+from muscriptor.events import NoteStartEvent, NoteEndEvent
+from muscriptor.transcription_model import TranscriptionModel
+from muscriptor.tokenizer.notes import Note
+
+__all__ = ["TranscriptionModel", "Note", "NoteStartEvent", "NoteEndEvent"]

muscriptor/__main__.py

@@ -0,0 +1,3 @@
+from muscriptor.main import main
+
+main()

muscriptor/events.py

@@ -0,0 +1,208 @@
+"""Public streaming events and per-chunk event builder.
+
+`TranscriptionModel.transcribe` is a generator that yields these dataclasses
+one at a time. Every :class:`NoteStartEvent` is guaranteed to be followed by
+exactly one matching :class:`NoteEndEvent` (same `index`) later in the stream.
+"""
+
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+
+from muscriptor.tokenizer.notes import (
+    MINIMUM_NOTE_DURATION_SEC,
+    Event,
+)
+
+
+_DRUM_INSTRUMENT = "drums"
+
+
+@dataclass
+class NoteStartEvent:
+    pitch: int
+    start_time: float
+    index: int
+    instrument: str
+
+
+@dataclass
+class NoteEndEvent:
+    end_time: float
+    start_event: NoteStartEvent
+
+    @property
+    def start_event_index(self) -> int:
+        return self.start_event.index
+
+
+@dataclass
+class ProgressEvent:
+    """A coarse transcription-progress signal, woven into the event stream.
+
+    Marks that ``completed`` of ``total`` fixed-size audio chunks have been
+    transcribed (``completed == 0`` is emitted once up front so consumers learn
+    ``total`` and get a timing baseline; ``completed == total`` marks the end).
+    These are deliberately coarse anchors — the frontend smooths between them
+    and derives an ETA, since wall-clock time per chunk is only observable
+    there. Advisory only: consumers that build notes/MIDI ignore them.
+    """
+
+    completed: int
+    total: int
+
+
+@dataclass
+class ChunkBoundary:
+    """Marks the start of a new model-output chunk in the token stream.
+
+    ``seek_time`` is the chunk's start time in seconds; ``next_seek_time`` is
+    the following chunk's start (``None`` for the last chunk), used to drop
+    events the model emits past its window.
+    """
+
+    seek_time: float
+    next_seek_time: float | None
+
+
+def decode_model_tokens(
+    stream: Iterator[int | ChunkBoundary | ProgressEvent],
+    vocab: list[Event],
+    instrument_for_program: Callable[[int], str],
+    frame_rate: int = 100,
+) -> Iterator[NoteStartEvent | NoteEndEvent | ProgressEvent]:
+    """Stream model token indices straight into NoteStart/NoteEnd events.
+
+    ``stream`` interleaves :class:`ChunkBoundary` markers with token indices:
+    each boundary starts a new chunk, followed by that chunk's tokens (EOS and
+    anything after it already stripped). All state — open notes, the running
+    note index, the per-chunk decode state — lives in this generator's frame
+    and persists across chunks.
+
+    Tokens are consumed strictly in order: no buffering, no end-of-chunk sort.
+    Each chunk begins with a *tie prologue* — ``(program, pitch)`` pairs for
+    notes sustained from the previous chunk, terminated by a ``tie`` token —
+    after which any prior open note not in that tie set is closed at the chunk
+    boundary. The rest of the chunk drives note onsets/offsets directly.
+    """
+    open_notes: dict[tuple[int, int], NoteStartEvent] = {}
+    next_index = 0
+
+    def mint(pitch: int, start_time: float, instrument: str) -> NoteStartEvent:
+        nonlocal next_index
+        ev = NoteStartEvent(
+            pitch=pitch, start_time=start_time, index=next_index, instrument=instrument
+        )
+        next_index += 1
+        return ev
+
+    # Per-chunk state (reset at every ChunkBoundary).
+    seek_time = 0.0
+    next_seek_time: float | None = None
+    start_tick = 0
+    tick_state = 0
+    program_state: int | None = None
+    velocity_state: int | None = None
+    in_prologue = True
+    skip_rest = False
+    tie_set: set[tuple[int, int]] = set()
+    chunk_started = False
+
+    for item in stream:
+        if isinstance(item, ProgressEvent):
+            # Advisory progress signal — pass straight through, untouched by the
+            # decode state machine.
+            yield item
+            continue
+        if isinstance(item, ChunkBoundary):
+            # If the previous chunk never closed its tie prologue (malformed:
+            # no `tie` token before it ended), treat its tie set as empty so
+            # every still-open note ends at that chunk's boundary.
+            if chunk_started and in_prologue:
+                for key in list(open_notes):
+                    yield NoteEndEvent(
+                        end_time=seek_time, start_event=open_notes.pop(key)
+                    )
+            seek_time = item.seek_time
+            next_seek_time = item.next_seek_time
+            start_tick = round(seek_time * frame_rate)
+            tick_state = start_tick
+            program_state = None
+            velocity_state = None
+            in_prologue = True
+            skip_rest = False
+            tie_set = set()
+            chunk_started = True
+            continue
+
+        event = vocab[item]
+        etype = event.type
+
+        if in_prologue:
+            if etype == "tie":
+                # End of the tie section: close prior notes not sustained here.
+                in_prologue = False
+                velocity_state = None
+                for key in list(open_notes):
+                    if key not in tie_set:
+                        yield NoteEndEvent(
+                            end_time=seek_time, start_event=open_notes.pop(key)
+                        )
+            elif etype == "shift":
+                # No tie token: the chunk is malformed. Close all open notes at
+                # the boundary and drop the rest of the chunk.
+                in_prologue = False
+                skip_rest = True
+                for key in list(open_notes):
+                    yield NoteEndEvent(
+                        end_time=seek_time, start_event=open_notes.pop(key)
+                    )
+            elif etype == "program":
+                program_state = event.value
+            elif etype == "pitch" and program_state is not None:
+                tie_set.add((program_state, event.value))
+            continue
+
+        if skip_rest:
+            continue
+
+        if etype == "shift":
+            if event.value > 0:
+                tick_state = start_tick + event.value
+        elif etype == "program":
+            program_state = event.value
+        elif etype == "velocity":
+            velocity_state = event.value
+        elif etype == "drum":
+            time = tick_state / frame_rate
+            if next_seek_time is None or time < next_seek_time:
+                start = mint(event.value, time, _DRUM_INSTRUMENT)
+                yield start
+                yield NoteEndEvent(
+                    end_time=time + MINIMUM_NOTE_DURATION_SEC, start_event=start
+                )
+        elif etype == "pitch":
+            if program_state is None or velocity_state is None:
+                continue
+            time = tick_state / frame_rate
+            if next_seek_time is not None and time >= next_seek_time:
+                continue
+            key = (program_state, event.value)
+            if key in open_notes:
+                yield NoteEndEvent(end_time=time, start_event=open_notes.pop(key))
+            if velocity_state > 0:
+                start = mint(event.value, time, instrument_for_program(program_state))
+                open_notes[key] = start
+                yield start
+
+    # End of stream: close anything still open. A well-formed final chunk uses
+    # the minimum-duration fallback; a chunk that ended mid-prologue closes at
+    # its boundary (matching the malformed-chunk behavior above).
+    if chunk_started and in_prologue:
+        for key in list(open_notes):
+            yield NoteEndEvent(end_time=seek_time, start_event=open_notes.pop(key))
+    else:
+        for ev in list(open_notes.values()):
+            yield NoteEndEvent(
+                end_time=ev.start_time + MINIMUM_NOTE_DURATION_SEC, start_event=ev
+            )
+        open_notes.clear()

muscriptor/main.py

@@ -0,0 +1,327 @@
+"""CLI for muscriptor: audio → MIDI transcription."""
+
+import dataclasses
+import json
+import sys
+from enum import Enum
+from pathlib import Path
+from typing import Annotated
+
+import typer
+
+from muscriptor.events import NoteEndEvent, NoteStartEvent, ProgressEvent
+from muscriptor.tokenizer.mt3 import (
+    MT3_FULL_PLUS_GROUP_NAMES,
+    resolve_instrument_names,
+)
+from muscriptor.transcription_model import TranscriptionModel
+
+app = typer.Typer(add_completion=False, help="muscriptor — audio-to-MIDI transcription")
+
+
+class OutputFormat(str, Enum):
+    midi = "midi"
+    json = "json"
+    jsonl = "jsonl"
+
+
+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,
+    }
+
+
+@app.command()
+def transcribe(
+    audio_file: Annotated[
+        Path, typer.Argument(help="Input audio file (wav, mp3, flac, …)")
+    ],
+    output: Annotated[
+        Path | None,
+        typer.Option(
+            "--output",
+            "-o",
+            help=(
+                "Output file path. Use '-' to write to stdout (all progress / "
+                "timing info is sent to stderr in that case). "
+                "Default: <audio_file>.<ext> where ext matches --format."
+            ),
+        ),
+    ] = None,
+    format: Annotated[
+        OutputFormat,
+        typer.Option(
+            "--format",
+            "-f",
+            help=(
+                "Output format: midi (default), json (single array of events), "
+                "or jsonl (one event per line, streamed as transcription progresses)"
+            ),
+            case_sensitive=False,
+        ),
+    ] = OutputFormat.midi,
+    notes: Annotated[
+        bool, typer.Option("--notes", help="Print decoded events to stdout")
+    ] = False,
+    sampling: Annotated[
+        bool,
+        typer.Option(
+            "--sampling", help="Use temperature sampling instead of greedy decoding"
+        ),
+    ] = False,
+    temperature: Annotated[
+        float,
+        typer.Option(
+            "--temperature", "-t", help="Sampling temperature (only with --sampling)"
+        ),
+    ] = 1.0,
+    cfg_coef: Annotated[
+        float, typer.Option("--cfg-coef", help="Classifier-free guidance coefficient")
+    ] = 1.0,  # todo: make it dynamic
+    model_path: Annotated[
+        str | None,
+        typer.Option(
+            "--model",
+            "-m",
+            help=(
+                "Model size ('small', 'medium', 'large'; default: medium), "
+                "a local safetensors path, or an hf:// / http(s):// URL"
+            ),
+        ),
+    ] = None,
+    device: Annotated[
+        str,
+        typer.Option(
+            "--device", "-d", help="Device: 'auto', 'cpu', 'cuda', 'cuda:0', …"
+        ),
+    ] = "auto",
+    batch_size: Annotated[
+        int | None,
+        typer.Option(
+            "--batch-size",
+            "-b",
+            help="Batch size for generation (default: 1 on CPU, 4 on GPU)",
+        ),
+    ] = None,
+    strict_eos: Annotated[
+        bool,
+        typer.Option(
+            "--strict-eos",
+            help="Raise an error if a chunk fails to emit EOS within the generation budget (default: downgrade to a warning)",
+        ),
+    ] = False,
+    beam_size: Annotated[
+        int,
+        typer.Option(
+            "--beam-size",
+            help="Beam search width (1 = greedy/sampling, ≥2 enables beam search)",
+        ),
+    ] = 1,
+    auralize: Annotated[
+        Path | None,
+        typer.Option(
+            "--auralize",
+            help=(
+                "Write a stereo auralization (L=original audio, R=MIDI synthesis) to "
+                "this path. Requires fluidsynth on PATH. Extension determines format: "
+                ".wav (default) or .mp3. Only valid with --format midi."
+            ),
+        ),
+    ] = None,
+    soundfont: Annotated[
+        Path | None,
+        typer.Option(
+            "--soundfont",
+            help=(
+                "Path to a .sf2 SoundFont for auralization. "
+                "Defaults to the bundled MuseScore_General.sf2."
+            ),
+        ),
+    ] = None,
+    instruments: Annotated[
+        str | None,
+        typer.Option(
+            "--instruments",
+            help=(
+                "Comma-separated list of expected instrument group names. "
+                "Case-insensitive; unambiguous abbreviations are accepted "
+                "(e.g. 'timp,cello,dist'). Run 'muscriptor list-instruments' "
+                "to see all available names."
+            ),
+        ),
+    ] = None,
+) -> None:
+    """Transcribe an audio file to MIDI."""
+    instrument_names: list[str] | None = None
+    if instruments is not None:
+        tokens = [n for n in instruments.split(",") if n.strip()]
+        try:
+            instrument_names = resolve_instrument_names(tokens)
+        except ValueError as e:
+            typer.echo(
+                f"Error: {e}. "
+                "Run 'muscriptor list-instruments' to see available names.",
+                err=True,
+            )
+            raise typer.Exit(1)
+        typer.echo(f"Instruments: {', '.join(instrument_names)}", err=True)
+
+    if not audio_file.exists():
+        typer.echo(f"Error: file not found: {audio_file}", err=True)
+        raise typer.Exit(1)
+
+    is_stdout = output is not None and str(output) == "-"
+
+    if output is None:
+        suffix = {
+            OutputFormat.midi: ".mid",
+            OutputFormat.json: ".json",
+            OutputFormat.jsonl: ".jsonl",
+        }[format]
+        output = audio_file.with_suffix(suffix)
+
+    _device = None if device == "auto" else device
+
+    # All chatty progress/timing info goes to stderr — stdout is reserved for
+    # the actual output when `-o -` is used.
+    typer.echo("Loading model…", err=True)
+    model = TranscriptionModel.load_model(
+        weights_path=model_path,
+        device=_device,
+    )
+    import torch
+
+    model._model = model._model.to(torch.float32)
+
+    typer.echo(f"Transcribing {audio_file} …", err=True)
+
+    if auralize is not None and format != OutputFormat.midi:
+        typer.echo("Error: --auralize requires --format midi", err=True)
+        raise typer.Exit(1)
+
+    kwargs = dict(
+        audio=audio_file,
+        use_sampling=sampling,
+        temperature=temperature,
+        cfg_coef=cfg_coef,
+        instruments=instrument_names,
+        batch_size=batch_size,
+        no_eos_is_ok=not strict_eos,
+        beam_size=beam_size,
+    )
+
+    if format == OutputFormat.midi:
+        midi_bytes = model.transcribe_to_midi(**kwargs)
+        if is_stdout:
+            sys.stdout.buffer.write(midi_bytes)
+            sys.stdout.buffer.flush()
+        else:
+            output.write_bytes(midi_bytes)
+            typer.echo(f"Saved MIDI to {output}", err=True)
+        if notes:
+            typer.echo(
+                "Re-run with --format json to inspect the event stream.", err=True
+            )
+        if auralize is not None and not is_stdout:
+            from muscriptor.utils.auralization import auralize as do_auralize
+
+            typer.echo(f"Auralizing → {auralize} …", err=True)
+            do_auralize(
+                midi_path=output,
+                original_audio_path=audio_file,
+                output_path=auralize,
+                soundfont_path=soundfont,
+            )
+            typer.echo(f"Saved auralization to {auralize}", err=True)
+    elif format == OutputFormat.jsonl:
+        # Stream one JSON object per line, flushing after each event so the
+        # file (or stdout pipe) can be consumed live.
+        if is_stdout:
+            sink = sys.stdout
+            close_after = False
+        else:
+            sink = output.open("w")
+            close_after = True
+        try:
+            for e in model.transcribe(**kwargs):
+                if isinstance(e, ProgressEvent):
+                    continue
+                sink.write(json.dumps(_event_to_dict(e)) + "\n")
+                sink.flush()
+                if notes:
+                    typer.echo(str(e), err=True)
+        finally:
+            if close_after:
+                sink.close()
+        if not is_stdout:
+            typer.echo(f"Saved JSONL to {output}", err=True)
+    else:  # json
+        events = [
+            e
+            for e in model.transcribe(**kwargs)
+            if not isinstance(e, ProgressEvent)
+        ]
+        payload = json.dumps([_event_to_dict(e) for e in events], indent=2)
+        if is_stdout:
+            sys.stdout.write(payload + "\n")
+            sys.stdout.flush()
+        else:
+            output.write_text(payload)
+            typer.echo(f"Saved JSON to {output}", err=True)
+        if notes:
+            for e in events:
+                typer.echo(str(e), err=True)
+
+
+@app.command()
+def serve(
+    host: Annotated[str, typer.Option("--host", help="Bind address")] = "127.0.0.1",
+    port: Annotated[int, typer.Option("--port", help="Port to listen on")] = 8222,
+    model_path: Annotated[
+        str | None,
+        typer.Option(
+            "--model",
+            "-m",
+            help=(
+                "Model size ('small', 'medium', 'large'; default: medium), "
+                "a local safetensors path, or an hf:// / http(s):// URL"
+            ),
+        ),
+    ] = None,
+    device: Annotated[
+        str,
+        typer.Option(
+            "--device", "-d", help="Device: 'auto', 'cpu', 'cuda', 'cuda:0', …"
+        ),
+    ] = "auto",
+):
+    """Run the HTTP transcription server (POST /transcribe → SSE event stream)."""
+    import uvicorn
+
+    from muscriptor.server import create_app
+
+    _device = None if device == "auto" else device
+    typer.echo("Loading model…")
+    model = TranscriptionModel.load_model(weights_path=model_path, device=_device)
+    web_dir = Path(__file__).resolve().parent.parent / "web" / "dist"
+    fastapi_app = create_app(model, web_dir=web_dir if web_dir.is_dir() else None)
+    uvicorn.run(fastapi_app, host=host, port=port)
+
+
+@app.command()
+def list_instruments():
+    """List the instrument group names accepted by --instruments."""
+    for name in MT3_FULL_PLUS_GROUP_NAMES:
+        typer.echo(name)
+
+
+def main():
+    app()
+
+
+if __name__ == "__main__":
+    main()