From 39b0f2bbed74f6487a41b82501ae3c6799e4b5c4 Mon Sep 17 00:00:00 2001 From: historia Date: Sun, 6 Sep 2026 20:55:27 -0400 Subject: feat: chunking, zipenhancer denoising --- lib/src/producer/cli.py | 95 +++++++- lib/src/producer/config.py | 34 ++- lib/src/producer/doctor.py | 23 +- lib/src/producer/dsp.py | 124 +++++----- lib/src/producer/engines/base.py | 13 +- lib/src/producer/engines/chunking.py | 166 +++++++++++++ lib/src/producer/engines/denoise_dfn.py | 51 ++-- lib/src/producer/engines/denoise_spectral.py | 186 ++++++++++++++ lib/src/producer/engines/denoise_zip.py | 56 ++++- lib/src/producer/engines/enhance_mossformer.py | 58 ++++- lib/src/producer/engines/enhance_resemble.py | 41 +++- lib/src/producer/engines/resemble_worker.py | 55 ++++- lib/src/producer/io.py | 2 +- lib/src/producer/lazy.py | 200 ++++++++++++++-- lib/src/producer/meters.py | 67 ++++-- lib/src/producer/pipeline.py | 132 ++++++++-- lib/src/producer/ui.py | 319 +++++++++++++++++++++++++ lib/src/producer/updates.py | 273 +++++++++++++++++++++ 18 files changed, 1719 insertions(+), 176 deletions(-) create mode 100644 lib/src/producer/engines/chunking.py create mode 100644 lib/src/producer/engines/denoise_spectral.py create mode 100644 lib/src/producer/ui.py create mode 100644 lib/src/producer/updates.py (limited to 'lib/src') diff --git a/lib/src/producer/cli.py b/lib/src/producer/cli.py index 476adf3..779de96 100644 --- a/lib/src/producer/cli.py +++ b/lib/src/producer/cli.py @@ -2,10 +2,11 @@ from __future__ import annotations import argparse import itertools +import os import sys from pathlib import Path -from . import __version__, pipeline +from . import __version__, pipeline, ui from . import config as cfgmod from . import dsp as pdsp from . import io as pio @@ -26,9 +27,19 @@ def build_parser() -> argparse.ArgumentParser: p.add_argument("-o", "--output", help="output file, or directory for batch") p.add_argument("--profile", choices=tuple(PROFILES), help="audiobook (default) or podcast") p.add_argument( - "--denoise", choices=["dfn3", "zipenhancer", "off"], help="denoise engine (default dfn3)" + "--denoise", + choices=["dfn3", "zipenhancer", "spectral", "off"], + help="denoise engine (default dfn3)", ) p.add_argument("--denoise-strength", type=float, help="0-1 blend of denoised signal") + p.add_argument( + "--denoise-pf", + action="store_true", + dest="denoise_pf", + default=None, + help="enable the DeepFilterNet post filter (extra noise reduction, " + "may eat soft speech on clean recordings)", + ) p.add_argument( "--enhance", choices=["off", "mossformer2", "resemble"], help="speech enhancement engine" ) @@ -52,10 +63,27 @@ def build_parser() -> argparse.ArgumentParser: "--format", choices=["wav", "flac", "mp3"], dest="out_format", help="output format" ) p.add_argument("--device", choices=["auto", "cuda", "cpu"], help="compute device") + p.add_argument( + "--engine-chunk", + type=float, + dest="engine_chunk_s", + metavar="SEC", + help="GPU engine chunk length in seconds (default 30, 0 = whole file)", + ) + p.add_argument( + "--engine-overlap", + type=float, + dest="engine_overlap_s", + metavar="SEC", + help="GPU engine chunk overlap in seconds (default 0.5)", + ) p.add_argument("--batch", action="store_true", help="inputs are directories/globs to expand") p.add_argument("--report", action="store_true", help="write .report.json") p.add_argument("--report-path", help="explicit report file path") p.add_argument("--dry-run", action="store_true", help="show the processing chain and exit") + p.add_argument( + "--no-update-check", action="store_true", help="skip the dependency update check" + ) p.add_argument("--config", help="config file (default config.toml in the project root)") p.add_argument("-v", "--verbose", action="count", default=0) p.add_argument("--version", action="store_true") @@ -67,6 +95,7 @@ def _apply_args(opts: Options, args: argparse.Namespace) -> None: "profile": "profile", "denoise": "denoise", "denoise_strength": "denoise_strength", + "denoise_pf": "denoise_pf", "enhance": "enhance", "enhance_strength": "enhance_strength", "output": "output", @@ -77,6 +106,8 @@ def _apply_args(opts: Options, args: argparse.Namespace) -> None: "bit_depth": "bit_depth", "out_format": "out_format", "device": "device", + "engine_chunk_s": "engine_chunk_s", + "engine_overlap_s": "engine_overlap_s", "report": "report", "report_path": "report_path", "dry_run": "dry_run", @@ -98,6 +129,10 @@ def _apply_args(opts: Options, args: argparse.Namespace) -> None: raise SystemExit(f"unknown profile: {opts.profile}") if not 0.0 <= opts.denoise_strength <= 1.0 or not 0.0 <= opts.enhance_strength <= 1.0: raise SystemExit("engine strengths must be within 0-1") + if opts.engine_chunk_s < 0.0 or opts.engine_overlap_s < 0.0: + raise SystemExit("--engine-chunk/--engine-overlap must be >= 0") + if opts.engine_chunk_s > 0.0 and opts.engine_overlap_s >= opts.engine_chunk_s: + raise SystemExit("--engine-overlap must be smaller than --engine-chunk") for k in DSP_KEYS: v = opts.strengths[k] if v is not None and not 0.0 <= float(v) <= 1.0: @@ -146,7 +181,7 @@ def _resolve_output_conflict(out_path: Path, interactive: bool) -> Path | None: return out_path alt = _next_free(out_path) if not interactive: - print(f"[producer] {out_path} exists, writing {alt.name} instead") + ui.log(f"[producer] {out_path} exists, writing {alt.name} instead") return alt while True: try: @@ -170,17 +205,25 @@ def _resolve_output_conflict(out_path: Path, interactive: bool) -> Path | None: print("[producer] please answer o, r, or c") -def process_one(inp: Path, opts: Options, single: bool, quiet: bool = False) -> Path | None: +def process_one( + inp: Path, + opts: Options, + single: bool, + quiet: bool = False, + pos: tuple[int, int] | None = None, +) -> Path | None: out_path = _resolve_output(inp, opts, single) out_path = _resolve_output_conflict(out_path, sys.stdin.isatty()) if out_path is None: - print(f"[producer] skipped {inp}") + ui.log(f"[producer] skipped {inp}") return None x, sr = pio.decode(inp) target_sr = opts.out_sample_rate() - if not quiet: - print(f"[producer] {inp} ({sr} Hz, {len(x) / sr:.1f}s)") - res = pipeline.run_pipeline(x, sr, opts) + pos_s = f"[{pos[0]}/{pos[1]}] " if pos else "" + ui.log(f"[producer] {pos_s}{inp} ({sr} Hz, {len(x) / sr:.1f}s)") + status = ui.Status(prefix=f"[producer] {pos_s}{inp.name} — ") + res = pipeline.run_pipeline(x, sr, opts, reporter=status) + status.finish() y = res.audio if target_sr != sr: y = pdsp.resample(y, sr, target_sr) @@ -195,7 +238,7 @@ def process_one(inp: Path, opts: Options, single: bool, quiet: bool = False) -> repmod.save(rep, rp) if not quiet: repmod.print_human(rep) - print(f"[producer] wrote {out_path}") + ui.log(f"[producer] {pos_s}wrote {out_path}") return out_path @@ -205,11 +248,26 @@ def main(argv: list[str] | None = None) -> int: from . import doctor return doctor.main(argv[1:]) + if argv and argv[0] == "update": + from . import updates + + return updates.run_update_command(argv[1:]) parser = build_parser() args = parser.parse_args(argv) if args.version: print(f"producer {__version__}") return 0 + if ( + not args.no_update_check + and os.environ.get("PRODUCER_NO_UPDATE_CHECK") != "1" + and not args.dry_run + ): + try: + from . import updates + + updates.check_and_prompt() + except Exception: + pass opts = Options() cfg_path = Path(args.config) if args.config else DATA_DIR.parent / "config.toml" if not args.config: @@ -219,6 +277,13 @@ def main(argv: list[str] | None = None) -> int: inputs = _expand_inputs(args) if not inputs: parser.error("no input files given") + if len(inputs) > 1 and opts.output: + o = Path(opts.output) + if o.suffix and not o.is_dir(): + parser.error( + "-o/--output must be a directory when processing multiple files " + "(one _processed. file is written per input)" + ) if opts.dry_run: stages = pipeline.build_stages(opts) print(f"producer dry run (profile={opts.profile}, format={opts.out_format})") @@ -231,10 +296,18 @@ def main(argv: list[str] | None = None) -> int: (DATA_DIR / "cache").mkdir(parents=True, exist_ok=True) failed = 0 single = len(inputs) == 1 - for inp in inputs: + total = len(inputs) + for idx, inp in enumerate(inputs, start=1): try: - process_one(inp, opts, single, quiet=len(inputs) > 1) + process_one( + inp, + opts, + single, + quiet=total > 1, + pos=(idx, total) if total > 1 else None, + ) except Exception as e: + ui.finish_live() failed += 1 print(f"[producer] ERROR {inp}: {e}", file=sys.stderr) if opts.verbose: diff --git a/lib/src/producer/config.py b/lib/src/producer/config.py index e88e80b..968df4a 100644 --- a/lib/src/producer/config.py +++ b/lib/src/producer/config.py @@ -27,6 +27,10 @@ class Profile: target: float = -20.0 ceiling_db: float = -3.0 sample_rate: int = 44100 + # RMS level the DSP chain is gain-staged to before compressors bite; the + # pre-gain lets the absolute EQ/comp thresholds stay meaningful regardless + # of how hot or quiet the (denoised) input arrives. + dsp_ref_db: float = -20.0 hpf_hz: float = 80.0 mud: tuple[float, float, float] = (300.0, -2.5, 1.0) warmth: tuple[float, float] = (150.0, 1.5) @@ -58,6 +62,7 @@ PODCAST = Profile( target=-16.0, ceiling_db=-1.5, sample_rate=48000, + dsp_ref_db=-18.0, mud=(300.0, -1.5, 1.0), warmth=(150.0, 1.0), comp1=(-18.0, 2.0, 15.0, 150.0), @@ -84,6 +89,7 @@ RADIO = Profile( target=-16.0, ceiling_db=-1.5, sample_rate=48000, + dsp_ref_db=-17.0, hpf_hz=70.0, mud=(280.0, -3.5, 1.3), warmth=(100.0, 3.0), @@ -112,7 +118,10 @@ PROFILES = {"audiobook": AUDIOBOOK, "podcast": PODCAST, "radio": RADIO} class Options: profile: str = "audiobook" denoise: str = "dfn3" - denoise_strength: float = 1.0 + # 0.9 flattens the residual gain wobble mask-based denoisers leave on + # speech; the dry blend trades a whisper of noise back for stability. + denoise_strength: float = 0.9 + denoise_pf: bool = False enhance: str = "off" enhance_strength: float = 1.0 dsp: bool = True @@ -120,11 +129,15 @@ class Options: target: float | None = None ceiling_db: float | None = None sample_rate: int | None = None - bit_depth: int = 16 + bit_depth: int = 32 out_format: str = "wav" device: str = "auto" strengths: dict[str, float | None] = field(default_factory=lambda: {k: None for k in DSP_KEYS}) hpf_hz: float | None = None + # 0 = feed engines the whole file (one warm model pass, no seam artifacts); + # on OOM the run automatically falls back to large chunks. + engine_chunk_s: float = 0.0 + engine_overlap_s: float = 0.5 output: str | None = None report: bool = False report_path: str | None = None @@ -157,6 +170,7 @@ class Options: "profile": self.profile, "denoise": self.denoise, "denoise_strength": self.denoise_strength, + "denoise_pf": self.denoise_pf, "enhance": self.enhance, "enhance_strength": self.enhance_strength, "dsp": self.dsp, @@ -169,6 +183,8 @@ class Options: "device": self.device, "strengths": {k: self.eff(k) for k in DSP_KEYS}, "hpf_hz": self.hpf_hz if self.hpf_hz is not None else self.prof().hpf_hz, + "engine_chunk_s": self.engine_chunk_s, + "engine_overlap_s": self.engine_overlap_s, } @@ -182,6 +198,8 @@ _CFG_FIELDS = { "target": "target", "ceiling": "ceiling_db", "sample_rate": "sample_rate", + "engine_chunk": "engine_chunk_s", + "engine_overlap": "engine_overlap_s", } @@ -208,6 +226,8 @@ def apply_config(opts: Options, cfg: dict[str, Any]) -> None: opts.denoise = eng if strength is not None: opts.denoise_strength = float(strength) + if isinstance(cfg["denoise"], dict) and "pf" in cfg["denoise"]: + opts.denoise_pf = bool(cfg["denoise"]["pf"]) if "enhance" in cfg: eng, strength = _engine_cfg(cfg["enhance"]) if eng: @@ -229,9 +249,17 @@ def write_default_config(path: Path) -> None: lines = [ 'profile = "audiobook"', "", + "# GPU engine processing: engines receive the whole file by default so", + "# models run one warm pass (no seam artifacts). engine_chunk > 0 forces", + "# chunked processing in seconds for very long files / low VRAM.", + "# engine_chunk = 0.0 # 0 = whole file (automatic fallback on OOM)", + "# engine_overlap = 0.5 # crossfade between chunks in seconds", + "", "[denoise]", 'engine = "dfn3"', - "strength = 1.0", + "# 0.9 flattens residual denoiser gain wobble; 1.0 = full suppression", + "strength = 0.9", + "# pf = false # DFN post filter: extra noise reduction, may eat soft speech", "", "[enhance]", 'engine = "off"', diff --git a/lib/src/producer/doctor.py b/lib/src/producer/doctor.py index 482a0aa..2d8997f 100644 --- a/lib/src/producer/doctor.py +++ b/lib/src/producer/doctor.py @@ -26,6 +26,15 @@ def _import(name: str): return None +def _chain_error(name: str) -> str: + """Import failure text for a module whose deeper import chain doesn't load.""" + try: + importlib.import_module(name) + except Exception as e: + return str(e) or type(e).__name__ + return "" + + def main(_argv: list[str] | None = None) -> int: print(f"producer {__version__} doctor") print(f" data dir: {DATA_DIR}") @@ -54,9 +63,17 @@ def main(_argv: list[str] | None = None) -> int: for mod in ("deepfilternet", "zipenhancer", "clearvoice"): m = _import(mod) or _import(mod.replace("-", "_")) - _check( - mod, m is not None, "installed" if m is not None else "lazy (installed on first use)" - ) + if m is None: + _check(mod, False, "lazy (installed on first use)") + continue + detail = "installed" + if mod == "zipenhancer": + # zipenhancer loads modelscope lazily on first use; probe the + # config module whose extras-only deps have broken it in the wild + broken = _chain_error("modelscope.utils.config") + if broken and "modelscope" not in broken: + detail += f"; modelscope chain broken: {broken} (will install on first use)" + _check(mod, True, detail) dfn3 = DATA_DIR / "models" / "DeepFilterNet3" / "config.ini" _check("dfn3 weights", dfn3.is_file(), str(dfn3)) diff --git a/lib/src/producer/dsp.py b/lib/src/producer/dsp.py index 60cdbc6..e65e595 100644 --- a/lib/src/producer/dsp.py +++ b/lib/src/producer/dsp.py @@ -5,7 +5,7 @@ from scipy import ndimage, signal from .meters import SILENCE_DB -_EPS = 1e-12 +_EPS = np.float32(1e-12) def _coef(time_ms: float, sr: int) -> float: @@ -13,7 +13,10 @@ def _coef(time_ms: float, sr: int) -> float: def _onepole(x: np.ndarray, a: float) -> np.ndarray: - return signal.lfilter([1.0 - a], [1.0, -a], x) + a = np.float32(a) + return signal.lfilter( + np.array([1.0 - a], dtype=np.float32), np.array([1.0, -a], dtype=np.float32), x + ) def rbj(kind: str, sr: int, freq: float, gain_db: float = 0.0, q: float = 0.7071) -> np.ndarray: @@ -61,13 +64,15 @@ def rbj(kind: str, sr: int, freq: float, gain_db: float = 0.0, q: float = 0.7071 return arr / arr[3] -def biquad(x: np.ndarray, sos: np.ndarray) -> np.ndarray: - y = signal.sosfilt(sos[None, :] if sos.ndim == 1 else sos, x.astype(np.float64)) - return y.astype(np.float32) +def _as_sos(sos: np.ndarray) -> np.ndarray: + sos = sos[None, :] if sos.ndim == 1 else sos + return sos.astype(np.float32, copy=False) -def _as_sos(sos: np.ndarray) -> np.ndarray: - return sos[None, :] if sos.ndim == 1 else sos +def biquad(x: np.ndarray, sos: np.ndarray) -> np.ndarray: + x32 = x.astype(np.float32, copy=False) + y = signal.sosfilt(_as_sos(sos), x32) + return y.astype(np.float32, copy=False) def hpf(x: np.ndarray, sr: int, hz: float) -> np.ndarray: @@ -95,16 +100,16 @@ def compressor( release_ms: float, knee_db: float = 6.0, ) -> np.ndarray: - x64 = x.astype(np.float64) + x32 = x.astype(np.float32, copy=False) a_att = _coef(attack_ms, sr) a_rel = _coef(release_ms, sr) - env = np.sqrt(np.clip(_onepole(np.square(x64), a_att), 0.0, None)) + env = np.sqrt(np.clip(_onepole(np.square(x32), a_att), 0.0, None)) level_db = 20.0 * np.log10(env + _EPS) - over = level_db - threshold_db - k = knee_db + over = level_db - np.float32(threshold_db) + k = np.float32(knee_db) gr = np.where( over <= -k / 2, - 0.0, + np.float32(0.0), np.where( over < k / 2, (1.0 - 1.0 / ratio) * np.square(over + k / 2) / (2.0 * k), @@ -113,7 +118,7 @@ def compressor( ) gr = _onepole(np.clip(gr, 0.0, None), a_rel) gain = 10.0 ** (-gr / 20.0) - return (x64 * gain).astype(np.float32) + return (x32 * gain).astype(np.float32, copy=False) def deesser( @@ -121,24 +126,28 @@ def deesser( ) -> np.ndarray: if max_reduction_db < 0.05: return x - x64 = x.astype(np.float64) - low = signal.sosfiltfilt(signal.butter(4, lo_hz, btype="lowpass", fs=sr, output="sos"), x64) - high = x64 - low + x32 = x.astype(np.float32, copy=False) + low = signal.sosfiltfilt( + signal.butter(4, lo_hz, btype="lowpass", fs=sr, output="sos").astype(np.float32), x32 + ) + high = x32 - low mid = signal.sosfiltfilt( - signal.butter(4, [lo_hz, hi_hz], btype="bandpass", fs=sr, output="sos"), high + signal.butter(4, [lo_hz, hi_hz], btype="bandpass", fs=sr, output="sos").astype(np.float32), + high, ) win = max(3, int(0.005 * sr) | 1) - env = signal.convolve(np.abs(mid), np.hanning(win) / np.sum(np.hanning(win)), mode="same") + kernel = (np.hanning(win) / np.sum(np.hanning(win))).astype(np.float32) + env = signal.convolve(np.abs(mid), kernel, mode="same") act = env[env > _EPS] if act.size == 0: return x - thr = float(np.percentile(act, 95)) * 10.0 ** (-3.0 / 20.0) + thr = np.float32(float(np.percentile(act, 95)) * 10.0 ** (-3.0 / 20.0)) over = np.clip(20.0 * np.log10((env + _EPS) / thr), 0.0, None) - gr = np.clip(over * 0.6, 0.0, max_reduction_db) + gr = np.clip(over * np.float32(0.6), 0.0, np.float32(max_reduction_db)) gr = _onepole(_onepole(gr, _coef(1.0, sr)), _coef(30.0, sr)) gain = 10.0 ** (-gr / 20.0) high_out = high - mid + mid * gain - return (low + high_out).astype(np.float32) + return (low + high_out).astype(np.float32, copy=False) def expander( @@ -156,9 +165,9 @@ def expander( nf = n // frame if nf < 2: return x - frames = x[: nf * frame].reshape(nf, frame).astype(np.float64) - frms = np.sqrt(np.mean(np.square(frames), axis=1)) - fdb = 20.0 * np.log10(frms + _EPS) + frames = x[: nf * frame].reshape(nf, frame).astype(np.float32, copy=False) + frms = np.sqrt(np.mean(np.square(frames), axis=1, dtype=np.float64)) + fdb = 20.0 * np.log10(frms + 1e-12) active = fdb[fdb > SILENCE_DB + 1.0] if active.size == 0: return x @@ -167,12 +176,12 @@ def expander( thr = 0.5 * (speech_ref + floor_ref) drop = np.where(fdb < thr, np.minimum((thr - fdb) * (1.0 - 1.0 / ratio), max_drop_db), 0.0) centers = (np.arange(nf) + 0.5) * frame - drop_s = np.interp(np.arange(n), centers, drop) + drop_s = np.interp(np.arange(n, dtype=np.float32), centers, drop).astype(np.float32) drop_s = np.maximum( _onepole(drop_s, _coef(release_ms, sr)), _onepole(drop_s, _coef(attack_ms, sr)) ) - gain = 10.0 ** (-np.clip(drop_s, 0.0, max_drop_db) / 20.0) - return (x.astype(np.float64) * gain).astype(np.float32) + gain = 10.0 ** (-np.clip(drop_s, 0.0, np.float32(max_drop_db)) / 20.0) + return (x.astype(np.float32, copy=False) * gain).astype(np.float32, copy=False) def limit( @@ -182,17 +191,17 @@ def limit( release_ms: float = 60.0, lookahead_ms: float = 1.0, ) -> np.ndarray: - x64 = x.astype(np.float64) - n = x64.size + x32 = x.astype(np.float32, copy=False) + n = x32.size if n == 0: return x lin = 10.0 ** (ceiling_db / 20.0) win = max(1, int(lookahead_ms * sr / 1000.0) | 1) - env = ndimage.maximum_filter1d(np.abs(x64), size=win, mode="nearest") + env = ndimage.maximum_filter1d(np.abs(x32), size=win, mode="nearest") over = np.clip(20.0 * np.log10(env + _EPS) - ceiling_db, 0.0, None) block = max(1, int(0.005 * sr)) nb = (n + block - 1) // block - padded = np.zeros(nb * block, dtype=np.float64) + padded = np.zeros(nb * block, dtype=np.float32) padded[:n] = over block_over = padded.reshape(nb, block).max(axis=1) decay = float(np.exp(-1000.0 * 0.005 / max(release_ms, 0.1))) @@ -201,14 +210,14 @@ def limit( for i in range(nb): prev = max(block_over[i], prev * decay) held[i] = prev - held_s = np.interp(np.arange(n), (np.arange(nb) + 0.5) * block, held) - held_s = _onepole(held_s, _coef(1.0, sr)) + held_s = np.interp(np.arange(n, dtype=np.float32), (np.arange(nb) + 0.5) * block, held) + held_s = _onepole(held_s.astype(np.float32), _coef(1.0, sr)) gain = 10.0 ** (-held_s / 20.0) - y = x64 * gain + y = x32 * gain bad = np.abs(y) > lin if np.any(bad): - y[bad] = np.sign(y[bad]) * lin - return y.astype(np.float32) + y[bad] = np.float32(lin) * np.sign(y[bad]) + return y.astype(np.float32, copy=False) def resample(x: np.ndarray, sr_in: int, sr_out: int) -> np.ndarray: @@ -217,25 +226,23 @@ def resample(x: np.ndarray, sr_in: int, sr_out: int) -> np.ndarray: from math import gcd g = gcd(int(sr_in), int(sr_out)) - y = signal.resample_poly( - x.astype(np.float64), int(sr_out) // g, int(sr_in) // g, window=("kaiser", 10.0) - ) - return y.astype(np.float32) + y = signal.resample_poly(x, int(sr_out) // g, int(sr_in) // g, window=("kaiser", 10.0)) + return y.astype(np.float32, copy=False) def tape(x: np.ndarray, sr: int, amount: float) -> np.ndarray: """Gentle asymmetric soft-clip saturation — analog-style even-harmonic warmth.""" if amount < 0.02: return x - s = float(np.clip(amount, 0.0, 1.0)) - drive = 1.0 + 2.5 * s - x64 = x.astype(np.float64) - curve = np.where(x64 < 0.0, x64 * (1.0 + 0.4 * s), x64) + s = np.float32(np.clip(amount, 0.0, 1.0)) + drive = np.float32(1.0 + 2.5 * float(s)) + x32 = x.astype(np.float32, copy=False) + curve = np.where(x32 < 0.0, x32 * (1.0 + 0.4 * s), x32) y = np.tanh(curve * drive) / np.tanh(drive) - rms_in = np.sqrt(np.mean(np.square(x64))) - rms_out = np.sqrt(np.mean(np.square(y))) - y *= rms_in / max(rms_out, _EPS) - return (x64 * (1.0 - s) + y * s).astype(np.float32) + rms_in = float(np.sqrt(np.mean(np.square(x32)))) + rms_out = float(np.sqrt(np.mean(np.square(y)))) + y *= np.float32(rms_in / max(rms_out, float(_EPS))) + return (x32 * (1.0 - s) + y * s).astype(np.float32, copy=False) def soothe(x: np.ndarray, sr: int, amount: float) -> np.ndarray: @@ -247,27 +254,30 @@ def soothe(x: np.ndarray, sr: int, amount: float) -> np.ndarray: """ if amount < 0.02: return x - s = float(np.clip(amount, 0.0, 1.0)) - y = x.astype(np.float64) + s = np.float32(np.clip(amount, 0.0, 1.0)) + y = x.astype(np.float32, copy=False) bands = ((200.0, 450.0, 3.5), (2500.0, 6000.0, 5.0)) win = max(3, int(0.010 * sr) | 1) - kernel = np.hanning(win) / np.sum(np.hanning(win)) + kernel = (np.hanning(win) / np.sum(np.hanning(win))).astype(np.float32) + changed = False for lo, hi, max_db in bands: - if max_db * s < 0.1: + if max_db * float(s) < 0.1: continue band = signal.sosfiltfilt( - signal.butter(4, [lo, hi], btype="bandpass", fs=sr, output="sos"), y + signal.butter(4, [lo, hi], btype="bandpass", fs=sr, output="sos").astype(np.float32), + y, ) env = signal.convolve(np.abs(band), kernel, mode="same") act = env[env > _EPS] if act.size == 0: continue - thr = float(np.percentile(act, 88)) * 10.0 ** (-6.0 / 20.0) + thr = np.float32(float(np.percentile(act, 88)) * 10.0 ** (-6.0 / 20.0)) over = np.clip(20.0 * np.log10((env + _EPS) / thr), 0.0, None) - gr = np.clip(over * 0.7, 0.0, max_db) * s + gr = np.clip(over * np.float32(0.7), 0.0, np.float32(max_db)) * s gr = _onepole(_onepole(gr, _coef(5.0, sr)), _coef(60.0, sr)) gain = 10.0 ** (-gr / 20.0) y = y - band + band * gain - if np.array_equal(y, x.astype(np.float64)): + changed = True + if not changed or np.array_equal(y, x.astype(np.float32, copy=False)): return x - return y.astype(np.float32) + return y.astype(np.float32, copy=False) diff --git a/lib/src/producer/engines/base.py b/lib/src/producer/engines/base.py index 2988357..9192a17 100644 --- a/lib/src/producer/engines/base.py +++ b/lib/src/producer/engines/base.py @@ -28,7 +28,16 @@ def device_name(device: str) -> str: def blend(x: np.ndarray, y: np.ndarray, strength: float) -> np.ndarray: + if y.size != x.size: + # a resample round trip can drift by a sample or two (resample_poly + # emits ceil(n * up/down) per hop); realign before the elementwise math + y = y[: x.size] if y.size > x.size else np.pad(y, (0, x.size - y.size)) s = float(np.clip(strength, 0.0, 1.0)) if s >= 0.999: - return y - return (x.astype(np.float64) * (1.0 - s) + y.astype(np.float64) * s).astype(np.float32) + return y.astype(np.float32, copy=False) + if s <= 0.001: + return x.astype(np.float32, copy=False) + out = np.asarray(y, dtype=np.float32) - np.asarray(x, dtype=np.float32) + out *= s + out += x + return out diff --git a/lib/src/producer/engines/chunking.py b/lib/src/producer/engines/chunking.py new file mode 100644 index 0000000..627e3b8 --- /dev/null +++ b/lib/src/producer/engines/chunking.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +import gc +from collections.abc import Callable, Iterator + +import numpy as np + +from .. import ui + +MIN_CHUNK_S = 2.0 + +_OOM_ERRORS: tuple[type[BaseException], ...] | None = None + + +def _oom_errors() -> tuple[type[BaseException], ...]: + global _OOM_ERRORS + if _OOM_ERRORS is None: + errs: list[type[BaseException]] = [MemoryError] + try: + import torch + + errs.append(torch.cuda.OutOfMemoryError) + except Exception: + pass + _OOM_ERRORS = tuple(errs) + return _OOM_ERRORS + + +def free_vram() -> None: + gc.collect() + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: + pass + + +def plan_chunks(n: int, chunk: int, overlap: int) -> list[tuple[int, int]]: + """Split n samples into spans of at most `chunk` samples sharing `overlap`. + + Consecutive spans advance by chunk - overlap; the last span always covers + the tail (its length is at least overlap + 1). + """ + n = int(n) + if n <= 0: + return [] + chunk = int(max(1, chunk)) + overlap = int(max(0, overlap)) + if overlap >= chunk: + raise ValueError("overlap must be smaller than chunk") + if chunk >= n: + return [(0, n)] + spans: list[tuple[int, int]] = [] + start = 0 + while True: + end = min(start + chunk, n) + spans.append((start, end)) + if end >= n: + return spans + start += chunk - overlap + + +def stitch(spans: list[tuple[int, int]], n: int, pieces) -> np.ndarray: + """Overlap-add engine outputs (aligned with spans) into one n-sample array. + + The first `overlap` samples of each piece crossfade linearly against the + tail already written by the previous piece. + """ + out = np.zeros(n, dtype=np.float32) + prev_end = -1 + for (start, end), piece in zip(spans, pieces, strict=True): + seg = np.asarray(piece, dtype=np.float32).reshape(-1) + want = end - start + if seg.size < want: + seg = np.pad(seg, (0, want - seg.size), mode="edge") + seg = seg[:want] + if start < 0 or end > n: + raise ValueError(f"span ({start}, {end}) outside output of length {n}") + ov = max(0, prev_end - start) + if ov > 0: + ramp = (np.arange(ov, dtype=np.float32) + 0.5) / float(ov) + out[start : start + ov] *= 1.0 - ramp + out[start : start + ov] += seg[:ov] * ramp + out[start + ov : end] = seg[ov:] + else: + out[start:end] = seg + prev_end = end + return out + + +def _tracked( + x: np.ndarray, + spans: list[tuple[int, int]], + ctx: int, + fn, + on_progress: Callable[[int, int], None] | None, +) -> Iterator[np.ndarray]: + total = len(spans) + if on_progress is not None: + on_progress(0, total) + for i, (a, b) in enumerate(spans): + lo = max(0, a - ctx) + hi = min(x.size, b + ctx) + piece = fn(x[lo:hi]) + if on_progress is not None: + on_progress(i + 1, total) + seg = np.asarray(piece, dtype=np.float32).reshape(-1) + want_full = hi - lo + if seg.size < want_full: + seg = np.pad(seg, (0, want_full - seg.size), mode="edge") + seg = seg[:want_full] + yield seg[a - lo : a - lo + (b - a)] + + +def apply_chunked( + x: np.ndarray, + sr: int, + chunk_s: float, + overlap_s: float, + fn, + min_chunk_s: float = MIN_CHUNK_S, + on_progress: Callable[[int, int], None] | None = None, + context_s: float = 0.0, +) -> np.ndarray: + """Run fn on the signal in overlapping chunks and stitch the results. + + fn receives a 1-D chunk and must return a sample-aligned array of the same + length. chunk_s <= 0 disables chunking (single whole-signal call). With + context_s > 0 each chunk is widened by up to that many seconds of + neighbouring audio on both sides before fn runs, and the extra context in + fn's output is trimmed away again — recurrent/stateful models then process + the region that survives stitching with warmed-up state instead of a cold + start. On CUDA or host OOM the whole run is retried with halved chunk + length down to min_chunk_s. on_progress(done, total_chunks) fires as chunks + complete. + """ + n = int(x.size) + chunk_s = float(chunk_s or 0.0) + overlap_s = max(0.0, float(overlap_s or 0.0)) + ctx = max(0, round(float(context_s or 0.0) * sr)) + if chunk_s <= 0.0: + try: + return fn(x) + except _oom_errors(): + # Whole-file passes can OOM on multi-hour files; fall back to + # large chunks and let the halving loop take it from there. + free_vram() + whole = 120.0 + ui.log( + f"[producer] whole-file pass ran out of memory; retrying with {whole:.0f}s chunks" + ) + return apply_chunked(x, sr, whole, overlap_s, fn, min_chunk_s, on_progress, context_s) + try: + spans = plan_chunks(n, round(chunk_s * sr), round(overlap_s * sr)) + return stitch(spans, n, _tracked(x, spans, ctx, fn, on_progress)) + except _oom_errors(): + free_vram() + smaller = chunk_s / 2.0 + if smaller < min_chunk_s: + raise + ui.log( + f"[producer] engine ran out of memory; retrying with {smaller:.0f}s chunks", + ) + return apply_chunked(x, sr, smaller, overlap_s, fn, min_chunk_s, on_progress, context_s) diff --git a/lib/src/producer/engines/denoise_dfn.py b/lib/src/producer/engines/denoise_dfn.py index 6b7d1d5..186ffd2 100644 --- a/lib/src/producer/engines/denoise_dfn.py +++ b/lib/src/producer/engines/denoise_dfn.py @@ -2,18 +2,22 @@ from __future__ import annotations import sys import types -import urllib.request import warnings import zipfile +from collections.abc import Callable from pathlib import Path import numpy as np -from .. import dsp, lazy +from .. import dsp, lazy, ui from .base import blend, device_name, pick_device +from .chunking import apply_chunked, free_vram MODELS_DIR = lazy.DATA_DIR / "models" TAG = "v0.5.6" +# Neighbouring audio fed to each chunk so the recurrent model and its feature +# normalizers run warm at chunk seams; trimmed away before stitching. +CONTEXT_S = 2.0 MODEL_ZIPS = { "DeepFilterNet3": "models/DeepFilterNet3.zip", "DeepFilterNet2": "models/DeepFilterNet2.zip", @@ -57,8 +61,8 @@ def ensure_model(model: str) -> Path: url = f"{BASE_URL}/{MODEL_ZIPS[model]}" MODELS_DIR.mkdir(parents=True, exist_ok=True) zpath = MODELS_DIR / f"{model}.zip" - print(f"[producer] downloading {model} weights...", flush=True) - urllib.request.urlretrieve(url, zpath) + ui.download(url, zpath, f"downloading {model} weights") + ui.log(f"[producer] extracting {model} weights...") with zipfile.ZipFile(zpath) as z: z.extractall(target) zpath.unlink(missing_ok=True) @@ -74,10 +78,25 @@ def ensure_model(model: str) -> Path: def denoise( - x: np.ndarray, sr: int, strength: float, device_pref: str = "auto" + x: np.ndarray, + sr: int, + strength: float, + device_pref: str = "auto", + chunk_s: float = 30.0, + overlap_s: float = 0.5, + on_progress: Callable[[int, int], None] | None = None, + post_filter: bool = False, ) -> tuple[np.ndarray, str, str]: - lazy.ensure(["deepfilternet==0.5.6"], purpose="DeepFilterNet") + """Denoise with DeepFilterNet. + + post_filter opts into DFN's extra noise-reduction post filter; it + over-attenuates and can eat soft speech on clean recordings, so it stays + off unless requested. + """ + # torch first: deepfilternet's declared torch dependency would otherwise + # resolve to the newest (multi-GB CUDA) build before we pin our tested one. lazy.ensure_torch() + lazy.ensure(["deepfilternet==0.5.6"], purpose="DeepFilterNet") _shim_torchaudio_backend() _shim_df_git() model_name = "DeepFilterNet3" @@ -92,7 +111,7 @@ def denoise( device = pick_device(device_pref) model, df_state, _ = init_df( model_base_dir=str(model_dir), - post_filter=strength >= 0.95, + post_filter=post_filter, log_level="error", log_file=None, ) @@ -105,13 +124,17 @@ def denoise( xin = dsp.resample(x, sr, sr_df) import torch - xin_t = torch.from_numpy(xin.astype(np.float32)).unsqueeze(0) - y = df_enhance(model, df_state, xin_t) - if isinstance(y, torch.Tensor): - y = y.detach().cpu().numpy() - y = np.asarray(y, dtype=np.float32) - if y.ndim > 1: - y = y.reshape(-1) + def run(chunk: np.ndarray) -> np.ndarray: + t = torch.from_numpy(np.ascontiguousarray(chunk, dtype=np.float32)).unsqueeze(0) + y = df_enhance(model, df_state, t) + if isinstance(y, torch.Tensor): + y = y.detach().cpu().numpy() + return np.asarray(y, dtype=np.float32).reshape(-1) + + y = apply_chunked( + xin, sr_df, chunk_s, overlap_s, run, on_progress=on_progress, context_s=CONTEXT_S + ) + free_vram() y = dsp.resample(y, sr_df, sr) y = blend(x, y, strength) return y, f"dfn ({model_name})", f"{device_name(dev)} ({dev})" diff --git a/lib/src/producer/engines/denoise_spectral.py b/lib/src/producer/engines/denoise_spectral.py new file mode 100644 index 0000000..e1bb256 --- /dev/null +++ b/lib/src/producer/engines/denoise_spectral.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +from collections.abc import Callable + +import numpy as np +from scipy import signal + +from .. import ui +from .base import blend +from .chunking import apply_chunked + +N_FFT = 2048 +HOP = 512 +DD_ALPHA = 0.97 +PROFILE_PCT = 35.0 +GATE_PCT = 30.0 +GAIN_SMOOTH_MS = 20.0 +PROFILE_BLOCK_S = 60.0 +RESERVOIR_ROWS = 4096 +MIN_CHUNK_S = 8.0 + + +def _frame_rms(x: np.ndarray, frame: int) -> np.ndarray: + nf = x.size // frame + if nf < 1: + return np.ones(1, dtype=np.float64) + return np.sqrt( + np.mean(np.square(x[: nf * frame].reshape(nf, frame).astype(np.float64)), axis=1) + ) + + +def _noise_profile(x: np.ndarray, sr: int) -> np.ndarray: + """File-global per-bin noise PSD estimate. + + Two cheap passes: a time-domain pass finds the quietest fraction of + frames (speech gaps), then an STFT pass subsamples those frames' PSDs + into a fixed-size per-bin reservoir. Because the profile comes from the + whole file instead of a per-chunk percentile, speech-heavy regions can + no longer leak into the estimate — the failure that made the old engine + under-suppress tape hiss between sentences. + """ + hop_t = _frame_rms(x, HOP) + thr = float(np.percentile(hop_t, GATE_PCT)) + nb_bins = N_FFT // 2 + 1 + reservoir = np.zeros((nb_bins, RESERVOIR_ROWS), dtype=np.float32) + count = 0 + rng = np.random.default_rng(0) + + def add_rows(rows: np.ndarray) -> None: + nonlocal count + k = rows.shape[0] + if k == 0: + return + if count < RESERVOIR_ROWS: + take = min(k, RESERVOIR_ROWS - count) + reservoir[:, count : count + take] = rows[:take].T + count += take + rows = rows[take:] + k = rows.shape[0] + if k > 0: + idx = rng.integers(0, RESERVOIR_ROWS, size=k) + reservoir[:, idx] = rows.T + + block = int(PROFILE_BLOCK_S * sr) + collected = 0 + for start in range(0, x.size, block): + seg = x[start : min(start + block, x.size)] + if seg.size < N_FFT * 2: + break + _f, _t, S = signal.stft( + seg, window="hann", nperseg=N_FFT, noverlap=N_FFT - HOP, boundary="zeros", padded=True + ) + psd = (S.real * S.real + S.imag * S.imag).astype(np.float32) + # STFT frames sit ~head/HOP before the block start due to zero padding + offset = start // HOP - (N_FFT // 2) // HOP + quiet = np.zeros(psd.shape[1], dtype=bool) + vals = hop_t[max(0, offset) : max(0, offset) + psd.shape[1]] + quiet[: vals.size] = vals < thr + add_rows(psd[:, quiet].T) + collected += int(quiet.sum()) + if collected == 0: + # no quiet frames at all (continuous dense speech): blind subsample + for start in range(0, x.size, block): + seg = x[start : min(start + block, x.size)] + if seg.size < N_FFT * 2: + break + _f, _t, S = signal.stft( + seg, + window="hann", + nperseg=N_FFT, + noverlap=N_FFT - HOP, + boundary="zeros", + padded=True, + ) + psd = (S.real * S.real + S.imag * S.imag).astype(np.float32) + step = max(1, psd.shape[1] // (RESERVOIR_ROWS // 8)) + add_rows(psd[:, ::step].T) + if count == 0: + return np.zeros(nb_bins, dtype=np.float32) + return np.percentile(reservoir[:, :count], PROFILE_PCT, axis=1).astype(np.float32) + + +def _denoise_chunk(chunk: np.ndarray, sr: int, strength: float, noise: np.ndarray) -> np.ndarray: + n = chunk.size + if n < N_FFT * 2: + return chunk + head = N_FFT // 2 + tail = head + (-(n - N_FFT)) % HOP + padded = np.pad(chunk, (head, tail), mode="edge") + _f, _t, S = signal.stft( + padded, window="hann", nperseg=N_FFT, noverlap=N_FFT - HOP, boundary="zeros", padded=True + ) + psd = (S.real * S.real + S.imag * S.imag).astype(np.float32) + if not np.any(noise > 0.0): + return chunk[:n].astype(np.float32) + eps = np.float32(1e-20) + over = 1.0 + 3.0 * float(strength) + snr_post = psd / (noise[:, None] * np.float32(over) + eps) + floor_gain = np.float32(10.0 ** (-(10.0 + 28.0 * float(strength)) / 20.0)) + gains = np.empty_like(psd) + prev_g2 = np.ones(psd.shape[0], dtype=np.float32) + prev_post = snr_post[:, 0] + for i in range(psd.shape[1]): + post = snr_post[:, i] + snr_prio = ( + DD_ALPHA * prev_g2 * prev_post + + (1.0 - DD_ALPHA) * np.maximum(post - 1.0, np.float32(0.0)) + ).astype(np.float32) + g = snr_prio / (1.0 + snr_prio) + np.maximum(g, floor_gain, out=g) + gains[:, i] = g + prev_g2 = g * g + prev_post = post + a = np.float32(np.exp(-1000.0 * HOP / sr / GAIN_SMOOTH_MS)) + gains = signal.lfilter([1.0 - a], [1.0, -a], gains, axis=1) + _t2, rec = signal.istft( + S * gains, window="hann", nperseg=N_FFT, noverlap=N_FFT - HOP, boundary=True + ) + out = rec[head : head + n] + if out.size < n: + out = np.pad(out, (0, n - out.size), mode="edge") + return out.astype(np.float32) + + +def denoise( + x: np.ndarray, + sr: int, + strength: float, + device_pref: str = "auto", + chunk_s: float = 30.0, + overlap_s: float = 0.5, + on_progress: Callable[[int, int], None] | None = None, +) -> tuple[np.ndarray, str, str]: + """Denoise with decision-directed spectral subtraction on a global noise + profile. + + Pure DSP: one pass measures the per-bin noise PSD across the whole file + (from its quietest frames), then a Wiener-style gain with decision- + directed smoothing is applied in overlapping blocks. The global profile + keeps suppression uniform everywhere — steady tape hiss and breath noise + between sentences go down by ~20+ dB instead of the old capped-at-24 dB + per-chunk estimate. Deterministic, no model download, + no time-varying gain wobble. + """ + strength = float(np.clip(strength, 0.0, 1.0)) + if strength <= 0.001: + return x.copy(), "spectral", "cpu" + chunk_s = float(chunk_s or 0.0) + if chunk_s <= 0.0: + # a whole-file STFT would need ~10 GB per hour of audio; block instead + ui.log("[producer] spectral denoiser processes in blocks; clamping --engine-chunk") + chunk_s = 30.0 + chunk_s = max(chunk_s, MIN_CHUNK_S) + overlap_s = max(0.0, min(float(overlap_s or 0.0), chunk_s / 2.0)) + noise = _noise_profile(x, sr) + y = apply_chunked( + x, + sr, + chunk_s, + overlap_s, + lambda chunk: _denoise_chunk(chunk, sr, strength, noise), + on_progress=on_progress, + context_s=2.0, + ) + y = blend(x, y, strength) + return y, "spectral (dd-wiener, global profile)", "cpu" diff --git a/lib/src/producer/engines/denoise_zip.py b/lib/src/producer/engines/denoise_zip.py index 1794a59..3f37f61 100644 --- a/lib/src/producer/engines/denoise_zip.py +++ b/lib/src/producer/engines/denoise_zip.py @@ -1,22 +1,68 @@ from __future__ import annotations +from collections.abc import Callable + import numpy as np from .. import dsp, lazy from .base import blend +from .chunking import apply_chunked + + +def _peak_normalize(y: np.ndarray) -> np.ndarray: + peak = float(np.max(np.abs(y))) if y.size else 0.0 + if peak > 1e-10: + y *= 10.0 ** (-3.0 / 20.0) / peak + if y.size: + peak = float(np.max(np.abs(y))) + if peak > 0.99: + y *= 0.95 / peak + return y def denoise( - x: np.ndarray, sr: int, strength: float, device_pref: str = "auto" + x: np.ndarray, + sr: int, + strength: float, + device_pref: str = "auto", + chunk_s: float = 30.0, + overlap_s: float = 0.5, + on_progress: Callable[[int, int], None] | None = None, ) -> tuple[np.ndarray, str, str]: - lazy.ensure(["zipenhancer==0.3.2"], purpose="ZipEnhancer") + # torch first: zipenhancer's declared torch>=2.0 dependency would otherwise + # resolve to the newest (multi-GB CUDA) build before we pin our tested one. + lazy.ensure_torch() + # zipenhancer's weight loader imports modelscope even for the default + # model, but doesn't declare it (it is only an optional extra on PyPI). + lazy.ensure(["zipenhancer==0.3.2", "modelscope"], purpose="ZipEnhancer") + # the loader runs lazily on the first denoise() call, and modelscope's + # config/hub code imports extras-only helpers (addict, simplejson, ...) + # on the way; the import and call probes install whatever trips over. + lazy.ensure_import("zipenhancer", purpose="ZipEnhancer") + # the library API takes the full modelscope repo id; its short-name + # mapping lives only in its CLI, and anything unrecognized is treated + # as a repo id (modelscope E3021) + from zipenhancer import MODEL_ZIPENHANCER as Z_MODEL_REPO from zipenhancer import denoise as z_denoise sr_z = 16000 + strength = float(np.clip(strength, 0.0, 1.0)) x16 = dsp.resample(x, sr, sr_z) - result = z_denoise(x16, sr_z, model="zipenhancer", strength=float(np.clip(strength, 0.0, 1.0))) - y = result[0] if isinstance(result, tuple) else result - y = np.asarray(y, dtype=np.float32) + + def run(chunk: np.ndarray) -> np.ndarray: + result = z_denoise(chunk, sr_z, model=Z_MODEL_REPO, normalize=False, strength=strength) + y = result[0] if isinstance(result, tuple) else result + return np.asarray(y, dtype=np.float32).reshape(-1) + + y = apply_chunked( + x16, + sr_z, + chunk_s, + overlap_s, + lambda chunk: lazy.ensure_call(lambda: run(chunk), purpose="ZipEnhancer"), + on_progress=on_progress, + ) + _peak_normalize(y) y = dsp.resample(y, sr_z, sr) y = blend(x, y, strength) return y, "zipenhancer (16 kHz SOTA, bandwidth restored)", device_pref diff --git a/lib/src/producer/engines/enhance_mossformer.py b/lib/src/producer/engines/enhance_mossformer.py index f881045..3298c09 100644 --- a/lib/src/producer/engines/enhance_mossformer.py +++ b/lib/src/producer/engines/enhance_mossformer.py @@ -1,33 +1,65 @@ from __future__ import annotations +from collections.abc import Callable + import numpy as np from .. import dsp, io, lazy from .base import blend, device_name, pick_device +from .chunking import apply_chunked, free_vram + +# ClearVoice's batch decode path crashes on inputs longer than +# one_time_decode_length (20 s); keep chunks below it. +MAX_CHUNK_S = 20.0 def enhance( - x: np.ndarray, sr: int, strength: float, device_pref: str = "auto" + x: np.ndarray, + sr: int, + strength: float, + device_pref: str = "auto", + chunk_s: float = 30.0, + overlap_s: float = 0.5, + on_progress: Callable[[int, int], None] | None = None, ) -> tuple[np.ndarray, str, str]: + # torch first: clearvoice's declared torch dependency would otherwise + # resolve to the newest (multi-GB CUDA) build before we pin our tested one. + lazy.ensure_torch() lazy.ensure(["clearvoice==0.1.2"], purpose="MossFormer2 (ClearVoice)") - import tempfile - from pathlib import Path - + # clearvoice is modelscope-based and imports undeclared helpers; the probe + # installs whatever the import chain actually trips over. + lazy.ensure_import("clearvoice", purpose="MossFormer2 (ClearVoice)") from clearvoice import ClearVoice device = pick_device(device_pref) sr_target = 48000 xin = dsp.resample(x, sr, sr_target) cv = ClearVoice(task="speech_enhancement", model_name="MossFormer2_SE_48K") - with tempfile.TemporaryDirectory(prefix="producer_mf2_") as td: - inp = Path(td) / "in.wav" - outp = Path(td) / "out.wav" - io.encode(xin, sr_target, inp, "wav", 16) - result = cv(input_path=str(inp), output_name=str(outp)) - if isinstance(result, tuple): - y, fs = result[0], int(result[1]) - else: - y, fs = io.decode(outp) + if chunk_s <= 0: + import tempfile + from pathlib import Path + + with tempfile.TemporaryDirectory(prefix="producer_mf2_") as td: + inp = Path(td) / "in.wav" + outp = Path(td) / "out.wav" + io.encode(xin, sr_target, inp, "wav", 16) + result = cv(input_path=str(inp), output_name=str(outp)) + if isinstance(result, tuple): + y, fs = result[0], int(result[1]) + else: + y, fs = io.decode(outp) + else: + effective = min(float(chunk_s), MAX_CHUNK_S) + + def run(chunk: np.ndarray) -> np.ndarray: + out = cv(chunk[None, :].astype(np.float32)) + if isinstance(out, tuple): + out = out[0] + return np.asarray(out, dtype=np.float32).reshape(-1) + + y = apply_chunked(xin, sr_target, effective, overlap_s, run, on_progress=on_progress) + fs = sr_target + free_vram() y = np.asarray(y, dtype=np.float32) if y.ndim > 1: y = y.mean(axis=-1) if y.shape[-1] <= 2 else y.reshape(-1) diff --git a/lib/src/producer/engines/enhance_resemble.py b/lib/src/producer/engines/enhance_resemble.py index 47f8a6c..8042215 100644 --- a/lib/src/producer/engines/enhance_resemble.py +++ b/lib/src/producer/engines/enhance_resemble.py @@ -1,13 +1,15 @@ from __future__ import annotations +import contextlib import os import subprocess import sys +from collections.abc import Callable from pathlib import Path import numpy as np -from .. import dsp, io, lazy +from .. import dsp, io, lazy, ui from .base import blend VENV_DIR = lazy.DATA_DIR / "venvs" / "resemble" @@ -24,21 +26,25 @@ def _ensure_venv() -> Path: return py uv = lazy.find_uv() if not py.is_file(): - print("[producer] creating resemble venv (CPython 3.11)...", flush=True) - subprocess.run([uv, "venv", str(VENV_DIR), "--python", "3.11"], check=True) - print( - "[producer] installing resemble-enhance into isolated venv (one-time, large)...", flush=True - ) - subprocess.run( + ui.log("[producer] creating resemble venv (CPython 3.11)...") + ui.run([uv, "venv", str(VENV_DIR), "--python", "3.11"], "creating resemble venv") + ui.log("[producer] installing resemble-enhance into isolated venv (one-time, large)...") + ui.run( [uv, "pip", "install", "--python", str(py), *REQS], - check=True, + "installing resemble-enhance dependencies", ) REQ_HASH_FILE.write_text(want) return py def enhance( - x: np.ndarray, sr: int, strength: float, device_pref: str = "auto" + x: np.ndarray, + sr: int, + strength: float, + device_pref: str = "auto", + chunk_s: float = 30.0, + overlap_s: float = 0.5, + on_progress: Callable[[int, int], None] | None = None, ) -> tuple[np.ndarray, str, str]: import tempfile @@ -61,8 +67,23 @@ def enhance( str(outp), device, str(float(np.clip(strength, 0.0, 1.0))), + str(float(chunk_s)), + str(float(overlap_s)), ] - subprocess.run(cmd, check=True, env=env) + proc = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, text=True) + assert proc.stdout is not None + for line in proc.stdout: + parts = line.split() + if len(parts) == 4 and parts[0] == "PROGRESS": + if on_progress is not None: + with contextlib.suppress(ValueError): + on_progress(int(parts[2]), int(parts[3])) + continue + if line.strip(): + sys.stderr.write(line if line.endswith("\n") else line + "\n") + rc = proc.wait() + if rc != 0: + raise subprocess.CalledProcessError(rc, cmd) y, fs = io.decode(outp) if fs != sr: y = dsp.resample(y, fs, sr) diff --git a/lib/src/producer/engines/resemble_worker.py b/lib/src/producer/engines/resemble_worker.py index cdc5e2b..493ceb3 100644 --- a/lib/src/producer/engines/resemble_worker.py +++ b/lib/src/producer/engines/resemble_worker.py @@ -2,34 +2,65 @@ from __future__ import annotations import contextlib import sys +from pathlib import Path import numpy as np import soundfile as sf +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from producer.engines.chunking import apply_chunked + def load(path: str) -> tuple[np.ndarray, int]: data, sr = sf.read(path, dtype="float32", always_2d=True) return data.mean(axis=1).astype(np.float32), int(sr) +def _report(stage: str): + def cb(done: int, total: int) -> None: + print(f"PROGRESS {stage} {done} {total}", flush=True) + + return cb + + def main() -> int: import torch from resemble_enhance.enhancer.inference import denoise as r_denoise from resemble_enhance.enhancer.inference import enhance as r_enhance - inp, outp, device, strength = sys.argv[1], sys.argv[2], sys.argv[3], float(sys.argv[4]) + inp, outp, device = sys.argv[1], sys.argv[2], sys.argv[3] + strength = float(sys.argv[4]) + chunk_s = float(sys.argv[5]) if len(sys.argv) > 5 else 30.0 + overlap_s = float(sys.argv[6]) if len(sys.argv) > 6 else 0.5 x, sr = load(inp) - t = torch.from_numpy(x) - with contextlib.suppress(Exception): - t, sr = r_denoise(t, sr, device) - try: - y, sr = r_enhance(t, sr, device, nfe=64, solver="midpoint", lambd=1.0 - 0.1 * strength) - except TypeError: - y, sr = r_enhance(t, sr, device) - if isinstance(y, torch.Tensor): - y = y.detach().cpu().numpy() - y = np.asarray(y, dtype=np.float32).reshape(-1) - sf.write(outp, y, int(sr), subtype="FLOAT") + state = {"sr": sr} + + def denoise_fn(chunk: np.ndarray) -> np.ndarray: + t = torch.from_numpy(np.ascontiguousarray(chunk)) + with contextlib.suppress(Exception): + t, state["sr"] = r_denoise(t, state["sr"], device) + if isinstance(t, torch.Tensor): + t = t.detach().cpu().numpy() + return np.asarray(t, dtype=np.float32).reshape(-1) + + def enhance_fn(chunk: np.ndarray) -> np.ndarray: + t = torch.from_numpy(np.ascontiguousarray(chunk)) + try: + y, _ = r_enhance( + t, state["sr"], device, nfe=64, solver="midpoint", lambd=1.0 - 0.1 * strength + ) + except TypeError: + y, _ = r_enhance(t, state["sr"], device) + if isinstance(y, torch.Tensor): + y = y.detach().cpu().numpy() + return np.asarray(y, dtype=np.float32).reshape(-1) + + x = apply_chunked(x, sr, chunk_s, overlap_s, denoise_fn, on_progress=_report("denoise")) + y = apply_chunked( + x, state["sr"], chunk_s, overlap_s, enhance_fn, on_progress=_report("enhance") + ) + sf.write(outp, y, int(state["sr"]), subtype="FLOAT") return 0 diff --git a/lib/src/producer/io.py b/lib/src/producer/io.py index 3e9e769..b6d38b6 100644 --- a/lib/src/producer/io.py +++ b/lib/src/producer/io.py @@ -38,7 +38,7 @@ def encode( sr: int, path: str | Path, out_format: str = "wav", - bit_depth: int = 16, + bit_depth: int = 32, ) -> Path: path = Path(path) fmt = out_format.lower() diff --git a/lib/src/producer/lazy.py b/lib/src/producer/lazy.py index 0ad664d..55bbf6c 100644 --- a/lib/src/producer/lazy.py +++ b/lib/src/producer/lazy.py @@ -1,18 +1,42 @@ from __future__ import annotations +import importlib import importlib.util +import json +import platform +import re import shutil -import subprocess import sys +import urllib.request +from collections.abc import Callable from importlib import metadata from pathlib import Path +from typing import TypeVar +from urllib.parse import unquote + +from . import ui + +_T = TypeVar("_T") DATA_DIR = Path(__file__).resolve().parents[2] +TORCH_VERSION = "2.7.1" TORCH_GPU_INDEX = "https://download.pytorch.org/whl/cu126" +WHEEL_CACHE = DATA_DIR / "cache" / "wheels" DIST_ALIASES = { "deepfilternet": "df", "deepfilterlib": "libdf", } +# the probe installs by the failing import name, but some import names differ +# from their pip name; only entries actually tripped by a chain get used +IMPORT_ALIASES = { + "PIL": "pillow", + "yaml": "pyyaml", + "dateutil": "python-dateutil", + "dotenv": "python-dotenv", + "cv2": "opencv-python", + "sklearn": "scikit-learn", + "skimage": "scikit-image", +} class EngineUnavailable(RuntimeError): @@ -50,18 +74,120 @@ def gpu_present() -> bool: return shutil.which("nvidia-smi") is not None +def _py_tag() -> str: + return f"cp{sys.version_info.major}{sys.version_info.minor}" + + +def _arches() -> tuple[str, ...]: + m = platform.machine().lower() + if m in ("amd64", "x86_64"): + return ("x86_64", "amd64") + if m in ("aarch64", "arm64"): + return ("aarch64", "arm64") + return (m,) + + +def _get_text(url: str) -> str: + req = urllib.request.Request(url, headers={"User-Agent": ui.UA}) + with urllib.request.urlopen(req, timeout=30) as resp: + return resp.read().decode("utf-8", "replace") + + +def _choose_pypi(data: dict) -> tuple[str, str, str | None, int | None] | None: + """Pick the linux wheel for this interpreter from a PyPI JSON release.""" + tag = _py_tag() + arches = _arches() + for u in data.get("urls", []): + fn = u.get("filename", "") + if fn.endswith(".whl") and tag in fn and "linux" in fn and any(a in fn for a in arches): + return fn, u["url"], (u.get("digests") or {}).get("sha256"), u.get("size") + return None + + +def _choose_gpu( + html: str, pkg: str, version: str +) -> tuple[str, str, str | None, int | None] | None: + """Pick a cu126 wheel from a download.pytorch.org index page. + + Hrefs may be relative or absolute and carry a #sha256= fragment, e.g. + https://download-r2.pytorch.org/whl/cu126/torch-2.7.1%2Bcu126-cp311-...whl#sha256=... + """ + prefix = f"{pkg}-{version}%2Bcu126-{_py_tag()}-" + arches = _arches() + for raw in re.findall(r'href="([^"]+\.whl[^"]*)"', html): + href, _, frag = raw.partition("#") + sha = frag[len("sha256=") :] if frag.startswith("sha256=") else None + encoded = href.rsplit("/", 1)[-1] + fname = unquote(encoded) + if not encoded.startswith(prefix) or not any(a in fname for a in arches): + continue + url = href if href.startswith("http") else f"{TORCH_GPU_INDEX}/{pkg}/{href}" + return fname, url, sha, None + return None + + +def _pypi_wheel(pkg: str, version: str) -> tuple[str, str, str | None, int | None] | None: + data = json.loads(_get_text(f"https://pypi.org/pypi/{pkg}/{version}/json")) + return _choose_pypi(data) + + +def _gpu_wheel(pkg: str, version: str) -> tuple[str, str, str | None, int | None] | None: + return _choose_gpu(_get_text(f"{TORCH_GPU_INDEX}/{pkg}/"), pkg, version) + + +def _resolve_torch_wheels(gpu: bool) -> list[tuple[str, str, str | None, int | None]]: + """(filename, url, sha256, size) for torch + torchaudio; [] when unresolvable.""" + pick = _gpu_wheel if gpu else _pypi_wheel + wheels = [] + for pkg in ("torch", "torchaudio"): + found = pick(pkg, TORCH_VERSION) + if found is None: + return [] + wheels.append(found) + return wheels + + def ensure_torch() -> None: if has_module("torch"): return - cmd = [find_uv(), "pip", "install", "--python", sys.executable] - if gpu_present(): - cmd += ["torch==2.7.1+cu126", "torchaudio==2.7.1+cu126", "--index-url", TORCH_GPU_INDEX] + gpu = gpu_present() + flavor = "CUDA" if gpu else "CPU" + ui.log( + f"[producer] installing torch {TORCH_VERSION} ({flavor})..." + + (" ~2.5 GB download" if gpu else "") + ) + wheels: list[tuple[str, str, str | None, int | None]] = [] + try: + wheels = _resolve_torch_wheels(gpu) + except Exception as e: + ui.log(f"[producer] wheel index lookup failed ({e}); using uv directly") + if wheels: + paths = [ + ui.download( + url, + WHEEL_CACHE / name, + f"downloading {name}", + expected_size=size, + sha256=sha, + ) + for name, url, sha, size in wheels + ] + ui.run( + [find_uv(), "pip", "install", "--python", sys.executable, *(str(p) for p in paths)], + "installing torch", + ) else: - cmd += ["torch==2.7.1", "torchaudio==2.7.1"] - print("[producer] installing torch (one-time, large download)...", flush=True) - subprocess.run(cmd, check=True) - import importlib - + cmd = [find_uv(), "pip", "install", "--python", sys.executable] + if gpu: + cmd += [ + f"torch=={TORCH_VERSION}+cu126", + f"torchaudio=={TORCH_VERSION}+cu126", + "--index-url", + TORCH_GPU_INDEX, + ] + else: + cmd += [f"torch=={TORCH_VERSION}", f"torchaudio=={TORCH_VERSION}"] + ui.run(cmd, "installing torch") importlib.invalidate_caches() @@ -74,12 +200,58 @@ def ensure(packages: list[str], purpose: str) -> None: missing = [m for m in missing if not m.startswith("torch") and not is_installed(m)] if not missing: return - print(f"[producer] installing {purpose} dependencies (one-time)...", flush=True) - cmd = [find_uv(), "pip", "install", "--python", sys.executable, *missing] - subprocess.run(cmd, check=True) - import importlib - + ui.log(f"[producer] installing {purpose} dependencies (one-time)...") + ui.run( + [find_uv(), "pip", "install", "--python", sys.executable, *missing], + f"installing {purpose} dependencies", + ) importlib.invalidate_caches() for spec in missing: if not is_installed(spec): raise EngineUnavailable(f"failed to install {dist_name(spec)} for {purpose}") + + +def ensure_import(module: str, purpose: str, max_rounds: int = 4) -> None: + """Import `module`, installing any undeclared dependency it trips over. + + Some model loaders import modules their wheels don't declare (ModelScope's + config/hub code needs addict/simplejson/sortedcontainers, for example). + The failing import name usually doubles as the pip spec; the ones where + it doesn't are translated via IMPORT_ALIASES. Runs until the import + succeeds; gives up after max_rounds with a clear error. + """ + missing = module + for _ in range(max_rounds): + try: + importlib.import_module(module) + return + except ModuleNotFoundError as e: + missing = e.name or module + if "." in missing: + # a submodule of an installed package is broken; installing + # "pkg.sub" from pip would be nonsense — surface the real error + raise + ensure([IMPORT_ALIASES.get(missing, missing)], purpose=purpose) + raise EngineUnavailable(f"cannot import {module} for {purpose} (still missing: {missing})") + + +def ensure_call(fn: Callable[[], _T], purpose: str, max_rounds: int = 8) -> _T: + """Run fn(), installing any undeclared dependency its import chain trips over. + + Engine weight loaders import lazily on the first call rather than at + import time, and those imports reach modules whose wheels don't declare + their deps (ModelScope's config/hub code needs addict/simplejson/ + sortedcontainers, for example, which are extras-only on PyPI). Same + recovery loop as ensure_import, around a call instead of an import; only + the first invocation can trip, later ones find the chain importable. + """ + missing = "" + for _ in range(max_rounds): + try: + return fn() + except ModuleNotFoundError as e: + missing = e.name or "" + if not missing or "." in missing: + raise + ensure([IMPORT_ALIASES.get(missing, missing)], purpose=purpose) + raise EngineUnavailable(f"cannot run {purpose} (still missing: {missing})") diff --git a/lib/src/producer/meters.py b/lib/src/producer/meters.py index 0cfc3e3..dc4c566 100644 --- a/lib/src/producer/meters.py +++ b/lib/src/producer/meters.py @@ -4,6 +4,7 @@ import numpy as np from scipy import signal SILENCE_DB = -120.0 +_METER_BLOCK = 1 << 20 def _db(v: float) -> float: @@ -13,46 +14,82 @@ def _db(v: float) -> float: def rms_db(x: np.ndarray) -> float: - x64 = x.astype(np.float64) - return _db(float(np.sqrt(np.mean(np.square(x64))))) + total = 0.0 + for i in range(0, x.size, _METER_BLOCK): + blk = x[i : i + _METER_BLOCK] + total += float(np.square(blk.astype(np.float64)).sum()) + mean = total / x.size if x.size else 0.0 + return _db(float(np.sqrt(mean))) def sample_peak_db(x: np.ndarray) -> float: - x64 = np.abs(x.astype(np.float64)) - return _db(float(np.max(x64)) if x64.size else 0.0) + peak = 0.0 + for i in range(0, x.size, _METER_BLOCK): + blk = x[i : i + _METER_BLOCK] + peak = max(peak, float(np.max(np.abs(blk))) if blk.size else 0.0) + return _db(peak) def true_peak_db(x: np.ndarray, sr: int, oversample: int = 4) -> float: if x.size < 2: return sample_peak_db(x) - y = signal.resample_poly(x.astype(np.float64), oversample, 1, window=("kaiser", 8.0)) - return sample_peak_db(y.astype(np.float32)) - - -def noise_floor_db(x: np.ndarray, sr: int, block_ms: float = 50.0, pct: float = 5.0) -> float: + half = 10 * oversample + h = signal.firwin(2 * half + 1, 1.0 / oversample, window=("kaiser", 8.0)).astype(np.float32) + pad = 4 * half + oversample + block = max(1, int(30.0 * sr)) + peak = 0.0 + start = 0 + while start < x.size: + end = min(start + block, x.size) + a, b = max(0, start - pad), min(x.size, end + pad) + y = signal.resample_poly(x[a:b].astype(np.float32, copy=False), oversample, 1, window=h) + lo = (start - a) * oversample + hi = (end - a) * oversample + peak = max(peak, float(np.max(np.abs(y[lo:hi])))) + start = end + return _db(peak) + + +def block_pct_db(x: np.ndarray, sr: int, block_ms: float = 50.0, pct: float = 5.0) -> float: n = max(1, int(sr * block_ms / 1000.0)) nb = x.size // n if nb < 1: return sample_peak_db(x) if x.size else SILENCE_DB - blocks = x[: nb * n].reshape(nb, n).astype(np.float64) - block_rms = np.sqrt(np.mean(np.square(blocks), axis=1)) + view = x[: nb * n].reshape(nb, n) + rows = max(1, _METER_BLOCK // n) + block_rms = np.empty(nb, dtype=np.float64) + for i in range(0, nb, rows): + sl = view[i : i + rows] + block_rms[i : i + rows] = np.sqrt(np.mean(np.square(sl.astype(np.float64)), axis=1)) active = block_rms[block_rms > 0.0] if active.size == 0: return SILENCE_DB return _db(float(np.percentile(active, pct))) +def noise_floor_db(x: np.ndarray, sr: int) -> float: + return block_pct_db(x, sr, block_ms=50.0, pct=5.0) + + +def speech_level_db(x: np.ndarray, sr: int) -> float: + """Robust loudness of the actual speech in x (95th pct of 50 ms frame RMS). + + Unlike plain RMS this barely moves with long pauses, so it is a stable + anchor for gain staging the DSP chain. + """ + return block_pct_db(x, sr, block_ms=50.0, pct=95.0) + + def lufs(x: np.ndarray, sr: int) -> float: import pyloudnorm as pyln - x64 = x.astype(np.float64) - if not np.any(x64): + if not np.any(x): return SILENCE_DB - if x64.size < int(sr * 0.2): + if x.size < int(sr * 0.2): return SILENCE_DB meter = pyln.Meter(sr) try: - val = meter.integrated_loudness(x64) + val = meter.integrated_loudness(x) except ValueError: return SILENCE_DB if not np.isfinite(val): diff --git a/lib/src/producer/pipeline.py b/lib/src/producer/pipeline.py index 0a22bf8..e7aa4ae 100644 --- a/lib/src/producer/pipeline.py +++ b/lib/src/producer/pipeline.py @@ -1,14 +1,26 @@ from __future__ import annotations +import math import time from collections.abc import Callable from dataclasses import dataclass, field +from typing import Protocol import numpy as np from . import dsp, loudness from .config import Options -from .meters import all_meters +from .meters import all_meters, speech_level_db + + +class Reporter(Protocol): + """Status sink for the pipeline; see producer.ui.Status.""" + + def stage(self, name: str) -> None: ... + + def tick(self, done: int, total: int) -> None: ... + + def stage_done(self, name: str, elapsed: float | None = None) -> None: ... @dataclass @@ -30,20 +42,62 @@ class RunResult: after: dict[str, float] = field(default_factory=dict) -def _denoise_fn(opts: Options, notes: list[str]): +def _chunk_note(opts: Options, notes: list[str], x: np.ndarray, sr: int) -> None: + if opts.engine_chunk_s <= 0 or x.size <= opts.engine_chunk_s * sr: + return + eff = max(0.01, opts.engine_chunk_s - opts.engine_overlap_s) + n_chunks = math.ceil(x.size / sr / eff) + notes.append( + f"engine chunking: {n_chunks} chunks x {opts.engine_chunk_s:g}s" + f" ({opts.engine_overlap_s:g}s overlap)" + ) + + +def _denoise_fn(opts: Options, notes: list[str], reporter: Reporter | None = None): engine = opts.denoise + prog = reporter.tick if reporter is not None else None def fn(x: np.ndarray, sr: int) -> np.ndarray: if engine == "off": return x + _chunk_note(opts, notes, x, sr) if engine == "dfn3": from .engines import denoise_dfn - y, eng, dev = denoise_dfn.denoise(x, sr, opts.denoise_strength, opts.device) + y, eng, dev = denoise_dfn.denoise( + x, + sr, + opts.denoise_strength, + opts.device, + chunk_s=opts.engine_chunk_s, + overlap_s=opts.engine_overlap_s, + on_progress=prog, + post_filter=opts.denoise_pf, + ) elif engine == "zipenhancer": from .engines import denoise_zip - y, eng, dev = denoise_zip.denoise(x, sr, opts.denoise_strength, opts.device) + y, eng, dev = denoise_zip.denoise( + x, + sr, + opts.denoise_strength, + opts.device, + chunk_s=opts.engine_chunk_s, + overlap_s=opts.engine_overlap_s, + on_progress=prog, + ) + elif engine == "spectral": + from .engines import denoise_spectral + + y, eng, dev = denoise_spectral.denoise( + x, + sr, + opts.denoise_strength, + opts.device, + chunk_s=opts.engine_chunk_s, + overlap_s=opts.engine_overlap_s, + on_progress=prog, + ) else: raise ValueError(f"unknown denoise engine: {engine}") notes.append(f"denoise: {eng} on {dev}") @@ -52,20 +106,38 @@ def _denoise_fn(opts: Options, notes: list[str]): return fn -def _enhance_fn(opts: Options, notes: list[str]): +def _enhance_fn(opts: Options, notes: list[str], reporter: Reporter | None = None): engine = opts.enhance + prog = reporter.tick if reporter is not None else None def fn(x: np.ndarray, sr: int) -> np.ndarray: if engine == "off": return x + _chunk_note(opts, notes, x, sr) if engine == "mossformer2": from .engines import enhance_mossformer - y, eng, dev = enhance_mossformer.enhance(x, sr, opts.enhance_strength, opts.device) + y, eng, dev = enhance_mossformer.enhance( + x, + sr, + opts.enhance_strength, + opts.device, + chunk_s=opts.engine_chunk_s, + overlap_s=opts.engine_overlap_s, + on_progress=prog, + ) elif engine == "resemble": from .engines import enhance_resemble - y, eng, dev = enhance_resemble.enhance(x, sr, opts.enhance_strength, opts.device) + y, eng, dev = enhance_resemble.enhance( + x, + sr, + opts.enhance_strength, + opts.device, + chunk_s=opts.engine_chunk_s, + overlap_s=opts.engine_overlap_s, + on_progress=prog, + ) else: raise ValueError(f"unknown enhance engine: {engine}") notes.append(f"enhance: {eng} on {dev}") @@ -78,6 +150,15 @@ def _dsp_fn(opts: Options, notes: list[str]): p = opts.prof() def fn(x: np.ndarray, sr: int) -> np.ndarray: + # Gain-stage the (denoised) signal to the level the voice chain was + # designed around. Denoisers shift the level distribution, and the + # absolute comp thresholds only make sense at a known input level; + # anchoring on measured speech level keeps the chain predictable and + # stops the compressors over-riding sparse denoised speech. + ref = speech_level_db(x, sr) + pregain_db = float(np.clip(p.dsp_ref_db - ref, -24.0, 24.0)) + notes.append(f"dsp pregain: {pregain_db:+.1f} dB (speech level {ref:.1f} dB)") + x = x * np.float32(10.0 ** (pregain_db / 20.0)) hpf_hz = opts.hpf_hz if opts.hpf_hz is not None else p.hpf_hz if opts.eff("hpf") > 0.001: x = dsp.hpf(x, sr, hpf_hz) @@ -99,7 +180,11 @@ def _dsp_fn(opts: Options, notes: list[str]): x = dsp.peak_eq(x, sr, pf, pg * opts.eff("presence")) af, ag = p.air x = dsp.shelf(x, sr, af, ag * opts.eff("air"), low=False) - x = dsp.expander(x, sr, max_drop_db=6.0 * opts.eff("breath")) + # Between sentences a good denoiser leaves near-silence but tape hiss + # and breaths survive; duck them for real instead of the old 6 dB cap. + breath = opts.eff("breath") + max_drop = (12.0 + 12.0 * breath) if breath > 0.02 else 0.0 + x = dsp.expander(x, sr, max_drop_db=max_drop) return x return fn @@ -108,8 +193,13 @@ def _dsp_fn(opts: Options, notes: list[str]): def _blend(x: np.ndarray, y: np.ndarray, s: float) -> np.ndarray: s = float(np.clip(s, 0.0, 1.0)) if s >= 0.999: - return y - return (x.astype(np.float64) * (1.0 - s) + y.astype(np.float64) * s).astype(np.float32) + return y.astype(np.float32, copy=False) + if s <= 0.001: + return x.astype(np.float32, copy=False) + out = np.asarray(y, dtype=np.float32) - np.asarray(x, dtype=np.float32) + out *= s + out += x + return out def _level_fn(opts: Options, notes: list[str]): @@ -119,15 +209,18 @@ def _level_fn(opts: Options, notes: list[str]): return fn -def build_stages(opts: Options, notes: list[str] | None = None) -> list[Stage]: +def build_stages( + opts: Options, notes: list[str] | None = None, reporter: Reporter | None = None +) -> list[Stage]: notes = notes if notes is not None else [] stages: list[Stage] = [] stages.append( Stage( "denoise", - f"engine={opts.denoise} strength={opts.denoise_strength:.2f}", + f"engine={opts.denoise} strength={opts.denoise_strength:.2f}" + f" pf={'on' if opts.denoise_pf else 'off'}", opts.denoise != "off", - _denoise_fn(opts, notes), + _denoise_fn(opts, notes, reporter), ) ) stages.append( @@ -135,7 +228,7 @@ def build_stages(opts: Options, notes: list[str] | None = None) -> list[Stage]: "enhance", f"engine={opts.enhance} strength={opts.enhance_strength:.2f}", opts.enhance != "off", - _enhance_fn(opts, notes), + _enhance_fn(opts, notes, reporter), ) ) dsp_desc = "hpf→mud→warmth→soothe→comp2x→tape→deess→presence→air→breath" @@ -151,7 +244,9 @@ def build_stages(opts: Options, notes: list[str] | None = None) -> list[Stage]: return stages -def run_pipeline(x: np.ndarray, sr: int, opts: Options) -> RunResult: +def run_pipeline( + x: np.ndarray, sr: int, opts: Options, reporter: Reporter | None = None +) -> RunResult: notes: list[str] = [] stages = build_stages(opts, notes) before = all_meters(x, sr) @@ -160,9 +255,14 @@ def run_pipeline(x: np.ndarray, sr: int, opts: Options) -> RunResult: for st in stages: if not st.enabled or st.fn is None: continue + if reporter is not None: + reporter.stage(st.name) t0 = time.perf_counter() y = st.fn(y, sr) - timings[st.name] = round(time.perf_counter() - t0, 3) + elapsed = time.perf_counter() - t0 + timings[st.name] = round(elapsed, 3) + if reporter is not None: + reporter.stage_done(st.name, elapsed) after = all_meters(y, sr) return RunResult( audio=y, sr=sr, stages=stages, timings=timings, notes=notes, before=before, after=after diff --git a/lib/src/producer/ui.py b/lib/src/producer/ui.py new file mode 100644 index 0000000..bc08756 --- /dev/null +++ b/lib/src/producer/ui.py @@ -0,0 +1,319 @@ +"""Terminal progress primitives: live status lines, download bars, subprocess runners. + +All live rendering shares one line on stdout: renderers erase whatever line is +currently live before drawing their own, and `log()` erases it before printing +a permanent line, so engine prints, downloads, and stage ticks never garble +each other. Without a TTY nothing is drawn in place; instead milestones are +logged periodically so piped output still shows regular progress. +""" + +from __future__ import annotations + +import hashlib +import shutil +import subprocess +import sys +import threading +import time +import urllib.request +from pathlib import Path + +_lock = threading.Lock() +_live = "" # contents of the in-place line, "" when none + +UA = "producer/0.1.0" + + +def is_tty() -> bool: + try: + return bool(sys.stdout and sys.stdout.isatty()) + except Exception: + return False + + +def _cols() -> int: + try: + return max(20, shutil.get_terminal_size().columns) + except Exception: + return 80 + + +def _erase() -> None: + global _live + if _live: + sys.stdout.write("\r" + " " * len(_live) + "\r") + _live = "" + + +def _render_live(text: str) -> None: + global _live + if not is_tty(): + return + with _lock: + text = text[: _cols() - 1] + _erase() + sys.stdout.write("\r" + text) + sys.stdout.flush() + _live = text + + +def finish_live() -> None: + """Clear the in-place line, if any.""" + with _lock: + if is_tty(): + _erase() + sys.stdout.flush() + + +def log(msg: str) -> None: + """Print a permanent line, clearing any in-place progress line first.""" + with _lock: + if is_tty(): + _erase() + print(msg, flush=True) + + +def fmt_bytes(n: float | None) -> str: + if n is None: + return "?" + n = float(n) + for unit in ("B", "KB", "MB", "GB", "TB"): + if n < 1024.0 or unit == "TB": + return f"{n:.1f} {unit}" if unit != "B" else f"{int(n)} B" + n /= 1024.0 + return f"{n:.1f} TB" + + +def fmt_secs(s: float | None) -> str: + if s is None or s < 0 or s != s: # None, negative, NaN + return "?" + s = int(s) + h, rem = divmod(s, 3600) + m, sec = divmod(rem, 60) + return f"{h}:{m:02d}:{sec:02d}" if h else f"{m}:{sec:02d}" + + +def _bar(frac: float, width: int = 22) -> str: + frac = min(1.0, max(0.0, frac)) + fill = round(frac * width) + return "[" + "#" * fill + "-" * (width - fill) + "]" + + +class Progress: + """Byte progress for one download: live bar on a TTY, milestones otherwise.""" + + def __init__(self, label: str, total: int | None = None) -> None: + self.label = label + self.total = total + self.done = 0 + self._t0 = time.monotonic() + self._last_draw = 0.0 + self._last_ms_t = 0.0 + self._last_ms_pct = -100 + + def update(self, done: int, total: int | None = None) -> None: + if total is not None: + self.total = total + self.done = done + now = time.monotonic() + if is_tty(): + if now - self._last_draw >= 0.1 or (self.total and done >= self.total): + self._draw(now) + else: + self._milestone(now) + + def _draw(self, now: float) -> None: + self._last_draw = now + el = max(1e-6, now - self._t0) + speed = self.done / el + text = f"[producer] {self.label} {_bar(0)}" + if self.total: + frac = self.done / self.total + rem = max(0.0, el * (self.total - self.done) / max(1, self.done)) + text = ( + f"[producer] {self.label} {_bar(frac)} {fmt_bytes(self.done)}/" + f"{fmt_bytes(self.total)} ({frac * 100:.0f}%) {fmt_bytes(speed)}/s" + f" ETA {fmt_secs(rem)}" + ) + else: + text = f"[producer] {self.label} {fmt_bytes(self.done)} {fmt_bytes(speed)}/s" + _render_live(text) + + def _milestone(self, now: float) -> None: + pct = 100.0 * self.done / self.total if self.total else 0.0 + if ( + now - self._last_ms_t >= 30.0 + or (self.total and pct - self._last_ms_pct >= 10.0) + or (self.total and self.done >= self.total) + ): + self._last_ms_t = now + self._last_ms_pct = pct + log(f"[producer] {self.label} {fmt_bytes(self.done)}/{fmt_bytes(self.total)}") + + def close(self, final: str | None = None) -> None: + el = time.monotonic() - self._t0 + msg = final or (f"[producer] {self.label} done ({fmt_bytes(self.done)} in {fmt_secs(el)})") + log(msg) + + +def _sha256_of(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + h.update(chunk) + return h.hexdigest() + + +def _cache_ok(dest: Path, expected_size: int | None, sha256: str | None) -> bool: + if not dest.is_file(): + return False + if sha256: + if expected_size is not None and dest.stat().st_size != expected_size: + return False + return _sha256_of(dest) == sha256 + if expected_size is not None: + return dest.stat().st_size == expected_size + return True + + +def download( + url: str, + dest: Path, + label: str | None = None, + expected_size: int | None = None, + sha256: str | None = None, + timeout: float = 60.0, +) -> Path: + """Download `url` to `dest` with a progress bar; skips if already cached.""" + dest = Path(dest) + label = label or f"downloading {dest.name}" + if _cache_ok(dest, expected_size, sha256): + log(f"[producer] {label}: already cached ({fmt_bytes(dest.stat().st_size)})") + return dest + dest.parent.mkdir(parents=True, exist_ok=True) + tmp = dest.with_name(dest.name + ".part") + req = urllib.request.Request(url, headers={"User-Agent": UA}) + with urllib.request.urlopen(req, timeout=timeout) as resp: + total = expected_size + if total is None: + try: + total = int(resp.headers.get("Content-Length") or 0) or None + except (TypeError, ValueError): + total = None + prog = Progress(label, total) + hasher = hashlib.sha256() if sha256 else None + tmp.parent.mkdir(parents=True, exist_ok=True) + with tmp.open("wb") as f: + while True: + chunk = resp.read(1 << 20) + if not chunk: + break + f.write(chunk) + if hasher is not None: + hasher.update(chunk) + prog.update(f.tell()) + prog.close() + if hasher is not None and hasher.hexdigest() != sha256: + tmp.unlink(missing_ok=True) + raise RuntimeError(f"checksum mismatch for {dest.name} ({url})") + tmp.replace(dest) + return dest + + +def run(cmd: list[str], label: str, check: bool = True) -> int: + """Run a subprocess, keeping output visible and adding heartbeats when piped. + + On a TTY the child inherits the terminal (uv draws its own progress bars); + when output is piped, the child's lines are forwarded and a heartbeat with + elapsed time is logged every 30 s so long installs never look frozen. + """ + if is_tty(): + proc = subprocess.run(cmd, check=check) + return proc.returncode + start = time.monotonic() + p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + + def beat() -> None: + while p.poll() is None: + time.sleep(30.0) + if p.poll() is None: + log(f"[producer] {label}... {int(time.monotonic() - start)}s elapsed") + + threading.Thread(target=beat, daemon=True).start() + assert p.stdout is not None + for line in p.stdout: + line = line.rstrip() + if line: + log(line) + rc = p.wait() + if check and rc != 0: + raise subprocess.CalledProcessError(rc, cmd) + return rc + + +class Status: + """Per-file pipeline reporter: one live line for the current stage.""" + + def __init__(self, prefix: str = "") -> None: + self.prefix = prefix + self.stage_name: str | None = None + self.total: int | None = None + self.done = 0 + self._t0 = 0.0 + self._last_draw = 0.0 + self._last_ms_t = 0.0 + self._last_ms_pct = -100 + + def stage(self, name: str) -> None: + self.stage_name = name + self.total = None + self.done = 0 + self._t0 = time.monotonic() + self._last_ms_t = self._t0 + self._last_ms_pct = -100 + self._draw(force=True) + + def tick(self, done: int, total: int) -> None: + """Progress within the current stage (engine chunks).""" + self.total = total + self.done = done + now = time.monotonic() + if is_tty(): + if now - self._last_draw >= 0.1 or done >= total: + self._draw(now) + elif total and ( + done >= total + or now - self._last_ms_t >= 30.0 + or 100.0 * done / total - self._last_ms_pct >= 20.0 + ): + self._last_ms_t = now + self._last_ms_pct = 100.0 * done / total + log(f"{self.prefix}{self.stage_name} {done}/{total} chunks ({self._last_ms_pct:.0f}%)") + + def stage_done(self, name: str, elapsed: float | None = None) -> None: + el = time.monotonic() - self._t0 if elapsed is None else elapsed + chunks = f" ({self.done}/{self.total} chunks)" if self.total else "" + log(f"{self.prefix}{name} done in {el:.1f}s{chunks}") + self.stage_name = None + finish_live() + + def _draw(self, force: bool = False) -> None: + if self.stage_name is None: + return + now = time.monotonic() + if not force and now - self._last_draw < 0.1: + return + self._last_draw = now + el = now - self._t0 + text = f"{self.prefix}{self.stage_name}" + if self.total: + pct = 100.0 * self.done / self.total + eta = el * (self.total - self.done) / self.done if self.done else None + text += f" {_bar(self.done / self.total, 18)} {self.done}/{self.total}" + text += f" chunks ({pct:.0f}%) ETA {fmt_secs(eta)}" + else: + text += f"... {el:.0f}s" + _render_live(text) + + def finish(self) -> None: + finish_live() diff --git a/lib/src/producer/updates.py b/lib/src/producer/updates.py new file mode 100644 index 0000000..690e955 --- /dev/null +++ b/lib/src/producer/updates.py @@ -0,0 +1,273 @@ +"""Dependency update checks: probe uv for newer allowed versions, prompt, apply. + +Probes are `uv pip install --dry-run --upgrade` resolutions against the live +environment, so "newer" always respects requirements-core.txt pins. Engine +probes pin `torch==` and include the core requirements file so +an engine upgrade can never re-resolve a CUDA torch onto a CPU wheel or drag +core pins. Everything here is best-effort: uv/network failures are skipped +silently and never block a run. +""" + +from __future__ import annotations + +import importlib +import re +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor +from importlib import metadata +from pathlib import Path +from typing import NamedTuple + +from . import lazy, ui + +CORE_FILE = lazy.DATA_DIR / "requirements-core.txt" +RESEMBLE_PY = lazy.DATA_DIR / "venvs" / "resemble" / "bin" / "python" +ENGINE_DISTS = ("deepfilternet", "zipenhancer", "clearvoice") +PROBE_TIMEOUT = 8.0 + + +class Update(NamedTuple): + """One promptable upgrade; `cmd` is the uv invocation that applies it.""" + + label: str + old: str | None + new: str + note: str + cmd: list[str] + + +def _installed(dist: str) -> str | None: + try: + return metadata.version(dist) + except metadata.PackageNotFoundError: + return None + + +def _installed_in(python: Path, dist: str) -> str | None: + try: + proc = subprocess.run( + [ + str(python), + "-c", + "import sys; from importlib.metadata import version; print(version(sys.argv[1]))", + dist, + ], + capture_output=True, + text=True, + timeout=15.0, + ) + except (OSError, subprocess.SubprocessError): + return None + if proc.returncode != 0: + return None + return proc.stdout.strip() or None + + +def _core_names() -> list[str]: + names = [] + for line in CORE_FILE.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if line and not line.startswith("#"): + names.append(re.split(r"[<>=!~;\[\s]", line)[0]) + return names + + +def _parse_would_install(out: str) -> dict[str, str]: + found: dict[str, str] = {} + for line in out.splitlines(): + line = line.strip() + if line.startswith(("+ ", "~ ")): + name, _, ver = line[2:].partition("==") + if name and ver: + found[name.strip()] = ver.strip() + return found + + +def _vkey(version: str) -> tuple: + """Orderable key for simple version comparison ("1.2.10" > "1.2.9").""" + parts = [] + for chunk in re.split(r"[._\-+!]", version): + if chunk: + parts.append((1, int(chunk), "") if chunk.isdigit() else (0, 0, chunk)) + return tuple(parts) + + +def _probe(args: list[str], python: str | None, keep: set[str]) -> dict[str, str]: + """`uv pip install --dry-run --upgrade ` -> {name: version} for `keep`. + + Best-effort: any uv failure (offline, unsatisfiable resolution, timeout) + yields {} so the update check silently skips that group. + """ + try: + cmd = [ + lazy.find_uv(), + "pip", + "install", + "--dry-run", + "--upgrade", + "--python", + python or sys.executable, + *args, + ] + proc = subprocess.run(cmd, capture_output=True, text=True, timeout=PROBE_TIMEOUT) + except Exception: + return {} + if proc.returncode != 0: + return {} + got = _parse_would_install(proc.stdout) + return {name: ver for name, ver in got.items() if name in keep} + + +def _jobs() -> list[tuple[str, list[str], str | None, set[str], str]]: + """(key, probe args, venv python or None, keep-set, display note) per probe.""" + jobs: list[tuple[str, list[str], str | None, set[str], str]] = [ + ("core", ["-U", "-r", str(CORE_FILE)], None, set(_core_names()), "") + ] + torch = _installed("torch") + gpu = lazy.gpu_present() + if torch and _installed("torchaudio"): + args = ["-U", "torch", "torchaudio"] + note = "CUDA, ~2.5 GB download" if gpu else "CPU build" + if gpu: + args += ["--index-url", lazy.TORCH_GPU_INDEX] + jobs.append(("torch+torchaudio", args, None, {"torch", "torchaudio"}, note)) + for name in ENGINE_DISTS: + if torch and _installed(name): + jobs.append( + ( + name, + ["-U", "-r", str(CORE_FILE), f"torch=={torch.split('+')[0]}", name], + None, + {name}, + "", + ) + ) + if RESEMBLE_PY.is_file() and _installed_in(RESEMBLE_PY, "resemble-enhance"): + jobs.append( + ( + "resemble-enhance", + ["-U", "resemble-enhance"], + str(RESEMBLE_PY), + {"resemble-enhance"}, + "isolated venv", + ) + ) + return jobs + + +def collect() -> list[Update]: + """Best-effort list of available updates (empty when up to date or offline).""" + jobs = _jobs() + if not jobs: + return [] + with ThreadPoolExecutor(max_workers=len(jobs)) as ex: + probes = list(ex.map(lambda j: (j[0], j[1], j[4], _probe(j[1], j[2], j[3])), jobs)) + updates: list[Update] = [] + uv = lazy.find_uv() + for key, _args, note, got in probes: + if not got: + continue + if key == "core": + for name in _core_names(): + new, old = got.get(name), _installed(name) + if new and old and _vkey(new) > _vkey(old): + updates.append( + Update( + name, + old, + new, + "", + [ + uv, + "pip", + "install", + "--python", + sys.executable, + "-U", + "-r", + str(CORE_FILE), + ], + ) + ) + elif key == "torch+torchaudio": + new, ta_new, old = got.get("torch"), got.get("torchaudio"), _installed("torch") + if new and ta_new and old and _vkey(new) > _vkey(old): + cmd = [ + uv, + "pip", + "install", + "--python", + sys.executable, + f"torch=={new}", + f"torchaudio=={ta_new}", + ] + if lazy.gpu_present(): + cmd += ["--index-url", lazy.TORCH_GPU_INDEX] + updates.append(Update("torch + torchaudio", old, new, note, cmd)) + else: + new = got.get(key) + old = _installed_in(RESEMBLE_PY, key) if key == "resemble-enhance" else _installed(key) + if new and old and _vkey(new) > _vkey(old): + if key == "resemble-enhance": + cmd = [uv, "pip", "install", "--python", str(RESEMBLE_PY), f"{key}=={new}"] + else: + torch = _installed("torch") + cmd = [ + uv, + "pip", + "install", + "--python", + sys.executable, + "-r", + str(CORE_FILE), + f"torch=={torch.split('+')[0]}", + f"{key}=={new}", + ] + updates.append(Update(key, old, new, note, cmd)) + return updates + + +def _confirm(prompt: str) -> bool: + try: + ans = input(f"[producer] {prompt} [y/N]: ").strip().lower() + except EOFError: + print() + return False + return ans in ("y", "yes") + + +def check_and_prompt(force: bool = False, assume_yes: bool = False) -> bool: + """Check once; prompt on interactive TTYs. Returns True when updates applied.""" + updates = collect() + if not updates: + if force: + ui.log("[producer] everything is up to date") + return False + ui.log("[producer] updates available:") + for u in updates: + line = f" {u.label} {u.old} -> {u.new}" + if u.note: + line += f" ({u.note})" + ui.log(line) + if not assume_yes: + if not sys.stdin.isatty(): + if force: + raise SystemExit("producer update: stdin is not interactive; use --yes") + ui.log("[producer] run ./producer update to install these") + return False + if not _confirm("install updates?"): + ui.log("[producer] skipped updates") + return False + for u in updates: + ui.run(u.cmd, f"updating {u.label}") + importlib.invalidate_caches() + ui.log("[producer] updates installed") + return True + + +def run_update_command(args: list[str]) -> int: + """`producer update [--yes]`: check for updates and install them now.""" + assume_yes = any(a in ("-y", "--yes") for a in args) + check_and_prompt(force=True, assume_yes=assume_yes) + return 0 -- cgit v1.2.3