1 files changed, 98 insertions, 0 deletions
diff --git a/lib/src/producer/engines/denoise_dfn.py b/lib/src/producer/engines/denoise_dfn.py
new file mode 100644
index 0000000..66a8cf6
--- /dev/null
+++ b/lib/src/producer/engines/denoise_dfn.py
@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import sys
+import types
+import urllib.request
+import zipfile
+from pathlib import Path
+
+import numpy as np
+
+from .. import dsp, lazy
+from .base import blend, device_name, pick_device
+
+MODELS_DIR = lazy.DATA_DIR / "models"
+TAG = "v0.5.6"
+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:
+    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 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"
+    print(f"[producer] downloading {model} weights...", flush=True)
+    urllib.request.urlretrieve(url, zpath)
+    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"
+) -> tuple[np.ndarray, str, str]:
+    lazy.ensure(["deepfilternet==0.5.6"], purpose="DeepFilterNet")
+    lazy.ensure_torch()
+    _shim_torchaudio_backend()
+    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=strength >= 0.95,
+        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
+
+    xin_t = torch.from_numpy(xin.astype(np.float32)).unsqueeze(0)
+    y = df_enhance(model, df_state, xin_t)
+    if isinstance(y, torch.Tensor):
+        y = y.detach().cpu().numpy()
+    y = np.asarray(y, dtype=np.float32)
+    if y.ndim > 1:
+        y = y.reshape(-1)
+    y = dsp.resample(y, sr_df, sr)
+    y = blend(x, y, strength)
+    return y, f"dfn ({model_name})", f"{device_name(dev)} ({dev})"