import gradio as gr
from gradio_midibridge import MIDIBridge
from processor import Processor

model_location_repo = "adricl/midi_single_instrument_mistral_transformer"
model_tokenizer_file = "HuggingFace_Mistral_Transformer_Single_Instrument.json"

# Create the Gradio interface
with gr.Blocks(title="MIDI Jam Session") as demo:
    gr.Markdown("""
    # 🎹 MIDI Jam Session
    
    This demo shows the MIDI Jam Session in action:
    
    1. **Select your MIDI input device** - Connect your MIDI keyboard or controller
    2. **Select your MIDI output device** - Choose where to send the processed MIDI
    3. **Play some notes** - The component will record your input
    4. **Wait 5 seconds** - After inactivity, the recording is sent for processing
    5. **""" + model_location_repo + """** - The MIDI is processed using a transformer model to generate new content based on your input
    6. **Listen to the result** - The generated MIDI content is played back
    
    """)
    
    midi_session = MIDIBridge(
        label="MIDI Jam Session",
        bpm=120,
        interactive=True,
    )

    with gr.Accordion("Generation Settings", open=False):
        max_new_tokens = gr.Slider(
            minimum=0, maximum=10000, value=2000, step=1,
            label="Max New Tokens",
        )
        temperature = gr.Slider(
            minimum=0.0, maximum=1.0, value=0.9, step=0.01,
            label="Temperature",
        )
        top_p = gr.Slider(
            minimum=0.0, maximum=1.0, value=0.95, step=0.01,
            label="Top P",
        )
        do_sample = gr.Checkbox(value=True, label="Do Sample")

    processor = Processor(model_location_repo=model_location_repo, model_tokenizer_file=model_tokenizer_file)
    
    # Connect the processing function
    midi_session.change(
        fn=processor.transpose_midi,
        inputs=[midi_session, max_new_tokens, temperature, top_p, do_sample],
        outputs=midi_session,
    )



if __name__ == "__main__":
    demo.launch()