app.py

import gradio as gr
import os
import shutil
import zipfile
import librosa
import numpy as np
from pydub import AudioSegment
from moviepy.editor import AudioFileClip, ImageClip, CompositeVideoClip
import subprocess
from pathlib import Path
import sys

# --- Configuration ---
OUTPUT_DIR = Path("nightpulse_output")
TEMP_DIR = Path("temp_processing")

# --- Core Logic Functions ---

def analyze_and_separate(audio_file):
    """Phase 1: Separate 6 Stems (Drums, Bass, Guitar, Piano, Vocals, Other)"""
    try:
        if not audio_file:
            raise ValueError("No audio file provided.")
            
        # Cleanup
        if OUTPUT_DIR.exists(): shutil.rmtree(OUTPUT_DIR)
        if TEMP_DIR.exists(): shutil.rmtree(TEMP_DIR)
        (OUTPUT_DIR / "Stems").mkdir(parents=True, exist_ok=True)
        (OUTPUT_DIR / "Loops").mkdir(parents=True, exist_ok=True)
        TEMP_DIR.mkdir(parents=True, exist_ok=True)

        filename = Path(audio_file).stem
        
        # 1. BPM Detection
        print(f"Analyzing {filename}...")
        y, sr = librosa.load(audio_file, duration=60, mono=True)
        tempo, _ = librosa.beat.beat_track(y=y, sr=sr)
        
        if np.ndim(tempo) > 0:
            bpm = int(round(tempo[0]))
        else:
            bpm = int(round(tempo))
        print(f"Detected BPM: {bpm}")

        # 2. Demucs Separation (Using 6-Stem Model)
        print("Separating stems (6-Stem Model)...")
        subprocess.run([
            sys.executable, "-m", "demucs", "-n", "htdemucs_6s", "--out", str(TEMP_DIR), audio_file
        ], check=True, capture_output=True)

        # 3. Locate Stems
        demucs_out = TEMP_DIR / "htdemucs_6s" 
        track_folder = next(demucs_out.iterdir(), None)
        if not track_folder: raise FileNotFoundError("Demucs failed to output files.")
        
        # Map all 6 stems
        drums = track_folder / "drums.wav"
        bass = track_folder / "bass.wav"
        guitar = track_folder / "guitar.wav"
        piano = track_folder / "piano.wav"
        vocals = track_folder / "vocals.wav"
        other = track_folder / "other.wav"
        
        # Return paths to the UI and the BPM
        return str(drums), str(bass), str(guitar), str(piano), str(other), str(vocals), bpm, str(track_folder)

    except Exception as e:
        raise gr.Error(f"Separation Failed: {str(e)}")

def package_and_export(track_folder_str, bpm, start_offset_sec, cover_art):
    """Phase 2: Chop Loops, Generate Video, Zip"""
    try:
        track_folder = Path(track_folder_str)
        # Re-map paths
        stems = {
            "Drums": track_folder / "drums.wav",
            "Bass": track_folder / "bass.wav",
            "Guitar": track_folder / "guitar.wav",
            "Piano": track_folder / "piano.wav",
            "Synths": track_folder / "other.wav",
            "Vocals": track_folder / "vocals.wav"
        }

        # 1. Save Full Stems (Copy to Stems folder)
        for name, path in stems.items():
            if path.exists():
                shutil.copy(path, OUTPUT_DIR / "Stems" / f"{bpm}BPM_Full_{name}.wav")

        # 2. Create Loops (With User Offset)
        ms_per_beat = (60 / bpm) * 1000
        eight_bars_ms = ms_per_beat * 4 * 8
        start_ms = start_offset_sec * 1000
        
        created_loops = {} 

        def make_loop(src, name):
            if not src.exists(): return None
            audio = AudioSegment.from_wav(str(src))
            
            end_ms = start_ms + eight_bars_ms
            if len(audio) < end_ms:
                s = 0
                e = min(len(audio), eight_bars_ms)
            else:
                s, e = start_ms, end_ms
                
            loop = audio[s:int(e)].fade_in(15).fade_out(15).normalize()
            out_path = OUTPUT_DIR / "Loops" / f"{bpm}BPM_{name}.wav"
            loop.export(out_path, format="wav")
            return out_path

        # Generate loops for all 6 stems
        created_loops['melody'] = make_loop(stems['Synths'], "SynthLoop")
        make_loop(stems['Drums'], "DrumLoop")
        make_loop(stems['Bass'], "BassLoop")
        make_loop(stems['Guitar'], "GuitarLoop")
        make_loop(stems['Piano'], "PianoLoop")
        make_loop(stems['Vocals'], "VocalChop")
        
        # 3. Generate Video
        video_path = None
        if cover_art and created_loops['melody']:
            print("Rendering Dynamic Video...")
            vid_out = OUTPUT_DIR / "Promo_Video.mp4"
            audio_clip = AudioFileClip(str(created_loops['melody']))
            duration = audio_clip.duration
            
            # Load and Resize Image
            img = ImageClip(cover_art).resize(width=1080)
            
            # Simple Zoom Animation
            img = img.resize(lambda t : 1 + 0.02*t) 
            img = img.set_position(('center', 'center'))
            img = img.set_duration(duration)
            img = img.set_audio(audio_clip)
            
            final_clip = CompositeVideoClip([img], size=(1080, 1920))
            final_clip.duration = duration
            final_clip.audio = audio_clip
            final_clip.fps = 24
            
            final_clip.write_videofile(str(vid_out), codec="libx264", audio_codec="aac", logger=None)
            video_path = str(vid_out)

        # 4. Zip
        zip_file = "NightPulse_Pack.zip"
        with zipfile.ZipFile(zip_file, 'w') as zf:
            for root, dirs, files in os.walk(OUTPUT_DIR):
                for file in files:
                    file_path = Path(root) / file
                    arcname = file_path.relative_to(OUTPUT_DIR)
                    zf.write(file_path, arcname)
                    
        return zip_file, video_path

    except Exception as e:
        raise gr.Error(f"Packaging Failed: {str(e)}")


# --- GUI (Blocks) ---
with gr.Blocks(title="Night Pulse | Command Center (6-Stem)") as app:
    gr.Markdown("# 🎛️ Night Pulse | 6-Stem Command Center")
    gr.Markdown("Deconstruct audio into 6 stems: Drums, Bass, Guitar, Piano, Vocals, Synths.")
    
    # State storage
    stored_folder = gr.State()
    stored_bpm = gr.State()
    
    with gr.Row():
        with gr.Column(scale=1):
            input_audio = gr.Audio(type="filepath", label="1. Upload Master Track")
            input_art = gr.Image(type="filepath", label="Cover Art (9:16)")
            btn_analyze = gr.Button("🔍 Phase 1: Separate (6 Stems)", variant="primary")
            
        with gr.Column(scale=1):
            gr.Markdown("### 2. Stem Preview")
            with gr.Row():
                p_drums = gr.Audio(label="Drums")
                p_bass = gr.Audio(label="Bass")
            with gr.Row():
                p_guitar = gr.Audio(label="Guitar")
                p_piano = gr.Audio(label="Piano")
            with gr.Row():
                p_other = gr.Audio(label="Synths/Other")
                p_vocals = gr.Audio(label="Vocals")
    
    gr.Markdown("---")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 3. Loop Logic")
            slider_start = gr.Slider(minimum=0, maximum=120, value=15, label="Loop Start Time (Seconds)", info="Select the start point for the 8-bar loop cut.")
            btn_package = gr.Button("📦 Phase 2: Package & Export", variant="primary")
        
        with gr.Column():
            gr.Markdown("### 4. Final Output")
            out_zip = gr.File(label="Download Pack (ZIP)")
            out_video = gr.Video(label="Promo Video")

    # Events
    btn_analyze.click(
        fn=analyze_and_separate,
        inputs=[input_audio],
        outputs=[p_drums, p_bass, p_guitar, p_piano, p_other, p_vocals, stored_bpm, stored_folder]
    )
    
    btn_package.click(
        fn=package_and_export,
        inputs=[stored_folder, stored_bpm, slider_start, input_art],
        outputs=[out_zip, out_video]
    )

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