Spaces:
Running
Running
| import spaces # imported first, as ZeroGPU requires | |
| import sys | |
| sys.stdout.reconfigure(line_buffering=True) | |
| from huggingface_hub import hf_hub_download | |
| from pyharp import ModelCard, build_endpoint | |
| from pyharp.labels import LabelList, AudioLabel | |
| from pyharp.tags import Channels, SampleRate, Subcategory | |
| import gradio as gr | |
| import soundfile as sf | |
| import torch | |
| from beat_this.inference import File2Beats | |
| model_card = ModelCard( | |
| name="Beat This!", | |
| description=( | |
| "Detects the beat and downbeat (start-of-bar) positions in a piece of " | |
| "music. The input is mixed to mono and resampled to 22.05 kHz internally. " | |
| "Repository: https://github.com/CPJKU/beat_this Paper: " | |
| "https://arxiv.org/abs/2407.21658" | |
| ), | |
| author="Francesco Foscarin, Jan Schlüter, Gerhard Widmer", | |
| tags=[Subcategory.MUSIC_ANALYSIS, SampleRate(22050), Channels(1)], | |
| ) | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Built and placed in the worker alone, which imports this file as __mp_main__ and runs process_fn, | |
| # so a broken checkpoint fails in the logs when the Space starts and no request builds the model | |
| # again. The checkpoint lives in a model repository rather than in this Space, pinned to a revision. | |
| if __name__ != "__main__": | |
| model = File2Beats( | |
| checkpoint_path=hf_hub_download( | |
| repo_id="teamup-tech/Beat-This-weights", | |
| filename="final0.ckpt", | |
| revision="4afbe376535aaab71b76ae2ca2d7c16b08877321", | |
| ), | |
| device=DEVICE, | |
| dbn=False, | |
| ) | |
| BEAT_COLOR = AudioLabel.rgb_color_to_int(120, 170, 255) # blue | |
| DOWNBEAT_COLOR = AudioLabel.rgb_color_to_int(255, 130, 40) # orange | |
| # downbeats stay in the overhead row (no amplitude); beats get an amplitude | |
| # near the top of the waveform, so they show up as a row just below downbeats | |
| BEAT_AMPLITUDE = 0.9 | |
| def process_fn(input_audio_path: str) -> tuple[str, LabelList]: | |
| """Finds beat and downbeat times in the input audio and returns them as labeled markers on the audio.""" | |
| # 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 music to track.") | |
| try: | |
| info = sf.info(input_audio_path) | |
| except Exception as exc: | |
| raise gr.Error(f"Could not read the audio file: {exc}") from exc | |
| if info.frames == 0: | |
| raise gr.Error("The audio file is empty.") | |
| return track_beats(input_audio_path) | |
| def track_beats(input_audio_path: str) -> tuple[str, LabelList]: | |
| beats, downbeats = model(input_audio_path) | |
| # downbeats are a subset of beats with matching float values | |
| downbeat_times = set(downbeats.tolist()) | |
| output_labels = LabelList() | |
| for t in beats: | |
| is_downbeat = float(t) in downbeat_times | |
| if is_downbeat: | |
| output_labels.append( | |
| AudioLabel(t=float(t), label="downbeat", color=DOWNBEAT_COLOR) | |
| ) | |
| else: | |
| output_labels.append( | |
| AudioLabel( | |
| t=float(t), label="beat", color=BEAT_COLOR, amplitude=BEAT_AMPLITUDE | |
| ) | |
| ) | |
| # HARP only attaches labels to an output track, so pass the audio through unchanged | |
| return input_audio_path, output_labels | |
| if __name__ == "__main__": | |
| with gr.Blocks() as demo: | |
| input_components = [ | |
| gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info(( | |
| "Music to track. Beats and downbeats come back as labels over the " | |
| "audio." | |
| )), | |
| ] | |
| output_components = [ | |
| gr.Audio(type="filepath", label="Audio with Beats").set_info( | |
| "Input audio, unchanged, with detected beats and downbeats marked." | |
| ), | |
| gr.JSON(label="Beats").set_info( | |
| "Detected beat and downbeat times, labeled \"beat\" or \"downbeat\"." | |
| ), | |
| ] | |
| build_endpoint( | |
| model_card=model_card, | |
| input_components=input_components, | |
| output_components=output_components, | |
| process_fn=process_fn, | |
| ) | |
| demo.queue().launch(show_error=True, pwa=True) | |