""" Gradio demo for CountEM Automatic Music Transcription. This demo allows users to upload audio files and transcribe them to MIDI using pre-trained models from Hugging Face Hub. """ import gradio as gr import spaces import tempfile import os from pathlib import Path import numpy as np import soundfile as sf import librosa import logging from onsets_and_frames.hf_model import CountEMModel from onsets_and_frames.constants import SAMPLE_RATE # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Cache for loaded models to avoid reloading model_cache = {} def load_model(model_name: str) -> CountEMModel: """Load model from cache or download from Hugging Face Hub.""" if model_name not in model_cache: logger.info(f"Loading model: {model_name}") model_cache[model_name] = CountEMModel.from_pretrained(model_name) logger.info(f"Model loaded successfully") return model_cache[model_name] @spaces.GPU def transcribe_audio( audio_input, model_choice: str, onset_threshold: float, frame_threshold: float, ) -> tuple: """ Transcribe audio to MIDI. Args: audio_input: Tuple of (sample_rate, audio_data) from Gradio Audio component model_choice: Model to use ("MusicNet" or "Synth") onset_threshold: Threshold for onset detection frame_threshold: Threshold for frame detection Returns: Tuple of (output_midi_path, status_message) """ try: # Handle empty input if audio_input is None: return None, "Error: Please upload an audio file" # Map model choice to HuggingFace repo ID model_map = { "MusicNet (Recommended)": "Yoni232/countem-musicnet", "Synth": "Yoni232/countem-synth", } model_name = model_map[model_choice] # Extract audio data # Gradio Audio component returns (sample_rate, audio_array) or audio file path input_filename = None if isinstance(audio_input, tuple): sr, audio = audio_input # Convert to float32 if needed if audio.dtype == np.int16: audio = audio.astype(np.float32) / 32768.0 elif audio.dtype == np.int32: audio = audio.astype(np.float32) / 2147483648.0 elif isinstance(audio_input, str): # Audio file path provided audio, sr = librosa.load(audio_input, sr=None, mono=True) # Extract filename for output naming input_filename = Path(audio_input).stem else: return None, f"Error: Unexpected audio input type: {type(audio_input)}" # Convert stereo to mono if needed if len(audio.shape) > 1: audio = audio.mean(axis=1) # Resample to 16kHz if needed if sr != SAMPLE_RATE: logger.info(f"Resampling from {sr}Hz to {SAMPLE_RATE}Hz") audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE) sr = SAMPLE_RATE # Check audio length duration = len(audio) / sr if duration < 0.5: return None, "Error: Audio is too short (minimum 0.5 seconds)" if duration > 600: # 10 minutes return ( None, f"Error: Audio is too long ({duration:.1f}s). Maximum is 10 minutes (600s).", ) # Load model status = f"Loading {model_choice} model..." logger.info(status) model = load_model(model_name) # Transcribe status = f"Transcribing {duration:.1f} seconds of audio..." logger.info(status) # Create temporary MIDI file with original filename if available if input_filename: temp_dir = tempfile.gettempdir() output_path = os.path.join(temp_dir, f"{input_filename}.mid") else: with tempfile.NamedTemporaryFile(suffix=".mid", delete=False) as tmp: output_path = tmp.name model.transcribe_to_midi( audio, output_path, onset_threshold=onset_threshold, frame_threshold=frame_threshold, ) # Success message success_msg = f""" ✓ Transcription complete! - Model: {model_choice} - Duration: {duration:.2f} seconds - Sample rate: {sr} Hz - Onset threshold: {onset_threshold} - Frame threshold: {frame_threshold} Download your MIDI file using the button below. """ return output_path, success_msg.strip() except Exception as e: error_msg = f"Error during transcription: {str(e)}" logger.error(error_msg) return None, error_msg # Build Gradio interface with gr.Blocks(title="CountEM - Music Transcription") as demo: gr.Markdown( """ # CountEM - Automatic Music Transcription Upload a piano/music recording and transcribe it to MIDI using a model that was trained using the CountEM framework on the MusicNet dataset. **Paper:** [Count the Notes: Histogram-Based Supervision for Automatic Music Transcription](https://arxiv.org/abs/2511.14250) (ISMIR 2025) **Models on Hugging Face:** - [countem-musicnet](https://huggingface.co/Yoni232/countem-musicnet) - Trained on MusicNet dataset - [countem-synth](https://huggingface.co/Yoni232/countem-synth) - Trained on synthetic data """ ) with gr.Row(): with gr.Column(): # Input section audio_input = gr.Audio( label="Upload Audio File", type="filepath", sources=["upload"], ) model_choice = gr.Radio( choices=["MusicNet (Recommended)", "Synth"], value="MusicNet (Recommended)", label="Model Selection", info="MusicNet model is trained on real piano recordings, Synth on synthetic data", ) with gr.Row(): onset_threshold = gr.Slider( minimum=0.1, maximum=0.9, value=0.5, step=0.05, label="Onset Threshold", info="Higher = fewer notes detected", ) frame_threshold = gr.Slider( minimum=0.1, maximum=0.9, value=0.5, step=0.05, label="Frame Threshold", info="Higher = shorter note durations", ) transcribe_btn = gr.Button("Transcribe to MIDI", variant="primary") with gr.Column(): # Output section output_midi = gr.File(label="Download MIDI", interactive=False) status_output = gr.Textbox( label="Status", lines=10, interactive=False, placeholder="Upload audio and click 'Transcribe to MIDI' to start...", ) # Example files gr.Markdown( """ ### Notes: - Audio will be automatically resampled to 16kHz if needed, and converted to mono - Supports common formats: WAV, FLAC, MP3, M4a - Maximum duration: 10 minutes - Best results with classical music - Processing time depends on audio length (typically a few seconds per minute of audio) """ ) # Connect button to function transcribe_btn.click( fn=transcribe_audio, inputs=[audio_input, model_choice, onset_threshold, frame_threshold], outputs=[output_midi, status_output], ) gr.Markdown( """ --- **Project Links:** - [GitHub Repository](https://github.com/Yoni-Yaffe/count-the-notes) - [Project Page](https://yoni-yaffe.github.io/count-the-notes/) - [ArXiv Paper](https://arxiv.org/abs/2511.14250) If you use this work, please cite: ``` @misc{yaffe2025countnoteshistogrambasedsupervision, title={Count The Notes: Histogram-Based Supervision for Automatic Music Transcription}, author={Jonathan Yaffe and Ben Maman and Meinard Müller and Amit H. Bermano}, year={2025}, eprint={2511.14250}, archivePrefix={arXiv}, primaryClass={cs.SD}, url={https://arxiv.org/abs/2511.14250}, } ``` """ ) if __name__ == "__main__": # Pre-load the default model to speed up first transcription logger.info("Pre-loading default model...") load_model("Yoni232/countem-musicnet") logger.info("Model pre-loaded. Starting Gradio interface...") # Launch the demo demo.launch( share=False, # Set to True to create a public link server_name="0.0.0.0", # Allow access from network server_port=7860, )