app.py

from __future__ import annotations

import os
import shutil
import subprocess
import sys
import textwrap
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import gradio as gr
import numpy as np
import soundfile as sf

try:
    import librosa
except Exception:
    librosa = None

APP_DIR = Path(__file__).resolve().parent
OUTPUT_DIR = APP_DIR / "outputs"
OUTPUT_DIR.mkdir(exist_ok=True)


@dataclass
class GeneratedAsset:
    label: str
    path: Path


def _summarize_runtime_error(message: str) -> str:
    lowered = message.lower()

    if "save_with_torchcodec" in lowered or "torchcodec is required" in lowered:
        return "A library tried to save audio through TorchCodec. Restart the app and run the separation again after this update."
    if "no module named 'demucs'" in lowered:
        return "Demucs is not installed. Run: pip install demucs"
    if "pyannote.audio is not installed" in lowered:
        return "Speaker separation is turned on, but `pyannote.audio` is not installed."
    if "speaker separation needs a hugging face access token" in lowered:
        return "Speaker separation is turned on, but no Hugging Face token was provided."
    if "noisereduce is not installed" in lowered:
        return "Noise estimation is turned on, but `noisereduce` is not installed."
    if "librosa is not installed correctly" in lowered:
        return "The audio-processing stack is incomplete because `librosa` is missing or broken."

    lines = [line.strip() for line in message.splitlines() if line.strip()]
    if not lines:
        return "An unknown error occurred."
    return lines[-1]


def _run_command(command: list[str], cwd: Path | None = None) -> None:
    completed = subprocess.run(
        command,
        cwd=str(cwd) if cwd else None,
        capture_output=True,
        text=True,
        check=False,
    )
    if completed.returncode != 0:
        stderr = completed.stderr.strip() or completed.stdout.strip() or "Unknown error"
        raise RuntimeError(f"Command failed: {' '.join(command)}\n{stderr}")


def _copy_uploaded_file(audio_path: str | Path, run_dir: Path) -> Path:
    source = Path(audio_path)
    target = run_dir / source.name
    shutil.copy2(source, target)
    return target


def _safe_stem(name: str) -> str:
    return "".join(ch if ch.isalnum() or ch in {"-", "_"} else "_" for ch in name).strip("_") or "audio"


def _latest_demucs_output(base_dir: Path, audio_stem: str) -> Path:
    matches = sorted(base_dir.rglob(f"{audio_stem}.wav"))
    if not matches:
        raise RuntimeError("Demucs finished without producing any WAV files.")
    return matches[0].parent


def _require_demucs_save_stack() -> None:
    try:
        import demucs  # noqa: F401
        import torch  # noqa: F401
    except Exception as exc:
        raise RuntimeError("Demucs or PyTorch is not installed in the current Python environment.") from exc

    # Demucs 4.x uses torchaudio / soundfile for I/O — torchcodec is NOT required.


def _prepare_demucs_audio(
    signal: np.ndarray,
    sample_rate: int,
    target_sample_rate: int,
    target_channels: int,
) -> np.ndarray:
    if signal.ndim == 1:
        signal = np.expand_dims(signal, axis=0)

    if sample_rate != target_sample_rate:
        signal = np.stack(
            [
                librosa.resample(channel, orig_sr=sample_rate, target_sr=target_sample_rate)
                for channel in signal
            ],
            axis=0,
        )

    if signal.shape[0] == target_channels:
        return signal.astype(np.float32)
    if target_channels == 1:
        return np.mean(signal, axis=0, keepdims=True).astype(np.float32)
    if signal.shape[0] == 1:
        return np.repeat(signal, target_channels, axis=0).astype(np.float32)
    if signal.shape[0] > target_channels:
        return signal[:target_channels].astype(np.float32)
    raise RuntimeError("The input audio channel layout is not supported for the selected Demucs model.")


def _patch_torchaudio_save() -> None:
    try:
        import torchaudio
    except Exception:
        return

    def _soundfile_save(filepath, src, sample_rate, *args, **kwargs):
        audio = src.detach().cpu().numpy()
        if audio.ndim == 2:
            audio = audio.T
        sf.write(filepath, audio, sample_rate)

    torchaudio.save = _soundfile_save


def separate_with_demucs(audio_path: Path, run_dir: Path, model_name: str) -> list[GeneratedAsset]:
    _require_demucs_save_stack()
    _require_audio_stack()
    _patch_torchaudio_save()
    import torch
    from demucs.apply import apply_model
    from demucs.pretrained import get_model

    model = get_model(name=model_name)
    model.cpu()
    model.eval()

    signal, sample_rate = librosa.load(str(audio_path), sr=None, mono=False)
    prepared_signal = _prepare_demucs_audio(signal, sample_rate, model.samplerate, model.audio_channels)
    wav = torch.tensor(prepared_signal, dtype=torch.float32)

    ref = wav.mean(0)
    wav = wav - ref.mean()
    ref_std = ref.std()
    if float(ref_std) > 0:
        wav = wav / ref_std

    sources = apply_model(
        model,
        wav[None],
        device="cpu",
        shifts=1,
        split=True,
        overlap=0.25,
        progress=False,
    )[0]

    if float(ref_std) > 0:
        sources = sources * ref_std
    sources = sources + ref.mean()

    stem_dir = run_dir / "demucs" / model_name / audio_path.stem
    stem_dir.mkdir(parents=True, exist_ok=True)
    assets: list[GeneratedAsset] = []
    for source_tensor, source_name in zip(sources, model.sources):
        stem_path = stem_dir / f"{source_name}.wav"
        stem_audio = source_tensor.detach().cpu().numpy().T
        sf.write(stem_path, stem_audio, model.samplerate)
        assets.append(GeneratedAsset(label=f"Stem: {source_name}", path=stem_path))
    return assets


def _mix_audio_files(paths: list[Path], output_path: Path) -> Path | None:
    if not paths:
        return None

    mixed_audio = None
    sample_rate = None

    for path in paths:
        audio, current_sr = sf.read(path, always_2d=True)
        if mixed_audio is None:
            mixed_audio = np.zeros_like(audio, dtype=np.float32)
            sample_rate = current_sr
        elif current_sr != sample_rate or audio.shape != mixed_audio.shape:
            raise RuntimeError("Audio stems do not share the same shape/sample rate, so they cannot be mixed.")

        mixed_audio += audio.astype(np.float32)

    mixed_audio = np.clip(mixed_audio, -1.0, 1.0)
    sf.write(output_path, mixed_audio, sample_rate)
    return output_path


def _zip_directory(directory: Path) -> Path | None:
    if not directory.exists() or not any(directory.rglob("*")):
        return None
    archive_base = directory.parent / directory.name
    archive_path = shutil.make_archive(str(archive_base), "zip", root_dir=directory)
    return Path(archive_path)


def build_background_music_asset(assets: list[GeneratedAsset], run_dir: Path) -> GeneratedAsset | None:
    stem_paths = {asset.path.stem.lower(): asset.path for asset in assets if asset.label.startswith("Stem:")}
    music_parts = [stem_paths[name] for name in ("bass", "drums", "other", "guitar", "piano") if name in stem_paths]
    if not music_parts:
        return None

    background_dir = run_dir / "background"
    background_dir.mkdir(parents=True, exist_ok=True)
    output_path = background_dir / "background_music_only.wav"
    mixed_path = _mix_audio_files(music_parts, output_path)
    if mixed_path is None:
        return None
    return GeneratedAsset(label="Background music only", path=mixed_path)


def _require_audio_stack() -> None:
    if librosa is None:
        raise RuntimeError(
            "librosa is not installed correctly. Reinstall dependencies before using speaker or noise extraction."
        )


def export_speaker_tracks(
    audio_path: Path,
    run_dir: Path,
    hf_token: str,
    min_segment_length: float,
) -> list[GeneratedAsset]:
    _require_audio_stack()
    try:
        from pyannote.audio import Pipeline
    except Exception as exc:
        raise RuntimeError(
            "pyannote.audio is not installed. Speaker separation needs the optional diarization stack."
        ) from exc

    if not hf_token.strip():
        raise RuntimeError(
            "Speaker separation needs a Hugging Face access token for the pyannote diarization model."
        )

    pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=hf_token.strip())
    diarization = pipeline(str(audio_path))

    signal, sample_rate = librosa.load(str(audio_path), sr=None, mono=False)
    if signal.ndim == 1:
        signal = np.expand_dims(signal, axis=0)

    speakers_dir = run_dir / "speakers"
    speakers_dir.mkdir(parents=True, exist_ok=True)

    speaker_buffers: dict[str, np.ndarray] = {}
    speaker_clip_counts: dict[str, int] = {}
    speaker_durations: dict[str, float] = {}
    assets: list[GeneratedAsset] = []

    for segment, _, speaker in diarization.itertracks(yield_label=True):
        duration = float(segment.end - segment.start)
        if duration < min_segment_length:
            continue

        start = max(0, int(segment.start * sample_rate))
        end = min(signal.shape[-1], int(segment.end * sample_rate))
        if end <= start:
            continue

        if speaker not in speaker_buffers:
            speaker_buffers[speaker] = np.zeros_like(signal)
            speaker_clip_counts[speaker] = 0
            speaker_durations[speaker] = 0.0

        speaker_buffers[speaker][:, start:end] = signal[:, start:end]
        speaker_clip_counts[speaker] += 1
        speaker_durations[speaker] += duration

        clip_path = speakers_dir / f"{speaker}_clip_{speaker_clip_counts[speaker]:03d}.wav"
        sf.write(clip_path, signal[:, start:end].T, sample_rate)
        assets.append(GeneratedAsset(label=f"{speaker} clip {speaker_clip_counts[speaker]}", path=clip_path))

    for speaker, buffer in speaker_buffers.items():
        speaker_path = speakers_dir / f"{speaker}_timeline.wav"
        sf.write(speaker_path, buffer.T, sample_rate)
        assets.insert(0, GeneratedAsset(label=f"{speaker} isolated timeline", path=speaker_path))

    # Treat the longest-running speaker as the main voice and fold the rest
    # together into a single "background voices" track.
    if len(speaker_buffers) >= 2:
        primary_speaker = max(speaker_durations, key=speaker_durations.get)
        background_mix = np.zeros_like(signal)
        for speaker, buffer in speaker_buffers.items():
            if speaker != primary_speaker:
                background_mix += buffer

        background_voice_path = speakers_dir / "background_voices_only.wav"
        sf.write(background_voice_path, np.clip(background_mix.T, -1.0, 1.0), sample_rate)
        assets.insert(0, GeneratedAsset(label="Background voices only", path=background_voice_path))

    if not assets:
        raise RuntimeError("No speaker segments were produced. Try lowering the minimum segment length.")

    return assets


def export_noise_estimates(audio_path: Path, run_dir: Path, prop_decrease: float) -> list[GeneratedAsset]:
    _require_audio_stack()
    try:
        import noisereduce as nr
    except Exception as exc:
        raise RuntimeError(
            "noisereduce is not installed. Noise extraction needs the optional denoising stack."
        ) from exc

    signal, sample_rate = librosa.load(str(audio_path), sr=None, mono=False)
    cleaned = nr.reduce_noise(y=signal, sr=sample_rate, prop_decrease=prop_decrease)
    estimated_noise = signal - cleaned

    noise_dir = run_dir / "noise"
    noise_dir.mkdir(parents=True, exist_ok=True)

    cleaned_path = noise_dir / "denoised.wav"
    noise_path = noise_dir / "estimated_noise.wav"
    sf.write(cleaned_path, cleaned.T if cleaned.ndim > 1 else cleaned, sample_rate)
    sf.write(noise_path, estimated_noise.T if estimated_noise.ndim > 1 else estimated_noise, sample_rate)

    return [
        GeneratedAsset(label="Denoised audio", path=cleaned_path),
        GeneratedAsset(label="Estimated noise residue", path=noise_path),
    ]


def _build_summary(run_dir: Path, steps: Iterable[str], warnings: Iterable[str], assets: list[GeneratedAsset]) -> str:
    file_lines = [f"- {asset.label}: `{asset.path.name}`" for asset in assets]
    warning_lines = [f"- {item}" for item in warnings if item]
    return textwrap.dedent(
        f"""
        Processing finished.

        Output folder:
        `{run_dir}`

        Completed steps:
        {os.linesep.join(f"- {step}" for step in steps)}

        Files:
        {os.linesep.join(file_lines) if file_lines else "- No files were generated."}

        Notes:
        {os.linesep.join(warning_lines) if warning_lines else "- None"}
        """
    ).strip()


def _build_failure_summary(run_dir: Path, warnings: Iterable[str]) -> str:
    warning_lines = [f"- {item}" for item in warnings if item]
    return textwrap.dedent(
        f"""
        Processing could not generate any separated tracks.

        Output folder:
        `{run_dir}`

        Most likely causes:
        - One or more required packages are not installed correctly.
        - Speaker separation was enabled without a valid Hugging Face token.

        Detailed errors:
        {os.linesep.join(warning_lines) if warning_lines else "- No extra details were captured."}
        """
    ).strip()


def process_audio(
    audio_file: str | None,
    video_file: str | None,
    demucs_model: str,
    run_source_separation: bool,
    run_speaker_diarization: bool,
    run_noise_estimation: bool,
    hf_token: str,
    min_segment_length: float,
    prop_decrease: float,
):
    if not audio_file and not video_file:
        raise gr.Error("Upload an audio or video file first.")

    input_file = video_file if video_file else audio_file
    original_path = Path(input_file)
    run_id = f"{_safe_stem(original_path.stem)}_{uuid.uuid4().hex[:8]}"
    run_dir = OUTPUT_DIR / run_id
    run_dir.mkdir(parents=True, exist_ok=True)

    steps: list[str] = []
    warnings: list[str] = []
    assets: list[GeneratedAsset] = []

    if video_file:
        uploaded_copy = run_dir / f"{original_path.stem}.wav"
        try:
            _run_command([
                "ffmpeg", "-i", str(original_path), 
                "-vn", "-acodec", "pcm_s16le", "-ar", "44100", "-ac", "2", "-y", 
                str(uploaded_copy)
            ])
            assets.append(GeneratedAsset(label="Extracted audio from video", path=uploaded_copy))
            steps.append("Extracted audio from video")
        except Exception as exc:
            raise gr.Error(f"Failed to extract audio from video: {exc}")
    else:
        uploaded_copy = _copy_uploaded_file(input_file, run_dir)
        assets.append(GeneratedAsset(label="Original upload", path=uploaded_copy))

    # Each processing path is optional on purpose. If one model is missing,
    # the app still returns anything that did succeed instead of failing hard.
    if run_source_separation:
        try:
            stem_assets = separate_with_demucs(uploaded_copy, run_dir, demucs_model)
            assets.extend(stem_assets)
            background_music_asset = build_background_music_asset(stem_assets, run_dir)
            if background_music_asset is not None:
                assets.append(background_music_asset)
            demucs_archive = _zip_directory(run_dir / "demucs")
            if demucs_archive is not None:
                assets.append(GeneratedAsset(label="Demucs folder archive", path=demucs_archive))
            steps.append(f"Source separation with Demucs ({demucs_model})")
        except Exception as exc:
            warnings.append(_summarize_runtime_error(str(exc)))

    if run_speaker_diarization:
        try:
            assets.extend(export_speaker_tracks(uploaded_copy, run_dir, hf_token, min_segment_length))
            steps.append("Speaker diarization and speaker track export")
        except Exception as exc:
            warnings.append(_summarize_runtime_error(str(exc)))

    if run_noise_estimation:
        try:
            assets.extend(export_noise_estimates(uploaded_copy, run_dir, prop_decrease))
            noise_archive = _zip_directory(run_dir / "noise")
            if noise_archive is not None:
                assets.append(GeneratedAsset(label="Noise folder archive", path=noise_archive))
            steps.append("Noise reduction and residual-noise estimation")
        except Exception as exc:
            warnings.append(_summarize_runtime_error(str(exc)))

    if len(assets) == 1:
        summary = _build_failure_summary(run_dir, warnings)
        return summary, [], []

    downloadable_files = [str(asset.path) for asset in assets]
    summary = _build_summary(run_dir, steps or ["No optional processing step succeeded."], warnings, assets)
    audio_list = [{"label": asset.label, "path": str(asset.path)} for asset in assets if str(asset.path).lower().endswith(".wav")]
    return summary, downloadable_files, audio_list


def build_demo() -> gr.Blocks:
    # The UI stays intentionally simple: upload first, then tune the optional
    # processing steps only if the recording needs them.
    intro = """
    # Audio Separator Studio

    Upload one audio file and the app will try to split it into separate outputs:
    vocals, accompaniment, drums/bass/other stems, individual speaker tracks, and an estimated noise layer.

    Perfectly separating every sound in a messy real-world recording is not possible yet, but this app gets you
    several practical layers from the same file.
    """

    guide = """
    ## What each option means

    **Upload audio or video**
    Add the song, podcast, interview, or video recording you want to split. If you upload a video, the audio will be automatically extracted.

    **Separation model**
    This controls how Demucs separates music parts.

    - `mdx_extra_q`: fastest and lightest. Good when you want quicker results.
    - `htdemucs`: balanced option. Better for most people and the default choice.
    - `htdemucs_ft`: slowest and heaviest. Usually the best quality when you want the strongest separation.
    - `htdemucs_6s`: adds extra `guitar` and `piano` stems on top of the usual stems. `[Default Selected]`

    **Separate music stems**
    Creates files such as vocals, drums, bass, other, and also a combined `background music only` file with vocals removed.

    **Separate speakers**
    Tries to detect different human speakers in the recording and exports separate files for each speaker.
    If there are multiple speakers, it also creates `background voices only`, which keeps the secondary speakers and removes the main speaker.

    **Estimate background noise**
    Creates a cleaner version of the audio and also a file containing the estimated leftover noise.

    **Hugging Face token for speaker diarization**
    This is only needed for `Separate speakers`.
    A Hugging Face token is a private access key from Hugging Face.
    `Speaker diarization` means detecting who spoke when, so the app can split one recording into speaker-wise tracks.
    If you are not using speaker separation, you can leave this box empty.

    **Minimum speaker segment length**
    Very short speech fragments can create messy results. Increase this to ignore tiny fragments. Lower it if you want more detailed speaker clips.

    **Noise reduction strength**
    Higher values remove more noise, but can also affect the original sound more strongly.
    """

    css = """
    .instructions-panel {
      padding: 14px !important;
      background: #2b2b31 !important;
      border: 2px dotted rgba(255, 255, 255, 0.8) !important;
      border-radius: 10px !important;
      box-shadow: none !important;
    }

    .instructions-panel > div {
      border: none !important;
      box-shadow: none !important;
      background: transparent !important;
      padding: 0 !important;
    }
    """

    with gr.Blocks(title="Audio Separator Studio", css=css) as demo:
        gr.Markdown(intro)
        with gr.Group(elem_classes=["instructions-panel"]):
            gr.Markdown(guide)
        with gr.Row():
            with gr.Column(scale=2):
                audio_input = gr.Audio(type="filepath", sources=["upload"], label="Upload audio")
                video_input = gr.Video(sources=["upload"], label="Or upload video (MP4)", include_audio=True)
                demucs_model = gr.Dropdown(
                    choices=["htdemucs", "htdemucs_ft", "htdemucs_6s", "mdx_extra_q"],
                    value="htdemucs_6s",
                    label="Source separation model",
                    info="Choose speed vs quality: mdx_extra_q = good and fast, htdemucs = better balanced, htdemucs_ft = best quality but slower, htdemucs_6s = adds guitar and piano stems.",
                )
                with gr.Row():
                    run_source_separation = gr.Checkbox(
                        value=True,
                        label="Separate music stems",
                        info="Splits the audio into stems like vocals, drums, bass, other, and background-music-only.",
                    )
                    run_speaker_diarization = gr.Checkbox(
                        value=True,
                        label="Separate speakers",
                        info="Splits different human speakers into separate files. Needs a Hugging Face token.",
                    )
                    run_noise_estimation = gr.Checkbox(
                        value=True,
                        label="Estimate background noise",
                        info="Creates a denoised version and a separate noise-only estimate.",
                    )
                hf_token = gr.Textbox(
                    value=os.environ.get("HF_TOKEN", ""),
                    type="password",
                    label="Hugging Face token for speaker diarization",
                    info="Only needed for speaker separation. It is your Hugging Face access key used to load the speaker-diarization model.",
                )
                min_segment_length = gr.Slider(
                    minimum=0.25,
                    maximum=3.0,
                    value=0.75,
                    step=0.05,
                    label="Minimum speaker segment length (seconds)",
                    info="Higher values ignore tiny speech fragments. Lower values keep more short speech clips.",
                )
                prop_decrease = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.9,
                    step=0.05,
                    label="Noise reduction strength",
                    info="Higher values remove more noise, but can also remove more detail from the original audio.",
                )
                run_button = gr.Button("Process Audio", variant="primary")
            with gr.Column(scale=3):
                summary = gr.Markdown(label="Summary")
                downloads = gr.File(label="Generated files", file_count="multiple")
                
                audio_players_state = gr.State([])
                
                @gr.render(inputs=audio_players_state)
                def render_audio_players(audio_list):
                    if audio_list:
                        gr.Markdown("### Listen to Generated Audio")
                        for audio_info in audio_list:
                            gr.Audio(value=audio_info["path"], label=audio_info["label"])

        run_button.click(
            fn=process_audio,
            inputs=[
                audio_input,
                video_input,
                demucs_model,
                run_source_separation,
                run_speaker_diarization,
                run_noise_estimation,
                hf_token,
                min_segment_length,
                prop_decrease,
            ],
            outputs=[summary, downloads, audio_players_state],
        )

    return demo


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