Update app.py

app.py

@@ -23,7 +23,7 @@ in_space = os.getenv("SYSTEM") == "spaces"
 # =================================================================================================
                        
 @spaces.GPU
-def GenerateDrums(input_midi, input_num_tokens):
+def GenerateDrums(input_midi, input_num_tokens, input_top_k_value):
     print('=' * 70)
     print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
     start_time = reqtime.time()
@@ -74,6 +74,7 @@ def GenerateDrums(input_midi, input_num_tokens):
     print('-' * 70)
     print('Input file name:', fn)
     print('Req num toks:', input_num_tokens)
+    print('Req top_k value:', input_top_k_value)
     print('-' * 70)
 
     #===============================================================================
@@ -138,6 +139,8 @@ def GenerateDrums(input_midi, input_num_tokens):
           with ctx:
             out = model.generate(x[-num_memory_tokens:],
                                 1,
+                                filter_logits_fn=top_k,
+                                filter_kwargs={'k': input_top_k_value},
                                 temperature=temperature,
                                 return_prime=False,
                                 verbose=False)
@@ -287,6 +290,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=1, step=1, label="Model sampling top_k value")
         
         run_btn = gr.Button("generate", variant="primary")
 
@@ -299,21 +303,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],
+        run_event = run_btn.click(GenerateDrums, [input_midi, input_num_tokens, input_top_k_value],
                                   [output_midi_title, output_midi_summary, output_midi, output_audio, output_plot])
 
         gr.Examples(
-            [["Ultimate-Drums-Transformer-Melody-Seed-1.mid", 128], 
-             ["Ultimate-Drums-Transformer-Melody-Seed-2.mid", 128], 
-             ["Ultimate-Drums-Transformer-Melody-Seed-3.mid", 128],
-             ["Ultimate-Drums-Transformer-Melody-Seed-4.mid", 128],
-             ["Ultimate-Drums-Transformer-Melody-Seed-5.mid", 128],
-             ["Ultimate-Drums-Transformer-Melody-Seed-6.mid", 128],
-             ["Ultimate-Drums-Transformer-MI-Seed-1.mid", 128],
-             ["Ultimate-Drums-Transformer-MI-Seed-2.mid", 128],
-             ["Ultimate-Drums-Transformer-MI-Seed-3.mid", 128],
-             ["Ultimate-Drums-Transformer-MI-Seed-4.mid", 128]],
-            [input_midi, input_num_tokens],
+            [["Ultimate-Drums-Transformer-Melody-Seed-1.mid", 128, 1], 
+             ["Ultimate-Drums-Transformer-Melody-Seed-2.mid", 128, 1], 
+             ["Ultimate-Drums-Transformer-Melody-Seed-3.mid", 128, 1],
+             ["Ultimate-Drums-Transformer-Melody-Seed-4.mid", 128, 1],
+             ["Ultimate-Drums-Transformer-Melody-Seed-5.mid", 128, 1],
+             ["Ultimate-Drums-Transformer-Melody-Seed-6.mid", 128, 1],
+             ["Ultimate-Drums-Transformer-MI-Seed-1.mid", 128, 1],
+             ["Ultimate-Drums-Transformer-MI-Seed-2.mid", 128, 1],
+             ["Ultimate-Drums-Transformer-MI-Seed-3.mid", 128, 1],
+             ["Ultimate-Drums-Transformer-MI-Seed-4.mid", 128, 1]],
+            [input_midi, input_num_tokens, input_top_k_value],
             [output_midi_title, output_midi_summary, output_midi, output_audio, output_plot],
             GenerateDrums,
             cache_examples=True,