Update app.py

app.py

@@ -111,6 +111,8 @@ def load_model(model_selector):
     print('Model will use', dtype, 'precision...')
     print('=' * 70)
 
+    return model
+
 #==================================================================================
 
 def load_midi(input_midi):
@@ -195,7 +197,8 @@ def generate_music(prime,
                    num_mem_tokens,
                    num_gen_batches,
                    model_temperature,
-                   # model_sampling_top_p
+                   # model_sampling_top_p,
+                   model_state
                   ):
 
     if not prime:
@@ -203,6 +206,8 @@ def generate_music(prime,
 
     else:
         inputs = prime[-num_mem_tokens:]
+
+    model = model_state
         
     model.cuda()
     model.eval()
@@ -240,7 +245,8 @@ def generate_callback(input_midi,
                       # model_sampling_top_p,
                       final_composition, 
                       generated_batches, 
-                      block_lines
+                      block_lines,
+                      model_state
                      ):
 
     generated_batches = []
@@ -255,7 +261,8 @@ def generate_callback(input_midi,
                                         num_mem_tokens,
                                         NUM_OUT_BATCHES,
                                         model_temperature,
-                                        # model_sampling_top_p
+                                        # model_sampling_top_p,
+                                        model_state
                                        )
     
     outputs = []
@@ -310,7 +317,8 @@ def generate_callback_wrapper(input_midi,
                               final_composition, 
                               generated_batches, 
                               block_lines,
-                              model_selector
+                              model_selector,
+                              model_state
                              ):
 
     print('=' * 70)
@@ -325,7 +333,8 @@ def generate_callback_wrapper(input_midi,
 
     print('Selected model type:', model_selector)
 
-    load_model(model_selector)
+    if not model_State:
+        model_state = load_model(model_selector)
     
     print('Num prime tokens:', num_prime_tokens)
     print('Num gen tokens:', num_gen_tokens)
@@ -343,7 +352,8 @@ def generate_callback_wrapper(input_midi,
                                 # model_sampling_top_p,
                                 final_composition,
                                 generated_batches,
-                                block_lines
+                                block_lines,
+                                model_state
                              )
     
     generated_batches = [sublist[-1] for sublist in result[0]]
@@ -481,6 +491,7 @@ with gr.Blocks() as demo:
     final_composition = gr.State([])
     generated_batches = gr.State([])
     block_lines = gr.State([])
+    model_state = gr.State([])
     
     #==================================================================================
     
@@ -529,7 +540,8 @@ with gr.Blocks() as demo:
                         final_composition,
                         generated_batches,
                         block_lines,
-                        model_selector
+                        model_selector,
+                        model_state
                        ], 
                        outputs
                       )