"""Automatic, isolated DeepFilterNet3 inference (Linux, Python 3.11 worker). Public progress callback: (stage: str, completed: float | None, total: float | None). Known totals are audio frames or download bytes (identified by stage); unknown totals use elapsed seconds as completed. First use downloads Python, wheels, and upstream model weights into the VoiceForge data directory (VOICEFORGE_HOME prefix or the XDG data location). """ from __future__ import annotations import ctypes import fcntl import math import os from pathlib import Path import platform import shutil import time from .setup import Progress, data_dir, ensure_uv, run_process, tool_env def _cuda_available() -> bool: # Probe the driver without importing Torch or installing CUDA in the host env. if os.environ.get("CUDA_VISIBLE_DEVICES") in {"", "-1"}: return False try: driver = ctypes.CDLL("libcuda.so.1") count = ctypes.c_int() return driver.cuInit(0) == 0 and driver.cuDeviceGetCount(ctypes.byref(count)) == 0 and count.value > 0 except (OSError, AttributeError): return False def _flavor(device: str) -> str: if device not in {"auto", "cpu", "cuda"}: raise ValueError("device must be 'auto', 'cpu', or 'cuda'") if device == "cuda" or (device == "auto" and _cuda_available()): if platform.machine() not in {"x86_64", "AMD64"}: raise RuntimeError("The pinned CUDA 12.1 backend requires Linux x86_64.") return "cu121" return "cpu" def _worker_env() -> dict[str, str]: # Same isolation as tool_env(); kept as a named seam for worker launches. return tool_env() def _interpreter_is_local(python: Path, root: Path) -> bool: """Whether the environment's interpreter belongs to this prefix. An installation duplicated by copying (while the original remained) keeps interpreter links that resolve into the original directory. Such an environment must be rebuilt, or the copy would silently depend on files outside this prefix and break when the original is removed. """ try: return python.resolve().is_relative_to(root) except OSError: return False def ensure_ai(device: str = "auto", progress: Progress | None = None, verify: bool = False) -> Path: """Install/validate the backend and model; return its Python executable. CPU and CUDA environments are separate. Explicit CUDA never falls back; auto uses Torch's availability check in the selected worker. Downloads trust PyPI, download.pytorch.org, Astral's Python distribution, and upstream DF. verify forces a real self-test even when the environment is marked ready ('setup' uses this); ordinary processing trusts the marker so a ready environment does not pay for a second model load per run. """ flavor = _flavor(device) try: libc = os.confstr("CS_GNU_LIBC_VERSION") or "" except (ValueError, OSError): libc = "" if platform.system() != "Linux" or not libc.startswith("glibc ") or tuple( int(part) for part in libc.split()[1].split(".")[:2] ) < (2, 28): raise RuntimeError("The pinned DeepFilterNet3 backend requires Linux with glibc >= 2.28 " "(for example Ubuntu 22.04+ or Debian 12+). " "Alpine/musl is not supported by its prebuilt wheels.") root = data_dir() root.mkdir(parents=True, exist_ok=True) environment = root / f"ai-df-0.5.6-torch-2.5.1-{flavor}-py311-v1" python = environment / "bin/python" marker = environment / ".voiceforge-ready" started = time.monotonic() # Serialize bootstrap/model downloads, including calls from separate CLIs. with (root / "ai.lock").open("a") as lock: while True: try: fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) break except BlockingIOError: if progress: progress("Waiting for AI setup", time.monotonic() - started, None) time.sleep(0.25) ready = (marker.is_file() and python.is_file() and _interpreter_is_local(python, root)) if not ready: # A moved, copied, or half-installed environment cannot be # repaired in place: interpreter links reference absolute paths. # Rebuild it; the shared wheel caches make this fast, usually # offline. if environment.exists(): shutil.rmtree(environment) uv = ensure_uv(progress) env = _worker_env() run_process([uv, "venv", "--python", "3.11", "--managed-python", "--relocatable", "--allow-existing", str(environment)], "Preparing AI Python 3.11", progress, env=env) run_process([uv, "pip", "install", "--python", str(python), "--index-url", f"https://download.pytorch.org/whl/{flavor}", "torch==2.5.1", "torchaudio==2.5.1"], f"Installing PyTorch ({flavor})", progress, env=env) run_process([uv, "pip", "install", "--python", str(python), "--only-binary", ":all:", "--index-url", "https://pypi.org/simple", "deepfilternet==0.5.6", "numpy==1.26.4", "soundfile==0.12.1"], "Installing DeepFilterNet3", progress, env=env) # Check a freshly built environment (and any explicit setup) with a # real model load and inference; this also downloads the weights. if verify or not ready: run_process([str(python), "-I", str(Path(__file__).with_name("ai_worker.py")), "--device", device, "--model-dir", str(root / "models/df3-0.5.6"), "--check"], "Checking DeepFilterNet3", progress, env=_worker_env()) marker.touch() return python def denoise(source: Path, target: Path, device: str = "auto", strength: float = 1.0, progress: Progress | None = None) -> None: """Denoise mono 48 kHz WAV to float WAV, preserving the exact sample count. strength is a dry/wet mix in [0, 1]. The worker publishes output atomically; errors/cancellation leave an existing target untouched. Chunked inference is bounded-memory, not bit-identical to inference over an entire recording. """ if not math.isfinite(strength) or not 0 <= strength <= 1: raise ValueError("strength must be finite and between 0 and 1") source, target = Path(source).resolve(), Path(target).resolve() if source == target: raise ValueError("source and target must be different files") if not source.is_file(): raise FileNotFoundError(source) if not target.parent.is_dir(): raise FileNotFoundError(target.parent) python = ensure_ai(device, progress) run_process([str(python), "-I", str(Path(__file__).with_name("ai_worker.py")), "--device", device, "--source", str(source), "--target", str(target), "--model-dir", str(data_dir() / "models/df3-0.5.6"), "--strength", str(strength)], "Denoising", progress, env=_worker_env())