Update app.py

app.py

@@ -3,12 +3,18 @@ import requests
 import os
 import time
 import json
+import uuid
+from datetime import datetime
 
 # Suno API key
 SUNO_KEY = os.environ.get("SunoKey", "")
 if not SUNO_KEY:
     print("⚠️ SunoKey not set!")
 
+# Store ongoing separation tasks (in-memory for demo)
+# In production, use a database or Redis
+separation_tasks = {}
+
 def get_audio_info(task_id):
     """Get audio information from Suno task ID"""
     if not SUNO_KEY:
@@ -83,30 +89,20 @@ def get_audio_info(task_id):
     except Exception as e:
         return f"❌ Error: {str(e)}", []
 
-def separate_vocals(task_id, audio_id, separation_type, audio_data=None):
-    """Separate vocals and instruments from Suno tracks"""
-    if not SUNO_KEY:
-        yield "❌ Error: SunoKey not configured in environment variables"
-        return
-    
-    if not task_id.strip() or not audio_id.strip():
-        yield "❌ Error: Please enter both Task ID and Audio ID"
-        return
-    
-    # Validate separation type
-    if separation_type not in ["separate_vocal", "split_stem"]:
-        yield "❌ Error: Invalid separation type"
-        return
-    
-    # Submit separation task
+def submit_separation_task(task_id, audio_id, separation_type):
+    """Submit stem separation task and return task ID"""
     try:
+        # Generate a unique callback URL for this request
+        callback_id = str(uuid.uuid4())
+        callback_url = f"https://1hit.no/callback.php?callback_id={callback_id}"
+        
         resp = requests.post(
             "https://api.sunoapi.org/api/v1/vocal-removal/generate",
             json={
                 "taskId": task_id,
                 "audioId": audio_id,
                 "type": separation_type,
-                "callBackUrl": "https://1hit.no/callback.php"  # Required but not used for polling
+                "callBackUrl": callback_url  # Using your callback endpoint
             },
             headers={
                 "Authorization": f"Bearer {SUNO_KEY}",
@@ -115,103 +111,236 @@ def separate_vocals(task_id, audio_id, separation_type, audio_data=None):
             timeout=30
         )
         
+        print(f"Separation request response: {resp.status_code}")
+        print(f"Response text: {resp.text}")
+        
         if resp.status_code != 200:
-            yield f"❌ Submission failed: HTTP {resp.status_code}"
-            return
+            return None, f"❌ Submission failed: HTTP {resp.status_code}"
         
         data = resp.json()
-        if data.get("code") != 200:
-            yield f"❌ API error: {data.get('msg', 'Unknown')}"
-            return
+        print(f"Parsed response: {data}")
+        
+        # Check for different response formats
+        if data.get("code") == 200:
+            # New format with data object
+            separation_task_id = data.get("data", {}).get("taskId")
+            music_id = data.get("data", {}).get("musicId")
+            
+            if separation_task_id:
+                # Store task info for polling
+                separation_tasks[separation_task_id] = {
+                    "task_id": task_id,
+                    "audio_id": audio_id,
+                    "separation_type": separation_type,
+                    "submitted_at": datetime.now().isoformat(),
+                    "callback_id": callback_id,
+                    "music_id": music_id
+                }
+                
+                return separation_task_id, None
+            else:
+                return None, "❌ No taskId in response"
+        
+        # Try alternative response format
+        elif "taskId" in data:
+            separation_task_id = data.get("taskId")
+            music_id = data.get("musicId")
+            
+            separation_tasks[separation_task_id] = {
+                "task_id": task_id,
+                "audio_id": audio_id,
+                "separation_type": separation_type,
+                "submitted_at": datetime.now().isoformat(),
+                "callback_id": callback_id,
+                "music_id": music_id
+            }
+            
+            return separation_task_id, None
+        
+        else:
+            error_msg = data.get("msg", data.get("error", "Unknown error"))
+            return None, f"❌ API error: {error_msg}"
+            
+    except Exception as e:
+        return None, f"❌ Error submitting task: {str(e)}"
+
+def poll_separation_result(separation_task_id):
+    """Poll for separation results using the vocal-removal/record-info endpoint"""
+    try:
+        check = requests.get(
+            "https://api.sunoapi.org/api/v1/vocal-removal/record-info",
+            headers={"Authorization": f"Bearer {SUNO_KEY}"},
+            params={"taskId": separation_task_id},
+            timeout=30
+        )
+        
+        if check.status_code == 200:
+            data = check.json()
+            print(f"Polling response: {data}")
+            
+            if data.get("code") == 200:
+                status = data.get("data", {}).get("status", "PENDING")
+                return status, data.get("data", {})
+            else:
+                return "ERROR", {"error": data.get("msg", "Unknown error")}
+        else:
+            return "ERROR", {"error": f"HTTP {check.status_code}"}
+            
+    except Exception as e:
+        return "ERROR", {"error": str(e)}
+
+def separate_vocals(task_id, audio_id, separation_type):
+    """Separate vocals and instruments from Suno tracks"""
+    if not SUNO_KEY:
+        yield "❌ Error: SunoKey not configured in environment variables"
+        return
+    
+    if not task_id.strip() or not audio_id.strip():
+        yield "❌ Error: Please enter both Task ID and Audio ID"
+        return
+    
+    # Validate separation type
+    if separation_type not in ["separate_vocal", "split_stem"]:
+        yield "❌ Error: Invalid separation type"
+        return
+    
+    # Step 1: Submit separation task
+    yield "⏳ **Submitting separation request...**\n"
+    
+    separation_task_id, error = submit_separation_task(task_id, audio_id, separation_type)
+    
+    if error:
+        yield error
+        return
+    
+    if not separation_task_id:
+        yield "❌ Failed to get separation task ID"
+        return
+    
+    # Store the separation task ID
+    separation_task_info = separation_tasks.get(separation_task_id, {})
+    
+    yield f"✅ **Separation Task Submitted!**\n\n"
+    yield f"**Separation Task ID:** `{separation_task_id}`\n"
+    if separation_task_info.get("music_id"):
+        yield f"**Music ID:** `{separation_task_info.get('music_id')}`\n"
+    yield f"**Callback URL:** `{separation_task_info.get('callback_url', 'https://1hit.no/callback.php')}`\n\n"
+    yield "⏳ **Processing separation (this may take 1-3 minutes)...**\n"
+    yield "_Polling API for results..._\n"
+    
+    # Step 2: Poll for results
+    for attempt in range(60):  # 60 attempts * 5 seconds = 5 minutes max
+        time.sleep(5)
         
-        separation_task_id = data["data"]["taskId"]
-        yield f"✅ **Submitted!**\nSeparation Task ID: `{separation_task_id}`\n\n⏳ Processing separation...\n"
+        status, result_data = poll_separation_result(separation_task_id)
         
-        # Poll for results
-        for attempt in range(40):  # 40 attempts * 5 seconds = 200 seconds max
-            time.sleep(5)
+        if status == "SUCCESS":
+            # Parse successful result
+            separation_info = result_data.get("vocal_removal_info", {})
+            
+            if not separation_info:
+                # Try alternative field names
+                separation_info = result_data
             
-            try:
-                check = requests.get(
-                    "https://api.sunoapi.org/api/v1/vocal-removal/record-info",
-                    headers={"Authorization": f"Bearer {SUNO_KEY}"},
-                    params={"taskId": separation_task_id},
-                    timeout=30
-                )
+            if not separation_info:
+                yield "✅ Separation completed but no download links found in response\n"
+                yield f"Check your callback endpoint: {separation_task_info.get('callback_url')}"
+                return
+            
+            # Format output based on separation type
+            output = "🎵 **Separation Complete!** 🎵\n\n"
+            output += f"**Separation Task ID:** `{separation_task_id}`\n"
+            output += f"**Original Task ID:** `{task_id}`\n"
+            output += f"**Audio ID:** `{audio_id}`\n\n"
+            
+            if separation_type == "separate_vocal":
+                output += "## 🎤 2-Stem Separation Results\n\n"
                 
-                if check.status_code == 200:
-                    check_data = check.json()
-                    status = check_data["data"].get("status", "PENDING")
-                    
-                    if status == "SUCCESS":
-                        # Success! Extract separation results
-                        separation_info = check_data["data"].get("vocal_removal_info", {})
-                        
-                        if not separation_info:
-                            yield "✅ Completed but no separation results found"
-                            return
-                        
-                        # Format output based on separation type
-                        output = "🎵 **Separation Complete!**\n\n"
-                        
-                        if separation_type == "separate_vocal":
-                            output += "## 2-Stem Separation\n\n"
-                            output += f"**Vocals:** [Download]({separation_info.get('vocal_url', 'N/A')})\n"
-                            output += f"**Instrumental:** [Download]({separation_info.get('instrumental_url', 'N/A')})\n"
-                            if separation_info.get('origin_url'):
-                                output += f"**Original:** [Download]({separation_info.get('origin_url')})\n"
-                            
-                        elif separation_type == "split_stem":
-                            output += "## 12-Stem Separation\n\n"
-                            
-                            stems = [
-                                ("Vocals", separation_info.get('vocal_url')),
-                                ("Backing Vocals", separation_info.get('backing_vocals_url')),
-                                ("Drums", separation_info.get('drums_url')),
-                                ("Bass", separation_info.get('bass_url')),
-                                ("Guitar", separation_info.get('guitar_url')),
-                                ("Keyboard", separation_info.get('keyboard_url')),
-                                ("Strings", separation_info.get('strings_url')),
-                                ("Brass", separation_info.get('brass_url')),
-                                ("Woodwinds", separation_info.get('woodwinds_url')),
-                                ("Percussion", separation_info.get('percussion_url')),
-                                ("Synth", separation_info.get('synth_url')),
-                                ("FX/Other", separation_info.get('fx_url')),
-                                ("Instrumental", separation_info.get('instrumental_url')),
-                                ("Original", separation_info.get('origin_url'))
-                            ]
-                            
-                            for stem_name, stem_url in stems:
-                                if stem_url:
-                                    output += f"**{stem_name}:** [Download]({stem_url})\n"
-                        
-                        output += f"\n⏱️ Processed in about {(attempt + 1) * 5} seconds\n"
-                        output += f"⚠️ **Note:** Download links expire in 14 days"
-                        
-                        yield output
-                        return
-                    
-                    elif status == "FAILED":
-                        error = check_data["data"].get("errorMessage", "Unknown error")
-                        yield f"❌ Separation failed: {error}"
-                        return
+                # Try different possible field names
+                vocal_url = (separation_info.get('vocal_url') or 
+                           separation_info.get('vocal') or 
+                           separation_info.get('vocals_url'))
+                instrumental_url = (separation_info.get('instrumental_url') or 
+                                  separation_info.get('instrumental') or 
+                                  separation_info.get('accompaniment_url'))
+                origin_url = separation_info.get('origin_url') or separation_info.get('original_url')
+                
+                if vocal_url:
+                    output += f"**🎤 Vocals:** [Download MP3]({vocal_url})\n"
+                if instrumental_url:
+                    output += f"**🎵 Instrumental:** [Download MP3]({instrumental_url})\n"
+                if origin_url:
+                    output += f"**📁 Original:** [Download MP3]({origin_url})\n"
+                
+            elif separation_type == "split_stem":
+                output += "## 🎛️ 12-Stem Separation Results\n\n"
+                
+                # Map of possible stem field names
+                stem_fields = {
+                    "Vocals": ["vocal_url", "vocals", "vocal", "vocals_url"],
+                    "Backing Vocals": ["backing_vocals_url", "backing_vocals", "backing"],
+                    "Drums": ["drums_url", "drums", "drum"],
+                    "Bass": ["bass_url", "bass"],
+                    "Guitar": ["guitar_url", "guitar"],
+                    "Keyboard": ["keyboard_url", "keyboard", "piano"],
+                    "Strings": ["strings_url", "strings"],
+                    "Brass": ["brass_url", "brass"],
+                    "Woodwinds": ["woodwinds_url", "woodwinds"],
+                    "Percussion": ["percussion_url", "percussion"],
+                    "Synth": ["synth_url", "synth", "synthesizer"],
+                    "FX/Other": ["fx_url", "fx", "other", "effects"],
+                    "Instrumental": ["instrumental_url", "instrumental", "accompaniment_url"],
+                    "Original": ["origin_url", "original_url", "original"]
+                }
+                
+                found_stems = 0
+                for stem_name, possible_fields in stem_fields.items():
+                    url = None
+                    for field in possible_fields:
+                        if field in separation_info and separation_info[field]:
+                            url = separation_info[field]
+                            break
                     
-                    else:
-                        # PENDING or PROCESSING
-                        yield (f"⏳ Status: {status}\n"
-                               f"Attempt: {attempt + 1}/40\n"
-                               f"Separation Task ID: `{separation_task_id}`\n\n"
-                               f"Processing... (Usually takes 30-120 seconds)")
-                else:
-                    yield f"⚠️ Check error: HTTP {check.status_code}"
+                    if url:
+                        found_stems += 1
+                        output += f"**{stem_name}:** [Download MP3]({url})\n"
+                
+                if found_stems == 0:
+                    # If no stems found, show raw data for debugging
+                    output += "⚠️ **No stem URLs found in response. Raw data:**\n"
+                    output += f"```json\n{json.dumps(separation_info, indent=2)}\n```\n"
             
-            except Exception as e:
-                yield f"⚠️ Error checking status: {str(e)}"
-        
-        yield "⏰ Timeout after 200 seconds. Try checking again later."
+            output += f"\n⏱️ **Processing time:** {(attempt + 1) * 5} seconds\n"
+            output += "⚠️ **Note:** Download links may expire after some time\n"
+            output += f"📋 **Callback ID:** `{separation_task_info.get('callback_id', 'N/A')}`\n"
+            
+            yield output
+            return
+            
+        elif status == "FAILED":
+            error_msg = result_data.get("error", result_data.get("errorMessage", "Unknown error"))
+            yield f"❌ **Separation failed:** {error_msg}\n"
+            yield f"**Task ID:** `{separation_task_id}`\n"
+            return
+            
+        elif status == "ERROR":
+            error_msg = result_data.get("error", "Unknown error")
+            yield f"⚠️ **Polling error:** {error_msg}\n"
+            yield f"**Task ID:** `{separation_task_id}`\n"
+            # Continue polling despite error
+        else:
+            # PENDING or PROCESSING
+            yield (f"⏳ **Status:** {status}\n"
+                   f"**Attempt:** {attempt + 1}/60\n"
+                   f"**Task ID:** `{separation_task_id}`\n\n"
+                   f"_Still processing... (usually takes 1-3 minutes)_\n")
     
-    except Exception as e:
-        yield f"❌ Error: {str(e)}"
+    yield "⏰ **Timeout after 5 minutes.**\n"
+    yield f"The separation task is still processing.\n"
+    yield f"**Task ID:** `{separation_task_id}`\n"
+    yield f"**Music ID:** `{separation_task_info.get('music_id', 'N/A')}`\n"
+    yield "You can check the status later using this Task ID.\n"
+    yield f"Or check your callback endpoint: {separation_task_info.get('callback_url')}"
 
 # Create the app
 with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
@@ -222,7 +351,7 @@ with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
         with gr.Column(scale=1):
             # Step 1: Enter Task ID
             task_id_input = gr.Textbox(
-                label="Task ID",
+                label="Original Task ID",
                 placeholder="Example: 5c79****be8e",
                 info="Enter your Suno generation task ID",
                 elem_id="task_id_input"
@@ -239,9 +368,6 @@ with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
                 interactive=True
             )
             
-            # Hidden store for audio data
-            audio_data_store = gr.State(value=None)
-            
             separation_type = gr.Radio(
                 label="Separation Type",
                 choices=[
@@ -255,27 +381,32 @@ with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
             
             separate_btn = gr.Button("🎵 Separate Stems", variant="primary", visible=False)
             
-            gr.Markdown("""**Workflow:**
-            1. Enter Task ID → Click "Find Audio Tracks"
-            2. Select which audio track to separate
-            3. Choose separation type
-            4. Click "Separate Stems"
-            5. Wait 30-120 seconds
-            6. Get download links for each stem
+            gr.Markdown("""
+            ### 📋 Workflow:
+            1. **Enter Task ID** → Click "Find Audio Tracks"
+            2. **Select audio track** from dropdown
+            3. **Choose separation type**
+            4. **Click Separate Stems**
+            5. **Wait 1-3 minutes** for processing
+            6. **Get download links** for each stem
             
-            **Separation types:**
-            - 🎤 **separate_vocal**: Vocals + Instrumental (2 stems, 1 credit)
-            - 🎛️ **split_stem**: 12 detailed stems (5 credits)
+            ### 💡 Tips:
+            - Task IDs are from your Suno generation history
+            - Each separation consumes API credits
+            - Links may expire after some time
+            - Processing time varies based on audio length
             
-            ⚠️ **Important:** 
-            - Each request consumes credits
-            - Links expire in 14 days
+            ### 🔍 Check Status:
+            You'll receive a Separation Task ID that you can use to:
+            - Check status via API
+            - Receive callback on your endpoint
+            - Retry if needed
             """)
         
         with gr.Column(scale=2):
             status_output = gr.Markdown(
                 label="Status",
-                value="Enter a Task ID and click 'Find Audio Tracks'..."
+                value="### 👋 Welcome!\nEnter a Task ID above to get started..."
             )
             
             results_output = gr.Markdown(
@@ -290,7 +421,7 @@ with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
             <p>Powered by <a href="https://suno.ai" target="_blank">Suno AI</a> • 
             <a href="https://sunoapi.org" target="_blank">Suno API Docs</a> • 
             <a href="https://docs.sunoapi.org/separate-vocals" target="_blank">Stem Separation Guide</a></p>
-            <p><small>This app uses the Suno API to separate tracks into individual stems.</small></p>
+            <p><small>Track separation powered by Suno API. Processing may take 1-3 minutes.</small></p>
         </div>
         """,
         elem_id="footer"
@@ -300,12 +431,11 @@ with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
     def on_get_audio(task_id):
         if not task_id:
             return (
-                "❌ Please enter a Task ID",
+                "❌ **Please enter a Task ID**",
                 gr.Dropdown(choices=[], visible=False),
                 gr.Radio(visible=False),
                 gr.Button(visible=False),
-                gr.Markdown(visible=False),
-                None
+                gr.Markdown(visible=False)
             )
         
         output, audio_options = get_audio_info(task_id)
@@ -317,43 +447,39 @@ with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
                 gr.Dropdown(choices=[], visible=False),
                 gr.Radio(visible=False),
                 gr.Button(visible=False),
-                gr.Markdown(visible=False),
-                None
+                gr.Markdown(visible=False)
             )
         
         # Create choices for dropdown
         choices = [(display_text, audio_id) for display_text, audio_id, _ in audio_options]
         
-        # Store full audio data for later use
-        audio_data_map = {audio_id: audio for _, audio_id, audio in audio_options}
-        
         return (
             output,
             gr.Dropdown(choices=choices, value=choices[0][1] if choices else None, visible=True),
             gr.Radio(visible=True),
             gr.Button(visible=True),
-            gr.Markdown(visible=False),
-            audio_data_map
+            gr.Markdown(visible=False)
         )
     
     # Step 2: Separate stems when button is clicked
-    def on_separate(task_id, audio_id, separation_type, audio_data_store):
-        # Use the stored task_id from the first step
-        yield from separate_vocals(task_id, audio_id, separation_type)
+    def on_separate_click(task_id, audio_id, separation_type):
+        # Clear previous results
+        yield gr.Markdown(value="", visible=False)
+        
+        # Start separation process
+        for update in separate_vocals(task_id, audio_id, separation_type):
+            yield gr.Markdown(value=update, visible=True)
     
     # Connect events
     get_audio_btn.click(
         fn=on_get_audio,
         inputs=[task_id_input],
-        outputs=[status_output, audio_dropdown, separation_type, separate_btn, results_output, audio_data_store]
+        outputs=[status_output, audio_dropdown, separation_type, separate_btn, results_output]
     )
     
     separate_btn.click(
-        fn=on_separate,
-        inputs=[task_id_input, audio_dropdown, separation_type, audio_data_store],
-        outputs=[results_output]
-    ).then(
-        lambda: gr.Markdown(visible=True),
+        fn=on_separate_click,
+        inputs=[task_id_input, audio_dropdown, separation_type],
         outputs=[results_output]
     )