Update app.py

app.py

@@ -21,13 +21,31 @@ from midi_to_colab_audio import midi_to_colab_audio
 # =================================================================================================
                        
 @spaces.GPU
-def Generate_Rock_Song(input_midi, input_melody_seed_number):
+def Generate_Rock_Song(input_midi, 
+                       input_freestyle_continuation, 
+                       input_number_prime_chords, 
+                       input_use_original_durations,
+                       input_match_original_pitches_counts
+                      ):
     
     print('=' * 70)
     print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
     start_time = reqtime.time()
     print('=' * 70)
 
+    fn = os.path.basename(input_midi)
+    fn1 = fn.split('.')[0]
+
+    print('=' * 70)
+    print('Requested settings:')
+    print('=' * 70)
+    print('Input MIDI file name:', fn)
+    print('Freestyle continuation:', input_freestyle_continuation)
+    print('Number of prime chords:', input_number_prime_chords)
+    print('Use original durations:', input_use_original_durations)
+    print('Match original pitches counts:', input_match_original_pitches_counts)
+    print('=' * 70)
+
     print('Loading model...')
 
     SEQ_LEN = 4096
@@ -65,16 +83,6 @@ def Generate_Rock_Song(input_midi, input_melody_seed_number):
 
     print('Done!')
     print('=' * 70)
-
-    #==================================================================
-
-    fn = os.path.basename(input_midi)
-    fn1 = fn.split('.')[0]
-
-    print('=' * 70)
-    print('Requested settings:')
-    print('=' * 70)
-    print('Input MIDI file name:', fn)
     
     #===============================================================================
     # Raw single-track ms score
@@ -174,6 +182,25 @@ def Generate_Rock_Song(input_midi, input_melody_seed_number):
 
     #==================================================================
 
+    def generate_continuation(num_prime_tokens, num_gen_tokens):
+
+        x = torch.tensor(prime_toks[:num_prime_tokens], dtype=torch.long, device=DEVICE)
+        
+        with ctx:
+          out = model.generate(x,
+                              num_gen_tokens,
+                              #filter_logits_fn=top_k,
+                              #filter_kwargs={'k': 5},
+                              temperature=0.9,
+                              return_prime=True,
+                              verbose=True)
+        
+        y = out.tolist()[0]
+
+        return y
+
+    #==================================================================
+
     def generate_tokens(seq, max_num_ptcs=10):
     
         input = copy.deepcopy(seq)
@@ -185,7 +212,7 @@ def Generate_Rock_Song(input_midi, input_melody_seed_number):
     
         while pcount < max_num_ptcs and y > 255:
     
-            x = torch.tensor(input, dtype=torch.long, device='cuda')
+            x = torch.tensor(input, dtype=torch.long, device=DEVICE)
         
             with ctx:
               out = model.generate(x,
@@ -371,7 +398,11 @@ if __name__ == "__main__":
         output_plot = gr.Plot(label="Output MIDI score plot")
         output_midi = gr.File(label="Output MIDI file", file_types=[".mid"])
 
-        run_event = run_btn.click(Generate_Rock_Song, [input_midi, input_melody_seed_number],
+        run_event = run_btn.click(Generate_Rock_Song, [input_freestyle_continuation, 
+                                                       input_number_prime_chords, 
+                                                       input_use_original_durations,
+                                                       input_match_original_pitches_counts
+                                                      ],
                                   [output_midi_title, output_midi_summary, output_midi, output_audio, output_plot])
 
         gr.Examples(