mtg-upf/audio-difficulty

1import gradio as gr2from get_difficulty import predict_difficulty3import tempfile4import os5from pydub import AudioSegment6import yt_dlp7import mimetypes8from huggingface_hub import hf_hub_download9import torch10import sys11import io12import spaces13 14REPO_ID = "pramoneda/audio"15CACHE_BASE = "models"16 17def download_model_checkpoints(model_name: str, num_checkpoints: int = 5):18    cache_dir = os.path.join(CACHE_BASE, model_name)19    os.makedirs(cache_dir, exist_ok=True)20    for checkpoint_id in range(num_checkpoints):21        filename = f"{model_name}/checkpoint_{checkpoint_id}.pth"22        local_path = os.path.join(cache_dir, f"checkpoint_{checkpoint_id}.pth")23        if not os.path.exists(local_path):24            path = hf_hub_download(repo_id=REPO_ID, filename=filename, cache_dir=cache_dir)25            if path != local_path:26                import shutil27                shutil.copy(path, local_path)28 29def download_youtube_audio(url, cookie_file=None):30    output_path = "yt_audio.%(ext)s"31    ydl_opts = {32        "format": "bestaudio/best",33        "outtmpl": output_path,34        "postprocessors": [{35            "key": "FFmpegExtractAudio",36            "preferredcodec": "mp3",37            "preferredquality": "192",38        }],39        "quiet": True,40        "no_warnings": True41    }42    if cookie_file:43        ydl_opts["cookiefile"] = cookie_file  # <-- usa el archivo de cookies44 45    with yt_dlp.YoutubeDL(ydl_opts) as ydl:46        ydl.download([url])47 48    return "yt_audio.mp3"49 50def convert_to_mp3(input_path):51    audio = AudioSegment.from_file(input_path)52    temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")53    audio.export(temp_audio.name, format="mp3")54    return temp_audio.name55 56@spaces.GPU57def process_input(input_file, youtube_url, cookie_file):58    # captura consola59    captured_output = io.StringIO()60    sys.stdout = captured_output61 62    # procesa audio/video63    if youtube_url:64        audio_path = download_youtube_audio(youtube_url, cookie_file)65        mp3_path = audio_path66    elif input_file:67        mime_type, _ = mimetypes.guess_type(input_file)68        audio_path = convert_to_mp3(input_file)69        mp3_path = audio_path70    else:71        sys.stdout = sys.__stdout__72        return "No audio or video provided.", None, None, None, ""73 74    # descarga checkpoints75    for model in ["audio_midi_cqt5_ps_v5", "audio_midi_pianoroll_ps_5_v4", "audio_midi_multi_ps_v5"]:76        download_model_checkpoints(model)77 78    # predicciones79    diff_cqt = predict_difficulty(audio_path, model_name="audio_midi_cqt5_ps_v5", rep="cqt5")80    diff_pr = predict_difficulty(audio_path, model_name="audio_midi_pianoroll_ps_5_v4", rep="pianoroll5")81    diff_multi = predict_difficulty(audio_path, model_name="audio_midi_multi_ps_v5", rep="multimodal5")82 83    sys.stdout = sys.__stdout__84    log_output = captured_output.getvalue()85 86    midi_path = "temp.mid"87    if not os.path.exists(midi_path):88        return "MIDI not generated.", None, None, None, log_output89 90    difficulty_text = (91        f"CQT difficulty: {diff_cqt}\n"92        f"Pianoroll difficulty: {diff_pr}\n"93        f"Multimodal difficulty: {diff_multi}"94    )95 96    return difficulty_text, midi_path, midi_path, mp3_path, log_output97 98demo = gr.Interface(99    fn=process_input,100    inputs=[101        gr.File(label="Upload MP3 or MP4", type="filepath"),102        gr.Textbox(label="YouTube URL"),103        gr.File(label="Upload cookies.txt (optional)", file_types=["text"], type="filepath")104    ],105    outputs=[106        gr.Textbox(label="Difficulty predictions"),107        gr.File(label="Generated MIDI"),108        gr.Audio(label="MIDI Playback", type="filepath"),109        gr.Audio(label="Extracted MP3 Preview", type="filepath"),110        gr.Textbox(label="Console Output")111    ],112    title="Music Difficulty Estimator",113    description=(114        "Upload an MP3/MP4 or provide a YouTube URL. "115        "If you want to predict the difficulty directly from youtube, export your YouTube cookies as a Netscape-format file "116        "and upload it here. Then the app can download and process the audio."117        "Related publication: [IEEE TASLP paper](https://ieeexplore.ieee.org/document/10878288)"118    )119)120 121if __name__ == "__main__":122    demo.launch(debug=True)123