Update app.py

app.py

@@ -23,7 +23,7 @@ in_space = os.getenv("SYSTEM") == "spaces"
 # =================================================================================================
                        
 @spaces.GPU
-def GenerateDrums(input_midi, input_num_tokens, input_top_k_value):
+def GenerateDrums(input_midi, input_num_tokens, input_top_k_value, input_max_drums_per_step):
     print('=' * 70)
     print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
     start_time = reqtime.time()
@@ -75,6 +75,7 @@ def GenerateDrums(input_midi, input_num_tokens, input_top_k_value):
     print('Input file name:', fn)
     print('Req num toks:', input_num_tokens)
     print('Req top_k value:', input_top_k_value)
+    print('Req max number of drums pitches:', input_max_drums_per_step)
     print('-' * 70)
 
     #===============================================================================
@@ -120,7 +121,7 @@ def GenerateDrums(input_midi, input_num_tokens, input_top_k_value):
     output = []
 
     temperature=0.9
-    max_drums_limit=5
+    max_drums_limit=input_max_drums_per_step
     num_memory_tokens=4096
 
     for c in comp_times[:input_num_tokens]:
@@ -291,6 +292,7 @@ if __name__ == "__main__":
         input_midi = gr.File(label="Input MIDI", file_types=[".midi", ".mid", ".kar"])
         input_num_tokens = gr.Slider(16, 2048, value=256, step=16, label="Number of composition chords to generate drums for")
         input_top_k_value = gr.Slider(1, 50, value=5, step=1, label="Model sampling top_k value")
+        input_max_drums_per_step = gr.Slider(1, 10, value=5, step=1, label="Maximum number of drums pitches per step")
         
         run_btn = gr.Button("generate", variant="primary")
 
@@ -303,21 +305,21 @@ if __name__ == "__main__":
         output_midi = gr.File(label="Output MIDI file", file_types=[".mid"])
 
 
-        run_event = run_btn.click(GenerateDrums, [input_midi, input_num_tokens, input_top_k_value],
+        run_event = run_btn.click(GenerateDrums, [input_midi, input_num_tokens, input_top_k_value, input_max_drums_per_step],
                                   [output_midi_title, output_midi_summary, output_midi, output_audio, output_plot])
 
         gr.Examples(
-            [["Ultimate-Drums-Transformer-Melody-Seed-1.mid", 128, 5], 
-             ["Ultimate-Drums-Transformer-Melody-Seed-2.mid", 128, 5], 
-             ["Ultimate-Drums-Transformer-Melody-Seed-3.mid", 128, 5],
-             ["Ultimate-Drums-Transformer-Melody-Seed-4.mid", 128, 5],
-             ["Ultimate-Drums-Transformer-Melody-Seed-5.mid", 128, 5],
-             ["Ultimate-Drums-Transformer-Melody-Seed-6.mid", 128, 5],
-             ["Ultimate-Drums-Transformer-MI-Seed-1.mid", 128, 5],
-             ["Ultimate-Drums-Transformer-MI-Seed-2.mid", 128, 5],
-             ["Ultimate-Drums-Transformer-MI-Seed-3.mid", 128, 5],
-             ["Ultimate-Drums-Transformer-MI-Seed-4.mid", 128, 5]],
-            [input_midi, input_num_tokens, input_top_k_value],
+            [["Ultimate-Drums-Transformer-Melody-Seed-1.mid", 128, 5, 5], 
+             ["Ultimate-Drums-Transformer-Melody-Seed-2.mid", 128, 5, 5], 
+             ["Ultimate-Drums-Transformer-Melody-Seed-3.mid", 128, 5, 5],
+             ["Ultimate-Drums-Transformer-Melody-Seed-4.mid", 128, 5, 5],
+             ["Ultimate-Drums-Transformer-Melody-Seed-5.mid", 128, 5, 5],
+             ["Ultimate-Drums-Transformer-Melody-Seed-6.mid", 128, 5, 5],
+             ["Ultimate-Drums-Transformer-MI-Seed-1.mid", 128, 5, 5],
+             ["Ultimate-Drums-Transformer-MI-Seed-2.mid", 128, 5, 5],
+             ["Ultimate-Drums-Transformer-MI-Seed-3.mid", 128, 5, 5],
+             ["Ultimate-Drums-Transformer-MI-Seed-4.mid", 128, 5, 5]],
+            [input_midi, input_num_tokens, input_top_k_value, input_max_drums_per_step],
             [output_midi_title, output_midi_summary, output_midi, output_audio, output_plot],
             GenerateDrums,
             cache_examples=True,