feat: pipeline runner with oracle modes
Co-Authored-By: Claude Sonnet 5.5 <[email protected]>
This commit is contained in:
Lucas Winther
co-authored by
Claude Sonnet 5.5
Claude Sonnet 5.5
parent
bf93a4a32e
commit
604d1a67af
@@ -0,0 +1,128 @@
|
||||
// End-to-end pipeline: mix -> stems -> per-role transcription -> tempo grid -> quantise -> tab -> Guitar Pro.
|
||||
// Pure: no file I/O; callers pass model bytes.
|
||||
import type { Audio, Grid, Role, ScoreNote, Song, TimedNote, Track, TruthTrack } from "../core/types.ts";
|
||||
import { BASS_STANDARD, GUITAR_STANDARD, TPQ } from "../core/types.ts";
|
||||
import { secondsToTicks } from "../core/time.ts";
|
||||
import type { OracleFlags, OutTrack, PipelineOutput } from "../eval/score.ts";
|
||||
import { songToGp } from "../score/gp.ts";
|
||||
import { createSeparator } from "./separate.ts";
|
||||
import { createTranscriber, type TranscribeOpts } from "./transcribe.ts";
|
||||
import { transcribeDrums } from "./drums.ts";
|
||||
import { estimateGrid, gridFromTruth } from "./tempo.ts";
|
||||
import { quantize } from "./quantize.ts";
|
||||
import { assignTab } from "./tab.ts";
|
||||
|
||||
export interface PipelineInput {
|
||||
mix: Audio;
|
||||
oracle?: { stems: Map<string, Audio>; truth: { bpm: number; timeSig: [number, number]; tracks: TruthTrack[] } };
|
||||
}
|
||||
export interface PipelineResult { output: PipelineOutput; gp: Uint8Array; warnings: string[] }
|
||||
|
||||
const PROGRAM: Record<Role, number> = { guitar: 27, bass: 33, keys: 0, vocals: 53, drums: 0 };
|
||||
const RANGE: Record<Exclude<Role, "drums">, TranscribeOpts> = {
|
||||
guitar: { minPitch: 40, maxPitch: 88 },
|
||||
bass: { minPitch: 28, maxPitch: 67 },
|
||||
keys: { minPitch: 21, maxPitch: 108 },
|
||||
vocals: { minPitch: 48, maxPitch: 84, monophonic: true },
|
||||
};
|
||||
const TUNING: Partial<Record<Role, number[]>> = { guitar: GUITAR_STANDARD, bass: BASS_STANDARD };
|
||||
const STEP = TPQ / 4;
|
||||
const DEFAULT_GRID: Grid = { bpm: 120, firstDownbeat: 0, beatsPerBar: 4 };
|
||||
|
||||
interface Job { name: string; role: Role; audio: Audio | null; truthNotes: TimedNote[] | null }
|
||||
|
||||
/** Copy string/fret from tab-assigned score notes onto the timed notes they came from (same pitch, nearest tick). */
|
||||
function copyFingering(timed: TimedNote[], assigned: ScoreNote[], grid: Grid): void {
|
||||
const used = new Set<number>();
|
||||
for (const n of timed) {
|
||||
const tick = secondsToTicks(n.onset - grid.firstDownbeat, grid.bpm);
|
||||
let best = -1, bestD = Infinity;
|
||||
assigned.forEach((a, i) => {
|
||||
if (used.has(i) || a.pitch !== n.pitch) return;
|
||||
const d = Math.abs(a.tick - tick);
|
||||
if (d < bestD) { bestD = d; best = i; }
|
||||
});
|
||||
if (best >= 0 && bestD <= STEP) {
|
||||
used.add(best);
|
||||
n.string = assigned[best]!.string;
|
||||
n.fret = assigned[best]!.fret;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export async function createPipeline(
|
||||
models: { demucs: Uint8Array; basicPitch: Uint8Array },
|
||||
): Promise<(input: PipelineInput, flags: OracleFlags) => Promise<PipelineResult>> {
|
||||
const separate = await createSeparator(models.demucs);
|
||||
const transcribe = await createTranscriber(models.basicPitch);
|
||||
|
||||
return async (input, flags) => {
|
||||
const warnings: string[] = [];
|
||||
const { oracle } = input;
|
||||
if ((flags.oracleStems || flags.oracleTempo || flags.oracleNotes) && !oracle) {
|
||||
throw new Error("createPipeline: oracle flags set but input.oracle is missing");
|
||||
}
|
||||
|
||||
// Stems -> jobs
|
||||
const jobs: Job[] = [];
|
||||
if (flags.oracleStems) {
|
||||
for (const t of oracle!.truth.tracks) {
|
||||
const audio = oracle!.stems.get(t.name);
|
||||
if (!audio) throw new Error(`createPipeline: no oracle stem for track "${t.name}"`);
|
||||
jobs.push({ name: t.name, role: t.role, audio, truthNotes: t.notes });
|
||||
}
|
||||
} else {
|
||||
const s = await separate(input.mix);
|
||||
warnings.push('demucs "other" stem unused');
|
||||
const mk = (name: string, role: Role, audio: Audio): Job => ({ name, role, audio, truthNotes: null });
|
||||
jobs.push(mk("Guitar", "guitar", s.guitar), mk("Bass", "bass", s.bass), mk("Keys", "keys", s.piano),
|
||||
mk("Vocals", "vocals", s.vocals), mk("Drums", "drums", s.drums));
|
||||
}
|
||||
|
||||
// Grid
|
||||
let grid: Grid;
|
||||
if (flags.oracleTempo) grid = gridFromTruth(oracle!.truth.bpm, oracle!.truth.timeSig[0]);
|
||||
else {
|
||||
try { grid = estimateGrid(input.mix); }
|
||||
catch (e) {
|
||||
grid = { ...DEFAULT_GRID };
|
||||
warnings.push(`estimateGrid failed (${(e as Error).message}); using default 120 bpm grid`);
|
||||
}
|
||||
}
|
||||
|
||||
// Transcribe + quantise + tab
|
||||
const tracks: OutTrack[] = [];
|
||||
const songTracks: Track[] = [];
|
||||
for (const job of jobs) {
|
||||
let notes: TimedNote[];
|
||||
if (job.role === "drums") notes = transcribeDrums(job.audio!);
|
||||
else if (flags.oracleNotes) {
|
||||
const src = job.truthNotes ?? oracle!.truth.tracks.filter((t) => t.role === job.role).flatMap((t) => t.notes);
|
||||
notes = src.map((n) => ({ onset: n.onset, offset: n.offset, pitch: n.pitch, velocity: n.velocity }));
|
||||
} else notes = await transcribe(job.audio!, RANGE[job.role]);
|
||||
|
||||
let scoreNotes = quantize(notes, grid);
|
||||
const tuning = TUNING[job.role];
|
||||
if (tuning) {
|
||||
const r = assignTab(scoreNotes, tuning);
|
||||
scoreNotes = r.notes;
|
||||
if (r.dropped > 0) warnings.push(`${job.name}: ${r.dropped} note(s) dropped by tab assignment`);
|
||||
copyFingering(notes, scoreNotes, grid);
|
||||
}
|
||||
tracks.push({ name: job.name, role: job.role, notes, scoreNotes });
|
||||
songTracks.push({
|
||||
name: job.name, role: job.role, program: PROGRAM[job.role], ...(tuning ? { tuning } : {}),
|
||||
pan: 0, fx: "none", gainDb: 0, notes: scoreNotes,
|
||||
});
|
||||
}
|
||||
|
||||
const barTicks = TPQ * 4;
|
||||
const endTick = Math.max(0, ...songTracks.flatMap((t) => t.notes.map((n) => n.tick + n.dur)));
|
||||
const song: Song = {
|
||||
title: "STSheets output", tier: "pipeline", seed: 0, bpm: grid.bpm, timeSig: [4, 4],
|
||||
bars: Math.max(1, Math.ceil(endTick / barTicks)), key: "C", tracks: songTracks,
|
||||
};
|
||||
const gp = songToGp(song); // throws on alphaTex diagnostics: indicates a quantiser/writer contract bug
|
||||
return { output: { grid, tracks }, gp, warnings };
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
import { test, expect, beforeAll } from "bun:test";
|
||||
import { join } from "node:path";
|
||||
import { buildSong, loadSongDir } from "../../src/bench/dataset.ts";
|
||||
import { createPipeline } from "../../src/pipeline/run.ts";
|
||||
import { truthFromGp } from "../../src/score/gp.ts";
|
||||
import { scoreSong } from "../../src/eval/score.ts";
|
||||
|
||||
const SLOW = 600_000;
|
||||
const root = join(import.meta.dir, "../..");
|
||||
let run: Awaited<ReturnType<typeof createPipeline>>;
|
||||
|
||||
beforeAll(async () => {
|
||||
const bytes = async (p: string) => new Uint8Array(await Bun.file(join(root, p)).arrayBuffer());
|
||||
run = await createPipeline({ demucs: await bytes("models/demucs/htdemucs_6s.onnx"), basicPitch: await bytes("models/basic-pitch/nmp.onnx") });
|
||||
await buildSong("t1a", 1);
|
||||
await buildSong("t1c", 1);
|
||||
}, SLOW);
|
||||
|
||||
test("t1a/1 oracle stems+tempo: guitar track, GP re-imports", async () => {
|
||||
const d = await loadSongDir(join(root, "data/t1a/1"));
|
||||
const flags = { oracleStems: true, oracleTempo: true, oracleNotes: false };
|
||||
const t0 = performance.now();
|
||||
const { output, gp } = await run({ mix: d.mix, oracle: { stems: d.stems, truth: d.truth } }, flags);
|
||||
const s = scoreSong(d.truth, output, flags);
|
||||
console.log(`t1a oracleNotes=false ${(performance.now() - t0).toFixed(0)}ms F1=${s.guitarMerged.f1} str=${s.stringAgreement}`);
|
||||
const g = output.tracks.filter((t) => t.role === "guitar");
|
||||
expect(g.length).toBe(1);
|
||||
expect(g[0]!.notes.length).toBeGreaterThan(0);
|
||||
expect(truthFromGp(gp).tracks.some((t) => t.role === "guitar")).toBe(true);
|
||||
}, SLOW);
|
||||
|
||||
test("t1a/1 oracle notes: merged guitar F1 === 1, stringAgreement is a number", async () => {
|
||||
const d = await loadSongDir(join(root, "data/t1a/1"));
|
||||
const flags = { oracleStems: true, oracleTempo: true, oracleNotes: true };
|
||||
const { output } = await run({ mix: d.mix, oracle: { stems: d.stems, truth: d.truth } }, flags);
|
||||
const s = scoreSong(d.truth, output, flags);
|
||||
console.log(`t1a oracleNotes=true F1=${s.guitarMerged.f1} str=${s.stringAgreement}`);
|
||||
expect(s.guitarMerged.f1).toBe(1);
|
||||
expect(typeof s.stringAgreement).toBe("number");
|
||||
}, SLOW);
|
||||
|
||||
test("t1c/1 full pipeline yields all roles", async () => {
|
||||
const d = await loadSongDir(join(root, "data/t1c/1"));
|
||||
const flags = { oracleStems: false, oracleTempo: false, oracleNotes: false };
|
||||
const t0 = performance.now();
|
||||
const { output, gp, warnings } = await run({ mix: d.mix }, flags);
|
||||
const s = scoreSong(d.truth, output, flags);
|
||||
console.log(`t1c full ${(performance.now() - t0).toFixed(0)}ms bpm=${output.grid.bpm} (true ${d.truth.bpm}) F1=${s.guitarMerged.f1} warnings=${JSON.stringify(warnings)} rhythm=${JSON.stringify(s.rhythm)}`);
|
||||
const roles = new Set(output.tracks.map((t) => t.role));
|
||||
for (const r of ["guitar", "bass", "keys", "vocals", "drums"]) expect(roles.has(r as never)).toBe(true);
|
||||
expect(gp.length).toBeGreaterThan(0);
|
||||
}, SLOW);
|
||||
Reference in New Issue
Block a user