Update app.py

app.py

@@ -303,18 +303,31 @@ def generate_music(prime, num_gen_tokens, num_gen_batches, model_temperature, mo
     
     print("Generating...")
     inp = torch.LongTensor([inputs] * num_gen_batches).cuda()
-    with ctx:
-        out = model.generate(
-            inp,
-            num_gen_tokens,
-            filter_logits_fn=top_p,
-            filter_kwargs={'thres': model_top_p},
-            temperature=model_temperature,
-            eos_token=18818,
-            return_prime=False,
-            verbose=False
-        )
-        
+
+    if model_top_p < 1:
+        with ctx:
+            out = model.generate(
+                inp,
+                num_gen_tokens,
+                filter_logits_fn=top_p,
+                filter_kwargs={'thres': model_top_p},
+                temperature=model_temperature,
+                eos_token=18818,
+                return_prime=False,
+                verbose=False
+            )
+
+    else:
+        with ctx:
+            out = model.generate(
+                inp,
+                num_gen_tokens,
+                temperature=model_temperature,
+                eos_token=18818,
+                return_prime=False,
+                verbose=False
+            )
+            
     print("Done!")
     print_sep()
     return out.tolist()
@@ -577,7 +590,7 @@ with gr.Blocks() as demo:
     num_prime_tokens = gr.Slider(16, 6656, value=6656, step=1, label="Number of prime tokens")
     num_gen_tokens = gr.Slider(16, 1024, value=512, step=1, label="Number of tokens to generate")
     model_temperature = gr.Slider(0.1, 1, value=0.9, step=0.01, label="Model temperature")
-    model_top_p = gr.Slider(0.1, 0.99, value=0.96, step=0.01, label="Model sampling top p value")
+    model_top_p = gr.Slider(0.1, 1.0, value=0.96, step=0.01, label="Model sampling top p value")
     add_drums = gr.Checkbox(value=False, label="Add drums")
     add_outro = gr.Checkbox(value=False, label="Add an outro")
     generate_btn = gr.Button("Generate", variant="primary")