v2.0.0 — Transkun transformer-only ONNX export + decode spec (from audio-claudio)

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+freq2mels.f32 filter=lfs diff=lfs merge=lfs -text
+ref3b_audio.f32 filter=lfs diff=lfs merge=lfs -text
+ref3b_features.f32 filter=lfs diff=lfs merge=lfs -text
+ref3c_S.f32 filter=lfs diff=lfs merge=lfs -text

DECODE_SPEC.md

@@ -0,0 +1,68 @@
+# Transkun — self-contained ONNX export (v2 Stage 4)
+
+A transformer-only ONNX export of **Transkun** (Yujia Yan's Neural Semi-CRF piano transcriber, 0.984
+MAESTRO F1), plus the frozen front-end buffers and the decode spec needed to run it **in-process with no
+Python/torch at runtime** behind `audio-claudio`'s `ITranscriber` port.
+
+- **Upstream:** <https://github.com/Yujia-Yan/Skipping-The-Frame-Level> — Yujia Yan, Frank Cwitkowitz,
+  Zhiyao Duan. Package `transkun` 2.0.1, checkpoint `pretrained/2.0.pt`.
+- **License:** MIT (© 2021 Yujia Yan) — see [`LICENSE.transkun`](LICENSE.transkun). audio-claudio is UNLICENSE;
+  MIT is compatible.
+- This is a **transformer-only export + a decode spec**, not a drop-in `.onnx` transcriber: the mel front
+  end and the semi-CRF Viterbi decode are reimplemented in C# (Stages 4b/4c), because `torch.fft.rfft` and
+  the custom semi-CRF backtracking are not ONNX-exportable.
+
+## What the ONNX computes
+
+`transkun.onnx` maps **`featuresBatch` → (`S`, `ctx`)**:
+
+- input `featuresBatch` `[nBatch, T, 229, 6]` — log-mel features (229 mel bins × 6 windows), `T` dynamic.
+- output `S` `[T, T, nBatch*90]` — the semi-CRF pairwise interval scores. `S[e, b, k]` scores a note on
+  track `k` spanning frames `b→e` (diagonal `e==b` = a single-frame note). The 90 tracks are
+  `symbols = [-64, -67, 21..108]`: index 0 = sustain pedal (CC64), 1 = soft pedal (CC67), 2–89 = MIDI
+  21–108. (The `S_skip` "no-event" score is provably 0 and is hardcoded in C#.)
+- output `ctx` `[90, T, 256]` (Stage 4e) — the backbone features, gathered at decoded interval endpoints to
+  drive the attribute heads. `S` is byte-for-byte the same as the S-only 4a export (corr 1.0).
+
+**`transkun-heads.onnx`** (Stage 4e) maps the gathered interval features **`attr` `[N, 768]`**
+(`[ctx_a, ctx_b, ctx_a·ctx_b]`) → **`velLogits` `[N, 128]`** (`velocityPredictor`; velocity = argmax) and
+**`ofRaw` `[N, 4]`** (`refinedOFPredictor`: two sub-frame onset/offset value logits → a ContinuousBernoulli
+mean in `[-0.5, 0.5]` frames, + two presence logits). This adds real velocity + sub-frame timing on top of
+the frame-level decode — validated note-identical to the native CLI (velocity exact, onsets ~1 ms).
+
+Two ops needed care in export (see `export_transkun.py`): the backbone's **5-D `scaled_dot_product_attention`**
+is reshaped to 4-D (a mathematical identity — SDPA batches all but the last two dims) because the ONNX
+exporter only supports 4-D SDPA. `diag_embed` exported cleanly on this stack (torch 2.13 / onnx 1.22 /
+onnxruntime 1.27, opset 17), contrary to an earlier assumption. Validated **`corr = 1.000000`,
+maxRelErr ≈ 5e-6** vs PyTorch on random and dynamic-`T` inputs.
+
+## Files
+
+| File | What |
+|---|---|
+| `transkun.onnx` | the main export `featuresBatch → (S, ctx)` (opset 17, weights inlined, ~53 MB) |
+| `transkun-heads.onnx` | the velocity + onset/offset attribute heads `attr → (velLogits, ofRaw)` (~3.4 MB) |
+| `export_transkun_heads.py` | Stage-4e regeneration (main graph with `ctx` + the heads) |
+| `freq2mels.f32` `[2049, 229]` | mel filterbank (`torchaudio.melscale_fbanks`, 30–8000 Hz) — Stage 4b |
+| `windows.f32` `[6, 4096]` | analysis windows (row 0 Hann, rows 1–5 learned Gaussian) — Stage 4b |
+| `symbols.i32` `[90]` | the track→symbol map `[-64, -67, 21..108]` |
+| `params.json` | fs 44100, windowSize 4096, hopSize 1024, nMels 229, eps 1e-5, segment 16 s / hop 8 s |
+| `ref3b_audio.f32`, `ref3b_features.f32` | Stage-4b TDD fixture: 1.5 s of `two-bar.wav` → its `featuresBatch` `[66,229,6]` |
+| `ref3c_S.f32`, `ref3c_intervals.json` | Stage-4c TDD fixture: a real model `S` `[66,66,90]` → Viterbi intervals |
+| `ref3c_syn_S.f32`, `ref3c_syn_intervals.json` | hand-built multi-track `S` `[6,6,90]` → known intervals |
+| `ref3c_forced_intervals.json` | the same synthetic `S` with a `forcedStartPos` (Stage-4d stitching) |
+| `manifest.json` | shape/dtype/file for every raw `.f32`/`.i32` array (raw little-endian) |
+| `export_transkun.py` | the regeneration script (needs the transkun venv; see below) |
+
+## Regenerating
+
+Not needed for the build (everything above is committed). To reproduce, in a venv with
+`transkun==2.0.1`, `torch`, `onnx`, `onnxruntime`, `numpy`:
+
+```
+python export_transkun.py <output-dir>
+```
+
+It loads the model **from `2.0.conf`** (not class defaults — `baseSize=64`, `nHead=8`), wraps
+`backbone + scorer`, exports, validates against PyTorch, extracts the buffers, and regenerates the ref
+fixtures. Deterministic (seed 0).

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Yujia Yan
+
+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.

README.md

@@ -0,0 +1,62 @@
+---
+license: mit
+library_name: onnxruntime
+tags:
+  - audio
+  - music
+  - piano-transcription
+  - amt
+  - onnx
+  - transkun
+  - semi-crf
+pipeline_tag: audio-to-audio
+---
+
+# Transkun — transformer-only ONNX export + decode spec
+
+A **self-contained ONNX export of [Transkun](https://github.com/Yujia-Yan/Skipping-The-Frame-Level)**
+(Yujia Yan's Neural Semi-CRF piano transcriber, 0.984 MAESTRO note F1) that runs the full model **in-process
+with no Python/PyTorch at runtime**. This is **not** a drop-in `.onnx` transcriber: `torch.fft.rfft` and the
+custom semi-CRF backtracking are not ONNX-exportable, so the mel front end and the Viterbi decode are provided
+as a **documented decode spec** + a **reference decoder**.
+
+> **Attribution.** The model, weights and architecture are the work of **Yujia Yan, Frank Cwitkowitz and
+> Zhiyao Duan** (*"Skipping the Frame-Level: Event-Based Piano Transcription with Neural Semi-CRFs"*, NeurIPS
+> 2021). This is an independent export + decode spec, not affiliated with or endorsed by the authors. Upstream:
+> <https://github.com/Yujia-Yan/Skipping-The-Frame-Level>. License: **MIT** (© 2021 Yujia Yan).
+
+## What's in the package
+
+| File | Role |
+|---|---|
+| `transkun.onnx` (~53 MB, opset 17) | `featuresBatch [1,T,229,6] → (S [T,T,90], ctx [90,T,256])` — the transformer + semi-CRF scorer + backbone features |
+| `transkun-heads.onnx` (~3.4 MB) | `attr [N,768] → (velLogits [N,128], ofRaw [N,4])` — velocity + sub-frame onset/offset heads |
+| `freq2mels.f32 [2049,229]`, `windows.f32 [6,4096]`, `symbols.i32 [90]`, `params.json` | frozen front-end constants |
+| `LICENSE.transkun` | upstream MIT license |
+| `export_transkun.py`, `export_transkun_heads.py` | regeneration scripts (need the `transkun` PyTorch package) |
+
+The **decode spec** (`DECODE_SPEC.md`) documents the mel front end, the `S` layout, the 90-track
+symbol map (`[-64, -67, 21..108]` = sustain/soft pedal + MIDI 21–108), the semi-CRF `viterbiBackward`, the
+16 s/8 s segment stitching, and the attribute heads (velocity = argmax; `ofValue` = ContinuousBernoulli mean).
+
+## Reference decoder + validation
+
+The reference decoder is the C# implementation in **[audio-claudio](https://github.com/TuesdayCrowd/audio-claudio)**
+(mel front end, `SemiCrfViterbi`, `TranskunTranscriber`). It is validated **note-identical to the native
+`transkun` CLI (PyTorch)**: on the test clips it reaches **100% note-level F1 at ±25 ms** with **exact
+velocity** on every note — the export + decode spec reproduce the reference implementation, not merely
+approximate it.
+
+## Pipeline (how to run)
+
+```
+audio (mono, 44.1 kHz)
+  → mel front end (framing 4096/1024, 6 windows, rfft ortho, freq2mels, log-norm) → featuresBatch
+  → transkun.onnx → (S, ctx)
+  → semi-CRF viterbiBackward(S) → per-track note intervals, over 16 s/8 s stitched segments
+  → gather ctx at interval endpoints → transkun-heads.onnx → velocity + sub-frame onset/offset
+  → notes (+ sustain/soft pedal from tracks 0/1)
+```
+
+See the repo's decode spec and `TranskunTranscriber` for the exact arithmetic (segment padding, `forcedStartPos`
+carry, merge, `resolveOverlapping`).

export_transkun.py

@@ -0,0 +1,198 @@
+#!/usr/bin/env python
+"""v2 Stage 4a — re-derive the Transkun ONNX export + frozen buffers + TDD fixtures.
+
+Boundary: featuresBatch -> S (backbone + inner-product CRF scorer). The mel front end (audio ->
+featuresBatch) and the Viterbi decode (S -> intervals) are ported to C# (4b/4c); this script also emits
+their reference fixtures. All outputs are written raw little-endian + a manifest.json for trivial C# reads.
+"""
+import json, math, os, sys
+import numpy as np
+import torch
+
+TK_DIR = "/private/tmp/claude-501/-Users-lawls-Development-TuesdayCrowd-Projects-audio-claudio/37748a9b-31e0-48a1-896e-7a25c1faf008/scratchpad/transkun-env/lib/python3.14/site-packages/transkun"
+OUT = sys.argv[1] if len(sys.argv) > 1 else os.path.dirname(os.path.abspath(__file__)) + "/artifacts"
+os.makedirs(OUT, exist_ok=True)
+torch.manual_seed(0)
+np.random.seed(0)
+
+manifest = {}
+def save(name, arr, dtype):
+    a = np.ascontiguousarray(arr).astype(dtype)
+    fn = name + (".f32" if dtype == "<f4" else ".i32")
+    a.tofile(os.path.join(OUT, fn))
+    manifest[name] = {"file": fn, "shape": list(a.shape), "dtype": "f32" if dtype == "<f4" else "i32"}
+    print(f"  saved {name:22} shape={list(a.shape)} dtype={manifest[name]['dtype']}")
+
+# ---------------------------------------------------------------- 1. load model from the .conf ----------
+print("[1] loading model from 2.0.conf + 2.0.pt")
+import moduleconf
+conf_path = os.path.join(TK_DIR, "pretrained", "2.0.conf")
+pt_path = os.path.join(TK_DIR, "pretrained", "2.0.pt")
+confManager = moduleconf.parseFromFile(conf_path)
+TransKun = confManager["Model"].module.TransKun
+conf = confManager["Model"].config
+checkpoint = torch.load(pt_path, map_location="cpu")
+model = TransKun(conf=conf)
+key = "best_state_dict" if "best_state_dict" in checkpoint else "state_dict"
+missing = model.load_state_dict(checkpoint[key], strict=False)
+model.eval(); torch.set_grad_enabled(False)
+print(f"  loaded ({key}); missing={len(missing.missing_keys)} unexpected={len(missing.unexpected_keys)}")
+print(f"  fs={model.fs} windowSize={model.windowSize} hopSize={model.hopSize} "
+      f"nSym={len(model.targetMIDIPitch)} params={sum(p.numel() for p in model.parameters())/1e6:.2f}M")
+
+# ---------------------------------------------------------------- 2. export wrapper: featuresBatch -> S -
+class ExportWrapper(torch.nn.Module):
+    def __init__(self, m):
+        super().__init__()
+        self.backbone = m.backbone
+        self.scorer = m.scorer
+        self.register_buffer("outputIndices", torch.tensor(m.targetMIDIPitch))
+    def forward(self, featuresBatch):            # [nBatch, T, 229, 6]
+        ctx = self.backbone(featuresBatch, outputIndices=self.outputIndices)
+        S_batch, _ = self.scorer(ctx)            # S_batch: [T, T, nBatch, 90]; skip is provably 0
+        return S_batch.flatten(-2, -1)           # [T, T, nBatch*90]
+
+wrapper = ExportWrapper(model).eval()
+onnx_path = os.path.join(OUT, "model.onnx")
+T0 = 64
+feat_example = torch.randn(1, T0, 229, 6)
+
+# The backbone's axial attention calls SDPA on 5-D q/k/v ([B, T', H, L, D]); the ONNX exporter's SDPA only
+# supports 4-D. SDPA treats every dim before the last two as batch, so collapsing the two leading dims to
+# one 4-D batch is a mathematical identity (validated by corr below). Patch for the duration of the export.
+import torch.nn.functional as F
+_orig_sdpa = F.scaled_dot_product_attention
+def _sdpa_4d(q, k, v, *a, **kw):
+    if q.ndim == 5:
+        B, X, H, L, D = q.shape
+        rs = lambda t: t.reshape(B * X, H, L, D)
+        return _orig_sdpa(rs(q), rs(k), rs(v), *a, **kw).reshape(B, X, H, L, D)
+    return _orig_sdpa(q, k, v, *a, **kw)
+F.scaled_dot_product_attention = _sdpa_4d
+torch.nn.functional.scaled_dot_product_attention = _sdpa_4d
+
+print(f"[2] exporting featuresBatch{list(feat_example.shape)} -> S, opset 17 (SDPA reshaped 5D->4D)")
+try:
+    torch.onnx.export(
+        wrapper, (feat_example,), onnx_path, opset_version=17,
+        input_names=["featuresBatch"], output_names=["S"],
+        dynamic_axes={"featuresBatch": {1: "T"}, "S": {0: "T", 1: "T"}})
+    print(f"  export OK: {os.path.getsize(onnx_path)/1e6:.1f} MB (+ external .data)")
+except Exception as e:
+    print(f"  STOCK EXPORT FAILED: {type(e).__name__}: {e}\n  (would apply eye-multiply patch)")
+    raise
+
+# Consolidate the external-data weights into ONE self-contained .onnx for committing.
+import onnx
+m_full = onnx.load(onnx_path)  # pulls in model.onnx.data
+single_path = os.path.join(OUT, "transkun.onnx")
+onnx.save(m_full, single_path, save_as_external_data=False)
+os.remove(onnx_path)
+if os.path.exists(onnx_path + ".data"):
+    os.remove(onnx_path + ".data")
+print(f"  consolidated -> transkun.onnx {os.path.getsize(single_path)/1e6:.1f} MB (single file)")
+
+# ---------------------------------------------------------------- 3. validate ONNX == PyTorch -----------
+print("[3] validating single-file ONNX vs PyTorch")
+import onnxruntime as ort
+sess = ort.InferenceSession(single_path, providers=["CPUExecutionProvider"])
+def check(feat, tag):
+    s_torch = wrapper(feat).numpy()
+    s_onnx = sess.run(None, {"featuresBatch": feat.numpy()})[0]
+    corr = np.corrcoef(s_torch.ravel(), s_onnx.ravel())[0, 1]
+    denom = np.abs(s_torch).max() + 1e-9
+    relerr = np.abs(s_torch - s_onnx).max() / denom
+    print(f"  {tag:14} shape={list(s_onnx.shape)} corr={corr:.6f} maxRelErr={relerr:.2e}")
+    return corr, relerr, s_onnx
+check(feat_example, "random T=64")
+check(torch.randn(1, 100, 229, 6), "random T=100")   # dynamic-T sanity
+
+# ---------------------------------------------------------------- 4. frozen buffers ---------------------
+print("[4] extracting frozen buffers")
+fe = model.framewiseFeatureExtractor
+freq2mels = fe.freq2mels.numpy()                                     # [2049, 229]
+win = fe.spectrogramExtractor.win                                    # [4096] Hann
+wins = torch.cat([win.unsqueeze(0), fe.spectrogramExtractor.winGen.get().t()], dim=0).numpy()  # [6, 4096]
+save("freq2mels", freq2mels, "<f4")
+save("windows", wins, "<f4")
+save("symbols", np.array(model.targetMIDIPitch), "<i4")
+params = {
+    "fs": int(model.fs), "windowSize": int(model.windowSize), "hopSize": int(model.hopSize),
+    "nMels": int(fe.outputDim), "nWindows": int(wins.shape[0]), "eps": float(fe.eps),
+    "fMin": 30.0, "fMax": 8000.0, "rfftBins": int(model.windowSize // 2 + 1),
+    "segmentSizeSeconds": 16.0, "segmentHopSeconds": 8.0, "nSymbols": len(model.targetMIDIPitch),
+}
+print("  params:", params)
+
+# ---------------------------------------------------------------- 5. ref3b: audio -> featuresBatch ------
+print("[5] ref3b (mel front end reference)")
+fs, hop, wsz, eps, nmel = model.fs, model.hopSize, model.windowSize, fe.eps, fe.outputDim
+# The first 1.5 s of the committed two-bar MeltySynth piano render (44100 Hz mono int16) — a real piano
+# signal (attack + timbre) that fires the MAESTRO-trained model, fully reproducible from a committed WAV.
+import wave
+wf = wave.open("/Users/lawls/Development/TuesdayCrowd/Projects/audio-claudio/fixtures/golden/two-bar.wav")
+assert wf.getframerate() == fs and wf.getnchannels() == 1 and wf.getsampwidth() == 2
+nread = int(1.5 * fs)
+audio = (np.frombuffer(wf.readframes(nread), dtype="<i2").astype(np.float32) / 32768.0)
+wf.close()
+
+def make_frame(x, hopSize, windowSize):  # mirrors Util.makeFrame (leftPaddingHalfFrame=True)
+    n = x.shape[-1]
+    nFrame = math.ceil(n / hopSize) + 1
+    lPad = windowSize // 2
+    rPad = (nFrame - 1) * hopSize + windowSize // 2 - n
+    xp = torch.nn.functional.pad(torch.tensor(x), (lPad, rPad))
+    return xp.unfold(-1, windowSize, hopSize)  # [nFrame, windowSize]
+
+frames = make_frame(audio, hop, wsz).unsqueeze(0).unsqueeze(0)  # [1,1,nFrame,windowSize] (nBatch,nChan,..)
+mean = frames.mean(dim=[1, 2, 3], keepdim=True)
+std = frames.std(dim=[1, 2, 3], keepdim=True)
+framesN = (frames - mean) / (std + 1e-8)
+features = fe(framesN).contiguous()                             # [1,1,nFrame,229,6]
+features = features.view(1, *features.shape[-3:])               # [1,nFrame,229,6]
+save("ref3b_audio", audio, "<f4")
+save("ref3b_features", features.squeeze(0).numpy(), "<f4")      # [nFrame,229,6]
+print(f"  audio={len(audio)} samples -> features {list(features.shape)}")
+
+# ---------------------------------------------------------------- 6. ref3c: S -> intervals --------------
+print("[6] ref3c (Viterbi decode reference)")
+from transkun.CRF.NeuralSemiCRFInterval import viterbiBackward
+S_real = wrapper(features).squeeze()          # [nFrame,nFrame,90]  (nBatch=1 -> flatten gives 90)
+T = S_real.shape[0]
+noise = torch.zeros(T - 1, S_real.shape[2])
+intervals_real = viterbiBackward(S_real, noise, None)         # per-track list of (begin,end)
+save("ref3c_S", S_real.numpy(), "<f4")
+json.dump({str(k): v for k, v in enumerate(intervals_real)},
+          open(os.path.join(OUT, "ref3c_intervals.json"), "w"))
+nnotes = sum(len(v) for v in intervals_real)
+print(f"  real S {list(S_real.shape)} -> {nnotes} intervals across 90 tracks")
+
+# Hand-checkable synthetic S, multi-track (score[end, begin, track]): track 5 = interval (1,3) + singleton
+# (5,5); track 10 = interval (0,4); track 20 = singleton (2,2). Decoded with default forcedStartPos.
+Tsyn, nSym = 6, 90
+Ssyn = torch.full((Tsyn, Tsyn, nSym), -5.0)
+Ssyn[3, 1, 5] = 4.0
+Ssyn[5, 5, 5] = 2.0
+Ssyn[4, 0, 10] = 6.0
+Ssyn[2, 2, 20] = 3.0
+noise_syn = torch.zeros(Tsyn - 1, nSym)
+intervals_syn = viterbiBackward(Ssyn, noise_syn, None)
+save("ref3c_syn_S", Ssyn.numpy(), "<f4")
+json.dump({str(k): v for k, v in enumerate(intervals_syn)},
+          open(os.path.join(OUT, "ref3c_syn_intervals.json"), "w"))
+print(f"  synthetic S {list(Ssyn.shape)} -> t5={intervals_syn[5]} t10={intervals_syn[10]} t20={intervals_syn[20]}")
+
+# A forcedStartPos case (used by 4d segment stitching): same S, but track 5 forced to start at frame 4,
+# so its interval (1,3) is skipped and only the singleton (5,5) survives.
+forced = [0] * nSym
+forced[5] = 4
+intervals_forced = viterbiBackward(Ssyn, noise_syn, forced)
+json.dump({"forcedStartPos": forced, "intervals": {str(k): v for k, v in enumerate(intervals_forced)}},
+          open(os.path.join(OUT, "ref3c_forced_intervals.json"), "w"))
+print(f"  forced(t5->4) -> t5={intervals_forced[5]} (interval (1,3) skipped)")
+
+# ---------------------------------------------------------------- 7. write manifest + params ------------
+json.dump(manifest, open(os.path.join(OUT, "manifest.json"), "w"), indent=2)
+json.dump(params, open(os.path.join(OUT, "params.json"), "w"), indent=2)
+print(f"[7] wrote manifest.json + params.json to {OUT}")
+print("DONE")

export_transkun_heads.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python
+"""v2 Stage 4e — re-export the main graph to ALSO output ctx, and export the two attribute heads
+(velocityPredictor, refinedOFPredictor) so the C# engine can add real velocity + sub-frame onset/offset
+refinement. Overwrites transkun.onnx (S is byte-for-byte the same; ctx is an added output) and adds
+transkun-heads.onnx. Deterministic (seed 0)."""
+import json, os, sys
+import numpy as np
+import torch
+
+TK_DIR = "/private/tmp/claude-501/-Users-lawls-Development-TuesdayCrowd-Projects-audio-claudio/37748a9b-31e0-48a1-896e-7a25c1faf008/scratchpad/transkun-env/lib/python3.14/site-packages/transkun"
+OUT = sys.argv[1] if len(sys.argv) > 1 else os.path.dirname(os.path.abspath(__file__)) + "/artifacts4e"
+os.makedirs(OUT, exist_ok=True)
+torch.manual_seed(0); np.random.seed(0)
+
+print("[1] load model from 2.0.conf")
+import moduleconf
+cm = moduleconf.parseFromFile(os.path.join(TK_DIR, "pretrained", "2.0.conf"))
+model = cm["Model"].module.TransKun(conf=cm["Model"].config)
+ckpt = torch.load(os.path.join(TK_DIR, "pretrained", "2.0.pt"), map_location="cpu")
+model.load_state_dict(ckpt["best_state_dict" if "best_state_dict" in ckpt else "state_dict"], strict=False)
+model.eval(); torch.set_grad_enabled(False)
+
+# SDPA 5D->4D (identity) so the backbone exports; same as the 4a export.
+import torch.nn.functional as F
+_sdpa = F.scaled_dot_product_attention
+def _sdpa4d(q, k, v, *a, **kw):
+    if q.ndim == 5:
+        B, X, H, L, D = q.shape
+        r = lambda t: t.reshape(B * X, H, L, D)
+        return _sdpa(r(q), r(k), r(v), *a, **kw).reshape(B, X, H, L, D)
+    return _sdpa(q, k, v, *a, **kw)
+F.scaled_dot_product_attention = _sdpa4d
+torch.nn.functional.scaled_dot_product_attention = _sdpa4d
+
+import onnx, onnxruntime as ort
+def consolidate(split_path, single_path):
+    m = onnx.load(split_path)
+    onnx.save(m, single_path, save_as_external_data=False)
+    os.remove(split_path)
+    if os.path.exists(split_path + ".data"):
+        os.remove(split_path + ".data")
+
+# ---- main graph: featuresBatch -> (S, ctx) --------------------------------------------------------------
+print("[2] export main graph featuresBatch -> (S, ctx)")
+class MainWrapper(torch.nn.Module):
+    def __init__(self, m):
+        super().__init__()
+        self.backbone = m.backbone
+        self.scorer = m.scorer
+        self.register_buffer("outputIndices", torch.tensor(m.targetMIDIPitch))
+    def forward(self, featuresBatch):
+        ctx = self.backbone(featuresBatch, outputIndices=self.outputIndices)  # [1,90,T,256]
+        S_batch, _ = self.scorer(ctx)
+        S = S_batch.flatten(-2, -1)                                            # [T,T,90]
+        return S, ctx.squeeze(0)                                              # ctx -> [90,T,256]
+
+main = MainWrapper(model).eval()
+feat = torch.randn(1, 64, 229, 6)
+split = os.path.join(OUT, "_main.onnx")
+torch.onnx.export(main, (feat,), split, opset_version=17,
+                  input_names=["featuresBatch"], output_names=["S", "ctx"],
+                  dynamic_axes={"featuresBatch": {1: "T"}, "S": {0: "T", 1: "T"}, "ctx": {1: "T"}})
+main_path = os.path.join(OUT, "transkun.onnx")
+consolidate(split, main_path)
+print(f"  transkun.onnx {os.path.getsize(main_path)/1e6:.1f} MB")
+
+sess = ort.InferenceSession(main_path, providers=["CPUExecutionProvider"])
+S_t, ctx_t = main(feat)
+S_o, ctx_o = sess.run(None, {"featuresBatch": feat.numpy()})
+print(f"  S   corr={np.corrcoef(S_t.numpy().ravel(), S_o.ravel())[0,1]:.6f} shape={list(S_o.shape)}")
+print(f"  ctx corr={np.corrcoef(ctx_t.numpy().ravel(), ctx_o.ravel())[0,1]:.6f} shape={list(ctx_o.shape)}")
+
+# ---- heads: attr[N,768] -> (velLogits[N,128], ofRaw[N,4]) -----------------------------------------------
+print("[3] export heads attr -> (velLogits, ofRaw)")
+class HeadWrapper(torch.nn.Module):
+    def __init__(self, m):
+        super().__init__()
+        self.vel = m.velocityPredictor
+        self.of = m.refinedOFPredictor
+    def forward(self, attr):
+        return self.vel(attr), self.of(attr)
+
+heads = HeadWrapper(model).eval()
+attr = torch.randn(7, 768)
+hsplit = os.path.join(OUT, "_heads.onnx")
+torch.onnx.export(heads, (attr,), hsplit, opset_version=17,
+                  input_names=["attr"], output_names=["velLogits", "ofRaw"],
+                  dynamic_axes={"attr": {0: "N"}, "velLogits": {0: "N"}, "ofRaw": {0: "N"}})
+heads_path = os.path.join(OUT, "transkun-heads.onnx")
+consolidate(hsplit, heads_path)
+print(f"  transkun-heads.onnx {os.path.getsize(heads_path)/1e3:.0f} KB")
+
+hsess = ort.InferenceSession(heads_path, providers=["CPUExecutionProvider"])
+v_t, o_t = heads(attr)
+v_o, o_o = hsess.run(None, {"attr": attr.numpy()})
+print(f"  velLogits corr={np.corrcoef(v_t.numpy().ravel(), v_o.ravel())[0,1]:.6f} shape={list(v_o.shape)}")
+print(f"  ofRaw     corr={np.corrcoef(o_t.numpy().ravel(), o_o.ravel())[0,1]:.6f} shape={list(o_o.shape)}")
+
+print(f"[4] wrote {main_path} + {heads_path}")
+print("DONE")

freq2mels.f32

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:449f2a5269cde41f6705cfde7157240ce7c8930f20b50909d57e633395bbf0f4
+size 1876884

manifest.json

@@ -0,0 +1,59 @@
+{
+  "freq2mels": {
+    "file": "freq2mels.f32",
+    "shape": [
+      2049,
+      229
+    ],
+    "dtype": "f32"
+  },
+  "windows": {
+    "file": "windows.f32",
+    "shape": [
+      6,
+      4096
+    ],
+    "dtype": "f32"
+  },
+  "symbols": {
+    "file": "symbols.i32",
+    "shape": [
+      90
+    ],
+    "dtype": "i32"
+  },
+  "ref3b_audio": {
+    "file": "ref3b_audio.f32",
+    "shape": [
+      66150
+    ],
+    "dtype": "f32"
+  },
+  "ref3b_features": {
+    "file": "ref3b_features.f32",
+    "shape": [
+      66,
+      229,
+      6
+    ],
+    "dtype": "f32"
+  },
+  "ref3c_S": {
+    "file": "ref3c_S.f32",
+    "shape": [
+      66,
+      66,
+      90
+    ],
+    "dtype": "f32"
+  },
+  "ref3c_syn_S": {
+    "file": "ref3c_syn_S.f32",
+    "shape": [
+      6,
+      6,
+      90
+    ],
+    "dtype": "f32"
+  }
+}

params.json

@@ -0,0 +1,14 @@
+{
+  "fs": 44100,
+  "windowSize": 4096,
+  "hopSize": 1024,
+  "nMels": 229,
+  "nWindows": 6,
+  "eps": 1e-05,
+  "fMin": 30.0,
+  "fMax": 8000.0,
+  "rfftBins": 2049,
+  "segmentSizeSeconds": 16.0,
+  "segmentHopSeconds": 8.0,
+  "nSymbols": 90
+}

ref3b_audio.f32

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:805fabae442dc3262a3d28fe670f541ad157d6c72d3a9ce5e9403661eca6d587
+size 264600

ref3b_features.f32

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca43f2677f53ba60d175e0c401b163be66ef70932e303c5d03e7e444c7c99620
+size 362736

ref3c_S.f32

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a2cbcdcb7d1567dbf29c48d2f9cef420a1de61d8d819b4ebfc921639bc9c9a0e
+size 1568160

ref3c_forced_intervals.json

@@ -0,0 +1 @@
+{"forcedStartPos": [0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], "intervals": {"0": [], "1": [], "2": [], "3": [], "4": [], "5": [[5, 5]], "6": [], "7": [], "8": [], "9": [], "10": [[0, 4]], "11": [], "12": [], "13": [], "14": [], "15": [], "16": [], "17": [], "18": [], "19": [], "20": [[2, 2]], "21": [], "22": [], "23": [], "24": [], "25": [], "26": [], "27": [], "28": [], "29": [], "30": [], "31": [], "32": [], "33": [], "34": [], "35": [], "36": [], "37": [], "38": [], "39": [], "40": [], "41": [], "42": [], "43": [], "44": [], "45": [], "46": [], "47": [], "48": [], "49": [], "50": [], "51": [], "52": [], "53": [], "54": [], "55": [], "56": [], "57": [], "58": [], "59": [], "60": [], "61": [], "62": [], "63": [], "64": [], "65": [], "66": [], "67": [], "68": [], "69": [], "70": [], "71": [], "72": [], "73": [], "74": [], "75": [], "76": [], "77": [], "78": [], "79": [], "80": [], "81": [], "82": [], "83": [], "84": [], "85": [], "86": [], "87": [], "88": [], "89": []}}

ref3c_intervals.json

@@ -0,0 +1 @@
+{"0": [], "1": [], "2": [], "3": [], "4": [], "5": [], "6": [], "7": [], "8": [], "9": [], "10": [], "11": [], "12": [], "13": [], "14": [], "15": [], "16": [], "17": [], "18": [], "19": [], "20": [], "21": [], "22": [], "23": [], "24": [], "25": [], "26": [], "27": [], "28": [], "29": [], "30": [], "31": [], "32": [], "33": [], "34": [], "35": [], "36": [], "37": [], "38": [], "39": [], "40": [], "41": [], "42": [], "43": [], "44": [], "45": [[43, 63]], "46": [], "47": [], "48": [], "49": [], "50": [], "51": [], "52": [], "53": [], "54": [], "55": [], "56": [], "57": [], "58": [], "59": [], "60": [], "61": [], "62": [], "63": [], "64": [], "65": [], "66": [], "67": [], "68": [], "69": [], "70": [], "71": [], "72": [], "73": [], "74": [], "75": [], "76": [], "77": [], "78": [], "79": [], "80": [], "81": [], "82": [], "83": [], "84": [], "85": [], "86": [], "87": [], "88": [], "89": []}

ref3c_syn_intervals.json

@@ -0,0 +1 @@
+{"0": [], "1": [], "2": [], "3": [], "4": [], "5": [[1, 3], [5, 5]], "6": [], "7": [], "8": [], "9": [], "10": [[0, 4]], "11": [], "12": [], "13": [], "14": [], "15": [], "16": [], "17": [], "18": [], "19": [], "20": [[2, 2]], "21": [], "22": [], "23": [], "24": [], "25": [], "26": [], "27": [], "28": [], "29": [], "30": [], "31": [], "32": [], "33": [], "34": [], "35": [], "36": [], "37": [], "38": [], "39": [], "40": [], "41": [], "42": [], "43": [], "44": [], "45": [], "46": [], "47": [], "48": [], "49": [], "50": [], "51": [], "52": [], "53": [], "54": [], "55": [], "56": [], "57": [], "58": [], "59": [], "60": [], "61": [], "62": [], "63": [], "64": [], "65": [], "66": [], "67": [], "68": [], "69": [], "70": [], "71": [], "72": [], "73": [], "74": [], "75": [], "76": [], "77": [], "78": [], "79": [], "80": [], "81": [], "82": [], "83": [], "84": [], "85": [], "86": [], "87": [], "88": [], "89": []}

transkun-heads.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f609f53fbb4df0a2165f941aff1d5a1c1063ba40aaa699365448dca5d9f3544d
+size 3422740

transkun.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d6960ecbce5c3d51051ebc1ebdeec4d3018ea6cfd360fb1ae96359aa7a075a3
+size 53391226