DeepFilterNet2 single-file denoiser CLI (single file or recursive folder, with inline uv deps)

denoise.py

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.9,<3.12"
# dependencies = [
#     "torch>=2.0,<2.2",
#     "torchaudio>=2.0,<2.2",
#     "deepfilternet==0.5.6",
#     "numpy<2",
#     "soundfile",
# ]
# ///
"""
DeepFilterNet2 batch / single-file denoiser.

Usage:
    # single file
    ./denoise.py input.wav -o output.wav

    # folder (recursive, mirrors structure into output dir)
    ./denoise.py /path/to/noisy_dir -o /path/to/clean_dir

    # multiple files (e.g. shell glob expansion)
    ./denoise.py *.flac -o /path/to/clean_dir

    # run with uv (handles deps via the inline PEP 723 metadata above):
    uv run denoise.py noisy_dir -o clean_dir

The DeepFilterNet2 model is downloaded on first run by `init_df()` and cached
locally (under the deepfilternet package data dir).
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import torch

from df.enhance import enhance, init_df, load_audio, save_audio


AUDIO_EXTS = {".wav", ".flac", ".ogg", ".mp3", ".m4a", ".aac", ".aif", ".aiff", ".opus"}


def collect_files(root: Path) -> list[Path]:
    return sorted(p for p in root.rglob("*") if p.is_file() and p.suffix.lower() in AUDIO_EXTS)


def denoise_file(in_path: Path, out_path: Path, model, df_state, sr: int, atten_lim_db: float | None) -> None:
    audio, meta = load_audio(str(in_path), sr=sr)
    device = next(model.parameters()).device
    audio = audio.to(device)
    kwargs = {}
    if atten_lim_db is not None:
        kwargs["atten_lim_db"] = atten_lim_db
    enhanced = enhance(model, df_state, audio, **kwargs)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    save_audio(str(out_path), enhanced.cpu(), sr)


def main() -> int:
    parser = argparse.ArgumentParser(description="Denoise audio with DeepFilterNet2.")
    parser.add_argument(
        "inputs", type=Path, nargs="+",
        help="Input audio file(s) or directory. Multiple files (e.g. via shell glob) are supported.",
    )
    parser.add_argument(
        "-o", "--output", type=Path, required=True,
        help="Output file (single-file input) or output directory (folder / multiple files).",
    )
    parser.add_argument(
        "--atten-lim-db", type=float, default=None,
        help="Limit max attenuation in dB (e.g. 25). Default: no limit.",
    )
    parser.add_argument(
        "--device", default=None, choices=[None, "cpu", "cuda", "mps"],
        help="Torch device. Default: auto.",
    )
    parser.add_argument(
        "--overwrite", action="store_true",
        help="Overwrite existing output files (default: skip).",
    )
    args = parser.parse_args()

    # CPU is the default. MPS is slower than CPU for this model on Apple Silicon
    # and has unimplemented ops; CUDA is fast but opt-in. Pass --device cuda/mps to override.
    device = torch.device(args.device) if args.device else torch.device("cpu")

    print(f"Loading DeepFilterNet2 (device={device}) ...", file=sys.stderr)
    model, df_state, _ = init_df()
    model = model.to(device=device).eval()
    sr = df_state.sr()

    inputs: list[Path] = args.inputs
    out_path: Path = args.output

    # Single-file input -> output may be a file or a directory.
    if len(inputs) == 1 and inputs[0].is_file():
        in_file = inputs[0]
        if out_path.exists() and out_path.is_dir():
            out_file = out_path / in_file.name
        else:
            out_file = out_path
        if out_file.exists() and not args.overwrite:
            print(f"Skipping (exists): {out_file}", file=sys.stderr)
            return 0
        print(f"Denoising {in_file} -> {out_file}", file=sys.stderr)
        denoise_file(in_file, out_file, model, df_state, sr, args.atten_lim_db)
        return 0

    # Single directory input -> recursive, mirror tree into out_path.
    if len(inputs) == 1 and inputs[0].is_dir():
        in_dir = inputs[0]
        files = collect_files(in_dir)
        if not files:
            print(f"No audio files found under {in_dir}", file=sys.stderr)
            return 1
        print(f"Found {len(files)} audio files. Writing to {out_path}", file=sys.stderr)
        for i, f in enumerate(files, 1):
            rel = f.relative_to(in_dir)
            out_file = out_path / rel
            if out_file.exists() and not args.overwrite:
                print(f"[{i}/{len(files)}] skip (exists): {rel}", file=sys.stderr)
                continue
            print(f"[{i}/{len(files)}] {rel}", file=sys.stderr)
            try:
                denoise_file(f, out_file, model, df_state, sr, args.atten_lim_db)
            except Exception as e:
                print(f"  ERROR: {e}", file=sys.stderr)
        return 0

    # Multiple inputs -> output must be a directory; flat output (filenames only).
    files = [p for p in inputs if p.is_file()]
    missing = [p for p in inputs if not p.exists()]
    for p in missing:
        print(f"Input not found: {p}", file=sys.stderr)
    if not files:
        return 1
    out_path.mkdir(parents=True, exist_ok=True)
    print(f"Denoising {len(files)} files into {out_path}", file=sys.stderr)
    for i, f in enumerate(files, 1):
        out_file = out_path / f.name
        if out_file.exists() and not args.overwrite:
            print(f"[{i}/{len(files)}] skip (exists): {f.name}", file=sys.stderr)
            continue
        print(f"[{i}/{len(files)}] {f.name}", file=sys.stderr)
        try:
            denoise_file(f, out_file, model, df_state, sr, args.atten_lim_db)
        except Exception as e:
            print(f"  ERROR: {e}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())