Update app.py

app.py

@@ -49,10 +49,10 @@ def generate_music(prompt, max_new_tokens: int = 128):
     return sampling_rate, audio
 
 # Combined Gradio function
-def main(user_input):
+def main(user_input, max_new_tokens):
     detailed_prompt = refine_prompt(user_input)
-    sampling_rate, audio = generate_music(detailed_prompt)
-    return detailed_prompt, (sampling_rate, audio)
+    sampling_rate, audio = generate_music(detailed_prompt, max_new_tokens)
+    return detailed_prompt, (sampling_rate, audio), "/tmp/output.wav"
 
 # Build Gradio UI
 with gr.Blocks() as demo:
@@ -65,13 +65,13 @@ Enter a music idea or mood and get a short AI-generated track. (CPU mode)""")
 
     refined_output = gr.Textbox(label="Enhanced Prompt by GPT")
     audio_output = gr.Audio(label="Generated Audio", type="numpy")
-    download_wav = gr.File(label="Download .wav file", file_name="generated.wav")
+    download_wav = gr.File(label="Download .wav file")
 
     generate_btn.click(
-        lambda inp, tok: (main(inp)[0], (main(inp)[1][0], main(inp)[1][1]), "/tmp/output.wav"),
+        main,
         inputs=[user_input, max_tokens],
         outputs=[refined_output, audio_output, download_wav]
     )
 
-# Launch in SSR mode
+# Launch in SSR mode with queue enabled
 demo.launch(server_name="0.0.0.0", server_port=7860, share=False, enable_queue=True)