app.py

import gradio as gr
import os
import subprocess
import spaces
from typing import Tuple, List, Dict
from pydub import AudioSegment
from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn
from rich.text import Text
import time

console = Console()

def fade_text(text, duration=0.5):
    for i in range(10):
        opacity = i / 10
        yield f"<div style='opacity: {opacity}; transition: opacity 0.05s;'>{text}</div>"
        time.sleep(duration / 10)

@spaces.GPU
def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass: bool, other: bool, mp3: bool, mp3_bitrate: int) -> Tuple[str, List[str], gr.HTML]:
    log_messages = []

    def stream_log(message, style=""):
        formatted_message = f"[{model_name}] {message}"
        log_messages.append(formatted_message)
        for frame in fade_text(f"<pre style='margin-bottom: 0;{style}'>{formatted_message}</pre>"):
            yield None, None, gr.HTML(frame)

    yield from stream_log("Initializing Demucs...", "color: #4CAF50; font-weight: bold;")
    time.sleep(1)  # Simulate initialization time

    yield from stream_log("Loading audio file...", "color: #2196F3;")
    time.sleep(0.5)  # Simulate loading time
    
    if audio_file is None:
        yield from stream_log("Error: No audio file provided", "color: #F44336;")
        raise gr.Error("Please upload an audio file")

    # Use absolute paths
    base_output_dir = os.path.abspath("separated")
    output_dir = os.path.join(base_output_dir, model_name, os.path.splitext(os.path.basename(audio_file))[0])
    os.makedirs(output_dir, exist_ok=True)

    # Construct the Demucs command with full paths
    cmd = [
        "python", "-m", "demucs", 
        "--out", base_output_dir,
        "-n", model_name,
        audio_file
    ]

    yield from stream_log("Preparing separation process...", "color: #FF9800;")
    time.sleep(0.5)  # Simulate preparation time

    try:
        # Run the Demucs command
        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
        
        # Simulate a loading animation
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            TimeRemainingColumn(),
            console=console
        ) as progress:
            task = progress.add_task("[cyan]Separating stems...", total=100)
            while process.poll() is None:
                progress.update(task, advance=1)
                time.sleep(0.1)
            progress.update(task, completed=100)

        if process.returncode != 0:
            error_output = process.stderr.read()
            yield from stream_log(f"Error: Separation failed", "color: #F44336;")
            raise gr.Error(f"Demucs separation failed. Check the logs for details.")

    except Exception as e:
        yield from stream_log(f"Unexpected error: {str(e)}", "color: #F44336;")
        raise gr.Error(f"An unexpected error occurred: {str(e)}")

    yield from stream_log("Separation completed successfully!", "color: #4CAF50; font-weight: bold;")
    time.sleep(0.5)  # Pause for effect

    yield from stream_log("Processing stems...", "color: #9C27B0;")
    time.sleep(0.5)  # Simulate processing time

    # Change the stem search directory using full path
    stem_search_dir = os.path.join(base_output_dir, model_name, os.path.splitext(os.path.basename(audio_file))[0])
    yield from stream_log(f"Searching for stems in: {stem_search_dir}")

    stems: Dict[str, str] = {}
    for stem in ["vocals", "drums", "bass", "other"]:
        stem_path = os.path.join(stem_search_dir, f"{stem}.wav")
        yield from stream_log(f"Checking for {stem} stem at: {stem_path}")
        if os.path.exists(stem_path):
            stems[stem] = stem_path
            yield from stream_log(f"Found {stem} stem")
        else:
            yield from stream_log(f"Warning: {stem} stem not found")

    if not stems:
        yield from stream_log("Error: No stems found. Checking alternative directory...")
        stem_search_dir = os.path.join(base_output_dir, model_name)
        for stem in ["vocals", "drums", "bass", "other"]:
            stem_path = os.path.join(stem_search_dir, f"{stem}.wav")
            yield from stream_log(f"Checking for {stem} stem at: {stem_path}")
            if os.path.exists(stem_path):
                stems[stem] = stem_path
                yield from stream_log(f"Found {stem} stem")
            else:
                yield from stream_log(f"Warning: {stem} stem not found")

    yield from stream_log(f"All found stems: {list(stems.keys())}")

    selected_stems: List[str] = []
    for stem, selected in zip(["vocals", "drums", "bass", "other"], [vocals, drums, bass, other]):
        if selected:
            yield from stream_log(f"{stem} is selected by user")
            if stem in stems:
                selected_stems.append(stems[stem])
                yield from stream_log(f"Selected {stem} stem for mixing")
            else:
                yield from stream_log(f"Warning: {stem} was selected but not found")

    yield from stream_log(f"Final selected stems: {selected_stems}")

    if not selected_stems:
        yield from stream_log("Error: No stems selected for mixing", "color: #F44336;")
        raise gr.Error("Please select at least one stem to mix and ensure it was successfully separated.")

    output_file: str = os.path.join(output_dir, "mixed.wav")
    yield from stream_log("Mixing selected stems...", "color: #FF5722;")
    time.sleep(0.5)  # Simulate mixing time
    
    mixed_audio: AudioSegment = AudioSegment.empty()
    for stem_path in selected_stems:
        mixed_audio += AudioSegment.from_wav(stem_path)
    mixed_audio.export(output_file, format="wav")

    if mp3:
        yield from stream_log(f"Converting to MP3...", "color: #795548;")
        time.sleep(0.5)  # Simulate conversion time
        mp3_output_file: str = os.path.splitext(output_file)[0] + ".mp3"
        mixed_audio.export(mp3_output_file, format="mp3", bitrate=str(mp3_bitrate) + "k")
        output_file = mp3_output_file

    yield from stream_log("Process completed successfully!", "color: #4CAF50; font-weight: bold;")
    final_message = Panel(
        Text("Separation and mixing completed successfully!", style="bold green"),
        title="Demucs Result",
        border_style="green"
    )
    yield output_file, list(stems.values()), gr.HTML(console.export_html(final_message))

# Define the Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# Demucs Music Source Separation and Mixing")
    gr.Markdown("Separate vocals, drums, bass, and other instruments from your music using Demucs and mix the selected stems.")
    
    with gr.Row():
        with gr.Column(scale=1):
            audio_input = gr.Audio(type="filepath", label="Upload Audio File")
            model_dropdown = gr.Dropdown(
                ["htdemucs", "htdemucs_ft", "htdemucs_6s", "hdemucs_mmi", "mdx", "mdx_extra", "mdx_q", "mdx_extra_q"],
                label="Model Name",
                value="htdemucs_ft"
            )
            with gr.Row():
                vocals_checkbox = gr.Checkbox(label="Vocals", value=True)
                drums_checkbox = gr.Checkbox(label="Drums", value=True)
            with gr.Row():
                bass_checkbox = gr.Checkbox(label="Bass", value=True)
                other_checkbox = gr.Checkbox(label="Other", value=True)
            mp3_checkbox = gr.Checkbox(label="Save as MP3", value=False)
            mp3_bitrate = gr.Slider(128, 320, step=32, label="MP3 Bitrate", visible=False)
            submit_btn = gr.Button("Process", variant="primary")

        with gr.Column(scale=1):
            output_audio = gr.Audio(type="filepath", label="Processed Audio (Mixed)")
            stems_output = gr.File(label="Individual Stems", file_count="multiple")
            separation_log = gr.HTML()

    submit_btn.click(
        fn=inference,
        inputs=[audio_input, model_dropdown, vocals_checkbox, drums_checkbox, bass_checkbox, other_checkbox, mp3_checkbox, mp3_bitrate],
        outputs=[output_audio, stems_output, separation_log]
    )
    
    mp3_checkbox.change(
        fn=lambda mp3: gr.update(visible=mp3),
        inputs=mp3_checkbox,
        outputs=mp3_bitrate
    )

# Launch the Gradio interface
iface.launch()