import gradio as gr
import torch
from unidecode import unidecode
from musc.model import PretrainedModel

# Function to transcribe the WAV file and generate the MIDI file
def transcribe_and_generate_midi(wav_file, model, batch_size=32, postprocessing='spotify'):
    # Save the uploaded WAV file to disk
    wav_file_path = "temp_audio.wav"
    with open(wav_file_path, "wb") as f:
        f.write(wav_file.read())
    
    # Transcribe the WAV file and generate MIDI
    midi, _, title = model.transcribe_wav(wav_file_path, batch_size=batch_size, postprocessing=postprocessing)

    # Write the MIDI file
    midi_file_name = unidecode(title) + '.mid'
    midi.write(midi_file_name)

    return midi_file_name

# Set up the Pretrained Model
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = PretrainedModel(instrument='violin').to(device)

# Gradio interface function
def convert_to_midi(wav_file):
    try:
        midi_file_name = transcribe_and_generate_midi(wav_file, model)
        return midi_file_name
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio interface
iface = gr.Interface(
    fn=convert_to_midi, 
    inputs=gr.inputs.File(label="Upload your WAV file"), 
    outputs=gr.outputs.File(label="Download MIDI file"),
    live=False,
    title="Violin to MIDI Converter",
    description="Upload a WAV file of a violin performance, and it will be transcribed into MIDI format."
)

# Launch the Gradio app
iface.launch()