Update app.py

app.py

@@ -1,58 +1,58 @@
-import random
 import argparse
+import base64
+import io
 import os
-import glob
-import json
-import rtmidi
-import gradio as gr
+import random
 import numpy as np
+import gradio as gr
+import rtmidi
 import onnxruntime as rt
 from huggingface_hub import hf_hub_download
 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"
+IN_SPACE = os.getenv("SYSTEM") == "spaces"
+MAX_LENGTH = 1024  # Maximum tokens for generation
 
+# MIDI Device Manager
 class MIDIDeviceManager:
     def __init__(self):
         self.midiout = rtmidi.MidiOut()
         self.midiin = rtmidi.MidiIn()
 
-    def get_output_devices(self):
-        return self.midiout.get_ports() or ["No MIDI output devices"]
-
-    def get_input_devices(self):
-        return self.midiin.get_ports() or ["No MIDI input devices"]
-
     def get_device_info(self):
-        out_devices = self.get_output_devices()
-        in_devices = self.get_input_devices()
-        out_info = "\n".join([f"Out Port {i}: {name}" for i, name in enumerate(out_devices)]) if out_devices else "No MIDI output devices detected"
-        in_info = "\n".join([f"In Port {i}: {name}" for i, name in enumerate(in_devices)]) if in_devices else "No MIDI input devices detected"
-        return f"Output Devices:\n{out_info}\n\nInput Devices:\n{in_info}"
+        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):
         if self.midiout.is_port_open():
             self.midiout.close_port()
         if self.midiin.is_port_open():
             self.midiin.close_port()
-        del self.midiout
-        del self.midiin
+        del self.midiout, self.midiin
 
+# MIDI Processor with ONNX Generation
 class MIDIManager:
     def __init__(self):
-        self.soundfont_path = hf_hub_download(repo_id="skytnt/midi-model", filename="soundfont.sf2")
-        self.synthesizer = MidiSynthesizer(self.soundfont_path)
-        self.loaded_midi = {}
-        self.modified_files = []
+        self.soundfont = hf_hub_download(repo_id="skytnt/midi-model", filename="soundfont.sf2")
+        self.synthesizer = MidiSynthesizer(self.soundfont)
+        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 = []  # Store base64-encoded MIDI data
         self.is_playing = False
-        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'])
 
-    def load_tokenizer(self, repo_id):
+    def _load_tokenizer(self, repo_id):
         config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
         with open(config_path, "r") as f:
             config = json.load(f)
@@ -61,190 +61,131 @@ class MIDIManager:
         return tokenizer
 
     def load_midi(self, file_path):
-        midi = MIDI.load(file_path)
-        midi_id = f"midi_{len(self.loaded_midi)}"
-        self.loaded_midi[midi_id] = (file_path, midi)
-        return midi_id
-
-    def extract_notes_and_instruments(self, midi):
-        notes = []
-        instruments = set()
-        for track in midi.tracks:
-            for event in track.events:
-                if event.type == 'note_on' and event.velocity > 0:
-                    notes.append((event.note, event.velocity, event.time))
-                if hasattr(event, 'program'):
-                    instruments.add(event.program)
-        return notes, list(instruments)
-
-    def generate_variation(self, midi_id, length_factor=10, variation=0.3):
-        if midi_id not in self.loaded_midi:
-            return None
-        _, midi = self.loaded_midi[midi_id]
-        notes, instruments = self.extract_notes_and_instruments(midi)
-        new_notes = []
-        for _ in range(int(length_factor)):  # Max length: 10x repetition
-            for note, vel, time in notes:
-                if random.random() < variation:
-                    new_note = min(127, max(0, note + random.randint(-2, 2)))
-                    new_vel = min(127, max(0, vel + random.randint(-10, 10)))
-                    new_notes.append((new_note, new_vel, time))
-                else:
-                    new_notes.append((note, vel, time))
-
-        new_midi = MIDI.MIDIFile(len(instruments) or 1)
-        for i, inst in enumerate(instruments or [0]):
-            new_midi.addTrack()
-            new_midi.addProgramChange(i, 0, 0, inst)
-            for note, vel, time in new_notes:
-                new_midi.addNote(i, 0, note, time, 100, vel)
-        
-        midi_output = io.BytesIO()
-        new_midi.writeFile(midi_output)
-        midi_data = base64.b64encode(midi_output.getvalue()).decode('utf-8')
-        self.modified_files.append(midi_data)
-        return midi_data
-
-    def generate_onnx(self, midi_id, max_len=1024, temp=1.0, top_p=0.98, top_k=20):
-        if midi_id not in self.loaded_midi:
-            return None
-        _, mid = self.loaded_midi[midi_id]
-        mid_seq = self.tokenizer.tokenize(MIDI.midi2score(mid))
-        mid = np.asarray([mid_seq], dtype=np.int64)
-        generator = np.random.RandomState(random.randint(0, MAX_SEED))
-        
-        # Simplified ONNX generation from app_onnx.py
-        input_tensor = mid
+        try:
+            return MIDI.load(file_path)
+        except Exception as e:
+            raise ValueError(f"Failed to load MIDI file: {e}")
+
+    def generate_variation(self, midi_data, temp=1.0, top_p=0.98, top_k=20):
+        # Tokenize input MIDI
+        mid_seq = self.tokenizer.tokenize(MIDI.midi2score(midi_data))
+        input_tensor = np.array([mid_seq], dtype=np.int64)
         cur_len = input_tensor.shape[1]
-        model = [self.model_base, self.model_token, self.tokenizer]
-        
-        while cur_len < max_len:
-            inputs = {"x": rt.OrtValue.ortvalue_from_numpy(input_tensor[:, -1:], device_type="cuda")}
-            outputs = {"hidden": rt.OrtValue.ortvalue_from_shape_and_type((1, 1, 1024), np.float32, device_type="cuda")}
-            io_binding = model[0].io_binding()
-            for name, val in inputs.items():
-                io_binding.bind_ortvalue_input(name, val)
-            for name in outputs:
-                io_binding.bind_ortvalue_output(name, outputs[name])
-            model[0].run_with_iobinding(io_binding)
-            hidden = outputs["hidden"].numpy()[:, -1:]
-
-            logits = model[1].run(None, {"hidden": hidden})[0]
-            scores = softmax(logits / temp, -1)
-            next_token = sample_top_p_k(scores, top_p, top_k, generator)
+        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)
             input_tensor = np.concatenate([input_tensor, next_token], axis=1)
             cur_len += 1
 
-        mid_seq = input_tensor.tolist()[0]
-        new_midi = self.tokenizer.detokenize(mid_seq)
+        # 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.modified_files.append(midi_data)
+        self.generated_files.append(midi_data)
         return midi_data
 
-    def play_with_loop(self, midi_data):
+    def play_midi(self, midi_data):
         self.is_playing = True
-        midi_file = MIDI.load(io.BytesIO(base64.b64decode(midi_data)))
-        while self.is_playing:
-            self.synthesizer.play_midi(midi_file)
-
-    def stop_playback(self):
+        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
+
+    def stop(self):
         self.is_playing = False
-        return "Playback stopped"
 
+# Helper Functions
 def softmax(x, axis):
-    x_max = np.max(x, axis=axis, keepdims=True)
-    exp_x_shifted = np.exp(x - x_max)
-    return exp_x_shifted / np.sum(exp_x_shifted, axis=axis, keepdims=True)
+    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=None):
-    if generator is None:
-        generator = np.random
+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, -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
-    mask = np.zeros(probs_sort.shape[-1])
-    mask[:k] = 1
-    probs_sort *= mask
-    probs_sort /= np.sum(probs_sort, axis=-1, keepdims=True)
-    shape = probs_sort.shape
-    probs_sort_flat = probs_sort.reshape(-1, shape[-1])
-    probs_idx_flat = probs_idx.reshape(-1, shape[-1])
-    next_token = np.stack([generator.choice(idxs, p=pvals) for pvals, idxs in zip(probs_sort_flat, probs_idx_flat)])
-    return next_token.reshape(*shape[:-1])
+    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):
+    if not files:
+        return None, "No file uploaded", ""
+    file = files[0]  # Process first file
+    try:
+        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():
-    html = "<h3>Downloads</h3><ul>"
-    for i, midi_data in enumerate(midi_processor.modified_files):
-        html += f'<li><a href="data:audio/midi;base64,{midi_data}" download="midi_{i}.mid">MIDI {i}</a></li>'
+    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
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
+    parser = argparse.ArgumentParser(description="MIDI Composer with ONNX Generation")
     parser.add_argument("--port", type=int, default=7860)
     parser.add_argument("--share", action="store_true")
-    opt = parser.parse_args()
+    args = parser.parse_args()
 
-    midi_manager = MIDIDeviceManager()
+    device_manager = MIDIDeviceManager()
     midi_processor = MIDIManager()
 
-    with gr.Blocks(theme=gr.themes.Soft()) as app:
-        gr.Markdown("<h1>🎵 MIDI Composer 🎵</h1>")
-        
+    with gr.Blocks(title="MIDI Composer", theme=gr.themes.Soft()) as app:
+        gr.Markdown("# 🎵 MIDI Composer 🎵")
+
         with gr.Tabs():
-            # Tab 1: MIDI Prompt (Main Tab)
+            # MIDI Prompt Tab
             with gr.Tab("MIDI Prompt"):
-                midi_upload = gr.File(label="Upload MIDI File", file_count="multiple")
-                output = gr.Audio(label="Generated MIDI", type="bytes", autoplay=True)
+                midi_upload = gr.File(label="Upload MIDI File", file_types=[".mid", ".midi"])
+                audio_output = gr.Audio(label="Generated MIDI", type="bytes", autoplay=True)
                 status = gr.Textbox(label="Status", value="Ready", interactive=False)
-                
-                def process_midi(files):
-                    if not files:
-                        return None, "No file uploaded"
-                    midi_data = None
-                    for file in files:
-                        midi_id = midi_processor.load_midi(file.name)
-                        # Use ONNX generation for advanced synthesis
-                        midi_data = midi_processor.generate_onnx(midi_id, max_len=1024)
-                        midi_processor.play_with_loop(midi_data)
-                    return io.BytesIO(base64.b64decode(midi_data)), "Playing", create_download_list()
-                
-                midi_upload.change(process_midi, inputs=[midi_upload], 
-                                 outputs=[output, status, "downloads"])
-
-            # Tab 2: Downloads
+                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"):
-                downloads = gr.HTML(value="No generated files yet")
+                gr.HTML(value=create_download_list())
 
-            # Tab 3: Devices
+            # Devices Tab
             with gr.Tab("Devices"):
-                device_info = gr.Textbox(label="Connected MIDI Devices", value=midi_manager.get_device_info(), interactive=False)
+                device_info = gr.Textbox(label="MIDI Devices", value=device_manager.get_device_info(), interactive=False)
                 refresh_btn = gr.Button("Refresh Devices")
                 stop_btn = gr.Button("Stop Playback")
-                
-                def refresh_devices():
-                    return midi_manager.get_device_info()
-                
-                refresh_btn.click(refresh_devices, inputs=None, outputs=[device_info])
-                stop_btn.click(midi_processor.stop_playback, inputs=None, outputs=[status])
-
-        gr.Markdown("""
-        <div style='text-align: center; margin-top: 20px;'>
-            <img src='https://huggingface.co/front/assets/huggingface_logo-noborder.svg' alt='Hugging Face Logo' style='width: 50px;'><br>
-            <strong>Hugging Face</strong><br>
-            <a href='https://huggingface.co/models'>Models</a> |
-            <a href='https://huggingface.co/datasets'>Datasets</a> |
-            <a href='https://huggingface.co/spaces'>Spaces</a> |
-            <a href='https://huggingface.co/posts'>Posts</a> |
-            <a href='https://huggingface.co/docs'>Docs</a> |
-            <a href='https://huggingface.co/enterprise'>Enterprise</a> |
-            <a href='https://huggingface.co/pricing'>Pricing</a>
-        </div>
-        """)
-
-    app.queue().launch(server_port=opt.port, share=opt.share, inbrowser=True)
-    midi_manager.close()
+                refresh_btn.click(refresh_devices, outputs=[device_info])
+                stop_btn.click(stop_playback, outputs=[status])
+
+    app.launch(server_port=args.port, share=args.share, inbrowser=True)
+    device_manager.close()