Update app.py

app.py

@@ -1,45 +1,46 @@
+import gradio as gr
+import json
+import rtmidi
+import os
 import argparse
 import base64
 import io
-import os
-import random
 import numpy as np
-import gradio as gr
-import rtmidi
-import onnxruntime as rt
 from huggingface_hub import hf_hub_download
+import onnxruntime as rt
 import MIDI
 from midi_synthesizer import MidiSynthesizer
 from midi_tokenizer import MIDITokenizer
 
-# Constants
-MAX_SEED = np.iinfo(np.int32).max
-IN_SPACE = os.getenv("SYSTEM") == "spaces"
-MAX_LENGTH = 1024  # Maximum tokens for generation
+# Match the JavaScript constant
+MIDI_OUTPUT_BATCH_SIZE = 4
 
-# MIDI Device Manager
 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
 
-# MIDI Processor with ONNX Generation
 class MIDIManager:
+    """Handles MIDI processing, generation, and playback."""
     def __init__(self):
-        self.soundfont = hf_hub_download(repo_id="skytnt/midi-model", filename="soundfont.sf2")
-        self.synthesizer = MidiSynthesizer(self.soundfont)
+        # 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"),
@@ -49,10 +50,10 @@ class MIDIManager:
             hf_hub_download(repo_id="skytnt/midi-model", filename="onnx/model_token.onnx"),
             providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
         )
-        self.generated_files = []  # Store base64-encoded MIDI data
-        self.is_playing = False
+        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)
@@ -61,39 +62,32 @@ class MIDIManager:
         return tokenizer
 
     def load_midi(self, file_path):
-        try:
-            return MIDI.load(file_path)
-        except Exception as e:
-            raise ValueError(f"Failed to load MIDI file: {e}")
+        """Loads a MIDI file from the given path."""
+        return MIDI.load(file_path)
 
-    def generate_variation(self, midi_data, temp=1.0, top_p=0.98, top_k=20):
-        # Tokenize input MIDI
+    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]
-        generator = np.random.RandomState(random.randint(0, MAX_SEED))
-
-        # Generate up to MAX_LENGTH
-        while cur_len < MAX_LENGTH:
-            inputs = {"x": input_tensor[:, -1:]}  # Last token
-            hidden = self.model_base.run(None, inputs)[0]  # Base model output
-            logits = self.model_token.run(None, {"hidden": hidden})[0]  # Token model output
-            probs = softmax(logits / temp, axis=-1)
-            next_token = sample_top_p_k(probs, top_p, top_k, generator)
+        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
-
-        # Detokenize and save as MIDI
         new_seq = input_tensor[0].tolist()
-        new_midi = self.tokenizer.detokenize(new_seq)
-        midi_output = io.BytesIO()
-        MIDI.score2midi(new_midi, midi_output)
-        midi_data = base64.b64encode(midi_output.getvalue()).decode('utf-8')
-        self.generated_files.append(midi_data)
-        return midi_data
+        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):
-        self.is_playing = True
+        """Renders MIDI data to audio bytes."""
         midi_bytes = base64.b64decode(midi_data)
         midi_file = MIDI.load(io.BytesIO(midi_bytes))
         audio = io.BytesIO()
@@ -101,91 +95,103 @@ class MIDIManager:
         audio.seek(0)
         return audio
 
-    def stop(self):
-        self.is_playing = False
-
-# Helper Functions
-def softmax(x, axis):
-    exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
-    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
-
-def sample_top_p_k(probs, p, k, generator):
-    probs_idx = np.argsort(-probs, axis=-1)
-    probs_sort = np.take_along_axis(probs, probs_idx, axis=-1)
-    probs_sum = np.cumsum(probs_sort, axis=-1)
-    mask = probs_sum - probs_sort > p
-    probs_sort[mask] = 0.0
-    probs_sort[:, k:] = 0.0  # Top-k filtering
-    probs_sort /= probs_sort.sum(axis=-1, keepdims=True)
-    next_token = generator.choice(probs.shape[-1], p=probs_sort[0])
-    return np.array([[next_token]])
-
-# UI Functions
-def process_midi_upload(files):
+    @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:
-        return None, "No file uploaded", ""
-    file = files[0]  # Process first file
-    try:
+        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_variation(midi_data)
-        audio = midi_processor.play_midi(generated_midi)
-        download_html = create_download_list()
-        return audio, "Generated and playing", download_html
-    except Exception as e:
-        return None, f"Error: {e}", ""
-
-def create_download_list():
-    if not midi_processor.generated_files:
-        return "<p>No generated files yet.</p>"
-    html = "<h3>Generated MIDI Files</h3><ul>"
-    for i, midi_data in enumerate(midi_processor.generated_files):
-        html += f'<li><a href="data:audio/midi;base64,{midi_data}" download="generated_{i}.mid">Download MIDI {i}</a></li>'
-    html += "</ul>"
-    return html
-
-def refresh_devices():
-    return device_manager.get_device_info()
-
-def stop_playback():
-    midi_processor.stop()
-    return "Playback stopped"
-
-# Main Application
+        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 with ONNX Generation")
-    parser.add_argument("--port", type=int, default=7860)
-    parser.add_argument("--share", action="store_true")
-    args = parser.parse_args()
+    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(title="MIDI Composer", theme=gr.themes.Soft()) as app:
-        gr.Markdown("# 🎵 MIDI Composer 🎵")
-
+    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", file_types=[".mid", ".midi"])
-                audio_output = gr.Audio(label="Generated MIDI", type="bytes", autoplay=True)
+                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)
-                midi_upload.change(
-                    process_midi_upload,
-                    inputs=[midi_upload],
-                    outputs=[audio_output, status, gr.HTML(elem_id="downloads")]
-                )
 
-            # Downloads Tab
-            with gr.Tab("Downloads", elem_id="downloads"):
-                gr.HTML(value=create_download_list())
+            # 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="MIDI Devices", value=device_manager.get_device_info(), interactive=False)
+                device_info = gr.Textbox(label="Connected MIDI Devices", value=device_manager.get_device_info(), interactive=False)
                 refresh_btn = gr.Button("Refresh Devices")
-                stop_btn = gr.Button("Stop Playback")
-                refresh_btn.click(refresh_devices, outputs=[device_info])
-                stop_btn.click(stop_playback, outputs=[status])
+                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)
 
-    app.launch(server_port=args.port, share=args.share, inbrowser=True)
+    # Launch the app
+    app.launch(server_port=opt.port, share=opt.share, inbrowser=True)
     device_manager.close()