Update app.py

app.py

@@ -115,7 +115,7 @@ def load_model(model_selector):
 
 #==================================================================================
 
-def load_midi(input_midi):
+def load_midi(input_midi, model_selector=''):
 
     raw_score = TMIDIX.midi2single_track_ms_score(input_midi.name)
 
@@ -143,7 +143,7 @@ def load_midi(input_midi):
 
 #==================================================================================
 
-def save_midi(tokens, batch_number=None):
+def save_midi(tokens, batch_number=None, model_selector=''):
 
     song = tokens
     song_f = []
@@ -253,7 +253,7 @@ def generate_callback(input_midi,
 
     if not final_composition and input_midi is not None:
         final_composition = load_midi(input_midi)[:num_prime_tokens]
-        midi_score = save_midi(final_composition)
+        midi_score = save_midi(final_composition, model_selector=model_state[2])
         block_lines.append(midi_score[-1][1] / 1000)
         
     batched_gen_tokens = generate_music(final_composition, 
@@ -275,7 +275,7 @@ def generate_callback(input_midi,
         tokens_preview = final_composition[-PREVIEW_LENGTH:]
         
         # Save MIDI to a temporary file
-        midi_score = save_midi(tokens_preview + tokens, i)
+        midi_score = save_midi(tokens_preview + tokens, i, model_selector=model_state[2])
 
         # MIDI plot
 
@@ -335,6 +335,7 @@ def generate_callback_wrapper(input_midi,
 
     if not model_state:
         model_state = load_model(model_selector)
+        model_state.append(model_selector)
     
     print('Num prime tokens:', num_prime_tokens)
     print('Num gen tokens:', num_gen_tokens)
@@ -368,13 +369,13 @@ def generate_callback_wrapper(input_midi,
 
 #==================================================================================
 
-def add_batch(batch_number, final_composition, generated_batches, block_lines):
+def add_batch(batch_number, final_composition, generated_batches, block_lines, model_state=[]):
 
     if generated_batches:
         final_composition.extend(generated_batches[batch_number])
 
         # Save MIDI to a temporary file
-        midi_score = save_midi(final_composition)
+        midi_score = save_midi(final_composition, model_selector=model_state[2])
     
         block_lines.append(midi_score[-1][1] / 1000)
 
@@ -404,7 +405,7 @@ def add_batch(batch_number, final_composition, generated_batches, block_lines):
 
 #==================================================================================
 
-def remove_batch(batch_number, num_tokens, final_composition, generated_batches, block_lines):
+def remove_batch(batch_number, num_tokens, final_composition, generated_batches, block_lines, model_state=[]):
 
     if final_composition:
 
@@ -413,7 +414,7 @@ def remove_batch(batch_number, num_tokens, final_composition, generated_batches,
             block_lines.pop()
     
         # Save MIDI to a temporary file
-        midi_score = save_midi(final_composition)
+        midi_score = save_midi(final_composition, model_selector=model_state[2])
     
         # MIDI plot
         midi_plot = TMIDIX.plot_ms_SONG(midi_score, 
@@ -441,13 +442,14 @@ def remove_batch(batch_number, num_tokens, final_composition, generated_batches,
 
 #==================================================================================
 
-def reset(final_composition=[], generated_batches=[], block_lines=[]):
+def reset(final_composition=[], generated_batches=[], block_lines=[], model_state = []):
     
     final_composition = []
     generated_batches = []
     block_lines = []
+    model_state = []
 
-    return final_composition, generated_batches, block_lines
+    return final_composition, generated_batches, block_lines, model_state
     
 #==================================================================================
 
@@ -456,6 +458,7 @@ def reset_demo(final_composition=[], generated_batches=[], block_lines=[]):
     final_composition = []
     generated_batches = []
     block_lines = []
+    model_state = []
 
 #==================================================================================
 
@@ -557,10 +560,10 @@ with gr.Blocks() as demo:
     final_plot_output = gr.Plot(label="Final MIDI plot")
     final_file_output = gr.File(label="Final MIDI file")
 
-    add_btn.click(add_batch, [batch_number, final_composition, generated_batches, block_lines],
+    add_btn.click(add_batch, [batch_number, final_composition, generated_batches, block_lines, model_state],
                   [final_audio_output, final_plot_output, final_file_output, final_composition, generated_batches, block_lines])       
 
-    remove_btn.click(remove_batch, [batch_number, num_gen_tokens, final_composition, generated_batches, block_lines], 
+    remove_btn.click(remove_batch, [batch_number, num_gen_tokens, final_composition, generated_batches, block_lines, model_state], 
                      [final_audio_output, final_plot_output, final_file_output, final_composition, generated_batches, block_lines])
 
     demo.unload(reset_demo)