Update main.py

main.py

@@ -9,6 +9,7 @@ from numpy import ndarray
 from note_seq.protobuf.music_pb2 import NoteSequence
 from note_seq.constants import STANDARD_PPQ
 import logging
+import math
 
 logging.basicConfig(level=logging.INFO)
 
@@ -145,6 +146,8 @@ GM_INSTRUMENTS = [
 ]
 tokenizer = None
 model = None
+
+
 def get_model_and_tokenizer() -> Tuple[AutoModelForCausalLM, AutoTokenizer]:
     logging.info("get_model_and_tokenizer: Starting to load model and tokenizer...")
     global model, tokenizer
@@ -311,9 +314,9 @@ def get_outputs_from_string(
     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
+    if not note_sequence.notes:
         logging.warning("get_outputs_from_string: Note sequence is empty, skipping plot.")
-        fig = None # Handle case where fig is None
+        fig = None
     else:
         fig = note_seq.plot_sequence(note_sequence, show_figure=False)
 
@@ -379,16 +382,18 @@ def generate_song(
     text_sequence: str = "",
     qpm: int = 120,
     prompt: str = "",
-    duration: int = 1
+    duration: int = 30
 ) -> Tuple[ndarray, str, Figure, str, str, str]:
-    logging.info(f"generate_song: Starting song generation. Genre: {genre}, Temperature: {temp}, QPM: {qpm}, Duration: {duration}, Prompt: '{prompt}'")
+    logging.info(f"generate_song: Starting song generation. Genre: {genre}, Temperature: {temp}, QPM: {qpm}, Duration: {duration} seconds, Prompt: '{prompt}'")
     if text_sequence == "":
         seed_string = create_seed_string(genre, prompt)
     else:
         seed_string = text_sequence
 
+    num_tracks = max(1, int(math.ceil(duration / 17)))
+
     generated_sequence = seed_string
-    for _ in range(duration):
+    for _ in range(num_tracks):
         instrument_sequence = generate_new_instrument(seed=generated_sequence, temp=temp)
         if instrument_sequence:
             generated_sequence = instrument_sequence
@@ -410,7 +415,7 @@ def run():
         with gr.Row():
             with gr.Column():
                 prompt_text = gr.Textbox(lines=2, placeholder="Enter text prompt here...", label="Text Prompt (Optional)")
-                duration_slider = gr.Slider(minimum=1, maximum=5, step=1, value=1, label="Duration (Tracks)")
+                duration_slider = gr.Slider(minimum=1, maximum=1000, step=1, value=30, label="Duration (Seconds)")
                 temp = gr.Slider(
                     minimum=0, maximum=1, step=0.05, value=0.85, label="Temperature"
                 )