import gradio as gr
import json
import rtmidi
import os
import argparse
import base64
import io
import numpy as np
from huggingface_hub import hf_hub_download
import onnxruntime as rt
import MIDI
from midi_synthesizer import MidiSynthesizer
from midi_tokenizer import MIDITokenizer
# Match the JavaScript constant
MIDI_OUTPUT_BATCH_SIZE = 4
class MIDIDeviceManager:
"""Manages MIDI input/output devices."""
def __init__(self):
self.midiout = rtmidi.MidiOut()
self.midiin = rtmidi.MidiIn()
def get_device_info(self):
"""Returns a string listing available MIDI devices."""
out_ports = self.midiout.get_ports() or ["No MIDI output devices"]
in_ports = self.midiin.get_ports() or ["No MIDI input devices"]
return f"Output Devices:\n{'\n'.join(out_ports)}\n\nInput Devices:\n{'\n'.join(in_ports)}"
def close(self):
"""Closes open MIDI ports."""
if self.midiout.is_port_open():
self.midiout.close_port()
if self.midiin.is_port_open():
self.midiin.close_port()
del self.midiout, self.midiin
class MIDIManager:
"""Handles MIDI processing, generation, and playback."""
def __init__(self):
# Load soundfont and models from Hugging Face
self.soundfont_path = hf_hub_download(repo_id="skytnt/midi-model", filename="soundfont.sf2")
self.synthesizer = MidiSynthesizer(self.soundfont_path)
self.tokenizer = self._load_tokenizer("skytnt/midi-model")
self.model_base = rt.InferenceSession(
hf_hub_download(repo_id="skytnt/midi-model", filename="onnx/model_base.onnx"),
providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
)
self.model_token = rt.InferenceSession(
hf_hub_download(repo_id="skytnt/midi-model", filename="onnx/model_token.onnx"),
providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
)
self.generated_files = []
def _load_tokenizer(self, repo_id):
"""Loads the MIDI tokenizer configuration."""
config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
with open(config_path, "r") as f:
config = json.load(f)
tokenizer = MIDITokenizer(config["tokenizer"]["version"])
tokenizer.set_optimise_midi(config["tokenizer"]["optimise_midi"])
return tokenizer
def load_midi(self, file_path):
"""Loads a MIDI file from the given path."""
return MIDI.load(file_path)
def generate_onnx(self, midi_data):
"""Generates a MIDI variation using ONNX models."""
mid_seq = self.tokenizer.tokenize(MIDI.midi2score(midi_data))
input_tensor = np.array([mid_seq], dtype=np.int64)
cur_len = input_tensor.shape[1]
max_len = 1024
while cur_len < max_len:
inputs = {"x": input_tensor[:, -1:]}
hidden = self.model_base.run(None, inputs)[0]
logits = self.model_token.run(None, {"hidden": hidden})[0]
probs = self._softmax(logits, axis=-1)
next_token = self._sample_top_p_k(probs, 0.98, 20)
input_tensor = np.concatenate([input_tensor, next_token], axis=1)
cur_len += 1
new_seq = input_tensor[0].tolist()
generated_midi = self.tokenizer.detokenize(new_seq)
# Store base64-encoded MIDI data for downloads
midi_bytes = MIDI.save(generated_midi)
self.generated_files.append(base64.b64encode(midi_bytes).decode('utf-8'))
return generated_midi
def play_midi(self, midi_data):
"""Renders MIDI data to audio bytes."""
midi_bytes = base64.b64decode(midi_data)
midi_file = MIDI.load(io.BytesIO(midi_bytes))
audio = io.BytesIO()
self.synthesizer.render_midi(midi_file, audio)
audio.seek(0)
return audio
@staticmethod
def _softmax(x, axis):
"""Computes softmax probabilities."""
exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
@staticmethod
def _sample_top_p_k(probs, p, k):
"""Samples a token using top-p and top-k sampling (simplified)."""
# Placeholder: replace with actual sampling logic if needed
return np.array([[np.random.choice(len(probs[0]))]])
def process_midi(files):
"""Processes uploaded MIDI files and yields updates for Gradio components."""
if not files:
yield [gr.update()] * (1 + 2 * MIDI_OUTPUT_BATCH_SIZE)
return
for idx, file in enumerate(files):
output_idx = idx % MIDI_OUTPUT_BATCH_SIZE
midi_data = midi_processor.load_midi(file.name)
generated_midi = midi_processor.generate_onnx(midi_data)
# Placeholder for MIDI events; in practice, extract from generated_midi
# Expected format: ["note", delta_time, track, channel, pitch, velocity, duration]
events = [
["note", 0, 0, 0, 60, 100, 1000], # Example event
# Add logic to convert generated_midi to events using tokenizer
]
# Prepare updates list: [js_msg, audio0, midi0, audio1, midi1, ...]
updates = [gr.update()] * (1 + 2 * MIDI_OUTPUT_BATCH_SIZE)
# Clear visualizer
updates[0] = js_msg.update(value=json.dumps([{"name": "visualizer_clear", "data": [output_idx, "v2"]}]))
yield updates
# Send MIDI events
updates[0] = js_msg.update(value=json.dumps([{"name": "visualizer_append", "data": [output_idx, events]}]))
yield updates
# Finalize visualizer and update audio/MIDI outputs
audio_update = midi_processor.play_midi(generated_midi)
midi_update = gr.File.update(value=generated_midi, label=f"Generated MIDI {output_idx}")
updates[0] = js_msg.update(value=json.dumps([{"name": "visualizer_end", "data": output_idx}]))
updates[1 + 2 * output_idx] = audio_update # Audio component
updates[2 + 2 * output_idx] = midi_update # MIDI file component
yield updates
# Final yield to ensure all components are in a stable state
yield [gr.update()] * (1 + 2 * MIDI_OUTPUT_BATCH_SIZE)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="MIDI Composer App")
parser.add_argument("--port", type=int, default=7860, help="Server port")
parser.add_argument("--share", action="store_true", help="Share the app publicly")
opt = parser.parse_args()
device_manager = MIDIDeviceManager()
midi_processor = MIDIManager()
with gr.Blocks(theme=gr.themes.Soft()) as app:
# Hidden textbox for sending messages to JS
js_msg = gr.Textbox(visible=False, elem_id="msg_receiver")
with gr.Tabs():
# MIDI Prompt Tab
with gr.Tab("MIDI Prompt"):
midi_upload = gr.File(label="Upload MIDI File(s)", file_count="multiple")
generate_btn = gr.Button("Generate")
status = gr.Textbox(label="Status", value="Ready", interactive=False)
# Outputs Tab
with gr.Tab("Outputs"):
output_audios = []
output_midis = []
for i in range(MIDI_OUTPUT_BATCH_SIZE):
with gr.Column():
gr.Markdown(f"## Output {i+1}")
gr.HTML(elem_id=f"midi_visualizer_container_{i}")
output_audio = gr.Audio(label="Generated Audio", type="bytes", autoplay=True, elem_id=f"midi_audio_{i}")
output_midi = gr.File(label="Generated MIDI", file_types=[".mid"])
output_audios.append(output_audio)
output_midis.append(output_midi)
# Devices Tab
with gr.Tab("Devices"):
device_info = gr.Textbox(label="Connected MIDI Devices", value=device_manager.get_device_info(), interactive=False)
refresh_btn = gr.Button("Refresh Devices")
refresh_btn.click(fn=lambda: device_manager.get_device_info(), outputs=[device_info])
# Define output components for event handling
outputs = [js_msg] + output_audios + output_midis
# Bind the generate button to the processing function
generate_btn.click(fn=process_midi, inputs=[midi_upload], outputs=outputs)
# Launch the app
app.launch(server_port=opt.port, share=opt.share, inbrowser=True)
device_manager.close()