1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
|
"""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()
|
