Update app.py

app.py

@@ -118,18 +118,19 @@ def split_lyrics(lyrics: str):
     structured_lyrics = [f"[{seg[0]}]\n{seg[1].strip()}\n\n" for seg in segments]
     return structured_lyrics
 
-@spaces.GPU(duration=175)
+@spaces.GPU(duration=178)
 def generate_music(
         genre_txt=None,
         lyrics_txt=None,
         run_n_segments=2,
-        max_new_tokens=30,
+        max_new_tokens=50,
         use_audio_prompt=False,
         audio_prompt_path="",
         prompt_start_time=0.0,
         prompt_end_time=30.0,
         cuda_idx=0,
         rescale=False,
+        progress=gr.Progress()
 ):
     """
     Generates music based on given genre and lyrics, optionally using an audio prompt.
@@ -139,7 +140,7 @@ def generate_music(
     if use_audio_prompt and not audio_prompt_path:
         raise FileNotFoundError("Please provide an audio prompt file when 'Use Audio Prompt' is enabled!")
     cuda_idx = cuda_idx
-    max_new_tokens = max_new_tokens * 100
+    max_new_tokens = int(max_new_tokens * (100/run_n_segments))
 
     with tempfile.TemporaryDirectory() as output_dir:
         stage1_output_dir = os.path.join(output_dir, f"stage1")
@@ -332,7 +333,7 @@ with gr.Blocks() as demo:
                 audio_prompt_input = gr.Audio(type="filepath", label="Audio Prompt (Optional)")
             with gr.Column():
                 num_segments = gr.Number(label="Number of Segments", value=2, interactive=True)
-                max_new_tokens = gr.Slider(label="Duration of song", minimum=1, maximum=30, step=1, value=15, interactive=True)
+                max_new_tokens = gr.Slider(label="Duration of song", minimum=1, maximum=50, step=1, value=30, interactive=True)
                 submit_btn = gr.Button("Submit")
                 music_out = gr.Audio(label="Mixed Audio Result")
                 with gr.Accordion(label="Vocal and Instrumental Result", open=False):