import gradio as gr
from musiclang_predict import MusicLangPredictor  # Assuming this is the correct import
from midi2audio import FluidSynth
import os


def musiclang():
    nb_tokens = 1024
    temperature = 0.9
    top_p = 1.0
    seed = 16

    # Initialize the MusicLangPredictor
    ml = MusicLangPredictor('musiclang/musiclang-v2')
    # Generate the score
    score = ml.predict(nb_tokens=nb_tokens, temperature=temperature, topp=top_p, rng_seed=seed)

    # Save the score as a MIDI file
    midi_path = 'test.mid'
    score.to_midi(midi_path)

    # Convert MIDI to WAV then WAV to MP3
    wav_path = 'result.wav'
    mp3_path = 'result.mp3'
    FluidSynth("/usr/share/sounds/sf2/FluidR3_GM.sf2").midi_to_audio(midi_path, wav_path)

    os.system(f'ffmpeg -i {wav_path} -acodec libmp3lame -y -loglevel quiet -stats {mp3_path}')

    # Return the path to the MP3 for Gradio to display
    return mp3_path


# Gradio interface
iface = gr.Interface(
    fn=musiclang,
    inputs=None,
    outputs=gr.Audio(label="Generated Music"),
    title="Music Generation with MusicLang",
    description="Click the button to generate music."
)

iface.launch()