Update app.py

app.py

@@ -65,6 +65,8 @@ SOUDFONT_PATH = 'SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2'
 
 MAX_MELODY_NOTES = 64
 
+MAX_GEN_TOKS = 3072
+
 #==================================================================================
 
 print('=' * 70)
@@ -184,7 +186,7 @@ def Generate_Accompaniment(input_midi,
 
     #===============================================================================
 
-    def generate_full_seq(input_seq, max_toks=3072, temperature=0.9, verbose=True):
+    def generate_full_seq(input_seq, max_toks=3072, temperature=0.9, top_k_value=15, verbose=True):
     
         seq_abs_run_time = sum([t for t in input_seq if t < 128])
     
@@ -205,6 +207,8 @@ def Generate_Accompaniment(input_midi,
             with ctx:
                 out = model.generate(x,
                                      1,
+                                     filter_logits_fn=top_k,
+                                     filter_kwargs={'thres': top_k_value},
                                      temperature=temperature,
                                      return_prime=False,
                                      verbose=False)
@@ -217,6 +221,9 @@ def Generate_Accompaniment(input_midi,
             full_seq.append(y)
     
             toks_counter += 1
+
+            if toks_counter == max_toks:
+                return full_seq
     
         return full_seq
 
@@ -281,7 +288,11 @@ def Generate_Accompaniment(input_midi,
 
     #==================================================================
 
-    input_seq = generate_full_seq(start_score_seq, temperature=model_temperature)
+    input_seq = generate_full_seq(start_score_seq,
+                                  max_toks=MAX_GEN_TOKS,
+                                  temperature=model_temperature,
+                                  top_k_value=model_sampling_top_k,
+                                 )
 
     final_song = input_seq[len(start_score_seq):]