| """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])) |
|
|
