Update main.py

main.py

@@ -6,6 +6,7 @@ from utils import (
     generate_song,
     remove_last_instrument,
     regenerate_last_instrument,
+    change_last_instrument,
     change_tempo,
 )
 
@@ -17,11 +18,16 @@ os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
 # Genrs
 genres = ["ROCK", "POP", "OTHER", "R&B/SOUL", "JAZZ", "ELECTRONIC", "RANDOM"]
 
-#Artists
+# Artists
 with open('artist_names.json', 'r') as f:
     artist_names = json.load(f)
 #print("Loaded Artists names:", artist_names)
 
+# Instruments
+with open('instruments.json', 'r') as f:
+    instruments = json.load(f)
+
+
 demo = gr.Blocks()
 
 
@@ -39,11 +45,15 @@ def run():
                 artist = gr.Dropdown(
                     choices=artist_names, value=artist_names[0], label="Select the artist style"
                 )
+                instrument = gr.Dropdown(
+                    choices=instruments, value=instruments[0], label="Select the instrument to be generated"
+                )
                 with gr.Row():
                     btn_from_scratch = gr.Button("🧹 Start from scratch")
-                    btn_continue = gr.Button("➡️ Continue Generation")
+                    btn_continue = gr.Button("➡️ Generate New Track")
                     btn_remove_last = gr.Button("↩️ Remove last instrument")
                     btn_regenerate_last = gr.Button("🔄 Regenerate last instrument")
+                    btn_change_last = gr.Button("🔄 Swap last instrument")
             with gr.Column():
                 with gr.Group():
                     audio_output = gr.Video(show_share_button=True)
@@ -65,7 +75,7 @@ def run():
             num_tokens = gr.Text(visible=False)
         btn_from_scratch.click(
             fn=generate_song,
-            inputs=[genre, artist, temp, empty_sequence, qpm],
+            inputs=[genre, artist, instrument, temp, empty_sequence, qpm],
             outputs=[
                 audio_output,
                 midi_file,
@@ -77,7 +87,7 @@ def run():
         )
         btn_continue.click(
             fn=generate_song,
-            inputs=[genre, artist, temp, text_sequence, qpm],
+            inputs=[genre, artist, instrument, temp, text_sequence, qpm],
             outputs=[
                 audio_output,
                 midi_file,
@@ -111,6 +121,18 @@ def run():
                 num_tokens,
             ],
         )
+        btn_change_last.click(
+            fn=change_last_instrument,
+            inputs=[text_sequence, instrument, temp, qpm],
+            outputs=[
+                audio_output,
+                midi_file,
+                plot_output,
+                instruments_output,
+                text_sequence,
+                num_tokens,
+            ],
+        )
         btn_qpm.click(
             fn=change_tempo,
             inputs=[text_sequence, qpm],