Update app.py

app.py

@@ -23,7 +23,7 @@ in_space = os.getenv("SYSTEM") == "spaces"
 # =================================================================================================
 
 @spaces.GPU
-def GenerateMIDI(num_tok, idrums, iinstr, input_top_k_value):
+def GenerateMIDI(num_tok, idrums, iinstr, input_top_k_ratio):
     print('=' * 70)
     print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
     start_time = time.time()
@@ -32,7 +32,7 @@ def GenerateMIDI(num_tok, idrums, iinstr, input_top_k_value):
     print('Req num tok:', num_tok)
     print('Req instr:', iinstr)
     print('Drums:', idrums)
-    print('top_k:', input_top_k_value)
+    print('top_k:', input_top_k_ratio)
     print('-' * 70)
 
     if idrums:
@@ -127,8 +127,8 @@ def GenerateMIDI(num_tok, idrums, iinstr, input_top_k_value):
             with torch.inference_mode():
                 out = model.module.generate(inp,
                                             1,
-                                            filter_logits_fn=top_k,
-                                            filter_kwargs={'k': input_top_k_value},
+                                            filter_logits_fn = top_k,
+                                            filter_thres = input_top_k_ratio,
                                             temperature=0.9,
                                             return_prime=False,
                                             verbose=False)
@@ -210,13 +210,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")
+        input_top_k_ratio = gr.Slider(0.1, 1, value=0.95, step=0.01, label="Model sampling top_k ratio")
         
         run_btn = gr.Button("generate", variant="primary")
         
         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, input_top_k_value],
+        run_event = run_btn.click(GenerateMIDI, [input_num_tokens, input_drums, input_instrument, input_top_k_ratio],
                                   [output_plot, output_midi, output_audio])
         app.queue().launch()