Update app.py

app.py

@@ -3,14 +3,16 @@ import requests
 import os
 import time
 import json
+import threading
 
 # Suno API key
 SUNO_KEY = os.environ.get("SunoKey", "")
 if not SUNO_KEY:
     print("⚠️ SunoKey not set!")
 
-# Store task info for auto-fill
+# Store task info for auto-fill and polling
 current_task_info = {}
+polling_active = {}
 
 def get_audio_files(task_id):
     """Get audio files from Suno task ID"""
@@ -99,6 +101,8 @@ def submit_separation_task(task_id, audio_id, separation_type):
                 current_task_info["original_task_id"] = task_id
                 current_task_info["audio_id"] = audio_id
                 current_task_info["separation_type"] = separation_type
+                current_task_info["start_time"] = time.time()
+                current_task_info["attempt"] = 0
                 
                 return f"✅ Task submitted!\n\n**Separation Task ID:** `{separation_task_id}`\n\n⏳ Starting auto-polling...", separation_task_id
             else:
@@ -184,8 +188,44 @@ def format_download_links(results, separation_task_id):
     
     return output
 
+# Manual polling function
+def manual_poll(task_id):
+    if not task_id:
+        return "❌ Enter Task ID"
+    
+    status, results, error = poll_separation_status(task_id)
+    
+    if error:
+        return f"❌ Error: {error}"
+    
+    if status == "SUCCESS":
+        if results:
+            output = f"✅ **Complete!**\n\n"
+            output += f"**Task ID:** `{task_id}`\n\n"
+            output += format_download_links(results, task_id)
+        else:
+            output = f"✅ **Complete** (no direct links)\n\n"
+            output += f"**Task ID:** `{task_id}`\n\n"
+            output += f"Check: [Viewer](https://1hit.no/viewer.php?task_id={task_id})"
+    
+    elif status in ["PENDING", "PROCESSING", "RUNNING"]:
+        output = f"⏳ **Status:** {status}\n\n"
+        output += f"**Task ID:** `{task_id}`\n\n"
+        output += "Still processing. Try again in 30 seconds."
+    
+    elif status == "FAILED":
+        output = f"❌ **Failed**\n\n"
+        output += f"**Task ID:** `{task_id}`\n"
+        output += f"**Status:** {status}"
+    
+    else:
+        output = f"🔄 **Status:** {status}\n\n"
+        output += f"**Task ID:** `{task_id}`"
+    
+    return output
+
 # Create the app
-with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
+with gr.Blocks() as app:
     gr.Markdown("# 🎵 Suno Stem Separator")
     
     with gr.Row():
@@ -241,6 +281,7 @@ with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
                     interactive=False,
                     visible=False
                 )
+                stop_poll_btn = gr.Button("⏹️ Stop Auto-Polling", variant="stop", visible=False)
             
             # Manual polling section
             with gr.Group():
@@ -265,8 +306,10 @@ with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
                     elem_id="viewer_link"
                 )
     
-    # Store current separation task ID
+    # Store current separation task ID for auto-polling
     current_separation_task_id = gr.State(value="")
+    polling_active_flag = gr.State(value=False)
+    polling_thread = None
     
     # Step 1: Get audio files
     def on_get_audio(task_id):
@@ -287,27 +330,29 @@ with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
         )
     
     # Step 2-3: Submit separation task
-    def on_submit(task_id, audio_id, sep_type):
+    def on_submit(task_id, audio_id, sep_type, active_flag):
         if not task_id or not audio_id:
-            return "❌ Missing Task ID or Audio ID", None, gr.Markdown(), gr.Slider(visible=False)
+            return "❌ Missing Task ID or Audio ID", "", "⏳ Waiting...", gr.Slider(visible=False), gr.Button(visible=False), active_flag
         
         status, separation_task_id = submit_separation_task(task_id, audio_id, sep_type)
         
         if not separation_task_id:
-            return status, None, gr.Markdown(), gr.Slider(visible=False)
+            return status, "", "⏳ Waiting...", gr.Slider(visible=False), gr.Button(visible=False), active_flag
         
         # Start auto-polling
         return (
             status,
             separation_task_id,
-            gr.Markdown("⏳ Polling started..."),
-            gr.Slider(visible=True, value=0)
+            "⏳ Polling started... (attempt 0/60)",
+            gr.Slider(visible=True, value=0),
+            gr.Button(visible=True),
+            True  # Set polling active
         )
     
-    # Auto-polling function
-    def auto_poll(separation_task_id, attempt):
-        if not separation_task_id:
-            return "⏳ Waiting for task...", 0, gr.Markdown(), False
+    # Auto-polling update function
+    def update_auto_poll(separation_task_id, attempt, active_flag):
+        if not active_flag or not separation_task_id:
+            return "⏳ Polling stopped", attempt, "", False
         
         attempt += 1
         
@@ -315,76 +360,58 @@ with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
         status, results, error = poll_separation_status(separation_task_id)
         
         if error:
-            return f"❌ Poll error: {error}", attempt, gr.Markdown(), attempt >= 60
+            return f"❌ Poll error: {error}", attempt, "", False
         
         if status == "SUCCESS":
             if results:
                 output = f"✅ **Separation Complete!**\n\n"
                 output += f"**Task ID:** `{separation_task_id}`\n\n"
                 output += format_download_links(results, separation_task_id)
-                return output, attempt, gr.Markdown(), True
+                return f"✅ Complete! (attempt {attempt})", attempt, output, False
             else:
                 output = f"✅ **Processing Complete**\n\n"
                 output += f"**Task ID:** `{separation_task_id}`\n\n"
                 output += "Check callback results: "
                 output += f"[Viewer](https://1hit.no/viewer.php?task_id={separation_task_id})"
-                return output, attempt, gr.Markdown(), True
+                return f"✅ Complete! (attempt {attempt})", attempt, output, False
         
         elif status in ["PENDING", "PROCESSING", "RUNNING"]:
-            output = f"⏳ **Status:** {status}\n\n"
-            output += f"**Task ID:** `{separation_task_id}`\n"
-            output += f"**Attempt:** {attempt}/60\n"
-            output += f"**Estimated wait:** {attempt * 5} seconds\n\n"
-            output += "Polling every 5 seconds..."
-            return output, attempt, gr.Markdown(), attempt >= 60
+            status_text = f"⏳ **Status:** {status}\n\n"
+            status_text += f"**Task ID:** `{separation_task_id}`\n"
+            status_text += f"**Attempt:** {attempt}/60\n"
+            status_text += f"**Elapsed:** {attempt * 5} seconds\n\n"
+            status_text += "Polling every 5 seconds..."
+            
+            if attempt >= 60:
+                status_text = f"⏰ **Timeout after 5 minutes**\n\n"
+                status_text += f"**Task ID:** `{separation_task_id}`\n"
+                status_text += "Check manually or wait for callback."
+                return status_text, attempt, "", False
+            
+            return status_text, attempt, "", True
         
         elif status == "FAILED":
-            output = f"❌ **Separation Failed**\n\n"
-            output += f"**Task ID:** `{separation_task_id}`\n"
-            output += f"**Status:** {status}\n"
-            return output, attempt, gr.Markdown(), True
+            status_text = f"❌ **Separation Failed**\n\n"
+            status_text += f"**Task ID:** `{separation_task_id}`\n"
+            status_text += f"**Status:** {status}"
+            return status_text, attempt, "", False
         
         else:
-            output = f"🔄 **Status:** {status}\n\n"
-            output += f"**Task ID:** `{separation_task_id}`\n"
-            output += f"**Attempt:** {attempt}/60\n"
-            return output, attempt, gr.Markdown(), attempt >= 60
+            status_text = f"🔄 **Status:** {status}\n\n"
+            status_text += f"**Task ID:** `{separation_task_id}`\n"
+            status_text += f"**Attempt:** {attempt}/60"
+            
+            if attempt >= 60:
+                status_text = f"⏰ **Timeout after 5 minutes**\n\n"
+                status_text += f"**Task ID:** `{separation_task_id}`\n"
+                status_text += "Check manually or wait for callback."
+                return status_text, attempt, "", False
+            
+            return status_text, attempt, "", True
     
-    # Manual polling function
-    def manual_poll(task_id):
-        if not task_id:
-            return "❌ Enter Task ID"
-        
-        status, results, error = poll_separation_status(task_id)
-        
-        if error:
-            return f"❌ Error: {error}"
-        
-        if status == "SUCCESS":
-            if results:
-                output = f"✅ **Complete!**\n\n"
-                output += f"**Task ID:** `{task_id}`\n\n"
-                output += format_download_links(results, task_id)
-            else:
-                output = f"✅ **Complete** (no direct links)\n\n"
-                output += f"**Task ID:** `{task_id}`\n\n"
-                output += f"Check: [Viewer](https://1hit.no/viewer.php?task_id={task_id})"
-        
-        elif status in ["PENDING", "PROCESSING", "RUNNING"]:
-            output = f"⏳ **Status:** {status}\n\n"
-            output += f"**Task ID:** `{task_id}`\n\n"
-            output += "Still processing. Try again in 30 seconds."
-        
-        elif status == "FAILED":
-            output = f"❌ **Failed**\n\n"
-            output += f"**Task ID:** `{task_id}`\n"
-            output += f"**Status:** {status}"
-        
-        else:
-            output = f"🔄 **Status:** {status}\n\n"
-            output += f"**Task ID:** `{task_id}`"
-        
-        return output
+    # Stop polling function
+    def stop_polling(active_flag):
+        return "⏹️ Polling stopped", 0, "", False, gr.Button(visible=False)
     
     # Connect events
     get_audio_btn.click(
@@ -393,23 +420,42 @@ with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
         outputs=[audio_status, audio_dropdown, separation_type, submit_btn, submission_output]
     )
     
+    # Submit button starts polling
     submit_btn.click(
         fn=on_submit,
-        inputs=[original_task_id, audio_dropdown, separation_type],
-        outputs=[submission_output, current_separation_task_id, auto_poll_status, auto_poll_progress]
-    ).then(
-        fn=auto_poll,
-        inputs=[current_separation_task_id, auto_poll_progress],
-        outputs=[auto_poll_status, auto_poll_progress, download_output, auto_poll_progress],
-        every=5  # Poll every 5 seconds
+        inputs=[original_task_id, audio_dropdown, separation_type, polling_active_flag],
+        outputs=[submission_output, current_separation_task_id, auto_poll_status, auto_poll_progress, stop_poll_btn, polling_active_flag]
     )
     
+    # Create a separate function for polling that runs independently
+    def start_polling_thread(separation_task_id, attempt, active_flag):
+        """Background polling thread"""
+        while active_flag and attempt < 60:
+            time.sleep(5)
+            attempt += 1
+            
+            status, results, error = poll_separation_status(separation_task_id)
+            
+            if status == "SUCCESS" or status == "FAILED" or error:
+                # Update UI via queue or state
+                break
+        
+        return attempt
+    
+    # Manual polling
     poll_btn.click(
         fn=manual_poll,
         inputs=[poll_task_id],
         outputs=[poll_output]
     )
     
+    # Stop polling button
+    stop_poll_btn.click(
+        fn=stop_polling,
+        inputs=[polling_active_flag],
+        outputs=[auto_poll_status, auto_poll_progress, download_output, polling_active_flag, stop_poll_btn]
+    )
+    
     # Auto-fill poll field when separation task is submitted
     def update_poll_field(separation_task_id):
         if separation_task_id:
@@ -421,9 +467,43 @@ with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
         inputs=[current_separation_task_id],
         outputs=[poll_task_id]
     )
+    
+    # Simple polling timer (manual refresh)
+    def manual_refresh(separation_task_id, attempt, active_flag):
+        if active_flag and separation_task_id:
+            return update_auto_poll(separation_task_id, attempt, active_flag)
+        return "⏳ Waiting for task...", 0, "", active_flag
+    
+    # Add a refresh button for manual polling updates
+    refresh_btn = gr.Button("🔄 Refresh Status", variant="secondary", visible=False)
+    
+    # Show/hide refresh button based on polling state
+    def toggle_refresh_button(active_flag):
+        return gr.Button(visible=active_flag)
+    
+    polling_active_flag.change(
+        fn=toggle_refresh_button,
+        inputs=[polling_active_flag],
+        outputs=[refresh_btn]
+    )
+    
+    refresh_btn.click(
+        fn=manual_refresh,
+        inputs=[current_separation_task_id, auto_poll_progress, polling_active_flag],
+        outputs=[auto_poll_status, auto_poll_progress, download_output, polling_active_flag]
+    )
 
+# Launch the app
 if __name__ == "__main__":
     print("🚀 Starting Suno Stem Separator")
     print(f"🔑 SunoKey: {'✅ Set' if SUNO_KEY else '❌ Not set'}")
     print("🌐 Open your browser to: http://localhost:7860")
-    app.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    
+    # Launch with theme parameter in launch() method
+    app.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        title="Suno Stem Separator",
+        theme="soft"
+    )