app.py

import gradio as gr
import torch
import os
import spaces
from pydub import AudioSegment
from typing import Tuple, Dict, List
from demucs.apply import apply_model
from demucs.separate import load_track
from demucs.pretrained import get_model
from demucs.audio import save_audio

device: str = "cuda" if torch.cuda.is_available() else "cpu"

# Define the inference function
@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, gr.HTML]:
    separator = get_model(name=model_name)
    log_messages = []

    def stream_log(message):
        log_messages.append(f"[{model_name}] {message}")
        return gr.HTML("<pre style='margin-bottom: 0;'>" + "<br>".join(log_messages) + "</pre>")

    yield None, stream_log("Starting separation process...")
    yield None, stream_log(f"Loading audio file: {audio_file}")
    
    # Check if audio_file is None
    if audio_file is None:
        yield None, stream_log("Error: No audio file provided")
        raise gr.Error("Please upload an audio file")

    # Load the audio file with the correct samplerate and audio channels
    try:
        wav, sr = load_track(audio_file, samplerate=separator.samplerate, audio_channels=2)
    except Exception as e:
        yield None, stream_log(f"Error loading audio file: {str(e)}")
        raise gr.Error(f"Failed to load audio file: {str(e)}")

    # Check the number of channels and adjust if necessary
    if wav.dim() == 1:
        wav = wav.unsqueeze(0)  # Add channel dimension if mono
    if wav.shape[0] == 1:
        wav = wav.repeat(2, 1)  # If mono, duplicate to stereo
    elif wav.shape[0] > 2:
        wav = wav[:2]  # If more than 2 channels, keep only the first two

    wav = wav.to(device)
    
    ref = wav.mean(0)
    wav = (wav - ref.view(1, -1))
    yield None, stream_log("Audio loaded successfully. Applying model...")

    # Use apply_model as a standalone function
    try:
        result = apply_model(separator, wav.to(device), device=device)
        yield None, stream_log(f"Model application result type: {type(result)}")
        yield None, stream_log(f"Model application result shape: {result.shape if hasattr(result, 'shape') else 'N/A'}")
        
        if isinstance(result, tuple) and len(result) == 2:
            sources, _ = result
        elif isinstance(result, torch.Tensor):
            sources = result
        else:
            raise ValueError(f"Unexpected result type from apply_model: {type(result)}")
        
        yield None, stream_log(f"Sources shape: {sources.shape}")
    except ValueError as e:
        yield None, stream_log(f"Error applying model: {str(e)}")
        yield None, stream_log(f"Separator sources: {separator.sources}")
        yield None, stream_log(f"WAV shape: {wav.shape}")
        yield None, stream_log(f"Separator model: {separator.__class__.__name__}")
        yield None, stream_log(f"Separator config: {separator.config}")
        raise gr.Error(f"Failed to apply model: {str(e)}. This might be due to incompatible audio format or model configuration.")
    except Exception as e:
        yield None, stream_log(f"Unexpected error applying model: {str(e)}")
        raise gr.Error(f"An unexpected error occurred while applying the model: {str(e)}")

    # Process the sources
    sources = [source * ref.view(1, -1) + ref.view(1, -1) for source in sources]
    
    yield None, stream_log("Model applied. Processing stems...")

    output_dir: str = os.path.join("separated", model_name, os.path.splitext(os.path.basename(audio_file))[0])
    os.makedirs(output_dir, exist_ok=True)
    stems: Dict[str, str] = {}
    for stem, source in zip(separator.sources, sources):
        stem_path: str = os.path.join(output_dir, f"{stem}.wav")
        save_audio(source, stem_path, separator.samplerate)
        stems[stem] = stem_path
        yield None, stream_log(f"Saved {stem} stem")

    selected_stems: List[str] = [stems[stem] for stem, include in zip(["vocals", "drums", "bass", "other"], [vocals, drums, bass, other]) if include]
    if not selected_stems:
        raise gr.Error("Please select at least one stem to mix.")

    output_file: str = os.path.join(output_dir, "mixed.wav")
    yield None, stream_log("Mixing selected stems...")
    if len(selected_stems) == 1:
        os.rename(selected_stems[0], output_file)
    else:
        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 None, stream_log(f"Converting to MP3 (bitrate: {mp3_bitrate}k)...")
        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 None, stream_log("Process completed successfully!")
    yield output_file, gr.HTML("<pre style='color: green;'>Separation and mixing completed successfully!</pre>")

# 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")
            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, separation_log]
    )
    
    mp3_checkbox.change(
        fn=lambda mp3: gr.update(visible=mp3),
        inputs=mp3_checkbox,
        outputs=mp3_bitrate
    )

# Launch the Gradio interface
iface.launch()