"""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()