app.py

import gradio as gr
import tempfile
import random
import json
import re
import pretty_midi
import subprocess
import os
from openai import OpenAI
# Audio playback support
try:
    import pygame
    pygame.mixer.init()
    PYGAME_AVAILABLE = True
except Exception as e:
    print(f"[WARNING] pygame mixer init failed: {e}")
    PYGAME_AVAILABLE = False

# === LLM APIs ===
def query_llm(prompt, model_name=None):
    if model_name and model_name != "OpenAI":
        import requests
        response = requests.post("http://localhost:11434/api/generate", json={"model": model_name, "prompt": prompt, "stream": False})
        return response.json().get("response", "")
    else:
        # Replace or load from environment
        client = OpenAI(
    base_url="https://api.studio.nebius.com/v1/",
    api_key=os.environ.get("NEBIUS_API_KEY")
)
        response = client.chat.completions.create(
    model="meta-llama/Llama-3.3-70B-Instruct",
    max_tokens=512,
    temperature=0.6,
    top_p=0.9,
    extra_body={
        "top_k": 50
    },
    messages=[]
)
    return response["choices"][0]["message"]["content"]

# === Step 1: Parse intent ===
def get_intent_from_prompt(prompt, model_name):
    system_prompt = f"""
    Extract the musical intent from this prompt.
    Return JSON with keys: tempo (int), key (A-G#), scale (major/minor), genre (e.g., lo-fi, trap), emotion, instrument.
    Prompt: '{prompt}'
    """
    response = query_llm(system_prompt, model_name)
    match = re.search(r'\{.*\}', response, re.DOTALL)
    if match:
        try:
            return json.loads(match.group(0))
        except json.JSONDecodeError:
            return {"tempo": 120, "key": "C", "scale": "major", "genre": "default", "emotion": "neutral", "instrument": "piano"}
    return {"tempo": 120, "key": "C", "scale": "major", "genre": "default", "emotion": "neutral", "instrument": "piano"}

# === Step 2: Melody planning ===
def get_melody_from_intent(intent, model_name):
    melody_prompt = f"""
    You are a music composer.
    Based on this musical intent:
    {json.dumps(intent)}

    Generate a melody plan using a list of 16 notes with pitch (A-G#), octave (3-6), and duration (0.25 to 1.0 seconds).
    Output ONLY valid JSON like:
    [
        {{"note": "D", "octave": 4, "duration": 0.5}},
        {{"note": "F", "octave": 4, "duration": 1.0}}
    ]
    """
    response = query_llm(melody_prompt, model_name)
    print(f"\n[DEBUG] LLM Response for melody:\n{response}\n")
    match = re.search(r'\[.*\]', response, re.DOTALL)
    if match:
        try:
            melody = json.loads(match.group(0))
            if isinstance(melody, list) and len(melody) > 0:
                return melody
        except json.JSONDecodeError as e:
            print(f"[ERROR] Melody JSON decode error: {e}")

    print("[WARNING] Using fallback melody.")
    return [
        {"note": "C", "octave": 4, "duration": 0.5},
        {"note": "E", "octave": 4, "duration": 0.5},
        {"note": "G", "octave": 4, "duration": 0.5},
        {"note": "B", "octave": 4, "duration": 0.5},
    ]

# === Step 3: MIDI generation ===
def midi_from_plan(melody, tempo):
    midi = pretty_midi.PrettyMIDI()
    instrument = pretty_midi.Instrument(program=0)
    time = 0.0
    seconds_per_beat = 60.0 / tempo

    note_map = {"C": 0, "C#": 1, "D": 2, "D#": 3, "E": 4, "F": 5, "F#": 6,
                "G": 7, "G#": 8, "A": 9, "A#": 10, "B": 11}

    for note_info in melody:
        try:
            pitch = 12 * (note_info["octave"] + 1) + note_map[note_info["note"].upper()]
            duration = float(note_info["duration"])
            start = time
            end = time + duration
            instrument.notes.append(pretty_midi.Note(
                velocity=100, pitch=pitch, start=start, end=end
            ))
            time = end
        except:
            continue

    midi.instruments.append(instrument)
    return midi

# === Generate audio preview from MIDI ===
def midi_to_wav(midi_path):
    try:
        import tempfile
        import subprocess
        import os

        # Convert MIDI to WAV using FluidSynth if installed, else fallback to empty
        wav_path = tempfile.NamedTemporaryFile(delete=False, suffix=".wav").name

        # Use fluidsynth if available, else skip audio preview
        fluidsynth_cmd = ["fluidsynth", "-ni", "/usr/share/sounds/sf2/FluidR3_GM.sf2", midi_path, "-F", wav_path, "-r", "44100"]
        result = subprocess.run(fluidsynth_cmd, capture_output=True)
        if result.returncode != 0:
            print("[WARNING] FluidSynth conversion failed or is not installed.")
            return None
        return wav_path
    except Exception as e:
        print(f"[ERROR] midi_to_wav failed: {e}")
        return None

# === Main function to generate MIDI and audio preview ===
def generate_midi_and_audio(prompt, model_name):
    intent = get_intent_from_prompt(prompt, model_name)
    melody = get_melody_from_intent(intent, model_name)
    midi = midi_from_plan(melody, intent.get("tempo", 120))

    with tempfile.NamedTemporaryFile(delete=False, suffix=".mid") as tmp:
        midi.write(tmp.name)
        midi_path = tmp.name

    audio_path = None
    if PYGAME_AVAILABLE:
        audio_path = midi_path  # We'll use pygame to play midi directly if possible

    return midi_path, audio_path

# === Get Ollama models ===
def get_ollama_models():
    try:
        result = subprocess.run(["ollama", "list"], capture_output=True, text=True)
        models = [line.split()[0] for line in result.stdout.strip().splitlines()[1:]]
        return ["OpenAI"] + models
    except Exception as e:
        return ["OpenAI"]

# === Gradio UI ===
models = get_ollama_models()

demo = gr.Interface(
    fn=generate_midi_and_audio,
    inputs=[
        gr.Textbox(label="Music Prompt"),
        gr.Dropdown(choices=models, label="LLM Model", value=models[0])
    ],
    outputs=[
        gr.File(label="🎵 Download MIDI File"),
        gr.Audio(label="🎧 Audio Preview (MIDI Playback, if supported)", type="filepath")
    ],
    title="🎼 Music Command Prompt (MCP Agent)",
    description="Describe your music idea and download a generated MIDI file. Choose from local or OpenAI LLMs."
)

demo.launch(mcp_server=True)