Update app.py

app.py

@@ -270,7 +270,7 @@ def save_midi(tokens):
 # MUSIC GENERATION FUNCTION (Combined)
 # -----------------------------
 @spaces.GPU
-def generate_music(prime, num_gen_tokens, num_mem_tokens, num_gen_batches, model_temperature):
+def generate_music(prime, num_gen_tokens, num_mem_tokens, num_gen_batches, model_temperature, model_top_p):
     """Generate music tokens given prime tokens and parameters."""
     inputs = prime[-num_mem_tokens:] if prime else [18816]
     print("Generating...")
@@ -280,8 +280,9 @@ def generate_music(prime, num_gen_tokens, num_mem_tokens, num_gen_batches, model
             inp,
             num_gen_tokens,
             filter_logits_fn=top_p,
-            filter_kwargs={'thres': 0.96},
+            filter_kwargs={'thres': model_top_p},
             temperature=model_temperature,
+            eos_token=18818,
             return_prime=False,
             verbose=False
         )
@@ -327,10 +328,11 @@ def generate_music_and_state(input_midi, num_prime_tokens, num_gen_tokens, num_m
         final_composition.append(18817) # Outro token
 
     if add_drums:
-        final_composition.extend([(128*128)+38+256]) # Drum token (38)
+        drum_pitch = random.choice([36, 38])
+        final_composition.extend([(128*128)+drum_pitch+256]) # Drum token
 
     batched_gen_tokens = generate_music(final_composition, num_gen_tokens, num_mem_tokens,
-                                        NUM_OUT_BATCHES, model_temperature)
+                                        NUM_OUT_BATCHES, model_temperature, model_top_p)
     
     output_batches = []
     for i, tokens in enumerate(batched_gen_tokens):