Update app.py

app.py

@@ -11,7 +11,6 @@ import gradio as gr
 import numpy as np
 import torch
 import torch.nn.functional as F
-import tqdm
 from huggingface_hub import hf_hub_download
 from transformers import DynamicCache
 
@@ -22,7 +21,6 @@ from midi_synthesizer import MidiSynthesizer
 MAX_SEED = np.iinfo(np.int32).max
 in_space = os.getenv("SYSTEM") == "spaces"
 
-
 @torch.inference_mode()
 def generate(model: MIDIModel, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98, top_k=20,
              disable_patch_change=False, disable_control_change=False, disable_channels=None, generator=None):
@@ -118,15 +116,12 @@ def generate(model: MIDIModel, prompt=None, batch_size=1, max_len=512, temp=1.0,
             if all(end):
                 break
 
-
 def create_msg(name, data):
     return {"name": name, "data": data}
 
-
 def send_msgs(msgs):
     return json.dumps(msgs)
 
-
 def get_duration(model_name, tab, mid_seq, continuation_state, continuation_select, instruments, drum_kit, bpm,
                  time_sig, key_sig, mid, midi_events, reduce_cc_st, remap_track_channel, add_default_instr,
                  remove_empty_channels, seed, seed_rand, gen_events, temp, top_p, top_k, allow_cc):
@@ -135,7 +130,6 @@ def get_duration(model_name, tab, mid_seq, continuation_state, continuation_sele
         t = gen_events // 14
     return t + 5
 
-
 @spaces.GPU(duration=get_duration)
 def run(model_name, tab, mid_seq, continuation_state, continuation_select, instruments, drum_kit, bpm, time_sig,
         key_sig, mid, midi_events, reduce_cc_st, remap_track_channel, add_default_instr, remove_empty_channels,
@@ -246,7 +240,6 @@ def run(model_name, tab, mid_seq, continuation_state, continuation_select, instr
             t = time.time()
     yield mid_seq, continuation_state, seed, send_msgs([])
 
-
 def finish_run(model_name, mid_seq):
     if mid_seq is None:
         outputs = [None] * OUTPUT_BATCH_SIZE
@@ -267,10 +260,11 @@ def finish_run(model_name, mid_seq):
                      create_msg("visualizer_end", i)]
     return *outputs, send_msgs(end_msgs)
 
-
 def synthesis_task(mid):
     return synthesizer.synthesis(MIDI.score2opus(mid))
 
+
+        
 def render_audio(model_name, mid_seq, should_render_audio):
     if (not should_render_audio) or mid_seq is None:
         outputs = [None] * OUTPUT_BATCH_SIZE
@@ -351,7 +345,67 @@ drum_kits2number = {v: k for k, v in number2drum_kits.items()}
 key_signatures = ['C♭', 'A♭m', 'G♭', 'E♭m', 'D♭', 'B♭m', 'A♭', 'Fm', 'E♭', 'Cm', 'B♭', 'Gm', 'F', 'Dm',
                   'C', 'Am', 'G', 'Em', 'D', 'Bm', 'A', 'F♯m', 'E', 'C♯m', 'B', 'G♯m', 'F♯', 'D♯m', 'C♯', 'A♯m']
 
-if __name__ == "__main__":
+
+
+
+
+
+mid = tokenizer.detokenize(mid_seq[i])
+        audio_future = thread_pool.submit(synthesis_task, mid)
+        audio_futures.append(audio_future)
+    for future in audio_futures:
+        outputs.append((44100, future.result()))
+    if OUTPUT_BATCH_SIZE == 1:
+        return outputs[0]
+    return tuple(outputs)
+
+def undo_continuation(model_name, mid_seq, continuation_state):
+    if mid_seq is None or len(continuation_state) < 2:
+        return mid_seq, continuation_state, send_msgs([])
+    tokenizer = models[model_name].tokenizer
+    if isinstance(continuation_state[-1], list):
+        mid_seq = continuation_state[-1]
+    else:
+        mid_seq = [ms[:continuation_state[-1]] for ms in mid_seq]
+    continuation_state = continuation_state[:-1]
+    end_msgs = [create_msg("progress", [0, 0])]
+    for i in range(OUTPUT_BATCH_SIZE):
+        events = [tokenizer.tokens2event(tokens) for tokens in mid_seq[i]]
+        end_msgs += [create_msg("visualizer_clear", [i, tokenizer.version]),
+                     create_msg("visualizer_append", [i, events]),
+                     create_msg("visualizer_end", i)]
+    return mid_seq, continuation_state, send_msgs(end_msgs)
+
+def create_arpeggio_events(chord, pattern, duration=480):
+    events = []
+    notes = {
+        'C': [60, 64, 67],
+        'D': [62, 66, 69],
+        'Am': [57, 60, 64],
+        'G': [55, 59, 62]
+    }
+    
+    for step in pattern:
+        note = notes[chord][step]
+        events.extend([
+            ['note_on', 0, 0, 0, 0, note, 80],
+            ['note_off', duration, 0, 0, 0, note, 0]
+        ])
+    
+    return events
+
+def add_arpeggio_sequence(tokenizer, mid_seq, sequence, pattern):
+    events = []
+    for chord in sequence:
+        events.extend(create_arpeggio_events(chord, pattern))
+    
+    tokens = [tokenizer.event2tokens(event) for event in events]
+    mid_seq[0].extend(tokens)
+    return mid_seq
+
+
+
+    if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--share", action="store_true", default=False, help="share gradio app")
     parser.add_argument("--port", type=int, default=7860, help="gradio server port")
@@ -402,17 +456,7 @@ if __name__ == "__main__":
     load_javascript()
     app = gr.Blocks(theme=gr.themes.Soft())
     with app:
-        gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Midi Composer</h1>")
-        gr.Markdown("![Visitors](https://api.visitorbadge.io/api/visitors?path=skytnt.midi-composer&style=flat)\n\n"
-                    "Midi event transformer for symbolic music generation\n\n"
-                    "Demo for [SkyTNT/midi-model](https://github.com/SkyTNT/midi-model)\n\n"
-                    "[Open In Colab]"
-                    "(https://colab.research.google.com/github/SkyTNT/midi-model/blob/main/demo.ipynb)"
-                    " or [download windows app](https://github.com/SkyTNT/midi-model/releases)"
-                    " for unlimited generation\n\n"
-                    "**Update v1.3**: MIDITokenizerV2 and new MidiVisualizer\n\n"
-                    "The current **best** model: generic pretrain model (tv2o-medium) by skytnt"
-                    )
+        gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Midi Composer with Arpeggios</h1>")
         js_msg = gr.Textbox(elem_id="msg_receiver", visible=False)
         js_msg.change(None, [js_msg], [], js="""
         (msg_json) =>{
@@ -431,18 +475,24 @@ if __name__ == "__main__":
                 input_drum_kit = gr.Dropdown(label="🥁drum kit", choices=list(drum_kits2number.keys()), type="value",
                                              value="None")
                 input_bpm = gr.Slider(label="BPM (beats per minute, auto if 0)", minimum=0, maximum=255,
-                                      step=1,
-                                      value=0)
+                                      step=1, value=0)
                 input_time_sig = gr.Radio(label="time signature (only for tv2 models)",
                                           value="auto",
                                           choices=["auto", "4/4", "2/4", "3/4", "6/4", "7/4",
-                                                   "2/2", "3/2", "4/2", "3/8", "5/8", "6/8", "7/8", "9/8", "12/8"]
-                                          )
+                                                   "2/2", "3/2", "4/2", "3/8", "5/8", "6/8", "7/8", "9/8", "12/8"])
                 input_key_sig = gr.Radio(label="key signature (only for tv2 models)",
                                          value="auto",
                                          choices=["auto"] + key_signatures,
-                                         type="index"
-                                         )
+                                         type="index")
+                
+                with gr.Row():
+                    arpeggio_intro = gr.Button("🎵 Intro Arpeggio", variant="primary")
+                    arpeggio_verse = gr.Button("🎸 Verse Arpeggio", variant="primary")
+                    arpeggio_chorus = gr.Button("🎹 Chorus Arpeggio", variant="primary")
+                    arpeggio_outro = gr.Button("🎷 Outro Arpeggio", variant="primary")
+
+
+                
                 example1 = gr.Examples([
                     [[], "None"],
                     [["Acoustic Grand"], "None"],
@@ -457,6 +507,7 @@ if __name__ == "__main__":
                     [["Electric Guitar(clean)", "Electric Guitar(muted)", "Overdriven Guitar", "Distortion Guitar",
                       "Electric Bass(finger)"], "Standard"]
                 ], [input_instruments, input_drum_kit])
+                
             with gr.TabItem("midi prompt") as tab2:
                 input_midi = gr.File(label="input midi", file_types=[".midi", ".mid"], type="binary")
                 input_midi_events = gr.Slider(label="use first n midi events as prompt", minimum=1, maximum=512,
@@ -470,6 +521,7 @@ if __name__ == "__main__":
                 input_remove_empty_channels = gr.Checkbox(label="remove channels without notes", value=False)
                 example2 = gr.Examples([[file, 128] for file in glob.glob("example/*.mid")],
                                        [input_midi, input_midi_events])
+                
             with gr.TabItem("last output prompt") as tab3:
                 gr.Markdown("Continue generating on the last output.")
                 input_continuation_select = gr.Radio(label="select output to continue generating", value="all",
@@ -530,5 +582,38 @@ if __name__ == "__main__":
         #                queue=False)
         undo_btn.click(undo_continuation, [input_model, output_midi_seq, output_continuation_state],
                        [output_midi_seq, output_continuation_state, js_msg], queue=False)
+
+
+        
+        def add_intro_arpeggio(model_name, mid_seq):
+            tokenizer = models[model_name].tokenizer
+            sequence = ['C', 'D', 'Am', 'G']
+            pattern = [0, 1, 2, 1]  # Root, Third, Fifth, Third
+            return add_arpeggio_sequence(tokenizer, mid_seq, sequence, pattern)
+
+        def add_verse_arpeggio(model_name, mid_seq):
+            tokenizer = models[model_name].tokenizer
+            sequence = ['D', 'C', 'Am', 'G']
+            pattern = [0, 2, 1, 2]  # Root, Fifth, Third, Fifth
+            return add_arpeggio_sequence(tokenizer, mid_seq, sequence, pattern)
+
+        def add_chorus_arpeggio(model_name, mid_seq):
+            tokenizer = models[model_name].tokenizer
+            sequence = ['G', 'D', 'Am', 'C']
+            pattern = [0, 1, 2, 1, 0, 2]  # Root, Third, Fifth, Third, Root, Fifth
+            return add_arpeggio_sequence(tokenizer, mid_seq, sequence, pattern)
+
+        def add_outro_arpeggio(model_name, mid_seq):
+            tokenizer = models[model_name].tokenizer
+            sequence = ['Am', 'G', 'D', 'C']
+            pattern = [2, 1, 0, 1]  # Fifth, Third, Root, Third
+            return add_arpeggio_sequence(tokenizer, mid_seq, sequence, pattern)
+
+        arpeggio_intro.click(add_intro_arpeggio, [input_model, output_midi_seq], output_midi_seq)
+        arpeggio_verse.click(add_verse_arpeggio, [input_model, output_midi_seq], output_midi_seq)
+        arpeggio_chorus.click(add_chorus_arpeggio, [input_model, output_midi_seq], output_midi_seq)
+        arpeggio_outro.click(add_outro_arpeggio, [input_model, output_midi_seq], output_midi_seq)
+
+    
     app.queue().launch(server_port=opt.port, share=opt.share, inbrowser=True, ssr_mode=False)
     thread_pool.shutdown()