import os

os.system("python3 -m pip install -e .")

import gradio as gr

import note_seq

from inferencemodel import InferenceModel
from utils import upload_audio

SAMPLE_RATE = 16000
SF2_PATH = 'SGM-v2.01-Sal-Guit-Bass-V1.3.sf2'

# Start inference model
inference_model = InferenceModel('/home/user/app/checkpoints/mt3/', 'mt3')

def inference(audio):
  with open(audio, 'rb') as fd:
      contents = fd.read()

  audio = upload_audio(contents,sample_rate=16000)

  est_ns = inference_model(audio)
  
  note_seq.sequence_proto_to_midi_file(est_ns, './transcribed.mid')
 
  return './transcribed.mid'
  
title = "MT3"
description = "Gradio demo for MT3: Multi-Task Multitrack Music Transcription. To use it, simply upload your audio file, or click one of the examples to load them. Read more at the links below."

article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2111.03017' target='_blank'>MT3: Multi-Task Multitrack Music Transcription</a> | <a href='https://github.com/magenta/mt3' target='_blank'>Github Repo</a></p>"

examples=[['download.wav']]

gr.Interface(
    inference, 
    gr.inputs.Audio(type="filepath", label="Input"), 
    [gr.outputs.File(label="Output")],
    title=title,
    description=description,
    article=article,
    examples=examples,
    ).launch().queue()