aboutsummaryrefslogtreecommitdiff
path: root/lib/src
diff options
context:
space:
mode:
authorhistoria <historiavg@proton.me>2026-09-06 20:55:27 -0400
committerhistoria <historiavg@proton.me>2026-09-06 20:55:27 -0400
commit39b0f2bbed74f6487a41b82501ae3c6799e4b5c4 (patch)
tree620ce0462d029ebca927283c62872e5c18bda818 /lib/src
parent13e15d78830ad61211d2cadb7d3a4dca8a29ab5c (diff)
downloadproducer-39b0f2bbed74f6487a41b82501ae3c6799e4b5c4.tar.gz
feat: chunking, zipenhancer denoising
Diffstat (limited to 'lib/src')
-rw-r--r--lib/src/producer/cli.py95
-rw-r--r--lib/src/producer/config.py34
-rw-r--r--lib/src/producer/doctor.py23
-rw-r--r--lib/src/producer/dsp.py124
-rw-r--r--lib/src/producer/engines/base.py13
-rw-r--r--lib/src/producer/engines/chunking.py166
-rw-r--r--lib/src/producer/engines/denoise_dfn.py51
-rw-r--r--lib/src/producer/engines/denoise_spectral.py186
-rw-r--r--lib/src/producer/engines/denoise_zip.py56
-rw-r--r--lib/src/producer/engines/enhance_mossformer.py58
-rw-r--r--lib/src/producer/engines/enhance_resemble.py41
-rw-r--r--lib/src/producer/engines/resemble_worker.py55
-rw-r--r--lib/src/producer/io.py2
-rw-r--r--lib/src/producer/lazy.py200
-rw-r--r--lib/src/producer/meters.py67
-rw-r--r--lib/src/producer/pipeline.py132
-rw-r--r--lib/src/producer/ui.py319
-rw-r--r--lib/src/producer/updates.py273
18 files changed, 1719 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