Update utils.py

utils.py

@@ -11,14 +11,19 @@ from constants import GM_INSTRUMENTS, SAMPLE_RATE
 from string_to_notes import token_sequence_to_note_sequence
 from model import get_model_and_tokenizer
 
+import json
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load the tokenizer and the model
 model, tokenizer = get_model_and_tokenizer()
 
+# Instruments
+with open('instruments.json', 'r') as f:
+    instruments = json.load(f)
 
-def create_seed_string(genre: str = "OTHER", artist: str = "OTHER") -> str:
+
+def create_seed_string(genre: str = "OTHER", artist: str = "OTHER", instrument:str="0") -> str:
     """
     Creates a seed string for generating a new piece.
 
@@ -29,13 +34,13 @@ def create_seed_string(genre: str = "OTHER", artist: str = "OTHER") -> str:
         str: The seed string.
     """
     if genre == "RANDOM" and artist == "RANDOM":
-        seed_string = "PIECE_START"
+        seed_string = f"PIECE_START GENRE=RANDOM ARTIST=RANDOM TRACK_START INST={instrument}"
     elif genre == "RANDOM" and artist != "RANDOM":
-        seed_string = f"PIECE_START GENRE=RANDOM ARTIST={artist} TRACK_START"
+        seed_string = f"PIECE_START GENRE=RANDOM ARTIST={artist} TRACK_START INST={instrument}"
     elif genre != "RANDOM" and artist == "RANDOM":
-        seed_string = f"PIECE_START GENRE={genre} ARTIST=RANDOM TRACK_START"
+        seed_string = f"PIECE_START GENRE={genre} ARTIST=RANDOM TRACK_START INST={instrument}"
     else:
-        seed_string = f"PIECE_START GENRE={genre} ARTIST={artist} TRACK_START"
+        seed_string = f"PIECE_START GENRE={genre} ARTIST={artist} TRACK_START INST={instrument}"
     return seed_string
 
 
@@ -61,6 +66,39 @@ def get_instruments(text_sequence: str) -> List[str]:
     return instruments
 
 
+def change_last_instrument( text_sequence: str, 
+                           instrument: str, 
+                           temp: float = 0.75, 
+                           qpm: int = 120
+                          ) -> Tuple[ndarray, str, Figure, str, str, str]:
+
+
+    instrument_idx = instruments.index(instrument)
+    #Drums
+    if instrument_idx == 0:
+        instrument_idx='DRUMS'
+    else:
+        instrument_idx = str(instrument_idx-1)
+    text_sequence = text_sequence.split()
+    for token_idx in reversed(range(len(text_sequence))):
+        if "INST=" in text_sequence[token_idx]:
+            text_sequence[token_idx] = f"INST={instrument_idx}"
+            break
+    text_sequence = (' ').join(text_sequence)
+    #print(text_sequence)
+
+    audio, midi_file, fig, instruments_str, num_tokens = get_outputs_from_string(
+        text_sequence, qpm
+    )
+    # print(type(audio),audio)
+    # print(type(midi_file),midi_file) 
+    # print(type(fig),fig)
+    # print(type(instruments_str),instruments_str)
+    # print(type(num_tokens),num_tokens)
+    return audio, midi_file, fig, instruments_str, text_sequence, num_tokens
+            
+
+
 def generate_new_instrument(seed: str, temp: float = 0.75) -> str:
     """
     Generates a new instrument sequence from a given seed and temperature.
@@ -165,6 +203,13 @@ def remove_last_instrument(
     return audio, midi_file, fig, instruments_str, new_song, num_tokens
 
 
+    genre: str = "OTHER",
+    artist: str = "KATE_BUSH",
+    instrument: str = "Acoustic Grand Piano",
+    temp: float = 0.75,
+    text_sequence: str = "",
+    qpm: int = 120
+    
 def regenerate_last_instrument(
     text_sequence: str, qpm: int = 120
 ) -> Tuple[ndarray, str, Figure, str, str, str]:
@@ -179,19 +224,32 @@ def regenerate_last_instrument(
         Tuple[ndarray, str, Figure, str, str, str]: The audio waveform, MIDI file name, plot figure,
                                                     instruments string, new song string, and number of tokens string.
     """
-    last_inst_index = text_sequence.rfind("INST=")
-    if last_inst_index == -1:
-        # No instrument so start from empty sequence
-        audio, midi_file, fig, instruments_str, new_song, num_tokens = generate_song(
-            text_sequence="", qpm=qpm
-        )
+
+    def remove_last_track(text_sequence):
+        tracks = text_sequence.split("TRACK_START")
+        # We keep all tracks except the last one
+        useful_tracks = tracks[:-1]
+        # We join the tracks back together, adding back the 'TRACK_START' that was removed by split
+        text_sequence = "TRACK_START".join(useful_tracks)
+        return text_sequence
+    
+    #last_inst_index = text_sequence.rfind("INST=")
+    
+    for token in reversed(text_sequence.split()):
+        if 'INST=' in token:
+            instrument_id = token.split('=')[1]
+            break
+    
+    if instrument_id=="DRUMS":
+        instrument="Drums"
     else:
-        # Take it from the last instrument and continue generation
-        next_space_index = text_sequence.find(" ", last_inst_index)
-        new_seed = text_sequence[:next_space_index]
-        audio, midi_file, fig, instruments_str, new_song, num_tokens = generate_song(
-            text_sequence=new_seed, qpm=qpm
-        )
+        instrument=instruments[int(instrument_id)+1]# Index 0 instrument is 'Acoustic Grand Piano' for rendering:https://soundprogramming.net/file-formats/general-midi-instrument-list/#google_vignette
+
+    new_seed = remove_last_track(text_sequence=text_sequence)
+    
+    audio, midi_file, fig, instruments_str, new_song, num_tokens = generate_song(
+        instrument=instrument,text_sequence=new_seed, qpm=qpm
+    )
     return audio, midi_file, fig, instruments_str, new_song, num_tokens
 
 
@@ -218,9 +276,10 @@ def change_tempo(
 def generate_song(
     genre: str = "OTHER",
     artist: str = "KATE_BUSH",
+    instrument: str = "Acoustic Grand Piano",
     temp: float = 0.75,
     text_sequence: str = "",
-    qpm: int = 120,
+    qpm: int = 120
 ) -> Tuple[ndarray, str, Figure, str, str, str]:
     """
     Generates a song given a genre, temperature, initial text sequence, and tempo.
@@ -238,13 +297,21 @@ def generate_song(
         Tuple[ndarray, str, Figure, str, str, str]: The audio waveform, MIDI file name, plot figure,
                                                     instruments string, generated song string, and number of tokens string.
     """
+    instrument = instruments.index(instrument)
+    #Drums
+    if instrument == 0:
+        instrument='DRUMS'
+    else:
+        instrument = str(instrument-1)
+        
     if text_sequence == "":
-        seed_string = create_seed_string(genre, artist)
+        seed_string = create_seed_string(genre, artist, instrument)
     else:
-        seed_string = text_sequence
+        seed_string = text_sequence + " TRACK_START INST=" + instrument
 
     generated_sequence = generate_new_instrument(seed=seed_string, temp=temp)
     audio, midi_file, fig, instruments_str, num_tokens = get_outputs_from_string(
         generated_sequence, qpm
     )
     return audio, midi_file, fig, instruments_str, generated_sequence, num_tokens
+