diff options
Diffstat (limited to 'lib/project/src/voiceforge/ai.py')
| -rw-r--r-- | lib/project/src/voiceforge/ai.py | 157 |
1 files changed, 0 insertions, 157 deletions
diff --git a/lib/project/src/voiceforge/ai.py b/lib/project/src/voiceforge/ai.py deleted file mode 100644 index 493be65..0000000 --- a/lib/project/src/voiceforge/ai.py +++ /dev/null @@ -1,157 +0,0 @@ -"""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()) |
