app.py

import gradio as gr
import requests
import os
import json

# Suno API key
SUNO_KEY = os.environ.get("SunoKey", "")
if not SUNO_KEY:
    print("⚠️ SunoKey not set!")

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)
            
            display = f"{i+1}. {title} ({duration:.1f}s)"
            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_vocal_separation(task_id, audio_id):
    """Submit vocal separation task (2-stem only)"""
    try:
        resp = requests.post(
            "https://api.sunoapi.org/api/v1/vocal-removal/generate",
            json={
                "taskId": task_id,
                "audioId": audio_id,
                "type": "separate_vocal",  # Always 2-stem
                "callBackUrl": "https://1hit.no/callback.php"
            },
            headers={
                "Authorization": f"Bearer {SUNO_KEY}",
                "Content-Type": "application/json"
            },
            timeout=30
        )
        
        print(f"Submission response: {resp.status_code}")
        print(f"Response: {resp.text}")
        
        if resp.status_code == 200:
            data = resp.json()
            print(f"Parsed data: {data}")
            
            # Get separation task ID from response
            separation_task_id = None
            
            # Try different response formats
            if data.get("taskId"):
                separation_task_id = data.get("taskId")
            elif data.get("data", {}).get("taskId"):
                separation_task_id = data["data"]["taskId"]
            elif data.get("code") == 200 and data.get("data", {}).get("taskId"):
                separation_task_id = data["data"]["taskId"]
            
            if separation_task_id:
                # Also get musicId if available
                music_id = data.get("musicId") or data.get("data", {}).get("musicId", "N/A")
                
                output = f"✅ **Task submitted!**\n\n"
                output += f"**Separation Task ID:** `{separation_task_id}`\n"
                output += f"**Music ID:** `{music_id}`\n\n"
                output += "Check status below using the Separation Task ID."
                return output, separation_task_id
            else:
                error_msg = data.get("msg", data.get("error", "Unknown error"))
                return f"❌ No task ID in response: {error_msg}\n\nRaw response:\n```json\n{json.dumps(data, indent=2)}\n```", None
        else:
            return f"❌ HTTP Error {resp.status_code}:\n{resp.text}", None
            
    except Exception as e:
        return f"❌ Error: {str(e)}", None

def check_separation_status(task_id):
    """Check vocal separation task status"""
    if not task_id:
        return "❌ Enter a Task ID to check"
    
    try:
        resp = requests.get(
            "https://api.sunoapi.org/api/v1/vocal-removal/record-info",
            headers={"Authorization": f"Bearer {SUNO_KEY}"},
            params={"taskId": task_id},
            timeout=30
        )
        
        print(f"Status check for {task_id}: {resp.status_code}")
        
        if resp.status_code == 200:
            data = resp.json()
            print(f"Status data: {json.dumps(data, indent=2)}")
            
            # Check API response code
            if data.get("code") != 200:
                error_msg = data.get("msg", "Unknown error")
                return f"❌ API Error: {error_msg}\n\n**Task ID:** `{task_id}`"
            
            # Get the main data object
            response_data = data.get("data", {})
            
            # Check success flag
            success_flag = response_data.get("successFlag", "UNKNOWN")
            
            if success_flag == "SUCCESS":
                # Get response with download URLs
                result_response = response_data.get("response", {})
                
                if result_response:
                    # Format results
                    output = "✅ **Vocal Separation Complete!**\n\n"
                    output += f"**Task ID:** `{task_id}`\n"
                    output += f"**Music ID:** `{response_data.get('musicId', 'N/A')}`\n\n"
                    output += "## 🎵 Download Links\n\n"
                    
                    # Check for URLs (case-insensitive)
                    urls_found = False
                    
                    # Vocal URL
                    vocal_url = (result_response.get("vocalUrl") or 
                                result_response.get("vocal_url") or
                                result_response.get("vocalURL"))
                    if vocal_url:
                        output += f"**🎤 Vocals:** [Download MP3]({vocal_url})\n"
                        urls_found = True
                    
                    # Instrumental URL  
                    instrumental_url = (result_response.get("instrumentalUrl") or 
                                      result_response.get("instrumental_url") or
                                      result_response.get("instrumentalURL"))
                    if instrumental_url:
                        output += f"**🎵 Instrumental:** [Download MP3]({instrumental_url})\n"
                        urls_found = True
                    
                    # Original URL
                    origin_url = (result_response.get("originUrl") or 
                                 result_response.get("origin_url") or
                                 result_response.get("originURL"))
                    if origin_url:
                        output += f"**📁 Original:** [Download MP3]({origin_url})\n"
                        urls_found = True
                    
                    if not urls_found:
                        output += "No download URLs found in response.\n"
                        output += f"**Raw response:**\n```json\n{json.dumps(result_response, indent=2)}\n```\n"
                    
                    output += f"\n**🔗 Viewer:** [Open in Viewer](https://1hit.no/viewer.php?task_id={task_id})"
                    
                    # Add other fields if they exist
                    other_fields = []
                    for key, value in result_response.items():
                        if value and key.endswith("Url") and key not in ["vocalUrl", "instrumentalUrl", "originUrl"]:
                            field_name = key.replace("Url", "").replace("_", " ").title()
                            other_fields.append(f"**{field_name}:** {value}")
                    
                    if other_fields:
                        output += "\n\n**Other URLs:**\n" + "\n".join(other_fields)
                    
                    return output
                else:
                    return f"✅ **Processing complete!**\n\n**Task ID:** `{task_id}`\n\nNo download URLs in response yet. Check your callback endpoint or try again in a moment."
            
            elif success_flag in ["PENDING", "PROCESSING", "RUNNING"]:
                return f"⏳ **Status:** {success_flag}\n\n**Task ID:** `{task_id}`\n\nStill processing. Check again in 30 seconds."
            
            elif success_flag == "FAILED":
                error_msg = response_data.get("errorMessage", "Unknown error")
                return f"❌ **Failed:** {error_msg}\n\n**Task ID:** `{task_id}`"
            
            else:
                return f"🔄 **Status:** {success_flag}\n\n**Task ID:** `{task_id}`\n\n**Full response:**\n```json\n{json.dumps(data, indent=2)}\n```"
        
        elif resp.status_code == 404:
            return f"❌ Task not found: `{task_id}`\n\nMake sure the Task ID is correct and the separation has started."
        else:
            return f"❌ HTTP Error {resp.status_code}:\n{resp.text}"
            
    except Exception as e:
        return f"❌ Error checking status: {str(e)}"

# Create the app
with gr.Blocks() as app:
    gr.Markdown("# 🎵 Suno Vocal Separator")
    gr.Markdown("Separate vocals from Suno tracks (2-stem only)")
    
    with gr.Row():
        # Left column: Input and submit
        with gr.Column(scale=1):
            # Step 1: Get audio files
            gr.Markdown("### 1. Get Audio Files")
            original_task_id = gr.Textbox(
                label="Original Task ID",
                placeholder="Enter Suno generation task ID",
                info="From your Suno history"
            )
            get_audio_btn = gr.Button("📥 Get Audio Files", variant="secondary")
            audio_status = gr.Markdown("Enter Task ID above")
            
            # Step 2: Select audio file
            gr.Markdown("### 2. Select Audio File")
            audio_dropdown = gr.Dropdown(
                label="Select Audio",
                choices=[],
                interactive=True,
                visible=False
            )
            
            # Step 3: Start separation
            gr.Markdown("### 3. Start Vocal Separation")
            submit_btn = gr.Button("🚀 Separate Vocals", variant="primary", visible=False)
            submission_output = gr.Markdown("Select audio file first", visible=False)
        
        # Right column: Check status and results
        with gr.Column(scale=1):
            # Step 4: Check status
            gr.Markdown("### 4. Check Separation Status")
            separation_task_id = gr.Textbox(
                label="Separation Task ID",
                placeholder="Will auto-fill after submission"
            )
            check_btn = gr.Button("🔍 Check Status", variant="secondary")
            status_output = gr.Markdown("Enter Task ID to check")
            
            # Step 5: Results
            gr.Markdown("### 5. Download Links")
            results_output = gr.Markdown("Results will appear here")
            
            # Step 6: Viewer link
            gr.Markdown("### 6. Viewer")
            gr.Markdown("[Open Viewer](https://1hit.no/viewer.php)")
    
    # 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.Button(visible=False), gr.Markdown(visible=False)
        
        status, options = get_audio_files(task_id)
        
        if not options:
            return status, gr.Dropdown(choices=[], visible=False), gr.Button(visible=False), gr.Markdown(visible=False)
        
        return (
            status,
            gr.Dropdown(choices=options, value=options[0][1] if options else None, visible=True),
            gr.Button(visible=True),
            gr.Markdown("Ready for vocal separation!", visible=True)
        )
    
    # Step 2-3: Submit vocal separation
    def on_submit(task_id, audio_id):
        if not task_id or not audio_id:
            return "❌ Missing Task ID or Audio ID", "", "⏳ Waiting..."
        
        status, sep_task_id = submit_vocal_separation(task_id, audio_id)
        
        if sep_task_id:
            return status, sep_task_id, "✅ Task submitted! Check status above."
        else:
            return status, "", "❌ Failed to submit task"
    
    # Step 4: Check status
    def on_check_status(task_id):
        return check_separation_status(task_id)
    
    # Connect events
    get_audio_btn.click(
        fn=on_get_audio,
        inputs=[original_task_id],
        outputs=[audio_status, audio_dropdown, submit_btn, submission_output]
    )
    
    submit_btn.click(
        fn=on_submit,
        inputs=[original_task_id, audio_dropdown],
        outputs=[submission_output, separation_task_id, status_output]
    )
    
    check_btn.click(
        fn=on_check_status,
        inputs=[separation_task_id],
        outputs=[results_output]
    )
    
    # Auto-fill separation task ID field after submission
    def auto_fill_sep_id(sep_task_id):
        return sep_task_id
    
    separation_task_id.change(
        fn=auto_fill_sep_id,
        inputs=[separation_task_id],
        outputs=[separation_task_id]
    )

# Launch the app
if __name__ == "__main__":
    print("🚀 Starting Suno Vocal 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)