import gradio as gr
from musc.model import PretrainedModel
from json import load as json_load
from mido import MidiFile,MidiTrack
from os import remove as os_remove
Model = PretrainedModel(json_load(open("violin.json")),"violin_model.pt").to("cpu")
def merge_violin_tracks(input_midi, output_midi):
    mid = MidiFile(input_midi)
    new_mid = MidiFile(ticks_per_beat=mid.ticks_per_beat)
    new_track = MidiTrack()
    new_mid.tracks.append(new_track)
    events = []
    for track in mid.tracks:
        current_time = 0
        for msg in track:
            current_time += msg.time  
            events.append((current_time, msg))
    events.sort(key=lambda x: x[0])
    last_time = 0
    for event_time, msg in events:
        delta_time = event_time - last_time
        new_track.append(msg.copy(time=delta_time))
        last_time = event_time  
    for track in mid.tracks:
        for msg in track:
            if msg.type == 'set_tempo':
                new_track.insert(0, msg)  
    new_mid.save(output_midi)

def transcribe_and_generate_midi(music_file_path, model=Model, batch_size=32):
    model.transcribe(music_file_path, batch_size=batch_size).write("output.mid")
    merge_violin_tracks("output.mid","output.mid")
    os_remove(music_file_path)
    return "output.mid"

gr.Interface(
    fn=transcribe_and_generate_midi, 
    inputs=gr.Audio(label="Upload your Audio file",type="filepath"),
    outputs=gr.File(label="Download MIDI file"),
    title="Audio2Violin",
    description="Upload a Audio file, and it will be transcribed into Violin MIDI format."
).launch()