Upload 2 files

app.py

@@ -0,0 +1,625 @@
+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()

requirements.txt

@@ -0,0 +1,10 @@
+# Best on Python 3.10 or 3.11.
+gradio>=5.25.0
+numpy>=1.26.0
+soundfile>=0.12.1
+librosa>=0.10.2
+noisereduce>=3.0.2
+demucs>=4.0.1
+torch>=2.2.0
+torchaudio>=2.2.0
+pyannote.audio>=3.3.2