| """Clearline-NC-GPU on a recording: background noise out, every voice kept. |
| |
| pip install -r requirements.txt |
| python example.py noisy.wav enhanced.wav [--live] [--device cuda] |
| |
| The model is Voice Isolator's released network (its code is in vi/) fine-tuned to keep every talker, and it is never |
| given an enrollment. It is causal (it looks 10 ms ahead) and remembers a whole call, so a recording is never cut into |
| pieces (60 s pieces with 10 s of context match it to only 20.6 dB): up to two minutes it runs in one pass, longer ones |
| through Voice Isolator's live engine in float16 (memory stays at a few GB; over 600 s it matches an exact float32 run |
| to 55.1 dB, three times faster than real time). --live writes what a live call would hear: the same engine, 20 ms at a |
| time, the output 10 ms behind the input. |
| """ |
| import argparse, copy, importlib, importlib.util, os, sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| HERE = Path(os.path.abspath(__file__)).parent |
| SR = 16000 |
| _PKG = "clearline_nc_gpu_vi" |
|
|
|
|
| def _vi(sub): |
| if _PKG not in sys.modules: |
| spec = importlib.util.spec_from_file_location(_PKG, HERE / "vi" / "__init__.py", |
| submodule_search_locations=[str(HERE / "vi")]) |
| mod = importlib.util.module_from_spec(spec) |
| sys.modules[_PKG] = mod |
| spec.loader.exec_module(mod) |
| return importlib.import_module(f"{_PKG}.{sub}") |
|
|
|
|
| def load(path=HERE, device="cuda"): |
| """The model from a copy of this repository (config.json and model.safetensors).""" |
| return _vi("infer").load_model(str(path), device=device) |
|
|
|
|
| @torch.no_grad() |
| def enhance(model, x, one_pass_s=120.0): |
| """x: float32 mono at 16 kHz -> the enhanced signal, the same length and aligned with x. One pass up to one_pass_s |
| (its causal WavLM's attention mask grows with the square of the length), else the live engine on a float16 copy.""" |
| dev = next(model.parameters()).device |
| x = np.ascontiguousarray(x, dtype=np.float32) |
| if len(x) <= one_pass_s * SR: |
| return model(torch.from_numpy(x)[None].to(dev))[0].float().cpu().numpy() |
| out = stream(copy.deepcopy(model), x, graphs=dev.type == "cuda") |
| k = out["delay"] |
| y = np.zeros_like(x) |
| y[: len(out["y"]) - k] = out["y"][k:] |
| return y |
|
|
|
|
| def stream(model, x, chunk_hops=2, dtype=torch.float16, graphs=True): |
| """x as one live call through Voice Isolator's engine (vi/stream.py), self-enrollment off: the model keeps its null |
| enrollment, as it was trained. Converts `model` to `dtype` in place (load a second copy for offline use). Returns |
| dict(y=the output of every whole 20 ms chunk, delay=the samples it lags the input).""" |
| dev = next(model.parameters()).device |
| s = _vi("stream").Stream(model, chunk_hops=chunk_hops, device=str(dev), dtype=dtype, graphs=graphs) |
| s.srv.self_enroll = False |
| return dict(y=s.push(np.asarray(x, np.float32)), delay=s.delay) |
|
|
|
|
| def main(): |
| import soundfile as sf |
| ap = argparse.ArgumentParser() |
| ap.add_argument("inp") |
| ap.add_argument("out") |
| ap.add_argument("--live", action="store_true") |
| ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| a = ap.parse_args() |
| x, sr = sf.read(a.inp, dtype="float32", always_2d=True) |
| x = x.mean(1) |
| if sr != SR: |
| import soxr |
| x = soxr.resample(x, sr, SR).astype(np.float32) |
| m = load(device=a.device) |
| y = stream(m, x, graphs=a.device != "cpu")["y"] if a.live else enhance(m, x) |
| sf.write(a.out, y, SR) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
