diff options
| author | historia <historiavg@proton.me> | 2026-09-07 06:47:47 -0400 |
|---|---|---|
| committer | historia <historiavg@proton.me> | 2026-09-07 06:47:47 -0400 |
| commit | 84dd2d068317998f6fb59400c534ef5be6b51b53 (patch) | |
| tree | 025293e9d9229e02960374771ae522d9de2628ce /lib/src/producer/engines/denoise_dfn.py | |
| parent | 39b0f2bbed74f6487a41b82501ae3c6799e4b5c4 (diff) | |
| download | producer-main.tar.gz | |
Diffstat (limited to 'lib/src/producer/engines/denoise_dfn.py')
| -rw-r--r-- | lib/src/producer/engines/denoise_dfn.py | 140 |
1 files changed, 0 insertions, 140 deletions
diff --git a/lib/src/producer/engines/denoise_dfn.py b/lib/src/producer/engines/denoise_dfn.py deleted file mode 100644 index 186ffd2..0000000 --- a/lib/src/producer/engines/denoise_dfn.py +++ /dev/null @@ -1,140 +0,0 @@ -from __future__ import annotations - -import sys -import types -import warnings -import zipfile -from collections.abc import Callable -from pathlib import Path - -import numpy as np - -from .. import dsp, lazy, ui -from .base import blend, device_name, pick_device -from .chunking import apply_chunked, free_vram - -MODELS_DIR = lazy.DATA_DIR / "models" -TAG = "v0.5.6" -# Neighbouring audio fed to each chunk so the recurrent model and its feature -# normalizers run warm at chunk seams; trimmed away before stitching. -CONTEXT_S = 2.0 -MODEL_ZIPS = { - "DeepFilterNet3": "models/DeepFilterNet3.zip", - "DeepFilterNet2": "models/DeepFilterNet2.zip", -} -BASE_URL = f"https://raw.githubusercontent.com/Rikorose/DeepFilterNet/{TAG}" - - -def _shim_torchaudio_backend() -> None: - warnings.filterwarnings( - "ignore", - message=r".*AudioMetaData.*has been moved.*", - category=UserWarning, - module=r"df[./]io", - ) - try: - import torchaudio.backend # noqa: F401 - except Exception: - pkg = sys.modules.get("torchaudio") - if pkg is not None and "torchaudio.backend" not in sys.modules: - stub = types.ModuleType("torchaudio.backend") - stub.__path__ = [] - sys.modules["torchaudio.backend"] = stub - pkg.backend = stub - - -def _shim_df_git() -> None: - import df.io - import df.logger - import df.utils - - for mod in (df.utils, df.logger, df.io): - for name in ("get_git_root", "get_commit_hash", "get_branch_name"): - if hasattr(mod, name): - setattr(mod, name, lambda: None) - - -def ensure_model(model: str) -> Path: - target = MODELS_DIR / model - if (target / "config.ini").is_file(): - return target - url = f"{BASE_URL}/{MODEL_ZIPS[model]}" - MODELS_DIR.mkdir(parents=True, exist_ok=True) - zpath = MODELS_DIR / f"{model}.zip" - ui.download(url, zpath, f"downloading {model} weights") - ui.log(f"[producer] extracting {model} weights...") - with zipfile.ZipFile(zpath) as z: - z.extractall(target) - zpath.unlink(missing_ok=True) - if not (target / "config.ini").is_file(): - inner = list(target.glob(f"**/{model}/config.ini")) - if inner: - src = inner[0].parent - for f in src.iterdir(): - f.rename(target / f.name) - if not (target / "config.ini").is_file(): - raise RuntimeError(f"{model} weights download failed") - return target - - -def denoise( - x: np.ndarray, - sr: int, - strength: float, - device_pref: str = "auto", - chunk_s: float = 30.0, - overlap_s: float = 0.5, - on_progress: Callable[[int, int], None] | None = None, - post_filter: bool = False, -) -> tuple[np.ndarray, str, str]: - """Denoise with DeepFilterNet. - - post_filter opts into DFN's extra noise-reduction post filter; it - over-attenuates and can eat soft speech on clean recordings, so it stays - off unless requested. - """ - # torch first: deepfilternet's declared torch dependency would otherwise - # resolve to the newest (multi-GB CUDA) build before we pin our tested one. - lazy.ensure_torch() - lazy.ensure(["deepfilternet==0.5.6"], purpose="DeepFilterNet") - _shim_torchaudio_backend() - _shim_df_git() - model_name = "DeepFilterNet3" - try: - model_dir = ensure_model(model_name) - except Exception: - model_name = "DeepFilterNet2" - model_dir = ensure_model(model_name) - from df.enhance import enhance as df_enhance - from df.enhance import init_df - - device = pick_device(device_pref) - model, df_state, _ = init_df( - model_base_dir=str(model_dir), - post_filter=post_filter, - log_level="error", - log_file=None, - ) - try: - model = model.to(device) - dev = device - except Exception: - dev = "cpu" - sr_df = int(df_state.sr()) - xin = dsp.resample(x, sr, sr_df) - import torch - - def run(chunk: np.ndarray) -> np.ndarray: - t = torch.from_numpy(np.ascontiguousarray(chunk, dtype=np.float32)).unsqueeze(0) - y = df_enhance(model, df_state, t) - if isinstance(y, torch.Tensor): - y = y.detach().cpu().numpy() - return np.asarray(y, dtype=np.float32).reshape(-1) - - y = apply_chunked( - xin, sr_df, chunk_s, overlap_s, run, on_progress=on_progress, context_s=CONTEXT_S - ) - free_vram() - y = dsp.resample(y, sr_df, sr) - y = blend(x, y, strength) - return y, f"dfn ({model_name})", f"{device_name(dev)} ({dev})" |
