Update app.py

app.py

@@ -1,102 +1,81 @@
-import time, requests, threading, io
-from flask import Flask, render_template, request, session, send_file, jsonify
+import time, requests, threading, os, io
+from flask import Flask, render_template, request, session, send_from_directory, jsonify
 from flask_socketio import SocketIO, emit
 
 app = Flask(__name__)
-app.config['SECRET_KEY'] = 'suno_studio_2026'
-# Optimized for Hugging Face proxying
-socketio = SocketIO(app, cors_allowed_origins="*", async_mode='threading')
+app.config['SECRET_KEY'] = 'server_direct_save_2026'
+socketio = SocketIO(app, cors_allowed_origins="*")
 
-BASE_URL = "https://api.sunoapi.org/api/v1"
+# Create a permanent folder on the server for the music
+MUSIC_FOLDER = 'saved_music'
+if not os.path.exists(MUSIC_FOLDER):
+    os.makedirs(MUSIC_FOLDER)
+
+# Map to link taskIds to User Session IDs (for the callback)
+task_to_sid = {}
 
 @app.route('/')
 def index():
     return render_template('index.html')
 
+# 1. THE CALLBACK (Saves the file to server storage)
 @app.route('/callback', methods=['POST'])
 def callback():
-    """Receives completion data from Suno."""
     data = request.json
     task_id = data.get('taskId') or data.get('task_id')
     audio_url = data.get('audio_url')
-    
-    if audio_url:
-        socketio.emit('status', {
-            'msg': '✅ Webhook: Track Ready!',
-            'url': audio_url,
-            'done': True
-        })
-    return jsonify({"status": "received"}), 200
 
-@app.route('/download_proxy')
-def download_proxy():
-    """Bypasses CORS to force download the MP3 file."""
-    url = request.args.get('url')
-    title = request.args.get('title', 'Suno_Track').replace(" ", "_")
-    try:
-        res = requests.get(url, stream=True)
-        res.raise_for_status()
-        return send_file(
-            io.BytesIO(res.content),
-            mimetype="audio/mpeg",
-            as_attachment=True,
-            download_name=f"{title}.mp3"
-        )
-    except Exception as e:
-        return f"Download Error: {str(e)}", 500
+    if audio_url and task_id:
+        # Download the file to the server immediately
+        filename = f"{task_id}.mp3"
+        filepath = os.path.join(MUSIC_FOLDER, filename)
+        
+        try:
+            r = requests.get(audio_url)
+            with open(filepath, 'wb') as f:
+                f.write(r.content)
+            
+            # Find the user who owns this task and tell them
+            if task_id in task_to_sid:
+                user_sid = task_to_sid[task_id]
+                socketio.emit('status', {
+                    'msg': '✅ Saved to Server!',
+                    'file_name': filename,
+                    'done': True
+                }, room=user_sid)
+        except Exception as e:
+            print(f"Callback save error: {e}")
+
+    return jsonify({"status": "ok"}), 200
+
+# 2. THE DIRECT DOWNLOAD (Serves the file from your server)
+@app.route('/get_music/<filename>')
+def get_music(filename):
+    return send_from_directory(MUSIC_FOLDER, filename, as_attachment=True)
 
 @socketio.on('save_token')
 def save_token(data):
-    token = data.get('token')
-    if token:
-        session['api_token'] = token
-        emit('status', {'msg': '🔑 Token linked successfully.'})
-        # Auto-check credits
-        try:
-            res = requests.get(f"{BASE_URL}/generate/credit", headers={"Authorization": f"Bearer {token}"})
-            emit('credits', {'val': res.json().get('data', 0)})
-        except: pass
+    session['api_token'] = data.get('token')
+    emit('status', {'msg': '🔑 Token Linked.'})
 
 @socketio.on('start_gen')
 def handle_gen(payload):
     token = session.get('api_token')
-    if not token:
-        emit('status', {'msg': '❌ Error: No Token Connected.'})
-        return
-
     headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
-    sid = request.sid
-
+    
+    # Set the callback URL dynamically
+    payload['callBackUrl'] = request.host_url + "callback"
+    
     try:
-        # Note: 'callBackUrl' is already in the payload from the frontend
-        res = requests.post(f"{BASE_URL}/generate/upload-cover", json=payload, headers=headers)
+        res = requests.post("https://api.sunoapi.org/api/v1/generate/upload-cover", json=payload, headers=headers)
         res_data = res.json()
+        task_id = res_data.get('data', {}).get('taskId')
         
-        if res.status_code == 200 and res_data.get('data'):
-            task_id = res_data['data'].get('taskId')
-            emit('status', {'msg': f'🚀 Task {task_id} initiated. Processing...'})
-            # Threaded backup polling in case callback fails
-            threading.Thread(target=poll_status, args=(task_id, sid, headers)).start()
-        else:
-            emit('status', {'msg': f"❌ API Error: {res_data.get('msg', 'Unknown')}"})
+        if task_id:
+            task_to_sid[task_id] = request.sid # Map user to task
+            emit('status', {'msg': f'🚀 Processing Task: {task_id}...'})
     except Exception as e:
-        emit('status', {'msg': f"❌ System Error: {str(e)}"})
-
-def poll_status(task_id, sid, headers):
-    """Checks status every 12 seconds as a backup."""
-    for _ in range(50): 
-        time.sleep(12)
-        try:
-            res = requests.get(f"{BASE_URL}/suno/cover/record-info?taskId={task_id}", headers=headers)
-            data = res.json().get('data', {})
-            if data.get('status') == 'complete':
-                socketio.emit('status', {
-                    'msg': '✨ Polling: Track Ready!',
-                    'url': data.get('audio_url'),
-                    'done': True
-                }, room=sid)
-                break
-        except: break
+        emit('status', {'msg': f'❌ Error: {str(e)}'})
 
 if __name__ == '__main__':
-    socketio.run(app, host="0.0.0.0", allow_unsafe_werkzeug=True)
+    socketio.run(app, host="0.0.0.0")