Update app.py

app.py

@@ -365,7 +365,7 @@ def generate_callback_wrapper(input_midi,
     print('Req execution time:', (reqtime.time() - start_time), 'sec')
     print('*' * 70)    
     
-    return tuple([result[1], generated_batches, result[3]] + [item for sublist in result[0] for item in sublist[:-1]])
+    return tuple([result[1], generated_batches, result[3]] + [item for sublist in result[0] for item in sublist[:-1]] + [model_state])
 
 #==================================================================================
 
@@ -442,7 +442,7 @@ def remove_batch(batch_number, num_tokens, final_composition, generated_batches,
 
 #==================================================================================
 
-def reset(final_composition=[], generated_batches=[], block_lines=[], model_state = []):
+def reset(final_composition=[], generated_batches=[], block_lines=[], model_state=[]):
     
     final_composition = []
     generated_batches = []
@@ -453,7 +453,7 @@ def reset(final_composition=[], generated_batches=[], block_lines=[], model_stat
     
 #==================================================================================
 
-def reset_demo(final_composition=[], generated_batches=[], block_lines=[]):
+def reset_demo(final_composition=[], generated_batches=[], block_lines=[], model_state=[]):
     
     final_composition = []
     generated_batches = []
@@ -533,6 +533,8 @@ with gr.Blocks() as demo:
             
             outputs.extend([audio_output, plot_output])
 
+    outputs.extend([[]])
+
     generate_btn.click(generate_callback_wrapper, 
                        [input_midi, 
                         num_prime_tokens,