Update app.py

app.py

@@ -278,7 +278,8 @@ def generate_music(prime, num_gen_tokens, num_gen_batches, model_temperature, mo
     if len(prime) >= 7168:
         prime = [18816] + prime[-7168:]
     
-    inputs = prime if prime else [18816]
+    inputs = prime
+    
     print("Generating...")
     inp = torch.LongTensor([inputs] * num_gen_batches).cuda()
     with ctx:
@@ -297,7 +298,7 @@ def generate_music(prime, num_gen_tokens, num_gen_batches, model_temperature, mo
     print_sep()
     return out.tolist()
 
-def generate_music_and_state(input_midi, num_prime_tokens, num_gen_tokens,
+def generate_music_and_state(input_midi, prime_instruments, num_prime_tokens, num_gen_tokens,
                              model_temperature, model_top_p, add_drums, add_outro, final_composition, generated_batches, block_lines):
     """
     Generate tokens using the model, update the composition state, and prepare outputs.
@@ -334,6 +335,14 @@ def generate_music_and_state(input_midi, num_prime_tokens, num_gen_tokens,
         # Use the last note's time as a marker.
         block_lines.append(midi_score[-1][1] / 1000 if final_composition else 0)
 
+    if not final_composition and input is None:
+        final_composition = [18816, 0]
+        
+        for i, instr in enumerate(prime_instruments):
+            instr_num = patch2number[instr]
+            final_composition.append((128*instr_num)+(72-(i*12))+256)
+            final_composition.append((8*16)+5+16768)
+
     if final_composition:
         if add_outro:
             final_composition.append(18817) # Outro token
@@ -446,6 +455,8 @@ def reset(final_composition=[], generated_batches=[], block_lines=[]):
     print_sep()
     return [], [], []
 
+patch2number = {v: k for k, v in TMIDIX.Number2patch.items()}
+
 # -----------------------------
 # GRADIO INTERFACE SETUP
 # -----------------------------
@@ -506,6 +517,8 @@ with gr.Blocks() as demo:
                       [final_composition, generated_batches, block_lines])
 
     gr.Markdown("## Generate")
+    prime_instruments = gr.Dropdown(label="Prime instruments (select up to 5)", choices=list(patch2number.keys()),
+                                    multiselect=True, max_choices=5, type="value")
     num_prime_tokens = gr.Slider(16, 7168, value=7168, step=1, label="Number of prime tokens")
     num_gen_tokens = gr.Slider(16, 1024, value=512, step=1, label="Number of tokens to generate")
     model_temperature = gr.Slider(0.1, 1, value=0.9, step=0.01, label="Model temperature")
@@ -524,7 +537,7 @@ with gr.Blocks() as demo:
             outputs.extend([audio_output, plot_output])
     generate_btn.click(
         generate_music_and_state,
-        [input_midi, num_prime_tokens, num_gen_tokens, model_temperature, model_top_p, add_drums, add_outro,
+        [input_midi, prime_instruments, num_prime_tokens, num_gen_tokens, model_temperature, model_top_p, add_drums, add_outro,
          final_composition, generated_batches, block_lines],
         outputs
     )