PiCoGen / picogen2 /mirtoolkit /sheetsage.py
Vansh Chugh
initial deploy
fe7e262
Raw
History Blame Contribute Delete
4.35 kB
"""
SheetSage: lead sheet transcription model
- reference: https://github.com/chrisdonahue/sheetsage
"""
import logging
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import Union
from sheetsage.infer import sheetsage as sheetsage_infer
from tqdm import tqdm as tqdm_fn
class SheetSage:
def __init__(self):
pass
def __call__(
self,
audio_path: Union[str, Path] = None,
audio_url: str = None,
segment_start_hint=None,
segment_end_hint=None,
use_jukebox=True,
measures_per_chunk=8,
dynamic_chunking=True,
segment_hints_are_downbeats=False,
beat_information=None,
beats_per_measure_hint=None,
beats_per_minute_hint=None,
detect_melody=True,
detect_harmony=True,
melody_threshold=None,
harmony_threshold=None,
beat_detection_padding=15.0,
avoid_chunking_if_possible=True,
legacy_behavior=False,
status_change_callback=lambda s: logging.info(s.name),
return_intermediaries=False,
tqdm=tqdm_fn,
return_dict=True,
):
# assert audio_path or audio_url should be provided but not both
assert (audio_path and not audio_url) or (audio_url and not audio_path), (
f"One of audio_path or audio_url should be provided but not both: {audio_path}, {audio_url}"
)
assert shutil.which("ffmpeg") is not None, "ffmpeg not found. Please install ffmpeg."
if audio_path: # if audio_path is provided
ext = "flac"
tmp_audio_file = tempfile.NamedTemporaryFile(suffix=f".{ext}")
subprocess.run(
[
"ffmpeg",
"-i",
str(audio_path),
"-vn",
"-f",
ext,
"-y",
tmp_audio_file.name,
]
)
audio_path = Path(tmp_audio_file.name)
assert audio_path.exists(), f"File not found: {audio_path}"
else: # if audio_url is provided
ext = "flac"
tmp_dir = tempfile.TemporaryDirectory()
tmp_audio_file = tmp_dir.name + f"/audio.{ext}"
subprocess.run(
[
"yt-dlp",
"-x",
"--audio-format",
ext,
"--audio-quality",
"0",
"-o",
tmp_dir.name + "/audio.%(ext)s",
audio_url,
],
check=True,
)
audio_path = Path(tmp_audio_file)
sheetsage_output = sheetsage_infer(
audio_path_bytes_or_url=audio_path,
segment_start_hint=segment_start_hint,
segment_end_hint=segment_end_hint,
use_jukebox=use_jukebox,
measures_per_chunk=measures_per_chunk,
dynamic_chunking=dynamic_chunking,
segment_hints_are_downbeats=segment_hints_are_downbeats,
beat_information=beat_information,
beats_per_measure_hint=beats_per_measure_hint,
beats_per_minute_hint=beats_per_minute_hint,
detect_melody=detect_melody,
detect_harmony=detect_harmony,
melody_threshold=melody_threshold,
harmony_threshold=harmony_threshold,
beat_detection_padding=beat_detection_padding,
avoid_chunking_if_possible=avoid_chunking_if_possible,
legacy_behavior=legacy_behavior,
status_change_callback=status_change_callback,
return_intermediaries=return_intermediaries,
tqdm=tqdm,
)
if return_dict:
return sheetsage_output
else:
return (
sheetsage_output["lead_sheet"],
sheetsage_output["segment_beats"],
sheetsage_output["segment_beats_times"],
sheetsage_output["chunks_tertiaries"],
sheetsage_output["melody_logits"],
sheetsage_output["harmony_logits"],
sheetsage_output["melody_last_hidden_state"],
sheetsage_output["harmony_last_hidden_state"],
)