app.py

import gradio as gr
import subprocess
import os
import librosa
import numpy as np
import soundfile as sf
from scipy.stats import pearsonr

def separate_audio(audio_filepath):
    # 1. Setup output directory
    output_base = "output"
    if not os.path.exists(output_base):
        os.makedirs(output_base)
    
    # 2. Run Demucs CLI
    print(f"Starting separation for: {audio_filepath}")
    subprocess.run([
        "python3", "-m", "demucs.separate", 
        "-n", "htdemucs_6s", 
        "-o", output_base, 
        audio_filepath
    ])
    
    # 3. Locate output files
    filename = os.path.splitext(os.path.basename(audio_filepath))[0]
    output_folder = os.path.join(output_base, "htdemucs_6s", filename)
    
    stems = ["vocals.wav", "drums.wav", "bass.wav", "other.wav", "piano.wav", "guitar.wav"]
    result_paths = []
    
    for stem in stems:
        path = os.path.abspath(os.path.join(output_folder, stem))
        if os.path.exists(path):
            result_paths.append(path)
        else:
            print(f"Warning: Could not find {stem} at {path}")
            result_paths.append(None)
            
    return result_paths

def analyze_advanced_metrics(audio_filepath):
    if audio_filepath is None:
        return None, {"error": "No audio file provided"}

    # Load audio
    y, sr = librosa.load(audio_filepath)
    
    # 1. Metronome Generation
    tempo, beat_frames = librosa.beat.beat_track(y=y, sr=sr)
    beat_times = librosa.frames_to_time(beat_frames, sr=sr)
    click_track = librosa.clicks(frames=beat_frames, sr=sr, length=len(y))
    
    output_base = "output"
    if not os.path.exists(output_base):
        os.makedirs(output_base)
    
    click_track_path = os.path.join(output_base, "click_track.wav")
    sf.write(click_track_path, click_track, sr)
    click_track_path = os.path.abspath(click_track_path)
    
    # 2. Structure Analysis
    # Heuristic: use novelty curve to find segment boundaries
    onset_env = librosa.onset.onset_strength(y=y, sr=sr)
    # Using a simple peak picking on the novelty curve for segment boundaries
    # A better way would be using librosa.segment, but let's keep it simple as requested
    # We want 4 to 5 major structural boundaries
    
    # Saliency-based segmentation
    hop_length = 512
    # Compute chroma features
    chroma = librosa.feature.chroma_cqt(y=y, sr=sr, hop_length=hop_length)
    # Use recurrence matrix for segmentation
    rec = librosa.segment.recurrence_matrix(chroma, mode='affinity', sym=True)
    # Compute the lag-similarity matrix
    lag_rec = librosa.segment.recurrence_to_lag(rec)
    
    # Instead of complex librosa.segment which might need more params, 
    # let's use a simpler approach to get 4-5 sections.
    duration = librosa.get_duration(y=y, sr=sr)
    
    # Find 4-5 major structural boundaries using novelty curve
    # Smooth the novelty curve
    novelty = librosa.util.normalize(onset_env)
    # We can use librosa.segment.subsegment or just pick top N peaks far apart
    # For simplicity, let's just divide the song into N chunks if novelty detection is too complex for this prompt
    # BUT the prompt says "Use librosa.segment or a novelty-curve heuristic to determine 4 to 5 major structural boundaries"
    
    boundaries = librosa.segment.agglomerative(chroma, 5) # Get 5 segments
    boundary_times = librosa.frames_to_time(boundaries, sr=sr, hop_length=hop_length)
    
    # Ensure start is 0 and end is duration
    boundary_times = np.unique(np.concatenate(([0.0], boundary_times, [duration])))
    
    structure_json = []
    for i in range(len(boundary_times) - 1):
        structure_json.append({
            "label": f"Section {i+1}",
            "start_time": round(float(boundary_times[i]), 2),
            "end_time": round(float(boundary_times[i+1]), 2)
        })
        
    return click_track_path, structure_json

def detect_key(audio_filepath):
    if audio_filepath is None:
        return {"error": "No audio file provided"}

    # Load audio
    y, sr = librosa.load(audio_filepath)
    
    # Extract chromagram
    chroma = librosa.feature.chroma_cqt(y=y, sr=sr)
    chroma_sum = np.sum(chroma, axis=1)
    
    # Krumhansl-Schmuckler profiles (Temperley)
    major_profile = [6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88]
    minor_profile = [6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17]
    
    notes = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
    
    results = []
    
    for i in range(12):
        # Rotate profiles to check each key
        shifted_major = np.roll(major_profile, i)
        shifted_minor = np.roll(minor_profile, i)
        
        # Pearson correlation
        corr_major, _ = pearsonr(chroma_sum, shifted_major)
        corr_minor, _ = pearsonr(chroma_sum, shifted_minor)
        
        results.append((corr_major, f"{notes[i]} Major", i))
        results.append((corr_minor, f"{notes[i]} Minor", i))
    
    # Find maximum correlation
    best_corr, best_key, root_idx = max(results, key=lambda x: x[0])
    
    return {
        "key_name": best_key,
        "root_index": int(root_idx)
    }

def extract_chords(audio_filepath):
    if audio_filepath is None:
        return {"error": "No audio file provided"}

    # Load audio
    y, sr = librosa.load(audio_filepath)
    
    # Extract chromagram
    chroma = librosa.feature.chroma_cqt(y=y, sr=sr)
    
    # Define chord templates (12 Major and 12 Minor)
    # C, C#, D, D#, E, F, F#, G, G#, A, A#, B
    maj_template = np.array([1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0])
    min_template = np.array([1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0])
    
    templates = []
    labels = []
    notes = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
    
    for i in range(12):
        templates.append(np.roll(maj_template, i))
        labels.append(f"{notes[i]} Major")
        templates.append(np.roll(min_template, i))
        labels.append(f"{notes[i]} Minor")
    
    templates = np.array(templates)
    
    # Analyze frame by frame
    chords_sequence = []
    times = librosa.frames_to_time(np.arange(chroma.shape[1]), sr=sr)
    
    for i in range(chroma.shape[1]):
        frame_chroma = chroma[:, i]
        if np.sum(frame_chroma) == 0:
            chords_sequence.append("N/A")
            continue
            
        correlations = np.dot(templates, frame_chroma)
        chord_idx = np.argmax(correlations)
        chords_sequence.append(labels[chord_idx])
        
    # Compress output: group consecutive identical chords
    compressed_chords = []
    if chords_sequence:
        current_chord = chords_sequence[0]
        compressed_chords.append({"time": round(float(times[0]), 2), "chord": current_chord})
        
        for i in range(1, len(chords_sequence)):
            if chords_sequence[i] != current_chord:
                current_chord = chords_sequence[i]
                compressed_chords.append({"time": round(float(times[i]), 2), "chord": current_chord})
                
    return compressed_chords

# Build Gradio UI with Blocks
with gr.Blocks(title="AI Stem Studio Backend") as demo:
    gr.Markdown("# AI Stem Studio Backend")
    
    with gr.Tab("Stem Separation"):
        sep_input = gr.Audio(type="filepath", label="Upload Audio")
        sep_btn = gr.Button("Separate Stems")
        with gr.Row():
            sep_vocals = gr.Audio(label="Vocals")
            sep_drums = gr.Audio(label="Drums")
            sep_bass = gr.Audio(label="Bass")
            sep_other = gr.Audio(label="Other")
            sep_piano = gr.Audio(label="Piano")
            sep_guitar = gr.Audio(label="Guitar")
        
        sep_btn.click(
            fn=separate_audio,
            inputs=sep_input,
            outputs=[sep_vocals, sep_drums, sep_bass, sep_other, sep_piano, sep_guitar],
            api_name="separate_audio"
        )

    with gr.Tab("Advanced Metrics"):
        metrics_input = gr.Audio(type="filepath", label="Upload Audio")
        metrics_btn = gr.Button("Analyze Metrics")
        with gr.Row():
            click_output = gr.Audio(label="Click Track")
            structure_output = gr.JSON(label="Song Structure")
        
        metrics_btn.click(
            fn=analyze_advanced_metrics,
            inputs=metrics_input,
            outputs=[click_output, structure_output],
            api_name="analyze_advanced_metrics"
        )
            
    with gr.Tab("Key Detection"):
        key_input = gr.Audio(type="filepath", label="Upload Audio")
        key_btn = gr.Button("Detect Key")
        key_output = gr.JSON(label="Key Analysis Result")
        
        key_btn.click(
            fn=detect_key,
            inputs=key_input,
            outputs=key_output,
            api_name="detect_key"
        )

    with gr.Tab("Chord Extraction"):
        chord_input = gr.Audio(type="filepath", label="Upload Audio")
        chord_btn = gr.Button("Extract Chords")
        chord_output = gr.JSON(label="Chord Progression")
        
        chord_btn.click(
            fn=extract_chords,
            inputs=chord_input,
            outputs=chord_output,
            api_name="extract_chords"
        )

if __name__ == "__main__":
    demo.launch()