Spaces:
Running on Zero
Running on Zero
| import spaces | |
| from transkun.ModelTransformer import TransKun, writeMidi | |
| import numpy as np | |
| import torch | |
| import librosa | |
| import tempfile | |
| import os | |
| import soundfile as sf | |
| # setuptools dropped this module in version 81, and it is only needed here | |
| # to locate data files inside the installed transkun package, which | |
| # importlib.resources does from the standard library. | |
| from importlib.resources import files | |
| import moduleconf | |
| from pyharp.core import ModelCard, build_endpoint | |
| from pyharp.media.audio import load_audio | |
| from pyharp.tags import Channels, SampleRate, Subcategory | |
| import gradio as gr | |
| model_card = ModelCard( | |
| name="Transkun", | |
| description=( | |
| "Transcribes solo piano performance into MIDI notation, with onsets, offsets, and " | |
| "velocities. The input is mixed to mono and resampled to 44.1 kHz internally. " | |
| "Repository: https://github.com/Yujia-Yan/Transkun Paper: " | |
| "https://arxiv.org/abs/2404.09466" | |
| ), | |
| author="Yujia Yan, Zhiyao Duan", | |
| tags=[Subcategory.MUSIC_ANALYSIS, SampleRate(44100), Channels(1), 'domain:music'] | |
| ) | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| def load_model(device: str): | |
| defaultWeight = str(files("transkun") / "pretrained" / "2.0.pt") | |
| defaultConf = str(files("transkun") / "pretrained" / "2.0.conf") | |
| confManager = moduleconf.parseFromFile(defaultConf) | |
| conf = confManager["Model"].config | |
| checkpoint = torch.load(defaultWeight, map_location="cpu") | |
| model = TransKun(conf=conf) | |
| # Checkpoints keep their weights under either key | |
| if "best_state_dict" in checkpoint: | |
| model.load_state_dict(checkpoint["best_state_dict"], strict=False) | |
| else: | |
| model.load_state_dict(checkpoint["state_dict"], strict=False) | |
| model.eval() | |
| return model.to(device) | |
| # pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the | |
| # server, so the model is built and placed there alone, rather than inside every GPU call | |
| if __name__ != "__main__": | |
| model = load_model(DEVICE) | |
| # The recording is transcribed in overlapping segments, so the cost grows with its length | |
| # PyHARP stops a job that runs past its timeout_s, 15 minutes by default, so a request whose | |
| # reservation would exceed that is refused rather than started. A shorter one is left to | |
| # ZeroGPU, which admits it when the caller's own allowance covers it. | |
| MAX_GPU_SECONDS = 900 | |
| def gpu_seconds(input_file): | |
| # Measured at about 2 to 8 s for 30 s of piano and 10 to 15 s for 300 s, upload included, | |
| # padded. The model is built when the worker starts, not inside the call. | |
| return 8 + 0.05 * sf.info(input_file).duration | |
| def transcribe(input_file): | |
| signal = load_audio(input_file) | |
| waveform = np.asarray(signal.audio_data) | |
| sr = int(signal.sample_rate) | |
| waveform = np.squeeze(waveform) | |
| # If 2D, assume (channels, samples). Make mono. | |
| if waveform.ndim == 2: | |
| if waveform.shape[0] > 1: | |
| waveform = waveform.mean(axis=0) | |
| else: | |
| waveform = waveform.reshape(-1) | |
| if sr != model.fs: | |
| waveform = librosa.resample(waveform.astype(np.float32), orig_sr=sr, target_sr=model.fs) | |
| sr = model.fs | |
| x = torch.from_numpy(waveform.reshape(-1, 1)).to(DEVICE) | |
| notesEst = model.transcribe(x, discardSecondHalf=False) | |
| outputMidi = writeMidi(notesEst) | |
| out_fd, out_path = tempfile.mkstemp(suffix=".mid") | |
| os.close(out_fd) | |
| outputMidi.write(out_path) | |
| return out_path | |
| def process_fn(input_audio_path: str) -> str: | |
| # Checked before the GPU is reserved, so a request the model cannot take spends none of the | |
| # caller's ZeroGPU allowance | |
| if not input_audio_path: | |
| raise gr.Error("Provide the piano recording to transcribe.") | |
| try: | |
| duration = sf.info(input_audio_path).duration | |
| except Exception as exc: | |
| raise gr.Error(f"Could not read the audio file: {exc}") from exc | |
| if duration == 0: | |
| raise gr.Error("The audio file is empty.") | |
| if gpu_seconds(input_audio_path) > MAX_GPU_SECONDS: | |
| raise gr.Error( | |
| f"The recording is {duration:.0f} s long, more than one request can transcribe in " | |
| f"the {MAX_GPU_SECONDS // 60} minutes it may run. Trim it." | |
| ) | |
| midi_path = transcribe(input_audio_path) | |
| return midi_path | |
| if __name__ == "__main__": | |
| with gr.Blocks() as demo: | |
| input_audio = gr.Audio(label="Upload Solo Piano Audio", type="filepath").harp_required(True).set_info(( | |
| "A solo piano recording. Trained on piano alone, so other " | |
| "instruments in the recording will confuse it." | |
| )) | |
| output_midi = gr.File( | |
| label="Transcription", | |
| file_types=[".mid", ".midi"], | |
| type="filepath" | |
| ).set_info("The transcribed notes, with onsets, offsets and velocities.") | |
| app = build_endpoint( | |
| model_card=model_card, | |
| input_components=[input_audio], | |
| output_components=[output_midi], | |
| process_fn=process_fn | |
| ) | |
| demo.queue().launch(pwa=True, show_error=True, ssr_mode=False) | |