diff options
Diffstat (limited to 'lib/project/src/voiceforge/ai_worker.py')
| -rw-r--r-- | lib/project/src/voiceforge/ai_worker.py | 161 |
1 files changed, 0 insertions, 161 deletions
diff --git a/lib/project/src/voiceforge/ai_worker.py b/lib/project/src/voiceforge/ai_worker.py deleted file mode 100644 index 99ec9f5..0000000 --- a/lib/project/src/voiceforge/ai_worker.py +++ /dev/null @@ -1,161 +0,0 @@ -"""Private script executed only by the isolated backend Python.""" -from __future__ import annotations - -import argparse -import hashlib -import json -import os -from pathlib import Path -import signal -import tempfile -import urllib.request -import zipfile - - -def report(stage: str, completed: float | None = None, total: float | None = None) -> None: - print("VOICEFORGE_PROGRESS " + json.dumps([stage, completed, total]), flush=True) - - -def ensure_model(directory: Path) -> Path: - """Install a release-pinned, hash-checked archive atomically under setup lock.""" - model = directory / "DeepFilterNet3" - if (model / "config.ini").is_file() and (model / "checkpoints/model_120.ckpt.best").is_file(): - return model - if directory.exists(): - raise RuntimeError(f"Incomplete model directory: {directory}. Remove it and retry AI setup.") - directory.parent.mkdir(parents=True, exist_ok=True) - url = "https://raw.githubusercontent.com/Rikorose/DeepFilterNet/v0.5.6/models/DeepFilterNet3.zip" - expected = "49c52edc8947ae1f9bf50d81530beaf3a2c3245aeaf34b6f31ff535cd22284d2" - # Release archive size is fixed alongside its hash, even without Content-Length. - expected_bytes = 7986207 - downloaded = 0 - report("Downloading DeepFilterNet3 model (bytes)", 0, expected_bytes) - with tempfile.TemporaryDirectory(prefix=".df3-", dir=directory.parent) as staging: - archive = Path(staging) / "model.zip" - digest = hashlib.sha256() - with urllib.request.urlopen(url, timeout=60) as response, archive.open("wb") as output: - while chunk := response.read(256 * 1024): - downloaded += len(chunk) - if downloaded > expected_bytes: - raise RuntimeError("DeepFilterNet3 model download exceeds the expected size") - digest.update(chunk) - output.write(chunk) - report("Downloading DeepFilterNet3 model (bytes)", downloaded, expected_bytes) - if downloaded != expected_bytes or digest.hexdigest() != expected: - raise RuntimeError("DeepFilterNet3 model checksum mismatch; refusing to load weights") - report("Verifying and extracting DeepFilterNet3 model") - unpacked = Path(staging) / "unpacked" - with zipfile.ZipFile(archive) as bundle: - for member in bundle.infolist(): - destination = (unpacked / member.filename).resolve() - if not destination.is_relative_to(unpacked.resolve()): - raise RuntimeError("Unsafe path in model archive") - bundle.extractall(unpacked) - os.replace(unpacked, directory) - return model - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--device", choices=("auto", "cpu", "cuda"), required=True) - parser.add_argument("--check", action="store_true") - parser.add_argument("--model-dir", type=Path, required=True) - parser.add_argument("--source", type=Path) - parser.add_argument("--target", type=Path) - parser.add_argument("--strength", type=float, default=1.0) - args = parser.parse_args() - - # SIGTERM from the parent takes the same cleanup path as Ctrl-C. - def terminate(signum, frame): - raise KeyboardInterrupt - - signal.signal(signal.SIGTERM, terminate) - - import numpy as np - import soundfile as sf - import torch - - torch.set_num_threads(min(4, os.cpu_count() or 1)) - available = torch.cuda.is_available() - if args.device == "cuda" and not available: - raise RuntimeError("CUDA was explicitly requested, but PyTorch cannot use it. " - "Check the NVIDIA driver and CUDA_VISIBLE_DEVICES, or select cpu.") - device = "cuda" if args.device != "cpu" and available else "cpu" - # df.utils.get_device() consults DEVICE on every call, including model creation. - os.environ["DEVICE"] = device - from df.enhance import enhance, init_df - from df.model import ModelParams - from libdf import DF - - model_path = ensure_model(args.model_dir) - report(f"Loading DeepFilterNet3 ({device})") - model, state, _ = init_df(str(model_path), log_level="ERROR", log_file=None) - if state.sr() != 48000: - raise RuntimeError("DeepFilterNet3 model must use 48 kHz") - if args.check: - # Exercise both model execution and the compiled libdf/NumPy ABI. - sample = enhance(model, state, torch.zeros(1, 4800), pad=True) - if sample.shape != (1, 4800) or not torch.isfinite(sample).all(): - raise RuntimeError("DeepFilterNet3 self-test failed") - report(f"DeepFilterNet3 ready ({device})", 1, 1) - return - if args.source is None or args.target is None: - parser.error("--source and --target are required unless --check is used") - if not np.isfinite(args.strength) or not 0 <= args.strength <= 1: - parser.error("--strength must be in [0, 1]") - - temporary = None - try: - with sf.SoundFile(args.source) as source: - if source.samplerate != 48000 or source.channels != 1 or source.format not in {"WAV", "WAVEX", "RF64"}: - raise ValueError("AI input must be a mono 48 kHz WAV; convert it before denoising") - frames = len(source) - fd, name = tempfile.mkstemp(prefix=".voiceforge-ai-", suffix=".wav", dir=args.target.parent) - os.close(fd) - temporary = Path(name) - # RF64 avoids RIFF's 4 GiB limit on very long recordings. - output_format = "RF64" if frames * 4 > 0xFFFFFFFF - 4096 else "WAV" - with sf.SoundFile(temporary, "w", samplerate=48000, channels=1, - format=output_format, subtype="FLOAT") as target: - block, context, overlap = 480000, 96000, 2400 - previous = None - params = ModelParams() - report(f"Denoising ({device})", 0, frames) - for start in range(0, frames, block): - end = min(start + block, frames) - left, right = max(0, start - context), min(frames, end + overlap + context) - source.seek(left) - audio = source.read(right - left, dtype="float32") - if not np.isfinite(audio).all(): - raise ValueError("Input contains non-finite audio samples") - if args.strength: - # Fresh STFT state per window; context warms normalization and - # recurrent layers. Pad to a whole hop before delay compensation. - state = DF(sr=params.sr, fft_size=params.fft_size, - hop_size=params.hop_size, nb_bands=params.nb_erb, - min_nb_erb_freqs=params.min_nb_freqs) - padded = np.pad(audio, (0, (-len(audio)) % params.hop_size)) - wet = enhance(model, state, torch.from_numpy(padded).unsqueeze(0), - pad=True).squeeze(0).numpy()[:len(audio)] - if len(wet) != len(audio) or not np.isfinite(wet).all(): - raise RuntimeError("DeepFilterNet3 returned invalid audio") - audio = audio * (1 - args.strength) + wet * args.strength - kept = audio[start - left:min(frames, end + overlap) - left].copy() - if previous is not None and args.strength: - n = len(previous) - fade = np.linspace(0, 1, n, dtype=np.float32) - kept[:n] = previous * (1 - fade) + kept[:n] * fade - target.write(kept[:end - start]) - previous = kept[end - start:].copy() - report(f"Denoising ({device})", end, frames) - if target.tell() != frames: - raise RuntimeError("Output frame count does not match input") - os.replace(temporary, args.target) - temporary = None - finally: - if temporary is not None: - temporary.unlink(missing_ok=True) - - -if __name__ == "__main__": - main() |