Update app.py

app.py

@@ -23,7 +23,7 @@ in_space = os.getenv("SYSTEM") == "spaces"
 # =================================================================================================
 
 @spaces.GPU
-def GenerateMIDI(num_tok, idrums, iinstr):
+def GenerateMIDI(num_tok, idrums, iinstr, input_top_k_value):
     print('=' * 70)
     print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
     start_time = time.time()
@@ -126,6 +126,8 @@ def GenerateMIDI(num_tok, idrums, iinstr):
             with torch.inference_mode():
                 out = model.module.generate(inp,
                                             1,
+                                            filter_logits_fn=top_k,
+                                            filter_kwargs={'k': input_top_k_value},
                                             temperature=0.9,
                                             return_prime=False,
                                             verbose=False)
@@ -207,12 +209,13 @@ if __name__ == "__main__":
             value="Piano", label="Lead Instrument Controls", info="Desired lead instrument")
         input_drums = gr.Checkbox(label="Add Drums", value=False, info="Add drums to the composition")
         input_num_tokens = gr.Slider(16, 1024, value=512, label="Number of Tokens", info="Number of tokens to generate")
+        input_top_k_value = gr.Slider(1, 100, value=15, label="Model sampling top_k value")
         
         run_btn = gr.Button("generate", variant="primary")
-
-        output_plot = gr.Plot(label='output plot')
+        
         output_audio = gr.Audio(label="output audio", format="mp3", elem_id="midi_audio")
+        output_plot = gr.Plot(label='output plot')
         output_midi = gr.File(label="output midi", file_types=[".mid"])
-        run_event = run_btn.click(GenerateMIDI, [input_num_tokens, input_drums, input_instrument],
+        run_event = run_btn.click(GenerateMIDI, [input_num_tokens, input_drums, input_instrument, input_top_k_value],
                                   [output_plot, output_midi, output_audio])
         app.queue().launch()