Update main.py

main.py

@@ -263,7 +263,7 @@ model, tokenizer = get_model_and_tokenizer()
 def create_seed_string(genre: str = "OTHER", prompt: str = "") -> str:
     logging.info(f"create_seed_string: Creating seed string. Genre: {genre}, Prompt: '{prompt}'")
     if prompt:
-        seed_string = f"PIECE_START PROMPT={prompt} GENRE={genre} TRACK_START" # Incorporate prompt more directly
+        seed_string = f"PIECE_START PROMPT={prompt} GENRE={genre} TRACK_START"
     elif genre == "RANDOM":
         seed_string = "PIECE_START"
     else:
@@ -281,7 +281,7 @@ def get_instruments(text_sequence: str) -> List[str]:
                 index = int(part[5:])
                 instruments.append(GM_INSTRUMENTS[index])
     return instruments
-def generate_new_instrument(seed: str, temp: float = 0.75, max_tokens=204) -> str: # Increased max_tokens
+def generate_new_instrument(seed: str, temp: float = 0.85, max_tokens=512) -> str:
     logging.info(f"generate_new_instrument: Starting instrument generation. Seed: '{seed}', Temperature: {temp}, Max Tokens: {max_tokens}")
     seed_length = len(tokenizer.encode(seed))
     input_ids = tokenizer.encode(seed, return_tensors="pt").to(model.device)
@@ -295,7 +295,7 @@ def generate_new_instrument(seed: str, temp: float = 0.75, max_tokens=204) -> st
     )
     generated_sequence = tokenizer.decode(generated_ids[0])
     new_generated_sequence = tokenizer.decode(generated_ids[0][seed_length:])
-    logging.info(f"generate_new_instrument: Generated sequence: '{new_generated_sequence}'") # Log generated sequence
+    logging.info(f"generate_new_instrument: Generated sequence: '{new_generated_sequence}'")
     if "NOTE_ON" in new_generated_sequence:
         logging.info("generate_new_instrument: New instrument generated successfully.")
         return generated_sequence
@@ -310,15 +310,22 @@ def get_outputs_from_string(
     instruments = get_instruments(generated_sequence)
     instruments_str = "\n".join(f"- {instrument}" for instrument in instruments)
     note_sequence = token_sequence_to_note_sequence(generated_sequence, qpm=qpm)
+
+    if not note_sequence.notes: # Check if note_sequence is empty
+        logging.warning("get_outputs_from_string: Note sequence is empty, skipping plot.")
+        fig = None # Handle case where fig is None
+    else:
+        fig = note_seq.plot_sequence(note_sequence, show_figure=False)
+
     synth = note_seq.fluidsynth
     array_of_floats = synth(note_sequence, sample_rate=SAMPLE_RATE)
     int16_data = note_seq.audio_io.float_samples_to_int16(array_of_floats)
-    fig = note_seq.plot_sequence(note_sequence, show_figure=False)
     num_tokens = str(len(generated_sequence.split()))
     audio = gr.make_waveform((SAMPLE_RATE, int16_data))
     note_seq.note_sequence_to_midi_file(note_sequence, "midi_ouput.mid")
     logging.info("get_outputs_from_string: Output generation complete.")
     return audio, "midi_ouput.mid", fig, instruments_str, num_tokens
+
 def remove_last_instrument(
     text_sequence: str, qpm: int = 120
 ) -> Tuple[ndarray, str, Figure, str, str, str]:
@@ -368,7 +375,7 @@ def change_tempo(
     return audio, midi_file, fig, instruments_str, text_sequence, num_tokens
 def generate_song(
     genre: str = "OTHER",
-    temp: float = 0.85, # Default temperature to 0.85 as in UI
+    temp: float = 0.85,
     text_sequence: str = "",
     qpm: int = 120,
     prompt: str = "",