fix genesis caching

decoder.py

@@ -178,7 +178,7 @@ class TextDecoder:
                 inst = 0
                 is_drum = 1
             if self.familized:
-                inst = Familizer(arbitrary=True).get_program_number(int(inst)) + 1
+                inst = Familizer(arbitrary=True).get_program_number(int(inst))
             instruments.append((int(inst), is_drum))
         return tuple(instruments)
 

generate.py

@@ -21,12 +21,12 @@ class GenerateMidiText:
         - self.process_prompt_for_next_bar()
         - self.generate_until_track_end()"""
 
-    def __init__(self, model, tokenizer):
+    def __init__(self, model, tokenizer, piece_by_track=[]):
         self.model = model
         self.tokenizer = tokenizer
         # default initialization
         self.initialize_default_parameters()
-        self.initialize_dictionaries()
+        self.initialize_dictionaries(piece_by_track)
 
     """Setters"""
 
@@ -38,8 +38,8 @@ class GenerateMidiText:
         self.set_nb_bars_generated()
         self.set_improvisation_level(0)
 
-    def initialize_dictionaries(self):
-        self.piece_by_track = []
+    def initialize_dictionaries(self, piece_by_track):
+        self.piece_by_track = piece_by_track
 
     def set_device(self, device="cpu"):
         self.device = ("cpu",)

generation_utils.py

@@ -2,14 +2,16 @@ import os
 import numpy as np
 import matplotlib.pyplot as plt
 import matplotlib
+
 from constants import INSTRUMENT_CLASSES
+from playback import get_music, show_piano_roll
 
 # matplotlib settings
 matplotlib.use("Agg")  # for server
 matplotlib.rcParams["xtick.major.size"] = 0
 matplotlib.rcParams["ytick.major.size"] = 0
-matplotlib.rcParams["axes.facecolor"] = "grey"
-matplotlib.rcParams["axes.edgecolor"] = "none"
+matplotlib.rcParams["axes.facecolor"] = "none"
+matplotlib.rcParams["axes.edgecolor"] = "grey"
 
 
 def define_generation_dir(model_repo_path):
@@ -93,7 +95,7 @@ def get_max_time(inst_midi):
 def plot_piano_roll(inst_midi):
     piano_roll_fig = plt.figure(figsize=(25, 3 * len(inst_midi.instruments)))
     piano_roll_fig.tight_layout()
-    piano_roll_fig.patch.set_alpha(0.1)
+    piano_roll_fig.patch.set_alpha(0)
     inst_count = 0
     beats_per_bar = 4
     sec_per_beat = 0.5
@@ -102,6 +104,14 @@ def plot_piano_roll(inst_midi):
         int
     )
     for inst in inst_midi.instruments:
+        # hardcoded for now
+        if inst.name == "Drums":
+            color = "purple"
+        elif inst.name == "Synth Bass 1":
+            color = "orange"
+        else:
+            color = "green"
+
         inst_count += 1
         plt.subplot(len(inst_midi.instruments), 1, inst_count)
 
@@ -118,24 +128,34 @@ def plot_piano_roll(inst_midi):
         for note in p_midi_note_list:
             note_time.append([note.start, note.end])
             note_pitch.append([note.pitch, note.pitch])
+        note_pitch = np.array(note_pitch)
+        note_time = np.array(note_time)
 
         plt.plot(
-            np.array(note_time).T,
-            np.array(note_pitch).T,
-            color="purple",
-            linewidth=3,
+            note_time.T,
+            note_pitch.T,
+            color=color,
+            linewidth=4,
             solid_capstyle="butt",
         )
         plt.ylim(0, 128)
         xticks = np.array(bars_time)[:-1]
         plt.tight_layout()
         plt.xlim(min(bars_time), max(bars_time))
-        # plt.xlabel("bars")
+        plt.ylim(max([note_pitch.min() - 5, 0]), note_pitch.max() + 5)
         plt.xticks(
             xticks + 0.5 * beats_per_bar * sec_per_beat,
             labels=xticks.argsort() + 1,
             visible=False,
         )
-        plt.title(inst.name, fontsize=10, color="white")
+        plt.text(
+            0.2,
+            note_pitch.max() + 4,
+            inst.name,
+            fontsize=20,
+            color=color,
+            horizontalalignment="left",
+            verticalalignment="top",
+        )
 
     return piano_roll_fig

playground.py

@@ -26,7 +26,6 @@ model, tokenizer = LoadModel(
     model_repo, from_huggingface=True, revision=revision
 ).load_model_and_tokenizer()
 
-
 miditok = get_miditok()
 decoder = TextDecoder(miditok)
 
@@ -40,32 +39,49 @@ def define_prompt(state, genesis):
 
 
 def generator(
-    regenerate, temp, density, instrument, state, add_bars=False, add_bar_count=1
+    label,
+    regenerate,
+    temp,
+    density,
+    instrument,
+    state,
+    piece_by_track,
+    add_bars=False,
+    add_bar_count=1,
 ):
 
+    genesis = GenerateMidiText(model, tokenizer, piece_by_track)
+    track = {"label": label}
     inst = next(
         (inst for inst in INSTRUMENT_CLASSES if inst["name"] == instrument),
         {"family_number": "DRUMS"},
     )["family_number"]
 
-    inst_index = index_has_substring(state, "INST=" + str(inst))
-
-    # Regenerate
-    if regenerate:
-        state.pop(inst_index)
-        genesis.delete_one_track(inst_index)
-        generated_text = (
-            genesis.get_whole_piece_from_bar_dict()
-        )  # maybe not useful here
-        inst_index = -1  # reset to last generated
+    inst_index = -1  # default to last generated
+    if state != []:
+        for index, instrum in enumerate(state):
+            if instrum["label"] == track["label"]:
+                inst_index = index  # changing if exists
 
     # Generate
     if not add_bars:
+        # Regenerate
+        if regenerate:
+            state.pop(inst_index)
+            genesis.delete_one_track(inst_index)
+
+            generated_text = (
+                genesis.get_whole_piece_from_bar_dict()
+            )  # maybe not useful here
+            inst_index = -1  # reset to last generated
+
         # NEW TRACK
         input_prompt = define_prompt(state, genesis)
         generated_text = genesis.generate_one_new_track(
             inst, density, temp, input_prompt=input_prompt
         )
+
+        regenerate = True  # set generate to true
     else:
         # NEW BARS
         genesis.generate_n_more_bars(add_bar_count)  # for all instruments
@@ -79,14 +95,23 @@ def generator(
     decoder.get_midi(inst_text, inst_midi_name)
     _, inst_audio = get_music(inst_midi_name)
     piano_roll = plot_piano_roll(mixed_inst_midi)
-    state.append(inst_text)
-
-    return inst_text, (44100, inst_audio), piano_roll, state, (44100, mixed_audio)
-
+    track["text"] = inst_text
+    state.append(track)
+
+    return (
+        inst_text,
+        (44100, inst_audio),
+        piano_roll,
+        state,
+        (44100, mixed_audio),
+        regenerate,
+        genesis.piece_by_track,
+    )
 
-def instrument_row(default_inst):
 
+def instrument_row(default_inst, row_id):
     with gr.Row():
+        row = gr.Variable(row_id)
         with gr.Column(scale=1, min_width=50):
             inst = gr.Dropdown(
                 [inst["name"] for inst in INSTRUMENT_CLASSES] + ["Drums"],
@@ -100,35 +125,33 @@ def instrument_row(default_inst):
             output_txt = gr.Textbox(label="output", lines=10, max_lines=10)
         with gr.Column(scale=1, min_width=100):
             inst_audio = gr.Audio(label="Audio")
-            regenerate = gr.Checkbox(value=False, label="Regenerate")
+            regenerate = gr.Checkbox(value=False, label="Regenerate", visible=False)
             # add_bars = gr.Checkbox(value=False, label="Add Bars")
             # add_bar_count = gr.Dropdown([1, 2, 4, 8], value=1, label="Add Bars")
             gen_btn = gr.Button("Generate")
             gen_btn.click(
                 fn=generator,
-                inputs=[
-                    regenerate,
-                    temp,
-                    density,
-                    inst,
+                inputs=[row, regenerate, temp, density, inst, state, piece_by_track],
+                outputs=[
+                    output_txt,
+                    inst_audio,
+                    piano_roll,
                     state,
+                    mixed_audio,
+                    regenerate,
+                    piece_by_track,
                 ],
-                outputs=[output_txt, inst_audio, piano_roll, state, mixed_audio],
             )
 
 
-with gr.Blocks(cache_examples=False) as demo:
-    genesis = GenerateMidiText(
-        model,
-        tokenizer,
-    )
-    genesis.set_nb_bars_generated(n_bars=n_bar_generated)
+with gr.Blocks() as demo:
+    piece_by_track = gr.State([])
     state = gr.State([])
     mixed_audio = gr.Audio(label="Mixed Audio")
     piano_roll = gr.Plot(label="Piano Roll")
-    instrument_row("Drums")
-    instrument_row("Bass")
-    instrument_row("Synth Lead")
+    instrument_row("Drums", 0)
+    instrument_row("Bass", 1)
+    instrument_row("Synth Lead", 2)
     # instrument_row("Piano")
 
 demo.launch(debug=True)
@@ -138,14 +161,9 @@ TODO: DEPLOY
 TODO: temp file situation
 TODO: clear cache situation
 TODO: reset button
-TODO: instrument mapping business
-TODO: Y lim axis of piano roll
 TODO: add a button to save the generated midi
 TODO: add improvise button
-TODO: making the piano roll fit on the horizontal scale
 TODO: set values for temperature as it is done for density
-TODO: set the color situation to be dark background
-TODO: make regeration default when an intrument has already been track has already been generated
 TODO: Add bar should be now set for the whole piece - regenerrate should regenerate the added bars only on all instruments
 TODO: row height to fix