Update app.py

app.py

@@ -2,25 +2,19 @@ import random
 import argparse
 import os
 import glob
+import json
 import rtmidi
 import gradio as gr
 import numpy as np
-import MIDI
-import base64
-import io
-import soundfile as sf  # Placeholder for audio rendering
+import onnxruntime as rt
 from huggingface_hub import hf_hub_download
+import MIDI
 from midi_synthesizer import MidiSynthesizer
+from midi_tokenizer import MIDITokenizer
 
 MAX_SEED = np.iinfo(np.int32).max
 in_space = os.getenv("SYSTEM") == "spaces"
 
-SONG_DATA = {
-    "title": "Do You Believe in Love",
-    "progression": ["G", "D", "Em", "C"],
-    "lyrics": ["I was walking down a one-way street", "Just a-looking for someone to meet"]
-}
-
 class MIDIDeviceManager:
     def __init__(self):
         self.midiout = rtmidi.MidiOut()
@@ -51,63 +45,24 @@ 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 = {}  # midi_id: (file_path, midi_obj)
-        self.modified_files = []  # Stores (midi_base64, audio_base64) tuples
-        self.example_variations = {}  # midi_id: (midi_base64, audio_base64) for pre-generated examples
+        self.loaded_midi = {}
+        self.modified_files = []
         self.is_playing = False
-        self.instruments = self.random_instrument_set()
-        self.drum_beat = self.create_drum_beat()
-        self.starter_midi = self.create_starter_midi()
-        self.loaded_midi["starter"] = ("Starter MIDI", self.starter_midi)
-        self.preload_default_midi()
-        self.load_example_midis()  # Pre-generate example variations
-
-    def random_instrument_set(self):
-        instrument_pool = [0, 24, 32, 48]  # Piano, Guitar, Bass, Strings
-        return random.sample(instrument_pool, 4)
-
-    def create_drum_beat(self):
-        return [(36, 100, 0), (42, 80, 50), (38, 90, 100), (42, 80, 150)]  # Kick, hi-hat, snare, hi-hat
-
-    def create_starter_midi(self):
-        midi = MIDI.MIDIFile(5)
-        for i, inst in enumerate(self.instruments):
-            midi.addTrack()
-            midi.addProgramChange(i, 0, 0, inst)
-            for t in range(0, 400, 100):
-                note = random.randint(60, 84)
-                midi.addNote(i, 0, note, t, 100, 100)
-        midi.addTrack()
-        for note, vel, time in self.drum_beat:
-            midi.addNote(4, 9, note, time, 100, vel)
-        return midi
-
-    def preload_default_midi(self):
-        default_path = "default.mid"
-        if os.path.exists(default_path):
-            midi_id = "default"
-            midi = MIDI.load(default_path)
-            self.loaded_midi[midi_id] = (default_path, midi)
-            midi_data, audio_data = self.generate_variation(midi_id)
-            self.play_with_loop(midi_data)
-
-    def load_example_midis(self):
-        examples = {}
-        for file_path in glob.glob("*.mid") + glob.glob("*.midi"):
-            if file_path == "default.mid":
-                continue
-            midi_id = f"example_{len(examples)}"
-            midi = MIDI.load(file_path)
-            self.loaded_midi[midi_id] = (file_path, midi)
-            # Pre-generate variation for each example
-            midi_data, audio_data = self.generate_variation(midi_id)
-            examples[midi_id] = (midi_data, audio_data)
-        self.example_variations = examples
-        return examples
+        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):
+        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):
         midi = MIDI.load(file_path)
-        midi_id = f"midi_{len(self.loaded_midi) - len(self.example_variations) - 1}"
+        midi_id = f"midi_{len(self.loaded_midi)}"
         self.loaded_midi[midi_id] = (file_path, midi)
         return midi_id
 
@@ -122,13 +77,13 @@ class MIDIManager:
                     instruments.add(event.program)
         return notes, list(instruments)
 
-    def generate_variation(self, midi_id, length_factor=2, variation=0.3):
+    def generate_variation(self, midi_id, length_factor=10, variation=0.3):
         if midi_id not in self.loaded_midi:
-            return None, None
+            return None
         _, midi = self.loaded_midi[midi_id]
         notes, instruments = self.extract_notes_and_instruments(midi)
         new_notes = []
-        for _ in range(int(length_factor)):
+        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)))
@@ -147,75 +102,91 @@ class MIDIManager:
         midi_output = io.BytesIO()
         new_midi.writeFile(midi_output)
         midi_data = base64.b64encode(midi_output.getvalue()).decode('utf-8')
-        
-        temp_midi = 'temp.mid'
-        with open(temp_midi, 'wb') as f:
-            f.write(midi_output.getvalue())
-        audio_output = io.BytesIO()
-        self.synthesizer.play_midi(new_midi)
-        audio_data = None  # Placeholder; see Notes
-        if os.path.exists(temp_midi):
-            os.remove(temp_midi)
-        
-        self.modified_files.append((midi_data, audio_data))
-        return midi_data, audio_data
+        self.modified_files.append(midi_data)
+        return midi_data
 
-    def apply_synth_effect(self, midi_data, effect, intensity):
-        midi = MIDI.load(io.BytesIO(base64.b64decode(midi_data)))
-        if effect == "tempo":
-            factor = 1 + (intensity - 0.5) * 0.4
-            for track in midi.tracks:
-                for event in track.events:
-                    event.time = int(event.time * factor)
+    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
+        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)
+            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)
         midi_output = io.BytesIO()
-        midi.writeFile(midi_output)
+        MIDI.score2midi(new_midi, midi_output)
         midi_data = base64.b64encode(midi_output.getvalue()).decode('utf-8')
-        
-        temp_midi = 'temp.mid'
-        with open(temp_midi, 'wb') as f:
-            f.write(midi_output.getvalue())
-        audio_output = io.BytesIO()
-        self.synthesizer.play_midi(midi)
-        audio_data = None  # Placeholder
-        if os.path.exists(temp_midi):
-            os.remove(temp_midi)
-        
-        self.modified_files.append((midi_data, audio_data))
-        return midi_data, audio_data
+        self.modified_files.append(midi_data)
+        return midi_data
 
     def play_with_loop(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)
-        return "Stopped"
 
     def stop_playback(self):
         self.is_playing = False
-        return "Stopping..."
+        return "Playback stopped"
+
+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)
+
+def sample_top_p_k(probs, p, k, generator=None):
+    if generator is None:
+        generator = np.random
+    probs_idx = np.argsort(-probs, axis=-1)
+    probs_sort = np.take_along_axis(probs, probs_idx, -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])
 
 def create_download_list():
     html = "<h3>Downloads</h3><ul>"
-    for i, (midi_data, audio_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>'
-        if audio_data:
-            html += f' | <a href="data:audio/wav;base64,{audio_data}" download="audio_{i}.wav">Audio {i}</a>'
-        html += '</li>'
+    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>'
     html += "</ul>"
     return html
 
-def get_midi_choices():
-    return [(os.path.basename(path), midi_id) for midi_id, (path, _) in midi_processor.loaded_midi.items()]
-
-def get_example_choices():
-    return [(os.path.basename(path), midi_id) for midi_id, (path, _) in midi_processor.loaded_midi.items() if midi_id.startswith("example")]
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--port", type=int, default=7860)
     parser.add_argument("--share", action="store_true")
-    parser.add_argument("--batch", type=int, default=1)
     opt = parser.parse_args()
 
     midi_manager = MIDIDeviceManager()
@@ -227,96 +198,50 @@ if __name__ == "__main__":
         with gr.Tabs():
             # Tab 1: MIDI Prompt (Main Tab)
             with gr.Tab("MIDI Prompt"):
-                midi_files = gr.File(label="Upload MIDI Files", file_count="multiple")
-                loaded_display = gr.HTML(value="No files loaded")
-                output = gr.Audio(label="Generated Preview", type="bytes", autoplay=True)
+                midi_upload = gr.File(label="Upload MIDI File", file_count="multiple")
+                output = gr.Audio(label="Generated MIDI", type="bytes", autoplay=True)
+                status = gr.Textbox(label="Status", value="Ready", interactive=False)
                 
-                def load_and_generate(files):
-                    html = "<h3>Loaded Files</h3>"
+                def process_midi(files):
+                    if not files:
+                        return None, "No file uploaded"
                     midi_data = None
-                    for file in files or []:
+                    for file in files:
                         midi_id = midi_processor.load_midi(file.name)
-                        html += f"<div>{file.name} <button onclick=\"remove_midi('{midi_id}')\">X</button></div>"
-                        midi_data, _ = midi_processor.generate_variation(midi_id)
-                    return html, (io.BytesIO(base64.b64decode(midi_data)) if midi_data else None), get_midi_choices(), create_download_list()
+                        # 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_files.change(load_and_generate, inputs=[midi_files], 
-                                outputs=[loaded_display, output, gr.State(get_midi_choices()), "downloads"])
+                midi_upload.change(process_midi, inputs=[midi_upload], 
+                                 outputs=[output, status, "downloads"])
 
-            # Tab 2: Examples
-            with gr.Tab("Examples"):
-                example_select = gr.Dropdown(label="Select Example", choices=get_example_choices(), value=None)
-                example_output = gr.Audio(label="Example Preview", type="bytes", autoplay=True)
-                
-                def load_example(midi_id):
-                    if not midi_id or midi_id not in midi_processor.example_variations:
-                        return None
-                    midi_data, audio_data = midi_processor.example_variations[midi_id]
-                    midi_processor.play_with_loop(midi_data)
-                    return io.BytesIO(base64.b64decode(midi_data))
-                
-                example_select.change(load_example, inputs=[example_select], outputs=[example_output])
+            # Tab 2: Downloads
+            with gr.Tab("Downloads", elem_id="downloads"):
+                downloads = gr.HTML(value="No generated files yet")
 
-            # Tab 3: Generate & Perform
-            with gr.Tab("Generate & Perform"):
-                midi_select = gr.Dropdown(label="Select MIDI", choices=get_midi_choices(), value="starter")
-                length_factor = gr.Slider(1, 5, value=2, step=1, label="Length Factor")
-                variation = gr.Slider(0, 1, value=0.3, label="Variation")
-                generate_btn = gr.Button("Generate")
-                effect = gr.Radio(["tempo"], label="Synth Effect", value="tempo")
-                intensity = gr.Slider(0, 1, value=0.5, label="Effect Intensity")
-                apply_btn = gr.Button("Apply Effect")
+            # Tab 3: Devices
+            with gr.Tab("Devices"):
+                device_info = gr.Textbox(label="Connected MIDI Devices", value=midi_manager.get_device_info(), interactive=False)
+                refresh_btn = gr.Button("Refresh Devices")
                 stop_btn = gr.Button("Stop Playback")
-                output = gr.Audio(label="Preview", type="bytes", autoplay=True)
-                status = gr.Textbox(label="Status", value="Ready")
-                midi_device = gr.Dropdown(label="MIDI Output Device", choices=midi_manager.get_output_devices(), type="index")
-                tempo = gr.Slider(label="Tempo (BPM)", minimum=40, maximum=200, value=120, step=1)
-                device_info = gr.Textbox(label="Connected MIDI Devices", value=midi_manager.get_device_info(), readonly=True)
-                refresh_btn = gr.Button("🔄 Refresh MIDI Devices")
-
-                def update_dropdown(choices):
-                    return gr.update(choices=choices)
-                
-                gr.State(get_midi_choices()).change(update_dropdown, inputs=[gr.State()], outputs=[midi_select])
-                
-                def generate(midi_id, length, var):
-                    if not midi_id:
-                        return None, "Select a MIDI file", create_download_list()
-                    midi_data, audio_data = midi_processor.generate_variation(midi_id, length, var)
-                    midi_processor.play_with_loop(midi_data)
-                    return io.BytesIO(base64.b64decode(midi_data)), "Playing", create_download_list()
-                
-                def apply_effect(midi_data, fx, inten):
-                    if not midi_data:
-                        return None, "Generate a MIDI first", create_download_list()
-                    new_midi_data, audio_data = midi_processor.apply_synth_effect(midi_data.decode('utf-8'), fx, inten)
-                    midi_processor.play_with_loop(new_midi_data)
-                    return io.BytesIO(base64.b64decode(new_midi_data)), "Playing", create_download_list()
                 
                 def refresh_devices():
-                    return midi_manager.get_output_devices(), midi_manager.get_device_info()
+                    return midi_manager.get_device_info()
                 
-                generate_btn.click(generate, inputs=[midi_select, length_factor, variation], 
-                                outputs=[output, status, "downloads"])
-                apply_btn.click(apply_effect, inputs=[output, effect, intensity], 
-                                outputs=[output, status, "downloads"])
+                refresh_btn.click(refresh_devices, inputs=None, outputs=[device_info])
                 stop_btn.click(midi_processor.stop_playback, inputs=None, outputs=[status])
-                refresh_btn.click(refresh_devices, inputs=None, outputs=[midi_device, device_info])
-
-            # Tab 4: Downloads
-            with gr.Tab("Downloads", elem_id="downloads"):
-                downloads = gr.HTML(value=create_download_list())
 
         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/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>
         """)