Integrate MusicGPT API generation into booth-musicgpt v0.7.0

app.py

@@ -33,6 +33,7 @@ import os
 import hmac
 import shutil
 import time
+import json
 from pathlib import Path
 
 # Flask and related imports
@@ -56,6 +57,11 @@ from helpers import (
     get_email_config, get_refresh_seconds, get_kiosk_cycle_seconds, get_kiosk_mode,
     get_max_revisions, get_callback_expiry_hours,
     get_hermes_api_key, set_hermes_api_key, generate_hermes_api_key, mask_api_key,
+    get_musicgpt_api_key, get_musicgpt_default_model, get_musicgpt_models,
+    build_musicgpt_webhook_url,
+    musicgpt_generate_request, musicgpt_queue_stems, musicgpt_poll_status,
+    download_musicgpt_outputs, format_musicgpt_cost, get_musicgpt_cost_totals,
+    download_album_cover,
     get_ntfy_config, send_ntfy,
     sign_prompt_callback, verify_prompt_callback, build_prompt_callback_url,
     send_email, build_signature_images,
@@ -272,6 +278,130 @@ def api_update_prompt(rid):
     return jsonify({'ok': True, 'request_id': rid, 'status': new_status}), 200
 
 
+@app.route('/api/musicgpt/webhook', methods=['POST'])
+def musicgpt_webhook():
+    """
+    Public webhook endpoint for MusicGPT async job completion.
+
+    Handles both Music AI generation webhooks and Extraction (stems) webhooks.
+    Generation webhooks may arrive multiple times per task (one per version,
+    plus lyrics, plus streaming URL, plus album cover). We update the request
+    incrementally and download files when a COMPLETED conversion payload arrives.
+    """
+    data = request.get_json(silent=True) or {}
+    task_id = data.get('task_id') or data.get('taskId')
+    conversion_type = data.get('conversion_type') or data.get('conversionType') or ''
+    status = data.get('status') or data.get('conversion_status') or 'COMPLETED'
+
+    if not task_id:
+        return jsonify({'ok': False, 'reason': 'missing_task_id'}), 400
+
+    conversion_id = data.get('conversion_id')
+    if not conversion_id and 'conversion_id_1' in data:
+        conversion_id = data.get('conversion_id_1')
+
+    # Music AI generation webhook
+    if 'Music' in conversion_type or 'Music AI' in conversion_type:
+        req = None
+        db = get_db()
+        row = db.execute('SELECT * FROM requests WHERE musicgpt_task_id = ?', (task_id,)).fetchone()
+        if row:
+            req = dict(row)
+        if not req:
+            return jsonify({'ok': False, 'reason': 'request_not_found'}), 404
+
+        new_status = status.upper()
+        update_fields = {'musicgpt_status': new_status}
+        if data.get('is_flagged'):
+            update_fields['musicgpt_error'] = data.get('reason') or 'Flagged by MusicGPT'
+        if new_status in ('COMPLETED', 'FINISHED'):
+            download_musicgpt_outputs(req, data)
+            # Re-fetch to get updated paths.
+            row = db.execute('SELECT * FROM requests WHERE id = ?', (req['id'],)).fetchone()
+            req = dict(row)
+            if req.get('song_a_path') and req.get('song_b_path'):
+                update_fields['status'] = 'songs_uploaded'
+        elif new_status in ('FAILED', 'ERROR'):
+            update_fields['musicgpt_error'] = data.get('reason') or data.get('error') or 'MusicGPT reported failure'
+        update_request(req['id'], **update_fields)
+        return jsonify({'ok': True, 'request_id': req['id']}), 200
+
+    # Extraction / stems webhook
+    if 'Extraction' in conversion_type:
+        req = None
+        db = get_db()
+        row = db.execute('SELECT * FROM requests WHERE stems_task_id = ?', (task_id,)).fetchone()
+        if row:
+            req = dict(row)
+        if not req:
+            return jsonify({'ok': False, 'reason': 'request_not_found'}), 404
+
+        new_status = status.upper()
+        update_fields = {'stems_status': new_status}
+        if new_status in ('COMPLETED', 'FINISHED'):
+            audio_url_map = data.get('audio_url')
+            if isinstance(audio_url_map, str):
+                try:
+                    audio_url_map = json.loads(audio_url_map)
+                except Exception:
+                    audio_url_map = {}
+            # Prefer a zip-style bundle URL if present, otherwise join stem URLs.
+            stems_url = data.get('bundle_url') or data.get('download_url')
+            if not stems_url and audio_url_map:
+                stems_url = '; '.join(f"{k}: {v}" for k, v in audio_url_map.items())
+            update_fields['stems_url'] = stems_url
+            update_fields['stems_link'] = stems_url
+            try:
+                update_fields['stems_cost'] = float(data.get('conversion_cost') or 0)
+            except (ValueError, TypeError):
+                pass
+        elif new_status in ('FAILED', 'ERROR'):
+            update_fields['stems_error'] = data.get('reason') or data.get('error') or 'Extraction failed'
+        update_request(req['id'], **update_fields)
+        return jsonify({'ok': True, 'request_id': req['id']}), 200
+
+    return jsonify({'ok': False, 'reason': 'unknown_conversion_type'}), 400
+
+
+@app.route('/admin/musicgpt/refresh', methods=['POST'])
+def admin_musicgpt_refresh():
+    """Manual dashboard action: poll all in-flight MusicGPT tasks and update statuses."""
+    redir = require_admin()
+    if redir:
+        return redir
+    db = get_db()
+    rows = db.execute(
+        "SELECT id, musicgpt_task_id, musicgpt_status FROM requests WHERE musicgpt_status IN ('IN_QUEUE', 'IN_PROGRESS')"
+    ).fetchall()
+    updated = 0
+    failed = 0
+    for row in rows:
+        task_id = row['musicgpt_task_id']
+        if not task_id:
+            continue
+        result = musicgpt_poll_status(task_id)
+        conversion = result.get('conversion') or {}
+        status = (conversion.get('status') or result.get('status') or '').upper()
+        req = get_request_by_id(row['id'])
+        if not req:
+            continue
+        if status in ('COMPLETED', 'FINISHED'):
+            download_musicgpt_outputs(req, conversion)
+            fields = {'musicgpt_status': 'COMPLETED'}
+            if req.get('song_a_path') and req.get('song_b_path'):
+                fields['status'] = 'songs_uploaded'
+            update_request(row['id'], **fields)
+            updated += 1
+        elif status in ('FAILED', 'ERROR'):
+            update_request(row['id'], musicgpt_status='FAILED', musicgpt_error=conversion.get('status_msg') or 'Polling reported failure')
+            failed += 1
+        elif status:
+            update_request(row['id'], musicgpt_status=status)
+            updated += 1
+    flash(f'MusicGPT refresh: {updated} updated, {failed} failed.', 'success' if not failed else 'warning')
+    return redirect(url_for('admin_dashboard'))
+
+
 @app.route('/thanks/<int:rid>')
 def thanks(rid):
     """Confirmation page shown after a customer submits a request."""
@@ -573,7 +703,8 @@ def admin_pricing():
             'name': entry.get('name', '') if isinstance(entry, dict) else '',
             'price': entry.get('price', '') if isinstance(entry, dict) else ''
         })
-    return render_template('admin/pricing.html', fixed=fixed, customs=customs)
+    cost_totals = get_musicgpt_cost_totals()
+    return render_template('admin/pricing.html', fixed=fixed, customs=customs, cost_totals=cost_totals)
 
 
 @app.route('/kiosk')
@@ -743,6 +874,94 @@ def admin_request(rid):
             )
             flash('Prompt saved.', 'success')
 
+        elif action == 'generate_musicgpt':
+            # Queue a MusicGPT generation job.
+            api_key = get_musicgpt_api_key()
+            if not api_key:
+                flash('MusicGPT API key is not configured.', 'error')
+                return redirect(url_for('admin_request', rid=rid))
+
+            if req.get('musicgpt_status') in ('IN_QUEUE', 'IN_PROGRESS'):
+                flash('A MusicGPT generation is already in progress for this request.', 'error')
+                return redirect(url_for('admin_request', rid=rid))
+
+            title = request.form.get('suno_title', '').strip()
+            style = request.form.get('suno_style', '').strip()
+            lyrics = request.form.get('suno_lyrics', '').strip()
+            if not (title and style and lyrics):
+                flash('Title, style, and lyrics are required to generate music.', 'error')
+                return redirect(url_for('admin_request', rid=rid))
+
+            model = request.form.get('musicgpt_model', get_musicgpt_default_model()).strip()
+            if model not in get_musicgpt_models():
+                flash('Invalid MusicGPT model selected.', 'error')
+                return redirect(url_for('admin_request', rid=rid))
+
+            gender = request.form.get('vocal_gender', req.get('vocal_gender') or '').strip()
+            task_id, conv1, conv2, estimate, error = musicgpt_generate_request(
+                rid, title, style, lyrics, gender=gender, model=model
+            )
+            if error:
+                update_request(rid, musicgpt_status='ERROR', musicgpt_error=error)
+                flash(f'Failed to queue MusicGPT generation: {error}', 'error')
+            else:
+                update_request(rid,
+                    suno_title=title,
+                    suno_style=style,
+                    suno_lyrics=lyrics,
+                    musicgpt_task_id=task_id,
+                    musicgpt_status='IN_QUEUE',
+                    musicgpt_error=None,
+                    musicgpt_cost=estimate,
+                    status='prompt_ready',
+                    deliver_wav=1 if request.form.get('deliver_wav') else 0
+                )
+                flash(f'MusicGPT generation queued (task {task_id}). Estimated cost: {estimate}', 'success')
+            return redirect(url_for('admin_request', rid=rid))
+
+        elif action == 'save_delivery_options':
+            update_request(rid, deliver_wav=1 if request.form.get('deliver_wav') else 0)
+            flash('Delivery options saved.', 'success')
+            return redirect(url_for('admin_request', rid=rid))
+
+        elif action == 'cancel_musicgpt':
+            # Mark an in-flight MusicGPT job as cancelled locally.
+            if req.get('musicgpt_status') in ('IN_QUEUE', 'IN_PROGRESS'):
+                update_request(rid, musicgpt_status='CANCELLED')
+                flash('MusicGPT generation marked as cancelled.', 'success')
+            else:
+                flash('No in-progress MusicGPT generation to cancel.', 'error')
+            return redirect(url_for('admin_request', rid=rid))
+
+        elif action == 'generate_stems':
+            # Queue a stem extraction job for Version A MP3.
+            api_key = get_musicgpt_api_key()
+            if not api_key:
+                flash('MusicGPT API key is not configured.', 'error')
+                return redirect(url_for('admin_request', rid=rid))
+            if req.get('stems_status') in ('IN_QUEUE', 'IN_PROGRESS'):
+                flash('A stems extraction is already in progress.', 'error')
+                return redirect(url_for('admin_request', rid=rid))
+            audio_url = request.form.get('stems_source_url', '').strip()
+            if not audio_url:
+                # Fall back to using the locally stored Version A MP3 path; upload not supported by extraction API.
+                a_path = req.get('song_a_path')
+                if a_path:
+                    # Extraction API supports audio_url; we need a public URL. Local path won't work.
+                    flash('Please provide a public audio URL for stems extraction.', 'error')
+                    return redirect(url_for('admin_request', rid=rid))
+                flash('No Version A audio available to extract stems from.', 'error')
+                return redirect(url_for('admin_request', rid=rid))
+            stems = request.form.getlist('stems') or ['vocals', 'instrumental']
+            task_id, conv_id, estimate, error = musicgpt_queue_stems(rid, audio_url, stems=stems)
+            if error:
+                update_request(rid, stems_status='ERROR', stems_error=error)
+                flash(f'Failed to queue stems extraction: {error}', 'error')
+            else:
+                update_request(rid, stems_task_id=task_id, stems_status='IN_QUEUE', stems_error=None, stems_cost=estimate)
+                flash(f'Stems extraction queued (task {task_id}). Estimated cost: {estimate}', 'success')
+            return redirect(url_for('admin_request', rid=rid))
+
         elif action == 'delete_songs':
             # Remove selected uploaded MP3s and reset status so new songs can be uploaded.
             to_delete = request.form.getlist('delete_song')
@@ -754,20 +973,34 @@ def admin_request(rid):
             for version in to_delete:
                 if version == 'a':
                     path = req.get('song_a_path')
+                    wav = req.get('song_a_wav_path')
                     if path and Path(path).exists():
                         try:
                             Path(path).unlink()
                         except OSError:
                             pass
+                    if wav and Path(wav).exists():
+                        try:
+                            Path(wav).unlink()
+                        except OSError:
+                            pass
                     update_fields['song_a_path'] = None
+                    update_fields['song_a_wav_path'] = None
                 elif version == 'b':
                     path = req.get('song_b_path')
+                    wav = req.get('song_b_wav_path')
                     if path and Path(path).exists():
                         try:
                             Path(path).unlink()
                         except OSError:
                             pass
+                    if wav and Path(wav).exists():
+                        try:
+                            Path(wav).unlink()
+                        except OSError:
+                            pass
                     update_fields['song_b_path'] = None
+                    update_fields['song_b_wav_path'] = None
 
             if update_fields:
                 # Wipe any prior customer approval because the old files are gone.
@@ -881,6 +1114,20 @@ def admin_request(rid):
                 body_lines.insert(5, f"This share link expires on {expiry_date} (3 months from today). Please download before then.")
                 body_lines.insert(6, "")
             body = '\n'.join(body_lines)
+
+            # Optionally attach WAV files and album cover.
+            if req.get('deliver_wav'):
+                for wav_field in ('song_a_wav_path', 'song_b_wav_path'):
+                    wav_path = req.get(wav_field)
+                    if wav_path and Path(wav_path).exists():
+                        p = Path(wav_path)
+                        attachments.append((str(p), p.name))
+            cfg = load_booth_settings()
+            if cfg.get('deliver_album_cover') and req.get('album_cover_url'):
+                cover_path = download_album_cover(rid, req['album_cover_url'])
+                if cover_path:
+                    attachments.append((cover_path, f"cover{Path(cover_path).suffix}"))
+
             try:
                 send_email(req['email'], 'Your theme song files are here!', body, attachments=attachments, inline_images=build_signature_images())
                 update_request(rid, delivery_sent_at=now_utc())
@@ -901,7 +1148,11 @@ def admin_request(rid):
         basename=basename,
         extra_files=extra_files,
         callback_url=build_prompt_callback_url(rid),
-        revision_history=list_revision_history(rid)
+        revision_history=list_revision_history(rid),
+        musicgpt_models=get_musicgpt_models(),
+        musicgpt_default_model=get_musicgpt_default_model(),
+        format_musicgpt_cost=format_musicgpt_cost,
+        deliver_album_cover=load_booth_settings().get('deliver_album_cover', False)
     )
 
 
@@ -1010,7 +1261,10 @@ def admin_settings():
         'style_genre', 'extra_requests', 'vocal_gender', 'status', 'suno_title', 'suno_style',
         'suno_lyrics', 'song_a_path', 'song_b_path', 'customer_approved',
         'approval_notified_at', 'preview_sent_at', 'delivery_sent_at',
-        'square_payment_ref', 'admin_alert_email', 'player_token', 'revision_note', 'revision_count', 'operator_notes', 'stems_link', 'stems_interest'
+        'square_payment_ref', 'admin_alert_email', 'player_token', 'revision_note', 'revision_count', 'operator_notes', 'stems_link', 'stems_interest',
+        'musicgpt_task_id', 'musicgpt_status', 'musicgpt_cost', 'musicgpt_error',
+        'album_cover_url', 'song_a_wav_path', 'song_b_wav_path', 'deliver_wav',
+        'stems_task_id', 'stems_status', 'stems_cost', 'stems_url', 'stems_error'
     }
     expected_tables = {'requests', 'revision_history'}
     health = {'ok': True, 'missing_columns': [], 'missing_tables': [], 'message': 'Database schema looks good.'}
@@ -1260,6 +1514,13 @@ def admin_settings():
                 flash('Failed to send ntfy test notification. Check server, topic, and access token.', 'error')
             return redirect(url_for('admin_settings'))
 
+        elif action == 'save_album_cover':
+            cfg = load_booth_settings()
+            cfg['deliver_album_cover'] = request.form.get('deliver_album_cover', '0') == '1'
+            save_booth_settings(cfg)
+            flash('Album cover delivery setting saved.', 'success')
+            return redirect(url_for('admin_settings'))
+
     return render_template(
         'admin/settings.html',
         health=health,
@@ -1289,6 +1550,9 @@ def admin_settings():
         version=current_app.config['VERSION'],
         current_callback_expiry_hours=get_callback_expiry_hours(),
         ntfy=runtime_settings,
+        deliver_album_cover=runtime_settings.get('deliver_album_cover', False),
+        musicgpt_api_key_set=bool(get_musicgpt_api_key()),
+        musicgpt_webhook_url=build_musicgpt_webhook_url(),
     )