Update app.py

app.py

@@ -109,7 +109,7 @@ def get_last_mp3_file(output_dir):
     # Return the most recent .mp3 file
     return mp3_files_with_path[0]
 
-def infer(genre_txt_content, lyrics_txt_content):
+def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
     # Create temporary files
     genre_txt_path = create_temp_file(genre_txt_content, prefix="genre_")
     lyrics_txt_path = create_temp_file(lyrics_txt_content, prefix="lyrics_")
@@ -131,12 +131,12 @@ def infer(genre_txt_content, lyrics_txt_content):
         "--stage2_model", "m-a-p/YuE-s2-1B-general",
         "--genre_txt", f"{genre_txt_path}",
         "--lyrics_txt", f"{lyrics_txt_path}",
-        "--run_n_segments", "2",
+        "--run_n_segments", f"{num_segments}",
         "--stage2_batch_size", "8",  # Increased from 4 to 8
         "--output_dir", f"{output_dir}",
         "--cuda_idx", "0",
-        "--max_new_tokens", "3000",
-        #"--disable_offload_model"
+        "--max_new_tokens", f"{max_new_tokens}",
+        "--disable_offload_model"
     ]
 
     # Set up environment variables for CUDA with optimized settings
@@ -192,13 +192,30 @@ with gr.Blocks() as demo:
             with gr.Column():
                 genre_txt = gr.Textbox(label="Genre")
                 lyrics_txt = gr.Textbox(label="Lyrics")
-                submit_btn = gr.Button("Submit")
+                gr.Examples(
+                    examples = [
+                        [
+                            "female blues airy vocal bright vocal piano sad romantic guitar jazz",
+                            """
+                            [chorus]
+                            Don't let this moment fade, hold me close tonight
+                            With you here beside me, everything's alright
+                            Can't imagine life alone, don't want to let you go
+                            Stay with me forever, let our love just flow
+                            """
+                        ]
+                    ], 
+                    inputs = [genre_txt, lyrics_txt]
+                )
             with gr.Column():
+                num_segments = gr.Number(label="Number of Song Segments", info="number of paragraphs", value=1, interactive=False)
+                max_new_tokens = gr.Slider(label="Max New Tokens / Duration", info="1000 token = 10 seconds", minimum=500, maximum="24000", step=500, value=1500, interactive=False)
+                submit_btn = gr.Button("Submit")
                 music_out = gr.Audio(label="Audio Result")
     
     submit_btn.click(
         fn = infer, 
-        inputs = [genre_txt, lyrics_txt],
+        inputs = [genre_txt, lyrics_txt, num_segments, max_new_tokens],
         outputs = [music_out]
     )
 demo.queue().launch(show_api=False, show_error=True)