diff options
Diffstat (limited to 'lib')
26 files changed, 3233 insertions, 176 deletions
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,10 +27,20 @@ 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" ) p.add_argument("--enhance-strength", type=float, help="0-1 blend of enhanced signal") @@ -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 <output>.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 <stem>_processed.<ext> 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=<hex> 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==<installed base>` 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 <args>` -> {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 diff --git a/lib/tests/test_chunking.py b/lib/tests/test_chunking.py new file mode 100644 index 0000000..f00c04e --- /dev/null +++ b/lib/tests/test_chunking.py @@ -0,0 +1,250 @@ +import numpy as np +import pytest + +from producer.engines.base import blend +from producer.engines.chunking import apply_chunked, plan_chunks, stitch + + +def test_plan_chunks_basic(): + assert plan_chunks(0, 5, 2) == [] + assert plan_chunks(4, 10, 2) == [(0, 4)] + assert plan_chunks(10, 5, 0) == [(0, 5), (5, 10)] + assert plan_chunks(10, 5, 2) == [(0, 5), (3, 8), (6, 10)] + assert plan_chunks(7, 5, 2) == [(0, 5), (3, 7)] + + +def test_plan_chunks_last_span_covers_tail(): + spans = plan_chunks(11, 5, 4) + assert spans[-1][1] == 11 + assert spans[-1][1] - spans[-1][0] >= 5 + + +def test_plan_chunks_rejects_overlap_ge_chunk(): + with pytest.raises(ValueError): + plan_chunks(100, 5, 5) + with pytest.raises(ValueError): + plan_chunks(100, 5, 6) + + +def test_stitch_identity_pieces_reconstruct_signal(): + rng = np.random.default_rng(0) + x = rng.standard_normal(10_000).astype(np.float32) + spans = plan_chunks(x.size, 2500, 300) + y = stitch(spans, x.size, (x[a:b] for a, b in spans)) + assert y.dtype == np.float32 + assert y.shape == x.shape + np.testing.assert_allclose(y, x, atol=1e-5) + + +def test_stitch_single_span_passthrough(): + x = np.arange(100, dtype=np.float32) + y = stitch([(0, 100)], 100, [x]) + np.testing.assert_array_equal(y, x) + + +def test_stitch_crossfade_of_complementary_pieces(): + # piece A is silence, piece B is a full-scale ramp; the crossfade zone + # must be a smooth blend, not a jump. + spans = [(0, 10), (5, 15)] + pieces = [np.zeros(10, dtype=np.float32), np.ones(10, dtype=np.float32)] + y = stitch(spans, 15, pieces) + mid = y[7] # 50% through the overlap + assert 0.4 < mid < 0.6 + assert y[0] == 0.0 + assert y[14] == 1.0 + + +def test_stitch_pads_short_pieces(): + spans = [(0, 10), (5, 15)] + pieces = [np.arange(10, dtype=np.float32), np.arange(8, dtype=np.float32)] + y = stitch(spans, 15, pieces) + assert y.shape == (15,) + assert np.all(np.isfinite(y)) + + +def test_apply_chunked_identity_matches_whole_file(): + rng = np.random.default_rng(1) + x = rng.standard_normal(50_000).astype(np.float32) + sr = 8000 + calls = [] + + def fn(chunk): + calls.append(chunk.size) + return chunk + + y = apply_chunked(x, sr, 7.0, 0.05, fn) + assert calls == [50_000] + np.testing.assert_array_equal(y, x) + + calls.clear() + y = apply_chunked(x, sr, 0.5, 0.05, fn) + assert len(calls) >= 10 + assert all(c <= 4000 for c in calls) + np.testing.assert_allclose(y, x, atol=1e-5) + + +def test_apply_chunked_zero_disables_chunking(): + x = np.zeros(10, dtype=np.float32) + calls = [] + + def fn(chunk): + calls.append(chunk.size) + return chunk + + apply_chunked(x, 8000, 0.0, 0.5, fn) + assert calls == [10] + + +def test_apply_chunked_context_reconstructs_identity(): + rng = np.random.default_rng(3) + x = rng.standard_normal(30_000).astype(np.float32) + y = apply_chunked(x, 8000, 1.5, 0.25, lambda c: c, context_s=0.5) + np.testing.assert_allclose(y, x, atol=1e-5) + + +def test_apply_chunked_context_feeds_padded_chunks(): + x = np.ones(5_000, dtype=np.float32) + sizes = [] + + def fn(chunk): + sizes.append(chunk.size) + return chunk + + apply_chunked(x, 1000, 2.0, 0.0, fn, context_s=0.5) + # plan_chunks(5000, 2000, 0) -> (0,2000), (2000,4000), (4000,5000) + assert sizes == [2500, 3000, 1500] + + +def test_apply_chunked_context_trims_cold_start_artifacts(): + x = np.ones(6_000, dtype=np.float32) + + def fn(chunk): + out = chunk.copy() + out[0] = 0.0 # cold-start artifact at the start of every model call + return out + + broken = apply_chunked(x, 1000, 2.0, 0.0, fn) + assert int(np.sum(broken == 0.0)) > 1 + + y = apply_chunked(x, 1000, 2.0, 0.0, fn, context_s=0.5) + assert y[0] == 0.0 # only the true file start stays degraded + assert np.all(y[1:] == 1.0) + + +def test_apply_chunked_context_survives_oom_retry(): + x = np.ones(8_000, dtype=np.float32) + attempts = [] + + def fn(chunk): + attempts.append(chunk.size) + if chunk.size > 3000: + raise MemoryError("simulated oom") + out = chunk.copy() + out[0] = 0.0 + return out + + y = apply_chunked(x, 1000, 6.0, 0.0, fn, min_chunk_s=1.0, context_s=0.5) + assert max(attempts) > 3000 + assert y[0] == 0.0 + assert np.all(y[1:] == 1.0) + + +def test_apply_chunked_retries_smaller_on_oom(): + x = np.random.default_rng(2).standard_normal(40_000).astype(np.float32) + sr = 8000 + attempts = [] + + def fn(chunk): + attempts.append(chunk.size) + if chunk.size > 3000: + raise MemoryError("simulated oom") + return chunk + + y = apply_chunked(x, sr, 5.0, 0.0, fn, min_chunk_s=0.25) + assert attempts[0] == 40_000 + assert attempts[-1] == 2500 + assert max(a for a in attempts if a <= 3000) == 2500 + np.testing.assert_allclose(y, x, atol=1e-5) + + +def test_apply_chunked_oom_gives_up_at_min_chunk(): + x = np.zeros(40_000, dtype=np.float32) + + def fn(chunk): + raise MemoryError("always") + + with pytest.raises(MemoryError): + apply_chunked(x, 8000, 1.0, 0.0, fn, min_chunk_s=0.5) + + +def test_apply_chunked_progress_callback(): + x = np.random.default_rng(3).standard_normal(20_000).astype(np.float32) + sr = 8000 + seen: list[tuple[int, int]] = [] + y = apply_chunked(x, sr, 0.5, 0.0, lambda c: c, on_progress=lambda d, t: seen.append((d, t))) + total = len(plan_chunks(x.size, round(0.5 * sr), 0)) + assert seen[0] == (0, total) + assert seen[-1] == (total, total) + assert [d for d, _t in seen[1:]] == list(range(1, total + 1)) + np.testing.assert_allclose(y, x, atol=1e-6) + + +def test_apply_chunked_progress_resets_after_oom_retry(): + x = np.random.default_rng(4).standard_normal(40_000).astype(np.float32) + sr = 8000 + seen: list[tuple[int, int]] = [] + + def fn(chunk): + if chunk.size > 3000: + raise MemoryError("simulated oom") + return chunk + + y = apply_chunked( + x, sr, 5.0, 0.0, fn, min_chunk_s=0.25, on_progress=lambda d, t: seen.append((d, t)) + ) + assert seen[0] == (0, 1) # first attempt is one whole-file chunk + done, total = seen[-1] + assert done == total and total > 2 # retried into smaller chunks + np.testing.assert_allclose(y, x, atol=1e-5) + + +def test_apply_chunked_progress_absent_when_disabled(): + x = np.zeros(10_000, dtype=np.float32) + y = apply_chunked(x, 8000, 0.5, 0.0, lambda c: c, on_progress=None) + assert y.shape == x.shape + + +def test_apply_chunked_whole_file_mode_falls_back_on_oom(): + x = np.random.default_rng(5).standard_normal(10_000).astype(np.float32) + sr = 1000 + + def fn(chunk): + if chunk.size > 1500: + raise MemoryError("simulated oom") + return chunk * 2 + + y = apply_chunked(x, sr, 0.0, 0.0, fn, min_chunk_s=0.5) + np.testing.assert_allclose(y, x * 2, atol=1e-6) + + +def test_blend_float32_no_float64_temporaries(): + x = np.full(1000, 0.25, dtype=np.float32) + y = np.full(1000, 0.75, dtype=np.float32) + out = blend(x, y, 0.5) + assert out.dtype == np.float32 + np.testing.assert_allclose(out, 0.5, atol=1e-7) + assert blend(x, y, 1.0) is y + np.testing.assert_allclose(blend(x, y, 0.0), x, atol=1e-7) + + +def test_blend_realigns_resample_drift(): + # a 48k -> 16k -> 48k round trip can come back a sample or two long + # (resample_poly emits ceil(n * up/down) per hop); blend must cope + x = np.ones(5, dtype=np.float32) + long_y = np.full(7, 0.5, dtype=np.float32) + np.testing.assert_allclose(blend(x, long_y, 0.5), np.full(5, 0.75), atol=1e-7) + + short_y = np.full(4, 0.5, dtype=np.float32) + out = blend(x, short_y, 0.5) + assert out.shape == x.shape + np.testing.assert_allclose(out, [0.75, 0.75, 0.75, 0.75, 0.5], atol=1e-7) diff --git a/lib/tests/test_cli.py b/lib/tests/test_cli.py index 1a157f7..76952ba 100644 --- a/lib/tests/test_cli.py +++ b/lib/tests/test_cli.py @@ -2,6 +2,7 @@ import sys from types import SimpleNamespace import numpy as np +import pytest import soundfile as sf from conftest import speechish @@ -165,6 +166,74 @@ def test_default_output_suffix_processed(tmp_path): assert _resolve_output(inp, opts, single=False) == odir / "song_processed.wav" +def test_engine_chunk_flags(): + args = build_parser().parse_args(["in.wav", "--engine-chunk", "15", "--engine-overlap", "1.0"]) + opts = Options() + _apply_args(opts, args) + assert opts.engine_chunk_s == 15.0 + assert opts.engine_overlap_s == 1.0 + + +def test_default_options_whole_file_and_soft_denoise(): + opts = Options() + assert opts.engine_chunk_s == 0.0 + assert opts.denoise_strength == 0.9 + + +def test_engine_chunk_validation(): + with pytest.raises(SystemExit): + args = build_parser().parse_args(["in.wav", "--engine-chunk", "5", "--engine-overlap", "5"]) + _apply_args(Options(), args) + with pytest.raises(SystemExit): + args = build_parser().parse_args(["in.wav", "--engine-chunk", "-1"]) + _apply_args(Options(), args) + + +def test_engine_chunk_config_override(tmp_path): + from producer import config as cfgmod + + cfg = tmp_path / "config.toml" + cfg.write_text("engine_chunk = 10.0\nengine_overlap = 1.0\n") + opts = Options() + cfgmod.apply_config(opts, cfgmod.load_config(cfg)) + assert opts.engine_chunk_s == 10.0 + assert opts.engine_overlap_s == 1.0 + + +def test_denoise_pf_flag_and_config_plumbing(tmp_path): + from producer import config as cfgmod + + opts = Options() + _apply_args(opts, build_parser().parse_args(["in.wav"])) + assert opts.denoise_pf is False + + opts = Options() + _apply_args(opts, build_parser().parse_args(["in.wav", "--denoise-pf"])) + assert opts.denoise_pf is True + + cfg = tmp_path / "config.toml" + cfg.write_text('[denoise]\nengine = "dfn3"\nstrength = 0.8\npf = true\n') + opts = Options() + cfgmod.apply_config(opts, cfgmod.load_config(cfg)) + assert opts.denoise_strength == 0.8 + assert opts.denoise_pf is True + + +def test_spectral_denoise_choice(tmp_path): + from producer import config as cfgmod + + opts = Options() + _apply_args(opts, build_parser().parse_args(["in.wav", "--denoise", "spectral"])) + assert opts.denoise == "spectral" + + cfg = tmp_path / "config.toml" + cfg.write_text('[denoise]\nengine = "spectral"\nstrength = 0.8\n') + opts = Options() + cfgmod.apply_config(opts, cfgmod.load_config(cfg)) + assert opts.denoise == "spectral" + assert opts.denoise_strength == 0.8 + + def test_conflict_auto_renames_when_not_a_tty(tmp_path): inp = _mk_wav(tmp_path, "in.wav") first = process_one(inp, _opts(), single=True) @@ -228,3 +297,61 @@ def test_conflict_prompt_eof_cancels(tmp_path, monkeypatch): monkeypatch.setattr("builtins.input", _eof) assert process_one(inp, _opts(), single=True) is None + + +def test_process_one_reports_stage_status(tmp_path, capsys): + inp = _mk_wav(tmp_path, "status.wav") + out_path = process_one(inp, _opts(), single=True) + assert out_path is not None + out = capsys.readouterr().out + assert f"[producer] {inp} (" in out + assert "dsp done in" in out + assert "levelling done in" in out + assert f"[producer] wrote {out_path}" in out + + +def test_multi_file_run_gets_position_prefixes(tmp_path, capsys): + from producer.cli import main + + a = _mk_wav(tmp_path, "a.wav") + b = _mk_wav(tmp_path, "b.wav") + rc = main([str(a), str(b), "--denoise", "off"]) + assert rc == 0 + out = capsys.readouterr().out + assert "[1/2]" in out and "[2/2]" in out + assert (tmp_path / "a_processed.wav").exists() + assert (tmp_path / "b_processed.wav").exists() + + +def test_output_file_rejected_for_multiple_inputs(tmp_path): + from producer.cli import main + + a = _mk_wav(tmp_path, "a.wav") + b = _mk_wav(tmp_path, "b.wav") + with pytest.raises(SystemExit): + main([str(a), str(b), "-o", str(tmp_path / "out.wav"), "--denoise", "off"]) + + +def test_output_dir_accepted_for_multiple_inputs(tmp_path): + from producer.cli import main + + a = _mk_wav(tmp_path, "a.wav") + b = _mk_wav(tmp_path, "b.wav") + outdir = tmp_path / "masters" + rc = main([str(a), str(b), "-o", str(outdir), "--denoise", "off"]) + assert rc == 0 + assert (outdir / "a_processed.wav").exists() + assert (outdir / "b_processed.wav").exists() + + +def test_batch_failure_continues_to_next_file(tmp_path, capsys): + from producer.cli import main + + good = _mk_wav(tmp_path, "good.wav") + missing = tmp_path / "missing.wav" + rc = main([str(missing), str(good), "--denoise", "off"]) + assert rc == 1 + captured = capsys.readouterr() + assert "[2/2]" in captured.out + assert "[producer] ERROR" in captured.err + assert (tmp_path / "good_processed.wav").exists() diff --git a/lib/tests/test_engines.py b/lib/tests/test_engines.py index 781b46e..9881859 100644 --- a/lib/tests/test_engines.py +++ b/lib/tests/test_engines.py @@ -12,6 +12,365 @@ AUDIO_META_MSG = ( ) +class _FakeTensor: + def __init__(self, arr): + self.arr = np.asarray(arr) + + def unsqueeze(self, axis): + return _FakeTensor(self.arr[None, ...]) + + def detach(self): + return self + + def cpu(self): + return self + + def numpy(self): + return self.arr + + +def _install_torch_stub(monkeypatch): + fake = types.ModuleType("torch") + fake.Tensor = _FakeTensor + fake.from_numpy = lambda a: _FakeTensor(a) + fake.cuda = types.SimpleNamespace(is_available=lambda: False) + monkeypatch.setitem(sys.modules, "torch", fake) + + +def _install_df_enhance_stub(monkeypatch, transform, calls, records=None): + pkg = types.ModuleType("df") + + class _State: + def sr(self): + return 44100 + + class _Model: + def to(self, dev): + return self + + def init_df(*_args, **kwargs): + if records is not None: + records["init_df"] = kwargs + return _Model(), _State(), "fake" + + def enhance(_model, _state, audio, **kwargs): + calls.append(audio.arr.shape[-1]) + if records is not None: + records["enhance"] = kwargs + return _FakeTensor(transform(audio.arr.copy())) + + pkg.enhance = types.ModuleType("df.enhance") + pkg.enhance.init_df = init_df + pkg.enhance.enhance = enhance + for sub in ("df.io", "df.logger", "df.utils"): + mod = types.ModuleType(sub) + setattr(pkg, sub.split(".")[1], mod) + monkeypatch.setitem(sys.modules, sub, mod) + monkeypatch.setitem(sys.modules, "df", pkg) + monkeypatch.setitem(sys.modules, "df.enhance", pkg.enhance) + + +def test_dfn3_chunked_engine_stitches_full_length(monkeypatch, sr, noisy_speech): + from producer.engines import denoise_dfn + + monkeypatch.setattr("producer.lazy.ensure", lambda *_a, **_k: None) + monkeypatch.setattr("producer.lazy.ensure_torch", lambda: None) + _install_torch_stub(monkeypatch) + calls: list[int] = [] + _install_df_enhance_stub(monkeypatch, lambda a: a * 0.5 + 0.001, calls) + + x = noisy_speech[: sr * 4] + y, eng, dev = denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=1.0, overlap_s=0.1) + assert "dfn" in eng + assert "cpu" in dev + assert len(calls) > 2 + assert y.shape == x.shape + np.testing.assert_allclose(y, x * 0.5 + 0.001, atol=1e-6) + + +def test_dfn3_whole_file_mode_single_call(monkeypatch, sr, noisy_speech): + from producer.engines import denoise_dfn + + monkeypatch.setattr("producer.lazy.ensure", lambda *_a, **_k: None) + monkeypatch.setattr("producer.lazy.ensure_torch", lambda: None) + _install_torch_stub(monkeypatch) + calls: list[int] = [] + _install_df_enhance_stub(monkeypatch, lambda a: a * 0.5, calls) + + x = noisy_speech[: sr * 2] + y, _eng, _dev = denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=0.0) + assert calls == [x.size] + np.testing.assert_allclose(y, x * 0.5, atol=1e-7) + + +def test_dfn3_post_filter_opt_in_no_atten_lim(monkeypatch, sr, noisy_speech): + from producer.engines import denoise_dfn + + monkeypatch.setattr("producer.lazy.ensure", lambda *_a, **_k: None) + monkeypatch.setattr("producer.lazy.ensure_torch", lambda: None) + _install_torch_stub(monkeypatch) + calls: list[int] = [] + records: dict = {} + _install_df_enhance_stub(monkeypatch, lambda a: a * 0.5, calls, records) + + x = noisy_speech[: sr * 2] + denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=0.0) + assert records["init_df"]["post_filter"] is False + assert records["enhance"] == {} # stock df_enhance call, no atten-lim override + + calls.clear() + records.clear() + denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=0.0, post_filter=True) + assert records["init_df"]["post_filter"] is True + + +def test_dfn3_chunked_feeds_context_padding(monkeypatch, sr, noisy_speech): + from producer.engines import denoise_dfn + + monkeypatch.setattr("producer.lazy.ensure", lambda *_a, **_k: None) + monkeypatch.setattr("producer.lazy.ensure_torch", lambda: None) + _install_torch_stub(monkeypatch) + calls: list[int] = [] + _install_df_enhance_stub(monkeypatch, lambda a: a * 0.5, calls) + + x = noisy_speech[: sr * 10] + y, _eng, _dev = denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=3.0, overlap_s=0.5) + # chunks are widened with context, then trimmed back to the spans + assert max(calls) > 3.0 * sr + assert min(calls) >= 3.0 * sr + np.testing.assert_allclose(y, x * 0.5, atol=1e-7) + + +_Z_MODEL_REPO = "iic/speech_zipenhancer_ans_multiloss_16k_base" + + +def _install_zipenhancer_stub(monkeypatch, calls): + fake = types.ModuleType("zipenhancer") + fake.MODEL_ZIPENHANCER = _Z_MODEL_REPO + + def denoise(chunk, sample_rate, model=_Z_MODEL_REPO, normalize=True, strength=1.0, **_kw): + calls.append( + {"n": chunk.size, "model": model, "normalize": normalize, "strength": strength} + ) + scale = 0.1 if len(calls) % 2 == 1 else 1.0 + return (chunk * scale, 0.0, chunk.size / sample_rate) + + fake.denoise = denoise + monkeypatch.setitem(sys.modules, "zipenhancer", fake) + + +def test_zipenhancer_chunked_normalizes_once(monkeypatch, sr): + from producer.engines import denoise_zip + + seq: list[str] = [] + monkeypatch.setattr( + "producer.lazy.ensure", + lambda pkgs, **_k: seq.append("ensure:" + ",".join(str(p) for p in pkgs)), + ) + monkeypatch.setattr("producer.lazy.ensure_torch", lambda: seq.append("torch")) + monkeypatch.setattr( + "producer.lazy.ensure_import", lambda mod, **_k: seq.append("import:" + mod) + ) + monkeypatch.setattr("producer.lazy.ensure_call", lambda fn, **_k: (seq.append("call"), fn())[1]) + calls: list[dict] = [] + _install_zipenhancer_stub(monkeypatch, calls) + + n = 4 * 44100 + t = np.arange(n) / sr + x = (0.5 * np.sin(2 * np.pi * 160.0 * t)).astype(np.float32) + y, eng, _dev = denoise_zip.denoise(x, sr, 1.0, "cpu", chunk_s=1.0, overlap_s=0.0) + assert "zipenhancer" in eng + # torch is pinned before package installs; the undeclared modelscope + # import ships with the engine; import and call probes run after both + assert seq == [ + "torch", + "ensure:zipenhancer==0.3.2,modelscope", + "import:zipenhancer", + "call", + "call", + "call", + "call", + ] + assert len(calls) == 4 + # the library API takes the full modelscope repo id; the short name + # would be treated as a repo id and fail with modelscope E3021 + assert all(c["model"] == _Z_MODEL_REPO for c in calls) + assert all(c["normalize"] is False for c in calls) + assert y.shape == x.shape + peak = float(np.max(np.abs(y))) + assert abs(peak - 10 ** (-3.0 / 20.0)) < 0.01 + even_rms = float(np.sqrt(np.mean(y[: n // 4].astype(np.float64) ** 2))) + odd_rms = float(np.sqrt(np.mean(y[n // 4 : n // 2].astype(np.float64) ** 2))) + assert 8.0 < odd_rms / even_rms < 12.0 + + +def test_zipenhancer_resample_roundtrip_length_realigned(monkeypatch): + # 120s @ 48k round-tripped through the 16k engine can come back a sample + # or two long (resample_poly emits ceil(n * up/down) per hop); the blend + # used to crash on the mismatch instead of realigning + from producer.engines import denoise_zip + + monkeypatch.setattr("producer.lazy.ensure_torch", lambda: None) + monkeypatch.setattr("producer.lazy.ensure", lambda *_a, **_k: None) + calls: list[dict] = [] + _install_zipenhancer_stub(monkeypatch, calls) + + sr = 48000 + n = 100001 # 48k -> 16k -> 48k drifts +1 for this length + x = (0.3 * np.sin(2 * np.pi * 220.0 * np.arange(n) / sr)).astype(np.float32) + y, _eng, _dev = denoise_zip.denoise(x, sr, 0.5, "cpu") + assert y.shape == x.shape + + +def test_mossformer_clearvoice_probe(monkeypatch, sr, noisy_speech): + from producer.engines import enhance_mossformer + + seq: list[str] = [] + monkeypatch.setattr("producer.lazy.ensure_torch", lambda: seq.append("torch")) + monkeypatch.setattr( + "producer.lazy.ensure", + lambda pkgs, **_k: seq.append("ensure:" + ",".join(str(p) for p in pkgs)), + ) + monkeypatch.setattr( + "producer.lazy.ensure_import", lambda mod, **_k: seq.append("import:" + mod) + ) + + fake = types.ModuleType("clearvoice") + + class _FakeCV: + def __init__(self, task=None, model_name=None): + pass + + def __call__(self, chunk): + return (chunk * 0.5, 48000) + + fake.ClearVoice = _FakeCV + monkeypatch.setitem(sys.modules, "clearvoice", fake) + + x = noisy_speech[: sr * 4] + y, eng, _dev = enhance_mossformer.enhance(x, 48000, 1.0, "cpu", chunk_s=2.0) + assert "mossformer2" in eng + assert seq == ["torch", "ensure:clearvoice==0.1.2", "import:clearvoice"] + assert y.shape == x.shape + np.testing.assert_allclose(y, x * 0.5, atol=1e-6) + + +def _hissy_speech(sr, dur=12.0, noise_db=-42.0, seed=3): + from conftest import speechish + + x = speechish(dur, sr, level_dbfs=-20.0, seed=seed) + rng = np.random.default_rng(seed) + noise = rng.standard_normal(x.size) + noise *= (10 ** (noise_db / 20.0)) / np.sqrt(np.mean(np.square(noise))) + return (x + noise).astype(np.float32) + + +def test_spectral_reduces_steady_noise(sr): + from producer import meters + from producer.engines import denoise_spectral + + x = _hissy_speech(sr) + y, eng, dev = denoise_spectral.denoise(x, sr, 1.0, "cpu") + assert "spectral" in eng and dev == "cpu" + assert y.shape == x.shape + assert np.all(np.isfinite(y)) + assert meters.noise_floor_db(y, sr) < meters.noise_floor_db(x, sr) - 8.0 + assert np.corrcoef(x.astype(np.float64), y.astype(np.float64))[0, 1] > 0.8 + + +def test_spectral_speech_level_flat(sr): + from producer.engines import denoise_spectral + + x = _hissy_speech(sr) + y, _eng, _dev = denoise_spectral.denoise(x, sr, 1.0, "cpu") + frame = int(0.03 * sr) + nf = x.size // frame + + def frms_db(z): + rms = np.sqrt(np.mean(z[: nf * frame].reshape(nf, frame).astype(np.float64) ** 2, axis=1)) + return 20.0 * np.log10(rms + 1e-12) + + xdb, ydb = frms_db(x), frms_db(y) + speech = xdb > np.percentile(xdb, 10) + 12.0 + assert speech.sum() > 20 + # the deterministic engine must not pump the speech level (dfn3's failure mode) + swing = np.abs(ydb[speech] - xdb[speech]) + assert np.percentile(swing, 95) < 2.0 + + +def test_spectral_strength_zero_is_identity(sr): + from producer.engines import denoise_spectral + + x = _hissy_speech(sr) + y, _eng, _dev = denoise_spectral.denoise(x, sr, 0.0, "cpu") + np.testing.assert_array_equal(y, x) + + +def test_spectral_sample_aligned(sr): + from producer.engines import denoise_spectral + + x = _hissy_speech(sr) + y, _eng, _dev = denoise_spectral.denoise(x, sr, 1.0, "cpu") + lo, hi = int(sr * 2.0), int(sr * 9.0) + corrs = {k: float(np.corrcoef(x[lo:hi], y[lo + k : hi + k])[0, 1]) for k in range(-3, 4)} + assert max(corrs, key=corrs.get) == 0 + assert corrs[0] > 0.8 + + +def test_spectral_suppression_capped(sr): + from producer import meters + from producer.engines import denoise_spectral + + noise = (0.008 * np.random.default_rng(5).standard_normal(sr * 6)).astype(np.float32) + y, _eng, _dev = denoise_spectral.denoise(noise, sr, 1.0, "cpu") + drop = meters.noise_floor_db(noise, sr) - meters.noise_floor_db(y, sr) + # bounded: deep enough to matter, never gated to digital silence + assert 20.0 < drop < 40.0 + + +def test_spectral_chunked_matches_whole_file(sr): + from producer.engines import denoise_spectral + + x = _hissy_speech(sr, dur=20.0) + y1, _e, _d = denoise_spectral.denoise(x, sr, 1.0, "cpu", chunk_s=60.0, overlap_s=0.5) + y2, _e, _d = denoise_spectral.denoise(x, sr, 1.0, "cpu", chunk_s=8.0, overlap_s=0.5) + assert y1.shape == y2.shape == x.shape + corr = np.corrcoef(y1.astype(np.float64), y2.astype(np.float64))[0, 1] + assert corr > 0.99 + + +def test_spectral_profile_global_not_per_chunk(sr): + """Tail hiss must get the same suppression with or without speech up front. + + The old per-chunk percentile leaked speech into the noise estimate and + under-suppressed exactly where it matters (between sentences). + """ + from conftest import speechish + + from producer import meters + from producer.engines import denoise_spectral + + speech = speechish(8.0, sr, level_dbfs=-20.0, seed=11) + rng = np.random.default_rng(9) + noise = rng.standard_normal(sr * 16).astype(np.float32) + noise *= (10 ** (-40.0 / 20.0)) / np.sqrt(np.mean(np.square(noise))) + tail_lo, tail_hi = sr * 10, sr * 16 + with_speech = np.concatenate([speech + noise[: speech.size], noise[speech.size :]]).astype( + np.float32 + ) + hiss_only = noise.copy() + y1, _e, _d = denoise_spectral.denoise(with_speech, sr, 1.0, "cpu") + y2, _e, _d = denoise_spectral.denoise(hiss_only, sr, 1.0, "cpu") + drop1 = meters.noise_floor_db(with_speech[tail_lo:tail_hi], sr) - meters.noise_floor_db( + y1[tail_lo:tail_hi], sr + ) + drop2 = meters.noise_floor_db(hiss_only[tail_lo:tail_hi], sr) - meters.noise_floor_db( + y2[tail_lo:tail_hi], sr + ) + assert drop1 > 14.0 + assert drop1 > drop2 - 4.0 + + def _stub_df_modules(calls: list[str]) -> tuple[types.ModuleType, list[types.ModuleType]]: pkg = types.ModuleType("df") mods = [] @@ -84,6 +443,43 @@ def test_dfn3_reduces_noise(sr, noisy_speech): @pytest.mark.slow +def test_dfn3_speech_level_flat(sr, noisy_speech): + """Guard against dfn3's reported failure mode: volume wobble in sentences.""" + pytest.importorskip("torch") + pytest.importorskip("df") + from producer.engines import denoise_dfn + + x = noisy_speech[: sr * 8] + y, _eng, _dev = denoise_dfn.denoise(x, sr, 0.9, "cpu") + frame = int(0.03 * sr) + nf = x.size // frame + + def frms_db(z): + rms = np.sqrt(np.mean(z[: nf * frame].reshape(nf, frame).astype(np.float64) ** 2, axis=1)) + return 20.0 * np.log10(rms + 1e-12) + + xdb, ydb = frms_db(x), frms_db(y) + speech = xdb > np.percentile(xdb, 10) + 12.0 + swing = np.abs(ydb[speech] - xdb[speech]) + assert np.percentile(swing, 95) < 2.5 + + +@pytest.mark.slow +def test_dfn3_chunked_matches_whole_file(sr, noisy_speech): + pytest.importorskip("torch") + pytest.importorskip("df") + from producer.engines import denoise_dfn + + x = noisy_speech[: sr * 60] + y_full, _, _ = denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=0.0) + y_chunk, _, _ = denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=15.0, overlap_s=0.5) + assert y_full.shape == y_chunk.shape == x.shape + corr = np.corrcoef(y_full.astype(np.float64), y_chunk.astype(np.float64))[0, 1] + assert corr > 0.99 + assert float(np.max(np.abs(y_full - y_chunk))) < 0.1 + + +@pytest.mark.slow def test_zipenhancer_reduces_noise(sr, noisy_speech): pytest.importorskip("torch") pytest.importorskip("zipenhancer") diff --git a/lib/tests/test_lazy.py b/lib/tests/test_lazy.py new file mode 100644 index 0000000..5ab5b95 --- /dev/null +++ b/lib/tests/test_lazy.py @@ -0,0 +1,286 @@ +import types + +import pytest + +from producer import lazy + +PYPI_JSON = { + "urls": [ + { + "filename": "torch-2.7.1-cp311-cp311-manylinux_2_28_x86_64.whl", + "url": "https://files.pythonhosted.org/packages/xx/torch-2.7.1-cp311-cp311-manylinux_2_28_x86_64.whl", + "digests": {"sha256": "abc123"}, + "size": 766_668_798, + }, + { + "filename": "torch-2.7.1-cp311-cp311-win_amd64.whl", + "url": "https://example.com/win.whl", + "digests": {}, + "size": 1, + }, + { + "filename": "torch-2.7.1.tar.gz", + "url": "https://example.com/src.tar.gz", + "digests": {}, + "size": 1, + }, + ] +} + +_SHA = "e1a846516570851234567890abcdef1234567890abcdef1234567890abcdef12" +GPU_HTML = ( + "<html><body>" + '<a href="https://download-r2.pytorch.org/whl/cu126/torch-2.6.0%2Bcu126-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=00' + f'00">old</a><a href="https://download-r2.pytorch.org/whl/cu126/torch-2.7.1%2Bcu126-cp311-cp311-manylinux_2_28_x86_64.whl#sha256={_SHA}">new</a>' + '<a href="https://download-r2.pytorch.org/whl/cu126/torch-2.7.1%2Bcu126-cp311-cp311-win_amd64.whl#sha256=1111">win</a>' + '<a href="torch-2.7.1%2Bcu126-cp311-cp311-manylinux_2_28_x86_64.whl">relative-dup</a>' + "</body></html>" +) + + +def test_choose_pypi_picks_linux_wheel(): + got = lazy._choose_pypi(PYPI_JSON) + assert got is not None + name, url, sha, size = got + assert name.endswith("manylinux_2_28_x86_64.whl") + assert "pythonhosted" in url + assert sha == "abc123" + assert size == 766_668_798 + + +def test_choose_gpu_matches_version_and_arch(): + got = lazy._choose_gpu(GPU_HTML, "torch", "2.7.1") + assert got is not None + name, url, sha, size = got + assert name == "torch-2.7.1+cu126-cp311-cp311-manylinux_2_28_x86_64.whl" + assert url == ( + "https://download-r2.pytorch.org/whl/cu126/" + "torch-2.7.1%2Bcu126-cp311-cp311-manylinux_2_28_x86_64.whl" + ) + assert sha == _SHA + assert size is None + + +def test_choose_gpu_no_match_returns_none(): + assert lazy._choose_gpu("<html></html>", "torch", "2.7.1") is None + + +def test_resolve_torch_wheels_needs_both_packages(monkeypatch): + def fake(pkg, ver): + return None if pkg == "torchaudio" else ("t.whl", "u", None, 1) + + monkeypatch.setattr(lazy, "_pypi_wheel", fake) + assert lazy._resolve_torch_wheels(gpu=False) == [] + + +def test_resolve_torch_wheels_returns_both(monkeypatch): + wheels = [("torch.whl", "u1", "s1", 1), ("torchaudio.whl", "u2", None, None)] + monkeypatch.setattr(lazy, "_pypi_wheel", lambda pkg, ver: wheels.pop(0)) + got = lazy._resolve_torch_wheels(gpu=False) + assert [w[0] for w in got] == ["torch.whl", "torchaudio.whl"] + + +def test_ensure_torch_downloads_wheels_then_installs(monkeypatch, tmp_path): + wheels = [ + ("torch-2.7.1-cp311.whl", "https://x/torch.whl", "sha", 10), + ("torchaudio-2.7.1-cp311.whl", "https://x/ta.whl", None, 5), + ] + monkeypatch.setattr(lazy, "has_module", lambda name: False) + monkeypatch.setattr(lazy, "_resolve_torch_wheels", lambda gpu: wheels) + monkeypatch.setattr(lazy, "find_uv", lambda: "uv") + monkeypatch.setattr(lazy, "WHEEL_CACHE", tmp_path / "wheels") + downloads: list[str] = [] + runs: list[tuple[list[str], str]] = [] + + def fake_download(url, dest, label=None, expected_size=None, sha256=None, timeout=60.0): + downloads.append(dest.name) + return dest + + monkeypatch.setattr(lazy.ui, "download", fake_download) + monkeypatch.setattr(lazy.ui, "run", lambda cmd, label, check=True: runs.append((cmd, label))) + lazy.ensure_torch() + assert downloads == ["torch-2.7.1-cp311.whl", "torchaudio-2.7.1-cp311.whl"] + assert len(runs) == 1 + cmd = runs[0][0] + assert cmd[0] == "uv" and cmd[1] == "pip" and cmd[2] == "install" + assert str(tmp_path / "wheels" / "torch-2.7.1-cp311.whl") in cmd + assert str(tmp_path / "wheels" / "torchaudio-2.7.1-cp311.whl") in cmd + + +def test_ensure_torch_falls_back_to_uv_index(monkeypatch): + monkeypatch.setattr(lazy, "has_module", lambda name: False) + monkeypatch.setattr(lazy, "_resolve_torch_wheels", lambda gpu: []) + monkeypatch.setattr(lazy, "find_uv", lambda: "uv") + monkeypatch.setattr(lazy, "gpu_present", lambda: True) + runs: list[list[str]] = [] + monkeypatch.setattr(lazy.ui, "run", lambda cmd, label, check=True: runs.append(cmd)) + lazy.ensure_torch() + cmd = runs[0] + assert "torch==2.7.1+cu126" in cmd + assert "--index-url" in cmd + assert lazy.TORCH_GPU_INDEX in cmd + + +def test_ensure_torch_falls_back_to_pypi_cpu(monkeypatch): + monkeypatch.setattr(lazy, "has_module", lambda name: False) + monkeypatch.setattr(lazy, "_resolve_torch_wheels", lambda gpu: []) + monkeypatch.setattr(lazy, "find_uv", lambda: "uv") + monkeypatch.setattr(lazy, "gpu_present", lambda: False) + runs: list[list[str]] = [] + monkeypatch.setattr(lazy.ui, "run", lambda cmd, label, check=True: runs.append(cmd)) + lazy.ensure_torch() + cmd = runs[0] + assert "torch==2.7.1" in cmd + assert "--index-url" not in cmd + + +def test_ensure_torch_noop_when_installed(monkeypatch): + monkeypatch.setattr(lazy, "has_module", lambda name: True) + runs: list = [] + monkeypatch.setattr(lazy.ui, "run", lambda cmd, label, check=True: runs.append(cmd)) + lazy.ensure_torch() + assert runs == [] + + +def test_ensure_import_noop_when_importable(monkeypatch): + installs: list[list[str]] = [] + monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) + lazy.ensure_import("numpy", purpose="test") + assert installs == [] + + +def test_ensure_import_installs_missing_module(monkeypatch): + installs: list[list[str]] = [] + monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) + state = {"round": 0} + + def fake_import(name): + if name == "fakeengine": + if state["round"] == 0: + state["round"] = 1 + raise ModuleNotFoundError("No module named 'addict'", name="addict") + return types.ModuleType("fakeengine") + raise ModuleNotFoundError(f"No module named {name!r}", name=name) + + monkeypatch.setattr(lazy.importlib, "import_module", fake_import) + lazy.ensure_import("fakeengine", purpose="test") + assert installs == [["addict"]] + + +def test_ensure_import_dotted_missing_reraises_without_install(monkeypatch): + installs: list[list[str]] = [] + monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) + + def fake_import(name): + raise ModuleNotFoundError("No module named 'pkg.sub'", name="pkg.sub") + + monkeypatch.setattr(lazy.importlib, "import_module", fake_import) + with pytest.raises(ModuleNotFoundError): + lazy.ensure_import("whatever", purpose="test") + assert installs == [] + + +def test_ensure_import_gives_up_after_rounds(monkeypatch): + installs: list[list[str]] = [] + monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) + + def fake_import(name): + raise ModuleNotFoundError("No module named 'ghost'", name="ghost") + + monkeypatch.setattr(lazy.importlib, "import_module", fake_import) + with pytest.raises(lazy.EngineUnavailable, match="ghost"): + lazy.ensure_import("whatever", purpose="test") + assert len(installs) == 4 + assert all(pkgs == ["ghost"] for pkgs in installs) + + +def test_ensure_import_maps_pil_to_pillow(monkeypatch): + installs: list[list[str]] = [] + monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) + state = {"round": 0} + + def fake_import(name): + if name == "fakeengine": + if state["round"] == 0: + state["round"] = 1 + raise ModuleNotFoundError("No module named 'PIL'", name="PIL") + return types.ModuleType("fakeengine") + raise ModuleNotFoundError(f"No module named {name!r}", name=name) + + monkeypatch.setattr(lazy.importlib, "import_module", fake_import) + lazy.ensure_import("fakeengine", purpose="test") + assert installs == [["pillow"]] + + +def test_ensure_call_passthrough_without_install(monkeypatch): + installs: list[list[str]] = [] + monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) + assert lazy.ensure_call(lambda: 42, purpose="test") == 42 + assert installs == [] + + +def test_ensure_call_installs_missing_module(monkeypatch): + installs: list[list[str]] = [] + monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) + state = {"round": 0} + + def fake_fn(): + if state["round"] == 0: + state["round"] = 1 + raise ModuleNotFoundError("No module named 'addict'", name="addict") + return "ok" + + assert lazy.ensure_call(fake_fn, purpose="test") == "ok" + assert installs == [["addict"]] + + +def test_ensure_call_dotted_missing_reraises_without_install(monkeypatch): + installs: list[list[str]] = [] + monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) + + def fake_fn(): + raise ModuleNotFoundError("No module named 'pkg.sub'", name="pkg.sub") + + with pytest.raises(ModuleNotFoundError): + lazy.ensure_call(fake_fn, purpose="test") + assert installs == [] + + +def test_ensure_call_unnamed_missing_reraises_without_install(monkeypatch): + installs: list[list[str]] = [] + monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) + + def fake_fn(): + raise ModuleNotFoundError("boom") + + with pytest.raises(ModuleNotFoundError): + lazy.ensure_call(fake_fn, purpose="test") + assert installs == [] + + +def test_ensure_call_maps_pil_to_pillow(monkeypatch): + installs: list[list[str]] = [] + monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) + state = {"round": 0} + + def fake_fn(): + if state["round"] == 0: + state["round"] = 1 + raise ModuleNotFoundError("No module named 'PIL'", name="PIL") + return "ok" + + assert lazy.ensure_call(fake_fn, purpose="test") == "ok" + assert installs == [["pillow"]] + + +def test_ensure_call_gives_up_after_rounds(monkeypatch): + installs: list[list[str]] = [] + monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) + + def fake_fn(): + raise ModuleNotFoundError("No module named 'ghost'", name="ghost") + + with pytest.raises(lazy.EngineUnavailable, match="ghost"): + lazy.ensure_call(fake_fn, purpose="test") + assert len(installs) == 8 + assert all(pkgs == ["ghost"] for pkgs in installs) diff --git a/lib/tests/test_meters.py b/lib/tests/test_meters.py index 73f3db6..b013ec9 100644 --- a/lib/tests/test_meters.py +++ b/lib/tests/test_meters.py @@ -17,6 +17,20 @@ def test_true_peak_bounds(sr): assert sp - 0.01 <= tp <= sp + 0.6 +def test_true_peak_blockwise_matches_whole_file(sr): + from scipy import signal + + rng = np.random.default_rng(3) + n = sr * 75 + x = ( + 0.5 * np.sin(2 * np.pi * 997.0 * np.arange(n) / sr) + 0.02 * rng.standard_normal(n) + ).astype(np.float32) + half = 10 * 4 + h = signal.firwin(2 * half + 1, 0.25, window=("kaiser", 8.0)).astype(np.float32) + ref = meters.sample_peak_db(signal.resample_poly(x, 4, 1, window=h)) + assert abs(meters.true_peak_db(x, sr) - ref) < 0.01 + + def test_lufs_of_sine(sr): x = sine(1000, 3.0, sr, peak_dbfs=-17.0) assert abs(meters.lufs(x, sr) + 20.05) < 0.2 @@ -32,6 +46,13 @@ def test_noise_floor_below_speech(sr): assert floor < meters.rms_db(x) - 4.0 +def test_speech_level_db_ignores_leading_silence(sr): + x = speechish(6.0, sr, level_dbfs=-20.0) + y = np.concatenate([np.zeros(sr * 4, dtype=np.float32), x]) + assert abs(meters.speech_level_db(x, sr) - meters.speech_level_db(y, sr)) < 1.0 + assert meters.speech_level_db(x, sr) > meters.rms_db(x) + + def test_all_meters_keys(sr): d = meters.all_meters(speechish(3.0, sr), sr) assert set(d) == { diff --git a/lib/tests/test_pipeline.py b/lib/tests/test_pipeline.py index da2e9da..b3f91f1 100644 --- a/lib/tests/test_pipeline.py +++ b/lib/tests/test_pipeline.py @@ -19,6 +19,16 @@ def test_stage_order_and_names(): assert by_name["levelling"].enabled is True +def test_denoise_stage_detail_shows_pf(): + opts = Options() + stages = pipeline.build_stages(opts) + assert "pf=off" in stages[0].detail + assert "leveling" not in stages[0].detail + opts.denoise_pf = True + stages = pipeline.build_stages(opts) + assert "pf=on" in stages[0].detail + + def test_full_chain_profile_bounds(sr): opts = Options() opts.denoise = "off" @@ -68,3 +78,37 @@ def test_podcast_profile_bounds(sr): res = pipeline.run_pipeline(x, sr, opts) assert abs(meters.lufs(res.audio, sr) + 16.0) < 0.8 assert meters.true_peak_db(res.audio, sr) <= -1.4 + + +def test_dsp_pregain_makes_chain_level_invariant(sr): + """The voice chain must treat a quiet and a hot take identically. + + Denoised files often arrive far below the chain's design level; without + the pre-gain the absolute comp thresholds would idle (or slam) depending + on input level alone. + """ + hot = Options() + hot.denoise = "off" + hot.enhance = "off" + hot.levelling = False + quiet = Options() + quiet.denoise = "off" + quiet.enhance = "off" + quiet.levelling = False + x_hot = speechish(6.0, sr, level_dbfs=-14.0) + x_quiet = speechish(6.0, sr, level_dbfs=-38.0) + y_hot = pipeline.run_pipeline(x_hot, sr, hot).audio + y_quiet = pipeline.run_pipeline(x_quiet, sr, quiet).audio + diff = abs(meters.rms_db(y_hot) - meters.rms_db(y_quiet)) + assert diff < 1.0 + + +def test_dsp_pregain_note_recorded(sr): + opts = Options() + opts.denoise = "off" + opts.enhance = "off" + notes: list[str] = [] + stages = pipeline.build_stages(opts, notes) + dsp_stage = next(s for s in stages if s.name == "dsp") + dsp_stage.fn(speechish(4.0, sr, level_dbfs=-30.0), sr) + assert any("pregain" in n for n in notes) diff --git a/lib/tests/test_ui.py b/lib/tests/test_ui.py new file mode 100644 index 0000000..156b3e8 --- /dev/null +++ b/lib/tests/test_ui.py @@ -0,0 +1,169 @@ +import hashlib +import subprocess +import sys + +import pytest + +from producer import ui + + +class _FakeResp: + def __init__(self, payload: bytes, length: str | None = None): + self._payload = payload + self.headers = {"Content-Length": length} if length else {} + + def read(self, n=-1): + if not self._payload: + return b"" + out, self._payload = self._payload[:n], self._payload[n:] + return out + + def __enter__(self): + return self + + def __exit__(self, *_a): + return False + + +def test_fmt_bytes(): + assert ui.fmt_bytes(512) == "512 B" + assert ui.fmt_bytes(2048) == "2.0 KB" + assert ui.fmt_bytes(8 << 20) == "8.0 MB" + assert ui.fmt_bytes(None) == "?" + + +def test_fmt_secs(): + assert ui.fmt_secs(0) == "0:00" + assert ui.fmt_secs(59) == "0:59" + assert ui.fmt_secs(61) == "1:01" + assert ui.fmt_secs(3700) == "1:01:40" + assert ui.fmt_secs(None) == "?" + + +def test_bar_fills(): + assert ui._bar(0.0) == "[" + "-" * 22 + "]" + assert ui._bar(1.0) == "[" + "#" * 22 + "]" + half = ui._bar(0.5) + assert half.count("#") == 11 and half.count("-") == 11 + + +def test_download_writes_file(tmp_path, monkeypatch, capsys): + payload = b"x" * (2 << 20) + sha = hashlib.sha256(payload).hexdigest() + monkeypatch.setattr( + "urllib.request.urlopen", lambda req, timeout=None: _FakeResp(payload, str(len(payload))) + ) + dest = tmp_path / "big.bin" + out = ui.download("https://example.com/big.bin", dest, sha256=sha) + assert out == dest + assert dest.read_bytes() == payload + assert "done" in capsys.readouterr().out + + +def test_download_checksum_mismatch(tmp_path, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", lambda req, timeout=None: _FakeResp(b"abc")) + dest = tmp_path / "f.bin" + with pytest.raises(RuntimeError): + ui.download("https://example.com/f.bin", dest, sha256="0" * 64) + assert not dest.exists() + assert not dest.with_name(dest.name + ".part").exists() + + +def test_download_cached_skips(tmp_path, monkeypatch, capsys): + dest = tmp_path / "cached.bin" + dest.write_bytes(b"hello") + + def boom(*_a, **_k): + raise AssertionError("should not download") + + monkeypatch.setattr("urllib.request.urlopen", boom) + ui.download("https://example.com/cached.bin", dest) + assert dest.read_bytes() == b"hello" + assert "cached" in capsys.readouterr().out + + +def test_download_size_mismatch_redownloads(tmp_path, monkeypatch): + dest = tmp_path / "short.bin" + dest.write_bytes(b"too short") + monkeypatch.setattr("urllib.request.urlopen", lambda req, timeout=None: _FakeResp(b"abcd")) + ui.download("https://example.com/short.bin", dest, expected_size=4) + assert dest.read_bytes() == b"abcd" + + +def test_download_milestones_non_tty(tmp_path, monkeypatch, capsys): + payload = b"y" * (8 << 20) + monkeypatch.setattr( + "urllib.request.urlopen", lambda req, timeout=None: _FakeResp(payload, str(len(payload))) + ) + ui.download("https://example.com/m.bin", tmp_path / "m.bin") + out = capsys.readouterr().out + assert "8.0 MB/8.0 MB" in out + assert "done" in out + + +def test_run_piped_forwards_output(monkeypatch, capsys): + monkeypatch.setattr(ui, "is_tty", lambda: False) + rc = ui.run([sys.executable, "-c", "print('hello-uv-output')"], "test run") + assert rc == 0 + assert "hello-uv-output" in capsys.readouterr().out + + +def test_run_piped_raises_on_failure(monkeypatch): + monkeypatch.setattr(ui, "is_tty", lambda: False) + with pytest.raises(subprocess.CalledProcessError): + ui.run([sys.executable, "-c", "raise SystemExit(3)"], "test run") + + +def test_status_non_tty_milestones_and_done(capsys): + s = ui.Status(prefix="[producer] in.wav — ") + s.stage("denoise") + for i in range(1, 11): + s.tick(i, 10) + s.stage_done("denoise", 2.5) + s.stage("dsp") + s.stage_done("dsp", 0.3) + s.finish() + out = capsys.readouterr().out + assert "denoise 10/10 chunks (100%)" in out + assert "denoise done in 2.5s (10/10 chunks)" in out + assert "dsp done in 0.3s" in out + + +def test_status_tty_live_line(monkeypatch, capsys): + monkeypatch.setattr(ui, "is_tty", lambda: True) + s = ui.Status(prefix="p — ") + s.stage("denoise") + s.tick(10, 10) # completion redraws immediately (bypasses the 0.1 s throttle) + out = capsys.readouterr().out + assert "denoise" in out and "10/10" in out + s.finish() + + +def test_progress_milestones_non_tty(capsys): + p = ui.Progress("downloading x", total=1000) + for i in (100, 300, 1000): + p.update(i) + p.close() + out = capsys.readouterr().out + assert "1000 B/1000 B" in out + assert "done (1000 B" in out + + +def test_progress_tty_bar(monkeypatch, capsys): + monkeypatch.setattr(ui, "is_tty", lambda: True) + p = ui.Progress("downloading torch", total=100) + p.update(50) + p.close() + out = capsys.readouterr().out + assert "50 B/100 B (50%)" in out + + +def test_log_clears_live_line(monkeypatch, capsys): + monkeypatch.setattr(ui, "is_tty", lambda: True) + p = ui.Progress("downloading", total=100) + p.update(50) + ui.log("[producer] some message") + out = capsys.readouterr().out + assert "some message" in out + assert out.index("some message") > out.index("downloading") + p.close() diff --git a/lib/tests/test_updates.py b/lib/tests/test_updates.py new file mode 100644 index 0000000..7ad1727 --- /dev/null +++ b/lib/tests/test_updates.py @@ -0,0 +1,221 @@ +import subprocess +import sys +from types import SimpleNamespace + +from producer import lazy, updates + + +def test_parse_would_install(): + out = "\n".join( + [ + "Using Python 3.11.16 environment at: /x", + "Resolved 2 packages in 1ms", + "Would install 2 packages", + " + scipy==1.14.1", + " - numpy==1.26.4", + " + numpy==2.2.6", + " ~ cffi==2.1.1", + "", + ] + ) + assert updates._parse_would_install(out) == { + "scipy": "1.14.1", + "numpy": "2.2.6", + "cffi": "2.1.1", + } + + +def test_vkey_ordering(): + assert updates._vkey("1.26.4") < updates._vkey("2.2.6") + assert updates._vkey("0.5.6") < updates._vkey("0.5.7") + assert updates._vkey("2.7.1") < updates._vkey("2.7.1+cu126") + assert updates._vkey("1.2.10") > updates._vkey("1.2.9") + assert not updates._vkey("0.5.6") > updates._vkey("0.5.6") + + +def test_core_names(tmp_path, monkeypatch): + reqs = tmp_path / "requirements-core.txt" + reqs.write_text("numpy==1.26.4\nscipy>=1.11,<1.15\n\n# comment\nsoundfile>=0.12,<0.13\n") + monkeypatch.setattr(updates, "CORE_FILE", reqs) + assert updates._core_names() == ["numpy", "scipy", "soundfile"] + + +def test_probe_filters_to_keep(monkeypatch): + calls = [] + + def fake_run(cmd, capture_output, text, timeout): + calls.append(cmd) + return SimpleNamespace( + returncode=0, + stdout="Resolved 2 packages\n + scipy==1.14.1\n + numpy==2.2.6\n", + ) + + monkeypatch.setattr(lazy, "find_uv", lambda: "uv") + monkeypatch.setattr(updates.subprocess, "run", fake_run) + got = updates._probe(["-U", "-r", "reqs.txt"], None, {"scipy"}) + assert got == {"scipy": "1.14.1"} + assert calls[0][1:4] == ["pip", "install", "--dry-run"] + + +def test_probe_failure_returns_empty(monkeypatch): + def boom(cmd, capture_output, text, timeout): + raise subprocess.TimeoutExpired(cmd, timeout) + + monkeypatch.setattr(lazy, "find_uv", lambda: "uv") + monkeypatch.setattr(updates.subprocess, "run", boom) + assert updates._probe(["torch"], None, {"torch"}) == {} + + +def test_collect_full_matrix(monkeypatch, tmp_path): + installed = { + "numpy": "1.26.4", + "scipy": "1.13.1", + "soundfile": "0.12.1", + "pyloudnorm": "0.2.0", + "torch": "2.7.1+cu126", + "torchaudio": "2.7.1+cu126", + "deepfilternet": "0.5.6", + "zipenhancer": None, + "clearvoice": None, + } + + def fake_probe(args, python, keep): + if "torch" in args and "torchaudio" in args and "--index-url" in args: + return {"torch": "2.9.0+cu126", "torchaudio": "2.9.0+cu126"} + if "deepfilternet" in args: + return {"deepfilternet": "0.5.7", "torch": "2.14.0"} + if "zipenhancer" in args: + return {} + return {"scipy": "1.14.1"} + + monkeypatch.setattr(updates, "_installed", lambda d: installed.get(d)) + monkeypatch.setattr(updates, "_probe", fake_probe) + monkeypatch.setattr(lazy, "gpu_present", lambda: True) + monkeypatch.setattr(lazy, "find_uv", lambda: "uv") + monkeypatch.setattr(updates, "RESEMBLE_PY", tmp_path / "missing" / "python") + + got = updates.collect() + labels = [u.label for u in got] + assert labels == ["scipy", "torch + torchaudio", "deepfilternet"] + + scipy = got[0] + assert (scipy.old, scipy.new) == ("1.13.1", "1.14.1") + assert "-U" in scipy.cmd and str(updates.CORE_FILE) in scipy.cmd + + torch_u = got[1] + assert (torch_u.old, torch_u.new) == ("2.7.1+cu126", "2.9.0+cu126") + assert "CUDA" in torch_u.note + assert "torch==2.9.0+cu126" in torch_u.cmd + assert "torchaudio==2.9.0+cu126" in torch_u.cmd + assert lazy.TORCH_GPU_INDEX in torch_u.cmd + + dfn = got[2] + assert (dfn.old, dfn.new) == ("0.5.6", "0.5.7") + assert "torch==2.7.1" in dfn.cmd + assert "deepfilternet==0.5.7" in dfn.cmd + assert "torch==2.14.0" not in dfn.cmd + + +def test_collect_skips_uninstalled_engines_and_missing_torch(monkeypatch, tmp_path): + monkeypatch.setattr(updates, "_installed", lambda d: None) + monkeypatch.setattr(lazy, "find_uv", lambda: "uv") + monkeypatch.setattr(updates, "RESEMBLE_PY", tmp_path / "missing" / "python") + + probed = [] + monkeypatch.setattr(updates, "_probe", lambda args, python, keep: probed.append(args) or {}) + assert updates.collect() == [] + assert len(probed) == 1 + assert "deepfilternet" not in probed[0] and "torch" not in probed[0] + + +def test_collect_resemble_isolated_venv(monkeypatch, tmp_path): + py = tmp_path / "resemble" / "bin" / "python" + py.parent.mkdir(parents=True) + py.write_text("") + monkeypatch.setattr(updates, "_installed", lambda d: "2.7.1+cu126" if d == "torch" else None) + monkeypatch.setattr(updates, "_installed_in", lambda p, d: "0.0.1") + monkeypatch.setattr(lazy, "find_uv", lambda: "uv") + monkeypatch.setattr(updates, "RESEMBLE_PY", py) + + def fake_probe(args, python, keep): + if "resemble-enhance" in args: + assert python == str(py) + return {"resemble-enhance": "0.0.2"} + return {} + + monkeypatch.setattr(updates, "_probe", fake_probe) + got = updates.collect() + assert len(got) == 1 + u = got[0] + assert u.label == "resemble-enhance" + assert "resemble-enhance==0.0.2" in u.cmd + assert str(py) in u.cmd + + +def test_confirm(monkeypatch): + answers = iter(["", "no", "y", "yes"]) + monkeypatch.setattr("builtins.input", lambda prompt: next(answers)) + assert not updates._confirm("install updates?") + assert not updates._confirm("install updates?") + assert updates._confirm("install updates?") + assert updates._confirm("install updates?") + + +def test_check_and_prompt_non_tty_notices_only(monkeypatch, capsys): + monkeypatch.setattr( + updates, + "collect", + lambda: [updates.Update("scipy", "1.13.1", "1.14.1", "", ["uv"])], + ) + monkeypatch.setattr(sys, "stdin", SimpleNamespace(isatty=lambda: False)) + assert updates.check_and_prompt() is False + out = capsys.readouterr().out + assert "updates available" in out + assert "./producer update" in out + + +def test_check_and_prompt_assume_yes_applies(monkeypatch, capsys): + u = updates.Update("scipy", "1.13.1", "1.14.1", "", ["uv", "pip", "install", "scipy"]) + monkeypatch.setattr(updates, "collect", lambda: [u]) + runs = [] + monkeypatch.setattr(updates.ui, "run", lambda cmd, label, check=True: runs.append(cmd)) + assert updates.check_and_prompt(assume_yes=True) is True + assert runs == [["uv", "pip", "install", "scipy"]] + assert "updates installed" in capsys.readouterr().out + + +def test_check_and_prompt_declined(monkeypatch, capsys): + u = updates.Update("scipy", "1.13.1", "1.14.1", "", ["uv"]) + monkeypatch.setattr(updates, "collect", lambda: [u]) + monkeypatch.setattr(sys, "stdin", SimpleNamespace(isatty=lambda: True)) + monkeypatch.setattr(updates, "_confirm", lambda prompt: False) + runs = [] + monkeypatch.setattr(updates.ui, "run", lambda cmd, label, check=True: runs.append(cmd)) + assert updates.check_and_prompt() is False + assert runs == [] + assert "skipped updates" in capsys.readouterr().out + + +def test_check_and_prompt_up_to_date(monkeypatch, capsys): + monkeypatch.setattr(updates, "collect", lambda: []) + assert updates.check_and_prompt(force=True) is False + assert "up to date" in capsys.readouterr().out + + +def test_update_command_requires_yes_when_not_interactive(monkeypatch): + monkeypatch.setattr(sys, "stdin", SimpleNamespace(isatty=lambda: False)) + monkeypatch.setattr( + updates, + "collect", + lambda: [updates.Update("scipy", "1.13.1", "1.14.1", "", ["uv"])], + ) + try: + updates.run_update_command([]) + except SystemExit as e: + assert "--yes" in str(e) + else: + raise AssertionError("expected SystemExit") + runs = [] + monkeypatch.setattr(updates.ui, "run", lambda cmd, label, check=True: runs.append(cmd)) + assert updates.run_update_command(["--yes"]) == 0 + assert len(runs) == 1 |