File size: 2,445 Bytes
dee8df0 6b1f7f2 dee8df0 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 | """One isolated SheetSage2 CUDA transcription per request."""
from __future__ import annotations
import importlib
import json
from pathlib import Path
import shutil
import sys
import numpy as np
import soundfile as sf
import torch
from transformers import AutoModel
SHEETSAGE_REV = "eab522a8168e8b8b8c4856bf8609cd86198f01fe"
def run(source: Path, output: Path) -> None:
if not torch.cuda.is_available():
raise RuntimeError("SheetSage2 requires the configured NVIDIA GPU")
waveform, rate = sf.read(source, dtype="float32", always_2d=True)
if rate != 24000 or waveform.shape[1] != 1 or not np.isfinite(waveform).all():
raise ValueError("Expected finite mono 24 kHz PCM")
torch.manual_seed(1024)
np.random.seed(1024)
model = AutoModel.from_pretrained(
"m-a-p/SheetSage2", revision=SHEETSAGE_REV, trust_remote_code=True,
torch_dtype=torch.float32, attn_implementation="sdpa",
).eval().to("cuda")
native = output / "native"
result = model.transcribe(waveform[:, 0], sampling_rate=24000,
output_dir=native, dtype="fp32", preset="default",
melody_only=False)
if result.get("abc_error") or not (native / "score.abc").is_file():
raise RuntimeError(f"SheetSage2 could not export ABC: {result.get('abc_error')}")
package = model.__class__.__module__.rsplit(".", 1)[0]
exporter = importlib.import_module(package + ".exports_sheetsage2").export_result
events = json.loads((native / "events.json").read_text())
melody = exporter(events, model.tokenizer, output / "melody-only",
duration=result["duration_seconds"], melody_only=True)
if melody.get("abc_error") or not (output / "melody-only/score.abc").is_file():
raise RuntimeError(f"SheetSage2 could not export melody ABC: {melody.get('abc_error')}")
shutil.copy2(native / "score.abc", output / "full.abc")
shutil.copy2(output / "melody-only/score.abc", output / "melody.abc")
shutil.copy2(native / "melody_vocal.mid", output / "notes.mid")
shutil.copy2(native / "beat.lab", output / "beat.lab")
(output / "sheet.json").write_text(json.dumps({
"model": "m-a-p/SheetSage2", "revision": SHEETSAGE_REV,
"duration_s": result["duration_seconds"],
"source_timeline": True,
}, indent=2) + "\n")
if __name__ == "__main__":
run(Path(sys.argv[1]), Path(sys.argv[2]))
|