Add model card, licenses, configs, and ONNX Runtime helper

LICENSE

@@ -0,0 +1,10 @@
+Creative Commons Attribution-NonCommercial 4.0 International
+
+The checkpoint weights in this repository are licensed by their original
+rightsholders under CC BY-NC 4.0. The complete legal code is available at:
+
+https://creativecommons.org/licenses/by-nc/4.0/legalcode.en
+
+This licence applies to the mirrored checkpoint weights and upstream model
+materials. It does not automatically apply to MIDI-ME source code or
+LUSKII-VSTS branding.

NOTICE.md

@@ -0,0 +1,31 @@
+# MuScriptor ONNX derivative attribution notice
+
+This repository contains derivative ONNX Runtime exports of MuScriptor Medium.
+
+- Original model: https://huggingface.co/MuScriptor/muscriptor-medium
+- Original project: https://github.com/muscriptor/muscriptor
+- Original developers: Mirelo × Kyutai
+- Original authors: Alexandre Rouard, Michael Krause, Axel Roebel,
+  Carl-Johann Simon-Gabriel, and Alexandre Défossez
+- ONNX conversion and quantization repository: happyme531/muscriptor-medium-onnx
+
+Changes from the original checkpoint:
+
+1. The model was exported to two ONNX graphs: a conditioner and a decoder with
+   explicit autoregressive KV-cache inputs and outputs.
+2. An FP16 ONNX variant was produced.
+3. A W4A16 variant was produced with 96 Transformer-backbone MatMul weights
+   quantized to UINT4 using HQQ and ONNX Runtime MatMulNBits (block size 128).
+4. Embeddings, the conditioner, and the LM head remain FP16 in the W4A16
+   variant. Log-mel audio preprocessing remains outside the ONNX graphs.
+
+The original and derivative model weights are licensed under CC BY-NC 4.0.
+Users must also comply with the original model card's supplemental conditions,
+including holding all necessary rights to input audio and avoiding illegal or
+unauthorized use. See `UPSTREAM_MODEL_CARD.md` for the complete upstream terms,
+intended uses, limitations, citation, and disclaimers.
+
+This derivative repository is not affiliated with or endorsed by Mirelo,
+Kyutai, or the original authors. The models and generated content are provided
+as-is without warranty.
+

README.md

@@ -0,0 +1,117 @@
+---
+license: cc-by-nc-4.0
+library_name: onnxruntime
+base_model: MuScriptor/muscriptor-medium
+tags:
+  - muscriptor
+  - music-transcription
+  - automatic-music-transcription
+  - audio-to-midi
+  - onnx
+  - onnxruntime
+  - int4
+  - hqq
+---
+
+# MuScriptor Medium — ONNX Runtime FP16 and W4A16
+
+This is an **unofficial derivative ONNX conversion** of
+[`MuScriptor/muscriptor-medium`](https://huggingface.co/MuScriptor/muscriptor-medium),
+developed by Mirelo × Kyutai. It contains an FP16 export and a smaller W4A16
+export for ONNX Runtime. This repository is not affiliated with or endorsed by
+the MuScriptor authors.
+
+The conversion uses upstream revision
+`f32236969308476e01fd3aae67357de5feb05a2d`. The source checkpoint used for the
+export has SHA-256
+`ac80adbdf85d87231735fd948af7013441c0afced316c4e9067fd5d8a7fb97ec`.
+
+## Variants
+
+| Directory | Transformer backbone | Embeddings | LM head | Activations | Weight files |
+|---|---:|---:|---:|---:|---:|
+| `onnx/fp16` | FP16 | FP16 | FP16 | FP16 | 613,150,680 bytes |
+| `onnx/w4a16` | HQQ UINT4 | FP16 | FP16 | FP16 | 169,649,866 bytes |
+
+The W4A16 graph quantizes only the 96 constant-weight MatMul operations in the
+24-layer Transformer backbone. It uses ONNX Runtime `MatMulNBits`, HQQ
+quantization, and block size 128. The token embedding, all other embedding
+tables, conditioner, and LM head deliberately remain FP16.
+
+`decoder.onnx.data` is external tensor data and **must remain next to**
+`onnx/w4a16/decoder.onnx`.
+
+## ONNX interface
+
+The conditioner consumes a 5-second, 16 kHz mono log-mel spectrogram plus
+instrument and dataset IDs. The decoder exposes an autoregressive interface
+with explicit `past_key`/`past_value` inputs and `present_key`/`present_value`
+outputs. Log-mel extraction is intentionally outside the ONNX graph.
+
+## Quick smoke test
+
+Download this public repository:
+
+```bash
+hf download happyme531/muscriptor-medium-onnx --local-dir muscriptor-medium-onnx
+cd muscriptor-medium-onnx
+uv venv
+source .venv/bin/activate
+uv pip install -r runtime/requirements.txt
+python runtime/test_onnxruntime.py --model-dir onnx/w4a16 --provider cpu
+```
+
+Use `--audio path/to/audio.wav` to replace the generated 440 Hz test tone. For
+CUDA, install a matching `onnxruntime-gpu` package instead of `onnxruntime`,
+and pass `--provider cuda`. PyTorch is optional; if installed, the helper
+imports it first so its CUDA/cuDNN libraries can be preloaded.
+
+The included script validates conditioner/decoder execution and greedy token
+generation. It is a low-level ONNX Runtime smoke test, not a complete
+audio-to-MIDI application.
+
+## Validation
+
+Both models pass `onnx.checker.check_model`. The FP16 and W4A16 exports were
+executed with ONNX Runtime on CUDA and x86-64 CPU; on a real audio sample their
+first eight greedy tokens matched. ARM64 execution was not validated and
+depends on whether the target ONNX Runtime build/provider implements the
+`MatMulNBits` graph used by the W4A16 decoder.
+
+| File | SHA-256 |
+|---|---|
+| `onnx/fp16/conditioner.onnx` | `cc94f58ba8023f339b4b8eba06a3f3e7623079c8234fa1be075612ee96efcf81` |
+| `onnx/fp16/decoder.onnx` | `6832149920078feb4562c1c211a0dcac52a4ea10d7b715459e2c6fd0aee83b25` |
+| `onnx/w4a16/conditioner.onnx` | `cc94f58ba8023f339b4b8eba06a3f3e7623079c8234fa1be075612ee96efcf81` |
+| `onnx/w4a16/decoder.onnx` | `2d3b8a35c353b581982a1567eea1629ce70d4b899de1d020f4fa9cb85757943e` |
+| `onnx/w4a16/decoder.onnx.data` | `8d92ff62bed61814593222a9fd8fcb731f0aa83cbd56591c56cc845b6e16b61d` |
+
+## License, attribution, and conditions
+
+The model weights and these derivative ONNX weights are distributed under
+[Creative Commons Attribution-NonCommercial 4.0 International](LICENSE).
+Commercial use is not permitted under that license. You must also comply with
+the upstream supplemental conditions summarized below, including
+having all necessary rights to any audio you transcribe and complying with all
+applicable laws.
+
+> **Summary of upstream supplemental conditions of use:** MuScriptor and any generated
+> content by MuScriptor are provided as is without any warranty of any kind,
+> including but not limited to any warranty of non-infringement. Use of
+> MuScriptor and its output must comply with all applicable laws and must not
+> result in, involve, or facilitate any illegal or unauthorized activity.
+> Prohibited uses include, without limitation, inputting music files and
+> transcribing them to MIDI/music sheet without having all the necessary
+> rights, including intellectual property rights, under applicable laws.
+> Accordingly, users undertake and warrant that they have all necessary rights,
+> including intellectual property rights, in connection with their use of
+> MuScriptor and its output. Mirelo and Kyutai disclaim liability for
+> non-compliant use; users shall indemnify, defend, and hold them harmless from
+> claims, damages, losses, liabilities, and expenses arising from failure to
+> comply with CC BY-NC 4.0 and/or these specific conditions. The exact original
+> wording is preserved in `UPSTREAM_MODEL_CARD.md`.
+
+See [NOTICE.md](NOTICE.md) for attribution and
+[UPSTREAM_MODEL_CARD.md](UPSTREAM_MODEL_CARD.md) for the original model card,
+intended uses, limitations, supplemental conditions, author list, and citation.
+The small runtime helpers are covered by [their MIT license](runtime/LICENSE).

UPSTREAM_MODEL_CARD.md

@@ -0,0 +1,170 @@
+---
+license: cc-by-nc-4.0
+library_name: muscriptor
+extra_gated_prompt: "MuScriptor is the result of a research collaboration between Mirelo and Kyutai whose purpose is to transcribe audio to MIDI/music sheet. It is provided primarily for research purposes under the CC BY-NC 4.0 licence supplemented by the below specific conditions of use.\nSpecific conditions of use: MuScriptor and any generated content by MuScriptor are provided as is without any warranty of any kind, including but not limited to any warranty of non-infringement. Use of MuScriptor and its output must comply with all applicable laws and must not result in, involve, or facilitate any illegal or unauthorized activity. Prohibited uses include, without limitation, inputting music files and transcribing them to MIDI/music sheet without having all the necessary rights, including intellectual property rights, under applicable laws. Accordingly, users of MuScriptor undertake and warrant to have all the necessary rights, including intellectual property rights, in connection with their use of MuScriptor and its output. We disclaim all liability for any non-compliant use and users of MuScriptor shall indemnify, defend, and hold harmless Mirelo and Kyutai from and against any and all claims, damages, losses, liabilities, and expenses (including reasonable attorneys' fees) incurred by Mirelo and/or Kyutai arising out of or resulting from their failure to comply with the terms of the CC BY-NC 4.0 licence and/or these specific conditions of use."
+extra_gated_fields:
+  Company or university if applicable: text
+  I am a:
+    type: select
+    options: 
+      - Musician
+      - AI Researcher
+      - Other
+tags:
+- music
+- music-transcription
+- automatic-music-transcription
+- amt
+- audio-to-midi
+- midi
+- music-information-retrieval
+- transformer
+- pytorch
+---
+
+# MuScriptor — medium (≈300M)
+
+**MuScriptor** is an open-weight model for **general-purpose, multi-instrument automatic music transcription (AMT)**: it converts a music recording (any genre, multiple simultaneous instruments) into a stream of notes played. This repository hosts the **medium** variant (≈300M parameters), the default checkpoint downloaded by the `muscriptor` library.
+
+`muscriptor-medium` balances quality and footprint. For the best transcription quality use [`muscriptor-large`](https://huggingface.co/MuScriptor/muscriptor-large) (≈1.3B); for the smallest/fastest option use [`muscriptor-small`](https://huggingface.co/MuScriptor/muscriptor-small) (≈100M).
+
+- Developed by [Mirelo](https://www.mirelo.ai/) x [kyutai](https://kyutai.org/)
+- 📄 Paper: *MuScriptor: An Open Model for Multi-Instrument Music Transcription* — Rouard, Krause, Roebel, Simon-Gabriel, Défossez (2026). _<!-- TODO: add arXiv link once public; it will auto-cross-link on the Hub -->_
+- 💻 Code: <https://github.com/muscriptor/muscriptor>
+- 🔊 Audio samples: <https://muscriptor.github.io>
+
+## Table of contents
+
+- [Quickstart](#quickstart)
+- [Model description](#model-description)
+- [Model variants](#model-variants)
+- [Intended uses & limitations](#intended-uses--limitations)
+- [Instrument conditioning](#instrument-conditioning)
+- [Training](#training)
+- [Evaluation](#evaluation)
+- [Citation](#citation)
+- [License](#license)
+
+## Quickstart
+
+Install the `muscriptor` package (it uses `huggingface_hub` to fetch weights automatically):
+
+```bash
+pip install git+https://github.com/muscriptor/muscriptor.git
+# TODO (PyPI release forthcoming: pip install muscriptor)
+```
+
+### Python
+
+```python
+from pathlib import Path
+from muscriptor import TranscriptionModel
+
+# "medium" resolves to hf://MuScriptor/muscriptor-medium and downloads on first use.
+model = TranscriptionModel.load_model("medium")
+
+# Get a MIDI file directly:
+Path("out.mid").write_bytes(model.transcribe_to_midi("audio.wav"))
+
+# Or stream note events as they are transcribed:
+for event in model.transcribe("audio.wav"):
+    print(event)  # NoteStartEvent / NoteEndEvent / ProgressEvent
+```
+
+`load_model` accepts a size keyword (`"small"`/`"medium"`/`"large"`), a local `.safetensors` path, or an `hf://` / `https://` URL. Weights loaded by size keyword (or any `hf://` URL) are cached in the standard Hugging Face cache (`~/.cache/huggingface/hub`, configurable via `HF_HOME`); weights fetched from a plain `http(s)://` URL are cached under `~/.cache/muscriptor/`. Input audio can be WAV or any format `libsndfile` reads (mp3, flac, ogg, m4a, …); it is resampled to 16 kHz mono internally.
+
+### CLI
+
+```bash
+muscriptor transcribe --model medium audio.wav -o out.mid
+```
+
+## Model description
+
+MuScriptor performs transcription by **autoregressively predicting a MIDI-like token sequence** given the mel-spectrogram of a short audio segment, following the sequence-to-sequence AMT paradigm (cf. MT3). It deliberately avoids complex architectural tweaks in favor of a simple, decoder-only Transformer.
+
+- **Architecture:** decoder-only Transformer (this variant: `dim=1024`, `num_heads=16`, `num_layers=24`).
+- **Input:** raw waveform (16 kHz, mono) of a 5-second segment → mel-spectrogram (STFT `n_fft=2048`, hop 160 → 100 Hz frame rate, 512 mel bins). The spectrogram is projected to the model dimension and used as a prefix condition.
+- **Output tokenization:** MT3-like note events; the 128 MIDI programs are mapped to **36 instrument subgroups** using the `MT3_FULL_PLUS` taxonomy. Decoding is greedy (argmax) by default, with optional classifier-free guidance (CFG).
+- **Inference:** audio is processed in 5-second chunks; note events are emitted in temporal order. Optional **instrument conditioning** stabilizes predictions across chunk boundaries and lets you restrict/customize the transcription (see below).
+
+**Note on the representation:** the tokenizer recovers onset/offset timing, pitch, and instrument, but **not velocity**. It also cannot represent two notes of the same pitch and instrument sounding at the same time. Drums are onset-only.
+
+## Model variants
+
+| Repo | Params | `dim` | heads | layers | Notes |
+|---|---|---|---|---|---|
+| [`muscriptor-small`](https://huggingface.co/MuScriptor/muscriptor-small) | ≈100M | 768 | 12 | 14 | smallest / fastest |
+| [`muscriptor-medium`](https://huggingface.co/MuScriptor/muscriptor-medium) | ≈300M | 1024 | 16 | 24 | **this model** · good trade-off |
+| [`muscriptor-large`](https://huggingface.co/MuScriptor/muscriptor-large) | ≈1.3B | 1536 | 24 | 48 | best quality |
+
+All variants share the same input pipeline, tokenizer, and training recipe; they differ only in latent dimension, attention heads, and depth.
+
+## Intended uses & limitations
+
+**Intended uses**
+- General-purpose transcription of real, multi-instrument music across genres (classical → heavy metal) into MIDI.
+- A building block for music information retrieval (chord/key recognition), musicological analysis, generative-modeling data pipelines, and tools for musicians.
+
+**Out of scope / use with care**
+- Not a substitute for a hand-annotated score; expect errors, especially on dense mixes, unusual timbres, and heavily processed audio.
+- Velocity/dynamics are **not** produced (see note above).
+- Onset/offset precision is lower for some styles (e.g. choral music), and exact offsets are inherently harder than onsets.
+
+**Limitations & biases**
+- Training data skews toward pop and Western classical music, and the instrument distribution is long-tailed (piano/guitar/bass/drums are most frequent). Rare instruments and underrepresented genres may be transcribed less reliably.
+- The fixed `MT3_FULL_PLUS` 36-group instrument taxonomy limits instrument granularity.
+- Simultaneous same-pitch/same-instrument notes cannot be represented by the tokenizer.
+
+## Instrument conditioning
+
+The model can be told which instrument groups are present in the track. Supplying the correct set improves quantitative scores and produces more coherent instrument assignments across segments.
+
+```python
+from muscriptor.tokenizer.mt3 import MT3_FULL_PLUS_GROUP_NAMES
+
+# `instrument_group` is a space-separated string of MT3_FULL_PLUS group IDs.
+# Convert readable group names to IDs:
+names = ["acoustic_piano", "acoustic_guitar", "acoustic_bass"]
+instrument_group = " ".join(str(MT3_FULL_PLUS_GROUP_NAMES[n]) for n in names)  # -> "0 4 7"
+
+# Only expect piano, acoustic guitar and bass in this track:
+model.transcribe_to_midi("audio.wav", instrument_group=instrument_group)
+```
+
+```bash
+muscriptor transcribe --model medium --instruments "acoustic_piano,acoustic_guitar,acoustic_bass" audio.wav -o out.mid
+muscriptor list-instruments   # show all available group names
+```
+
+## Evaluation
+
+Metrics are instrument-agnostic F1 scores computed with [`mir_eval`](https://github.com/craffel/mir_eval) on `D_Test`, the authors' held-out test set of 372 multi-instrument tracks.
+
+**Model-size comparison** (F1 ↑; from the paper's scaling study, models trained on `D_Real` only, CFG = 2):
+
+| Variant | Params | Onset | Frame | Offset | Drums | Multi |
+|---|---|---|---|---|---|---|
+| `muscriptor-small` | 100M | 51.2 | 67.2 | 38.7 | 41.5 | 38.2 |
+| **`muscriptor-medium`** | **300M** | **52.4** | **68.0** | **40.3** | **42.0** | **39.7** |
+| `muscriptor-large` | 1.3B | 53.2 | 68.7 | 41.0 | 42.5 | 40.5 |
+
+These numbers come from the model-size ablation, which trains on real audio **only**. The **released checkpoints additionally use synthetic pre-training and RL post-training**, which improve real-world quality substantially beyond these figures. See [`muscriptor-large`](https://huggingface.co/MuScriptor/muscriptor-large) and the paper for per-dataset results.
+
+## Citation
+
+```bibtex
+@inproceedings{muscriptor2026,
+  title     = {MuScriptor: An Open Model for Multi-Instrument Music Transcription},
+  author    = {Rouard, Simon and Krause, Michael and Roebel, Axel and
+               Simon-Gabriel, Carl-Johann and D{\'e}fossez, Alexandre},
+  year      = {2026},
+  note      = {Kyutai, Mirelo AI, IRCAM}
+}
+```
+
+<!-- TODO: replace with the final published citation (venue / arXiv id) once available. -->
+
+## License
+
+Code released under the [MIT License](https://github.com/muscriptor/muscriptor/blob/main/LICENSE). Weights released under CC-BY-NC.

onnx/fp16/config.json

@@ -0,0 +1,12 @@
+{
+  "source_model": "MuScriptor/muscriptor-medium@f32236969308476e01fd3aae67357de5feb05a2d",
+  "dtype": "float16",
+  "embedding_storage": "float16",
+  "num_layers": 24,
+  "num_heads": 16,
+  "head_dim": 64,
+  "card": 1395,
+  "initial_token_id": 1395,
+  "eos_token_id": 1,
+  "first_reserved_token_id": 1393
+}

onnx/w4a16/config.json

@@ -0,0 +1,19 @@
+{
+  "source_model": "MuScriptor/muscriptor-medium@f32236969308476e01fd3aae67357de5feb05a2d",
+  "dtype": "float16",
+  "embedding_storage": "float16",
+  "num_layers": 24,
+  "num_heads": 16,
+  "head_dim": 64,
+  "card": 1395,
+  "initial_token_id": 1395,
+  "eos_token_id": 1,
+  "first_reserved_token_id": 1393,
+  "backbone_weight_storage": "uint4-HQQ",
+  "lm_head_weight_storage": "float16",
+  "activation_dtype": "float16",
+  "quantization": "ORT MatMulNBits HQQ",
+  "block_size": 128,
+  "quantized_backbone_matmuls": 96,
+  "quantized_lm_head_matmuls": 0
+}

runtime/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Kyutai x Mirelo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

runtime/muscriptor_onnx/__init__.py

@@ -0,0 +1,6 @@
+"""ONNX export helpers for MuScriptor."""
+
+from .audio import log_mel_spectrogram, load_audio_16k
+
+__all__ = ["load_audio_16k", "log_mel_spectrogram"]
+

runtime/muscriptor_onnx/audio.py

@@ -0,0 +1,74 @@
+"""NumPy implementation of MuScriptor's 16 kHz log-mel frontend."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+from scipy.signal import resample_poly
+
+
+SAMPLE_RATE = 16_000
+N_FFT = 2_048
+HOP_LENGTH = 160
+N_MELS = 512
+SEGMENT_SAMPLES = 5 * SAMPLE_RATE
+
+
+def load_audio_16k(path: str | Path) -> np.ndarray:
+    """Load an audio file as mono float32 at 16 kHz."""
+    audio, sample_rate = sf.read(path, dtype="float32", always_2d=True)
+    audio = audio.mean(axis=1)
+    if sample_rate != SAMPLE_RATE:
+        divisor = np.gcd(sample_rate, SAMPLE_RATE)
+        audio = resample_poly(
+            audio, SAMPLE_RATE // divisor, sample_rate // divisor
+        ).astype(np.float32)
+    return audio.astype(np.float32, copy=False)
+
+
+def first_five_second_chunk(audio: np.ndarray) -> np.ndarray:
+    """Crop/pad audio to the five-second chunk used by MuScriptor."""
+    audio = np.asarray(audio, dtype=np.float32).reshape(-1)
+    if audio.size >= SEGMENT_SAMPLES:
+        return audio[:SEGMENT_SAMPLES]
+    return np.pad(audio, (0, SEGMENT_SAMPLES - audio.size))
+
+
+def _hz_to_mel_htk(freq: np.ndarray) -> np.ndarray:
+    return 2595.0 * np.log10(1.0 + freq / 700.0)
+
+
+def _mel_to_hz_htk(mel: np.ndarray) -> np.ndarray:
+    return 700.0 * (10.0 ** (mel / 2595.0) - 1.0)
+
+
+def mel_filterbank() -> np.ndarray:
+    """Match the pure-torch HTK filterbank bundled with MuScriptor."""
+    all_freqs = np.linspace(0, SAMPLE_RATE // 2, N_FFT // 2 + 1, dtype=np.float32)
+    mel_min = _hz_to_mel_htk(np.asarray(0.0, dtype=np.float32))
+    mel_max = _hz_to_mel_htk(np.asarray(SAMPLE_RATE / 2, dtype=np.float32))
+    mel_points = np.linspace(mel_min, mel_max, N_MELS + 2, dtype=np.float32)
+    freq_points = _mel_to_hz_htk(mel_points)
+    freq_diff = freq_points[1:] - freq_points[:-1]
+    slopes = freq_points[None, :] - all_freqs[:, None]
+    down = -slopes[:, :-2] / freq_diff[:-1]
+    up = slopes[:, 2:] / freq_diff[1:]
+    return np.maximum(0.0, np.minimum(down, up)).astype(np.float32)
+
+
+def log_mel_spectrogram(audio: np.ndarray) -> np.ndarray:
+    """Return `[1, 501, 512]` log-mel features for a five-second chunk.
+
+    The reflection padding, periodic Hann window, FFT, HTK mel bank, and log
+    epsilon mirror ``muscriptor.modules.mel_spectrogram``.
+    """
+    audio = first_five_second_chunk(audio)
+    padded = np.pad(audio, (N_FFT // 2, N_FFT // 2), mode="reflect")
+    frames = np.lib.stride_tricks.sliding_window_view(padded, N_FFT)[::HOP_LENGTH]
+    window = np.hanning(N_FFT + 1)[:-1].astype(np.float32)
+    spectrum = np.abs(np.fft.rfft(frames * window, n=N_FFT, axis=-1)).astype(np.float32)
+    mel = spectrum @ mel_filterbank()
+    return np.log(mel + np.float32(1e-6))[None, ...].astype(np.float32)
+

runtime/requirements.txt

@@ -0,0 +1,4 @@
+numpy>=1.26
+scipy>=1.11
+soundfile>=0.12
+onnxruntime>=1.26.0

runtime/test_onnxruntime.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+"""Run cached autoregressive MuScriptor inference with ONNX Runtime."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+# Importing PyTorch first preloads the CUDA/cuDNN shared libraries shipped in
+# the uv environment when available; CPU-only users do not need PyTorch.
+try:
+    import torch  # noqa: F401
+except ModuleNotFoundError:
+    torch = None  # type: ignore[assignment]
+import onnxruntime as ort
+
+from muscriptor_onnx.audio import SAMPLE_RATE, load_audio_16k, log_mel_spectrogram
+
+
+def providers(requested: str, device_id: int) -> list:
+    available = ort.get_available_providers()
+    if requested == "cuda" or (requested == "auto" and "CUDAExecutionProvider" in available):
+        if "CUDAExecutionProvider" not in available:
+            raise RuntimeError(f"CUDA EP unavailable; installed providers: {available}")
+        return [("CUDAExecutionProvider", {"device_id": device_id}), "CPUExecutionProvider"]
+    return ["CPUExecutionProvider"]
+
+
+def session(path: Path, selected_providers: list) -> ort.InferenceSession:
+    options = ort.SessionOptions()
+    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    result = ort.InferenceSession(path, sess_options=options, providers=selected_providers)
+    requested_cuda = bool(selected_providers) and (
+        selected_providers[0] == "CUDAExecutionProvider"
+        or (
+            isinstance(selected_providers[0], tuple)
+            and selected_providers[0][0] == "CUDAExecutionProvider"
+        )
+    )
+    if requested_cuda and "CUDAExecutionProvider" not in result.get_providers():
+        raise RuntimeError("CUDA EP was requested but session creation fell back to CPU")
+    return result
+
+
+def causal_mask(query_length: int, past_length: int) -> np.ndarray:
+    query_positions = past_length + np.arange(query_length)[:, None]
+    key_positions = np.arange(past_length + query_length)[None, :]
+    return np.where(key_positions <= query_positions, 0.0, -65504.0).astype(np.float16)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-dir", type=Path, required=True)
+    parser.add_argument("--audio", type=Path)
+    parser.add_argument("--provider", choices=("auto", "cuda", "cpu"), default="auto")
+    parser.add_argument("--device-id", type=int, default=0)
+    parser.add_argument("--max-new-tokens", type=int, default=8)
+    parser.add_argument(
+        "--instrument-id",
+        type=int,
+        action="append",
+        help="Optional MT3_FULL_PLUS group ID; repeat for multiple groups",
+    )
+    args = parser.parse_args()
+
+    metadata = json.loads((args.model_dir / "config.json").read_text())
+    selected = providers(args.provider, args.device_id)
+    conditioner = session(args.model_dir / "conditioner.onnx", selected)
+    decoder = session(args.model_dir / "decoder.onnx", selected)
+
+    if args.audio:
+        audio = load_audio_16k(args.audio)
+        source = str(args.audio)
+    else:
+        # Deterministic smoke-test input; it only tests execution, not quality.
+        time_axis = np.arange(5 * SAMPLE_RATE, dtype=np.float32) / SAMPLE_RATE
+        audio = (0.1 * np.sin(2 * np.pi * 440.0 * time_axis)).astype(np.float32)
+        source = "generated 440 Hz sine"
+    mel = log_mel_spectrogram(audio).astype(np.float16)
+    instrument_ids = np.asarray(
+        [args.instrument_id if args.instrument_id is not None else [-1]], dtype=np.int64
+    )
+    dataset_ids = np.asarray([[-1]], dtype=np.int64)
+
+    start = time.perf_counter()
+    condition = conditioner.run(
+        ["condition_embeddings"],
+        {
+            "log_mel": mel,
+            "instrument_ids": instrument_ids,
+            "dataset_ids": dataset_ids,
+        },
+    )[0]
+    condition_seconds = time.perf_counter() - start
+
+    layers = metadata["num_layers"]
+    heads = metadata["num_heads"]
+    head_dim = metadata["head_dim"]
+    past_key = np.zeros((layers, 1, heads, 0, head_dim), dtype=np.float16)
+    past_value = np.zeros_like(past_key)
+    input_ids = np.asarray([[metadata["initial_token_id"]]], dtype=np.int64)
+    generated: list[int] = []
+    decode_times: list[float] = []
+
+    for step in range(args.max_new_tokens):
+        prefix = condition if step == 0 else np.empty((1, 0, heads * head_dim), np.float16)
+        query_length = prefix.shape[1] + input_ids.shape[1]
+        mask = causal_mask(query_length, past_key.shape[3])
+        tick = time.perf_counter()
+        logits, past_key, past_value = decoder.run(
+            ("logits", "present_key", "present_value"),
+            {
+                "input_ids": input_ids,
+                "condition_embeddings": prefix,
+                "past_key": past_key,
+                "past_value": past_value,
+                "attention_mask": mask,
+            },
+        )
+        decode_times.append(time.perf_counter() - tick)
+        if not np.isfinite(logits).all():
+            raise RuntimeError("decoder produced non-finite logits")
+        logits[:, metadata["first_reserved_token_id"] :] = -np.inf
+        token = int(logits.argmax(axis=-1)[0])
+        generated.append(token)
+        input_ids = np.asarray([[token]], dtype=np.int64)
+        if token == metadata["eos_token_id"]:
+            break
+
+    active = decoder.get_providers()
+    sizes = {
+        path.name: path.stat().st_size
+        for path in args.model_dir.iterdir()
+        if path.is_file() and (path.suffix == ".onnx" or path.name.endswith(".onnx.data"))
+    }
+    sizes["total"] = sum(sizes.values())
+    print(f"model:             {args.model_dir}")
+    print(f"audio:             {source}")
+    print(f"providers:         {active}")
+    print(f"mel shape:         {mel.shape}")
+    print(f"condition shape:   {condition.shape}")
+    print(f"final KV shape:    {past_key.shape}")
+    print(f"generated tokens:  {generated}")
+    print(f"condition latency: {condition_seconds * 1000:.1f} ms")
+    print(f"decode latency:    {[round(x * 1000, 1) for x in decode_times]} ms")
+    print(f"ONNX file sizes:   {sizes}")
+
+
+if __name__ == "__main__":
+    main()