Update app.py

app.py

@@ -9,8 +9,64 @@ SUNO_KEY = os.environ.get("SunoKey", "")
 if not SUNO_KEY:
     print("⚠️ SunoKey not set!")
 
+# Store task info for auto-fill
+current_task_info = {}
+
+def get_audio_files(task_id):
+    """Get audio files from Suno task ID"""
+    if not SUNO_KEY:
+        return "❌ Error: SunoKey not configured", []
+    
+    if not task_id.strip():
+        return "❌ Please enter a Task ID", []
+    
+    try:
+        resp = requests.get(
+            "https://api.sunoapi.org/api/v1/generate/record-info",
+            headers={"Authorization": f"Bearer {SUNO_KEY}"},
+            params={"taskId": task_id.strip()},
+            timeout=30
+        )
+        
+        if resp.status_code != 200:
+            return f"❌ Request failed: HTTP {resp.status_code}", []
+        
+        data = resp.json()
+        
+        if data.get("code") != 200:
+            return f"❌ API error: {data.get('msg', 'Unknown')}", []
+        
+        # Check status
+        status = data.get("data", {}).get("status", "UNKNOWN")
+        if status != "SUCCESS":
+            return f"⏳ Task status: {status}. Wait for generation to complete.", []
+        
+        # Get audio files
+        suno_data = data.get("data", {}).get("response", {}).get("sunoData", [])
+        if not suno_data:
+            return "❌ No audio files found", []
+        
+        # Create dropdown options
+        audio_options = []
+        for i, audio in enumerate(suno_data):
+            audio_id = audio.get("id", f"audio_{i}")
+            title = audio.get("title", f"Track {i+1}")
+            duration = audio.get("duration", 0)
+            prompt = audio.get("prompt", "")[:50]
+            
+            display = f"{i+1}. {title} ({duration:.1f}s)"
+            if prompt:
+                display += f" - {prompt}..."
+            
+            audio_options.append((display, audio_id))
+        
+        return f"✅ Found {len(audio_options)} audio file(s)", audio_options
+        
+    except Exception as e:
+        return f"❌ Error: {str(e)}", []
+
 def submit_separation_task(task_id, audio_id, separation_type):
-    """Submit stem separation task - SIMPLE VERSION"""
+    """Submit separation task"""
     try:
         resp = requests.post(
             "https://api.sunoapi.org/api/v1/vocal-removal/generate",
@@ -30,239 +86,341 @@ def submit_separation_task(task_id, audio_id, separation_type):
         if resp.status_code == 200:
             data = resp.json()
             
-            # Try different response formats
+            # Get separation task ID
+            separation_task_id = None
             if "taskId" in data:
-                return data["taskId"], None  # New format
+                separation_task_id = data["taskId"]
             elif data.get("code") == 200 and "data" in data and "taskId" in data["data"]:
-                return data["data"]["taskId"], None  # Nested format
+                separation_task_id = data["data"]["taskId"]
+            
+            if separation_task_id:
+                # Store for auto-fill
+                current_task_info["separation_task_id"] = separation_task_id
+                current_task_info["original_task_id"] = task_id
+                current_task_info["audio_id"] = audio_id
+                current_task_info["separation_type"] = separation_type
+                
+                return f"✅ Task submitted!\n\n**Separation Task ID:** `{separation_task_id}`\n\n⏳ Starting auto-polling...", separation_task_id
             else:
-                return None, f"Unexpected response format: {data}"
+                return f"❌ No task ID in response:\n{json.dumps(data, indent=2)}", None
         else:
-            return None, f"HTTP Error: {resp.status_code}"
+            return f"❌ HTTP Error {resp.status_code}:\n{resp.text}", None
             
     except Exception as e:
-        return None, f"Error: {str(e)}"
+        return f"❌ Error: {str(e)}", None
 
-def poll_task_status(task_id):
-    """Poll separation task status - SIMPLE VERSION"""
+def poll_separation_status(separation_task_id):
+    """Poll separation task status"""
     try:
         resp = requests.get(
             "https://api.sunoapi.org/api/v1/vocal-removal/record-info",
             headers={"Authorization": f"Bearer {SUNO_KEY}"},
-            params={"taskId": task_id},
+            params={"taskId": separation_task_id},
             timeout=30
         )
         
         if resp.status_code == 200:
             data = resp.json()
             
-            # Simple status extraction
+            # Get status
             status = "UNKNOWN"
             if "status" in data:
                 status = data["status"]
             elif "data" in data and "status" in data["data"]:
                 status = data["data"]["status"]
             elif data.get("code") == 200:
-                status = "SUCCESS"  # Assume success if code is 200
+                status = "SUCCESS"
             
-            # Simple result extraction
+            # Get results
             results = {}
             if "vocal_removal_info" in data:
                 results = data["vocal_removal_info"]
             elif "data" in data and "vocal_removal_info" in data["data"]:
                 results = data["data"]["vocal_removal_info"]
-            elif data.get("code") == 200 and "data" in data:
-                # Maybe results are directly in data
-                results = {k: v for k, v in data["data"].items() if k.endswith("_url")}
             
             return status, results, None
         else:
-            return "ERROR", {}, f"HTTP {resp.status_code}"
+            return "ERROR", {}, f"HTTP {resp.status_code}: {resp.text}"
             
     except Exception as e:
         return "ERROR", {}, str(e)
 
-def format_results(results):
-    """Format results as simple markdown"""
+def format_download_links(results, separation_task_id):
+    """Format download links from results"""
     if not results:
-        return "No results found"
+        return "No download links found"
     
     output = "## 🎵 Download Links\n\n"
-    for key, url in results.items():
-        if url and "url" in key.lower():
-            stem_name = key.replace("_url", "").replace("_", " ").title()
-            output += f"**{stem_name}:** [Download]({url})\n"
+    
+    # 2-stem separation
+    if "vocal_url" in results or "instrumental_url" in results:
+        if results.get("vocal_url"):
+            output += f"**🎤 Vocals:** [Download MP3]({results['vocal_url']})\n"
+        if results.get("instrumental_url"):
+            output += f"**🎵 Instrumental:** [Download MP3]({results['instrumental_url']})\n"
+    
+    # 12-stem separation
+    stem_fields = [
+        ("backing_vocals_url", "🎤 Backing Vocals"),
+        ("bass_url", "🎸 Bass"),
+        ("brass_url", "🎺 Brass"),
+        ("drums_url", "🥁 Drums"),
+        ("fx_url", "🎛️ FX/Other"),
+        ("guitar_url", "🎸 Guitar"),
+        ("keyboard_url", "🎹 Keyboard"),
+        ("percussion_url", "🪘 Percussion"),
+        ("strings_url", "🎻 Strings"),
+        ("synth_url", "🎹 Synth"),
+        ("woodwinds_url", "🎷 Woodwinds"),
+        ("vocal_url", "🎤 Vocals"),
+        ("instrumental_url", "🎵 Instrumental"),
+    ]
+    
+    for field, name in stem_fields:
+        if results.get(field):
+            output += f"**{name}:** [Download MP3]({results[field]})\n"
+    
+    output += f"\n**🔗 Viewer:** [Open in Viewer](https://1hit.no/viewer.php?task_id={separation_task_id})"
     
     return output
 
 # Create the app
 with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
     gr.Markdown("# 🎵 Suno Stem Separator")
-    gr.Markdown("Separate Suno tracks into stems")
     
-    with gr.Tab("Submit New Task"):
-        with gr.Row():
-            with gr.Column():
-                task_id_input = gr.Textbox(
+    with gr.Row():
+        # Left column: Input and control
+        with gr.Column(scale=1):
+            # Step 1: Get audio files
+            with gr.Group():
+                gr.Markdown("### 1. Get Audio Files")
+                original_task_id = gr.Textbox(
                     label="Original Task ID",
-                    placeholder="From your Suno generation"
+                    placeholder="Enter Suno generation task ID",
+                    info="From your Suno history"
                 )
-                audio_id_input = gr.Textbox(
-                    label="Audio ID",
-                    placeholder="From your Suno generation"
+                get_audio_btn = gr.Button("📥 Get Audio Files", variant="secondary")
+                audio_status = gr.Markdown("Enter Task ID above")
+            
+            # Step 2: Select audio file
+            with gr.Group():
+                gr.Markdown("### 2. Select Audio File")
+                audio_dropdown = gr.Dropdown(
+                    label="Select Audio",
+                    choices=[],
+                    interactive=True,
+                    visible=False
                 )
+            
+            # Step 3: Start separation
+            with gr.Group():
+                gr.Markdown("### 3. Start Separation")
                 separation_type = gr.Radio(
-                    choices=["separate_vocal", "split_stem"],
                     label="Separation Type",
-                    value="separate_vocal"
+                    choices=[
+                        ("🎤 Vocals Only (1 credit)", "separate_vocal"),
+                        ("🎛️ Full Stems (5 credits)", "split_stem")
+                    ],
+                    value="separate_vocal",
+                    visible=False
                 )
-                submit_btn = gr.Button("Submit Separation", variant="primary")
-                
-            with gr.Column():
-                task_output = gr.Markdown("## Task Status")
-    
-    with gr.Tab("Poll Existing Task"):
-        with gr.Row():
-            with gr.Column():
+                submit_btn = gr.Button("🚀 Start Separation", variant="primary", visible=False)
+                submission_output = gr.Markdown("Select audio file first", visible=False)
+        
+        # Right column: Results and polling
+        with gr.Column(scale=2):
+            # Auto-polling section
+            with gr.Group():
+                gr.Markdown("### 4. Auto-Polling Status")
+                auto_poll_status = gr.Markdown("Waiting for task submission...")
+                auto_poll_progress = gr.Slider(
+                    minimum=0,
+                    maximum=60,
+                    value=0,
+                    label="Polling attempts",
+                    interactive=False,
+                    visible=False
+                )
+            
+            # Manual polling section
+            with gr.Group():
+                gr.Markdown("### 5. Manual Polling")
                 poll_task_id = gr.Textbox(
                     label="Separation Task ID",
-                    placeholder="Enter task ID to poll"
+                    placeholder="Will auto-fill from current task"
+                )
+                poll_btn = gr.Button("🔍 Check Status", variant="secondary")
+                poll_output = gr.Markdown("Enter Task ID to check")
+            
+            # Results section
+            with gr.Group():
+                gr.Markdown("### 6. Download Links")
+                download_output = gr.Markdown("Results will appear here")
+            
+            # Viewer link
+            with gr.Group():
+                gr.Markdown("### 7. Viewer")
+                viewer_link = gr.Markdown(
+                    "[Open Viewer](https://1hit.no/viewer.php)",
+                    elem_id="viewer_link"
                 )
-                poll_btn = gr.Button("Check Status", variant="secondary")
-                
-            with gr.Column():
-                poll_output = gr.Markdown("## Poll Results")
     
-    with gr.Tab("Instructions"):
-        gr.Markdown("""
-        ## 📋 How to Use
-        
-        1. **Get Task ID & Audio ID:**
-           - Go to your Suno generation history
-           - Copy the Task ID and Audio ID
-        
-        2. **Submit Separation:**
-           - Enter Task ID and Audio ID
-           - Choose separation type
-           - Click "Submit Separation"
-           - You'll get a new Separation Task ID
+    # Store current separation task ID
+    current_separation_task_id = gr.State(value="")
+    
+    # Step 1: Get audio files
+    def on_get_audio(task_id):
+        if not task_id:
+            return "❌ Enter Task ID", gr.Dropdown(choices=[], visible=False), gr.Radio(visible=False), gr.Button(visible=False), gr.Markdown(visible=False)
         
-        3. **Get Results:**
-           - Wait 1-3 minutes
-           - Use the Poll tab to check status
-           - OR wait for callback to https://1hit.no/callback.php
+        status, options = get_audio_files(task_id)
         
-        4. **View Results:**
-           - Visit: https://1hit.no/viewer.php?task_id=YOUR_TASK_ID
+        if not options:
+            return status, gr.Dropdown(choices=[], visible=False), gr.Radio(visible=False), gr.Button(visible=False), gr.Markdown(visible=False)
         
-        ## ⚠️ Notes
-        - Processing takes 1-3 minutes
-        - Each separation uses credits
-        - Links expire after some time
-        """)
+        return (
+            status,
+            gr.Dropdown(choices=options, value=options[0][1] if options else None, visible=True),
+            gr.Radio(visible=True),
+            gr.Button(visible=True),
+            gr.Markdown("Ready to separate!", visible=True)
+        )
     
-    # Submit task function
+    # Step 2-3: Submit separation task
     def on_submit(task_id, audio_id, sep_type):
         if not task_id or not audio_id:
-            return "❌ Please enter both Task ID and Audio ID"
+            return "❌ Missing Task ID or Audio ID", None, gr.Markdown(), gr.Slider(visible=False)
         
-        result = submit_separation_task(task_id, audio_id, sep_type)
+        status, separation_task_id = submit_separation_task(task_id, audio_id, sep_type)
         
-        if result[1]:  # Error
-            return f"❌ {result[1]}"
+        if not separation_task_id:
+            return status, None, gr.Markdown(), gr.Slider(visible=False)
         
-        separation_task_id = result[0]
+        # Start auto-polling
+        return (
+            status,
+            separation_task_id,
+            gr.Markdown("⏳ Polling started..."),
+            gr.Slider(visible=True, value=0)
+        )
+    
+    # Auto-polling function
+    def auto_poll(separation_task_id, attempt):
+        if not separation_task_id:
+            return "⏳ Waiting for task...", 0, gr.Markdown(), False
         
-        output = f"""
-        ## ✅ Task Submitted!
+        attempt += 1
         
-        **Separation Task ID:** `{separation_task_id}`
+        # Poll for status
+        status, results, error = poll_separation_status(separation_task_id)
         
-        **Next Steps:**
-        1. Wait 1-3 minutes for processing
-        2. Check status in the **Poll** tab using this Task ID
-        3. Results will also be sent to callback endpoint
+        if error:
+            return f"❌ Poll error: {error}", attempt, gr.Markdown(), attempt >= 60
         
-        **Poll this task:** Use ID: `{separation_task_id}`
-        """
+        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
+            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 output
+        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
+        
+        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
+        
+        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
     
-    # Poll task function  
-    def on_poll(task_id):
+    # Manual polling function
+    def manual_poll(task_id):
         if not task_id:
-            return "❌ Please enter a Task ID"
+            return "❌ Enter Task ID"
         
-        status, results, error = poll_task_status(task_id)
+        status, results, error = poll_separation_status(task_id)
         
         if error:
             return f"❌ Error: {error}"
         
         if status == "SUCCESS":
             if results:
-                return f"""
-                ## ✅ Separation Complete!
-                
-                **Status:** {status}
-                **Task ID:** `{task_id}`
-                
-                {format_results(results)}
-                
-                **Also check:** https://1hit.no/viewer.php?task_id={task_id}
-                """
+                output = f"✅ **Complete!**\n\n"
+                output += f"**Task ID:** `{task_id}`\n\n"
+                output += format_download_links(results, task_id)
             else:
-                return f"""
-                ## ✅ Separation Complete!
-                
-                **Status:** {status}
-                **Task ID:** `{task_id}`
-                
-                No direct links found in response.
-                Check callback results: https://1hit.no/viewer.php?task_id={task_id}
-                """
+                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"]:
-            return f"""
-            ## ⏳ Still Processing
-            
-            **Status:** {status}
-            **Task ID:** `{task_id}`
-            
-            Processing usually takes 1-3 minutes.
-            Check again in a minute.
-            """
+            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":
-            return f"""
-            ## ❌ Separation Failed
-            
-            **Status:** {status}
-            **Task ID:** `{task_id}`
-            
-            The separation task failed.
-            Please check your inputs and try again.
-            """
+            output = f"❌ **Failed**\n\n"
+            output += f"**Task ID:** `{task_id}`\n"
+            output += f"**Status:** {status}"
+        
         else:
-            return f"""
-            ## 🔄 Status Check
-            
-            **Status:** {status}
-            **Task ID:** `{task_id}`
-            
-            Current status: {status}
-            Keep checking every 30 seconds.
-            """
+            output = f"🔄 **Status:** {status}\n\n"
+            output += f"**Task ID:** `{task_id}`"
+        
+        return output
+    
+    # Connect events
+    get_audio_btn.click(
+        fn=on_get_audio,
+        inputs=[original_task_id],
+        outputs=[audio_status, audio_dropdown, separation_type, submit_btn, submission_output]
+    )
     
-    # Connect buttons
     submit_btn.click(
         fn=on_submit,
-        inputs=[task_id_input, audio_id_input, separation_type],
-        outputs=[task_output]
+        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
     )
     
     poll_btn.click(
-        fn=on_poll,
+        fn=manual_poll,
         inputs=[poll_task_id],
         outputs=[poll_output]
     )
+    
+    # Auto-fill poll field when separation task is submitted
+    def update_poll_field(separation_task_id):
+        if separation_task_id:
+            return separation_task_id
+        return ""
+    
+    current_separation_task_id.change(
+        fn=update_poll_field,
+        inputs=[current_separation_task_id],
+        outputs=[poll_task_id]
+    )
 
 if __name__ == "__main__":
     print("🚀 Starting Suno Stem Separator")