diff options
Diffstat (limited to 'lib/src')
23 files changed, 0 insertions, 3229 deletions
diff --git a/lib/src/__main__.py b/lib/src/__main__.py deleted file mode 100644 index cfe0046..0000000 --- a/lib/src/__main__.py +++ /dev/null @@ -1,4 +0,0 @@ -from producer.cli import main - -if __name__ == "__main__": - main() diff --git a/lib/src/producer/__init__.py b/lib/src/producer/__init__.py deleted file mode 100644 index 3dc1f76..0000000 --- a/lib/src/producer/__init__.py +++ /dev/null @@ -1 +0,0 @@ -__version__ = "0.1.0" diff --git a/lib/src/producer/cli.py b/lib/src/producer/cli.py deleted file mode 100644 index 779de96..0000000 --- a/lib/src/producer/cli.py +++ /dev/null @@ -1,321 +0,0 @@ -from __future__ import annotations - -import argparse -import itertools -import os -import sys -from pathlib import Path - -from . import __version__, pipeline, ui -from . import config as cfgmod -from . import dsp as pdsp -from . import io as pio -from . import report as repmod -from .config import DSP_KEYS, PROFILES, Options - -DATA_DIR = Path(__file__).resolve().parents[2] -EXTS = {".wav", ".flac", ".mp3", ".m4a", ".aac", ".ogg", ".opus", ".aif", ".aiff", ".wma"} - - -def build_parser() -> argparse.ArgumentParser: - p = argparse.ArgumentParser( - prog="producer", - description="One-click narration/podcast mastering: denoise, enhance, voice DSP, loudness.", - ) - p.add_argument("inputs", nargs="*", help="audio file(s) to process") - p.add_argument("doctor", nargs="?", help=argparse.SUPPRESS) - 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", "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") - p.add_argument( - "--no-dsp", action="store_false", dest="dsp", help="skip EQ/compression/de-ess chain" - ) - p.add_argument( - "--no-levelling", action="store_false", dest="levelling", help="skip loudness stage" - ) - for key in DSP_KEYS: - p.add_argument(f"--{key}", type=float, metavar="0-1", help=f"strength for {key} stage") - p.add_argument("--hpf-hz", type=float, help="high-pass corner (default 80)") - p.add_argument("--target", type=float, help="loudness target (RMS dB or LUFS per profile)") - p.add_argument("--ceiling", type=float, dest="ceiling_db", help="true-peak ceiling in dB") - p.add_argument("--sample-rate", type=int, help="output sample rate (44100/48000)") - p.add_argument( - "--bit-depth", type=int, choices=[16, 24, 32], default=None, help="wav/flac bit depth" - ) - p.add_argument( - "--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") - return p - - -def _apply_args(opts: Options, args: argparse.Namespace) -> None: - m = { - "profile": "profile", - "denoise": "denoise", - "denoise_strength": "denoise_strength", - "denoise_pf": "denoise_pf", - "enhance": "enhance", - "enhance_strength": "enhance_strength", - "output": "output", - "hpf_hz": "hpf_hz", - "target": "target", - "ceiling_db": "ceiling_db", - "sample_rate": "sample_rate", - "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", - "verbose": "verbose", - } - for arg_name, field_name in m.items(): - v = getattr(args, arg_name, None) - if v is not None: - setattr(opts, field_name, v) - if args.dsp is False: - opts.dsp = False - if args.levelling is False: - opts.levelling = False - for key in DSP_KEYS: - v = getattr(args, key, None) - if v is not None: - opts.strengths[key] = float(v) - if opts.profile not in PROFILES: - 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: - raise SystemExit(f"--{k} must be within 0-1") - - -def _expand_inputs(args: argparse.Namespace) -> list[Path]: - inputs: list[Path] = [] - for raw in args.inputs: - p = Path(raw) - if args.batch and p.is_dir(): - inputs.extend(sorted(q for q in p.iterdir() if q.suffix.lower() in EXTS)) - elif args.batch and not p.exists(): - import glob - - for q in sorted(glob.glob(raw)): - q = Path(q) - if q.suffix.lower() in EXTS: - inputs.append(q) - else: - inputs.append(p) - return inputs - - -def _resolve_output(inp: Path, opts: Options, single: bool) -> Path: - ext = opts.out_format.lower() - if opts.output: - o = Path(opts.output) - if single: - return o if o.suffix else o.with_suffix("." + ext) - o.mkdir(parents=True, exist_ok=True) - return o / f"{inp.stem}_processed.{ext}" - return inp.with_name(f"{inp.stem}_processed.{ext}") - - -def _next_free(path: Path) -> Path: - for i in itertools.count(1): - cand = path.with_name(f"{path.stem}_{i}{path.suffix}") - if not cand.exists(): - return cand - - -def _resolve_output_conflict(out_path: Path, interactive: bool) -> Path | None: - """Returns the path to write, or None when the user cancels.""" - if not out_path.exists(): - return out_path - alt = _next_free(out_path) - if not interactive: - ui.log(f"[producer] {out_path} exists, writing {alt.name} instead") - return alt - while True: - try: - ans = ( - input( - f"[producer] {out_path} exists " - f"([o]verwrite / [r]ename to {alt.name} / [c]ancel): " - ) - .strip() - .lower() - ) - except EOFError: - print() - return None - if ans in ("o", "overwrite"): - return out_path - if ans in ("r", "rename"): - return alt - if ans in ("c", "cancel"): - return None - print("[producer] please answer o, r, or c") - - -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: - ui.log(f"[producer] skipped {inp}") - return None - x, sr = pio.decode(inp) - target_sr = opts.out_sample_rate() - 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) - pio.encode(y, target_sr, out_path, opts.out_format, opts.bit_depth) - rep = repmod.build(str(inp), str(out_path), opts, res.before, res.after, res.timings, res.notes) - if opts.report or opts.report_path: - rp = ( - Path(opts.report_path) - if opts.report_path - else out_path.with_name(out_path.stem + ".report.json") - ) - repmod.save(rep, rp) - if not quiet: - repmod.print_human(rep) - ui.log(f"[producer] {pos_s}wrote {out_path}") - return out_path - - -def main(argv: list[str] | None = None) -> int: - argv = list(sys.argv[1:] if argv is None else argv) - if argv and argv[0] == "doctor": - 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: - cfgmod.write_default_config(cfg_path) - cfgmod.apply_config(opts, cfgmod.load_config(cfg_path)) - _apply_args(opts, args) - 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})") - for st in stages: - state = "on " if st.enabled else "off" - print(f" [{state}] {st.name:<10} {st.detail}") - print(f" output: {opts.out_format} @ {opts.out_sample_rate()} Hz, {opts.bit_depth}-bit") - return 0 - (DATA_DIR / "models").mkdir(parents=True, exist_ok=True) - (DATA_DIR / "cache").mkdir(parents=True, exist_ok=True) - failed = 0 - single = len(inputs) == 1 - total = len(inputs) - for idx, inp in enumerate(inputs, start=1): - try: - 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: - import traceback - - traceback.print_exc() - return 1 if failed else 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/lib/src/producer/config.py b/lib/src/producer/config.py deleted file mode 100644 index 968df4a..0000000 --- a/lib/src/producer/config.py +++ /dev/null @@ -1,286 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any - -import tomllib - -DSP_KEYS = ( - "hpf", - "mud", - "warmth", - "soothe", - "compress", - "tape", - "deess", - "presence", - "air", - "breath", -) - - -@dataclass -class Profile: - name: str - loudness_mode: str = "rms" - 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) - comp1: tuple[float, float, float, float] = (-20.0, 2.0, 15.0, 150.0) - comp2: tuple[float, float, float, float] = (-18.0, 3.0, 5.0, 100.0) - deess: tuple[float, float, float] = (5500.0, 8000.0, 5.0) - presence: tuple[float, float] = (3000.0, 1.5) - air: tuple[float, float] = (10000.0, 1.5) - strengths: dict[str, float] = field( - default_factory=lambda: { - "hpf": 1.0, - "mud": 0.8, - "warmth": 0.8, - "soothe": 0.3, - "compress": 0.8, - "tape": 0.0, - "deess": 0.6, - "presence": 0.8, - "air": 0.6, - "breath": 0.35, - } - ) - - -AUDIOBOOK = Profile("audiobook") -PODCAST = Profile( - "podcast", - loudness_mode="lufs", - 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), - comp2=(-16.0, 3.0, 5.0, 100.0), - deess=(5500.0, 8000.0, 6.0), - presence=(3000.0, 2.5), - air=(10000.0, 2.0), - strengths={ - "hpf": 1.0, - "mud": 0.6, - "warmth": 0.7, - "soothe": 0.5, - "compress": 0.9, - "tape": 0.2, - "deess": 0.7, - "presence": 1.0, - "air": 0.8, - "breath": 0.2, - }, -) -RADIO = Profile( - "radio", - loudness_mode="lufs", - 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), - comp1=(-16.0, 3.0, 10.0, 120.0), - comp2=(-14.0, 4.0, 3.0, 90.0), - deess=(5000.0, 7500.0, 4.0), - presence=(3500.0, 1.0), - air=(9000.0, 1.0), - strengths={ - "hpf": 1.0, - "mud": 0.9, - "warmth": 1.0, - "soothe": 0.7, - "compress": 1.0, - "tape": 0.55, - "deess": 0.5, - "presence": 0.7, - "air": 0.5, - "breath": 0.4, - }, -) -PROFILES = {"audiobook": AUDIOBOOK, "podcast": PODCAST, "radio": RADIO} - - -@dataclass -class Options: - profile: str = "audiobook" - denoise: str = "dfn3" - # 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 - levelling: bool = True - target: float | None = None - ceiling_db: float | None = None - sample_rate: int | None = None - 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 - dry_run: bool = False - verbose: int = 0 - - def prof(self) -> Profile: - return PROFILES[self.profile] - - def eff(self, key: str) -> float: - v = self.strengths.get(key) - if v is None: - return float(self.prof().strengths[key]) - return float(v) - - def target_value(self) -> float: - return self.prof().target if self.target is None else float(self.target) - - def ceiling(self) -> float: - return self.prof().ceiling_db if self.ceiling_db is None else float(self.ceiling_db) - - def out_sample_rate(self) -> int: - return self.prof().sample_rate if self.sample_rate is None else int(self.sample_rate) - - def loudness_mode(self) -> str: - return self.prof().loudness_mode - - def to_dict(self) -> dict[str, Any]: - return { - "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, - "levelling": self.levelling, - "target": self.target_value(), - "ceiling_db": self.ceiling(), - "sample_rate": self.out_sample_rate(), - "bit_depth": self.bit_depth, - "out_format": self.out_format, - "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, - } - - -_CFG_FIELDS = { - "profile": "profile", - "denoise_strength": "denoise_strength", - "enhance_strength": "enhance_strength", - "format": "out_format", - "bit_depth": "bit_depth", - "device": "device", - "target": "target", - "ceiling": "ceiling_db", - "sample_rate": "sample_rate", - "engine_chunk": "engine_chunk_s", - "engine_overlap": "engine_overlap_s", -} - - -def _engine_cfg(value: Any) -> tuple[Any, Any]: - if isinstance(value, dict): - return value.get("engine"), value.get("strength") - return value, None - - -def load_config(path: Path) -> dict[str, Any]: - if not path.exists(): - return {} - with path.open("rb") as f: - return tomllib.load(f) - - -def apply_config(opts: Options, cfg: dict[str, Any]) -> None: - for key, field_name in _CFG_FIELDS.items(): - if key in cfg: - setattr(opts, field_name, cfg[key]) - if "denoise" in cfg: - eng, strength = _engine_cfg(cfg["denoise"]) - if eng: - 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: - opts.enhance = eng - if strength is not None: - opts.enhance_strength = float(strength) - prof_cfg = cfg.get(opts.profile) - if isinstance(prof_cfg, dict): - for k in DSP_KEYS: - if k in prof_cfg: - opts.strengths[k] = float(prof_cfg[k]) - if "hpf_hz" in prof_cfg: - opts.hpf_hz = float(prof_cfg["hpf_hz"]) - - -def write_default_config(path: Path) -> None: - if path.exists(): - return - 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"', - "# 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"', - "", - "[audiobook]", - "mud = 0.8", - "warmth = 0.8", - "compress = 0.8", - "deess = 0.6", - "presence = 0.8", - "air = 0.6", - "breath = 0.35", - "", - "[podcast]", - "mud = 0.6", - "warmth = 0.7", - "compress = 0.9", - "deess = 0.7", - "presence = 1.0", - "air = 0.8", - "breath = 0.2", - "", - ] - path.write_text("\n".join(lines)) diff --git a/lib/src/producer/doctor.py b/lib/src/producer/doctor.py deleted file mode 100644 index 2d8997f..0000000 --- a/lib/src/producer/doctor.py +++ /dev/null @@ -1,95 +0,0 @@ -from __future__ import annotations - -import importlib.util -import shutil -import sys -from pathlib import Path - -from . import __version__ - -DATA_DIR = Path(__file__).resolve().parents[2] - - -def _check(name: str, ok: bool, detail: str = "") -> bool: - mark = "ok " if ok else "MISS" - line = f" [{mark}] {name}" - if detail: - line += f": {detail}" - print(line) - return ok - - -def _import(name: str): - try: - return importlib.import_module(name) - except Exception: - 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}") - critical_ok = True - - v = sys.version_info - critical_ok &= _check("python", v >= (3, 10), f"{v.major}.{v.minor}.{v.micro}") - - for mod in ("numpy", "scipy", "soundfile", "pyloudnorm"): - m = _import(mod) - ok = m is not None - critical_ok &= ok - _check(mod, ok, getattr(m, "__version__", "") if ok else "not installed") - - torch = _import("torch") - if torch is None: - _check("torch", False, "not installed (installs on first denoise/enhance run)") - else: - ver = getattr(torch, "__version__", "?") - cuda = bool(torch.cuda.is_available()) - dev = torch.cuda.get_device_name(0) if cuda else "" - _check("torch", True, f"{ver}, cuda={cuda}" + (f" ({dev})" if dev else "")) - - ffmpeg = shutil.which("ffmpeg") - _check("ffmpeg", ffmpeg is not None, ffmpeg or "missing (needed for mp3 + odd formats)") - - for mod in ("deepfilternet", "zipenhancer", "clearvoice"): - m = _import(mod) or _import(mod.replace("-", "_")) - 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)) - - res_venv = DATA_DIR / "venvs" / "resemble" / "bin" / "python" - _check("resemble venv", res_venv.is_file(), str(res_venv)) - - uv = DATA_DIR / "bin" / "uv" - _check("uv", uv.is_file() or shutil.which("uv") is not None, str(uv)) - - cfg = DATA_DIR.parent / "config.toml" - _check("config.toml", cfg.exists(), str(cfg)) - - print("doctor:", "core OK" if critical_ok else "core problems detected") - return 0 if critical_ok else 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/lib/src/producer/dsp.py b/lib/src/producer/dsp.py deleted file mode 100644 index e65e595..0000000 --- a/lib/src/producer/dsp.py +++ /dev/null @@ -1,283 +0,0 @@ -from __future__ import annotations - -import numpy as np -from scipy import ndimage, signal - -from .meters import SILENCE_DB - -_EPS = np.float32(1e-12) - - -def _coef(time_ms: float, sr: int) -> float: - return float(np.exp(-1000.0 / (max(time_ms, 1e-4) * sr))) - - -def _onepole(x: np.ndarray, a: float) -> np.ndarray: - 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: - a = 10.0 ** (gain_db / 40.0) - w0 = 2.0 * np.pi * freq / sr - cs = np.cos(w0) - sn = np.sin(w0) - alpha = sn / (2.0 * q) - if kind == "highpass": - b = [(1 + cs) / 2, -(1 + cs), (1 + cs) / 2] - aa = [1 + alpha, -2 * cs, 1 - alpha] - elif kind == "lowpass": - b = [(1 - cs) / 2, 1 - cs, (1 - cs) / 2] - aa = [1 + alpha, -2 * cs, 1 - alpha] - elif kind == "peaking": - b = [1 + alpha * a, -2 * cs, 1 - alpha * a] - aa = [1 + alpha / a, -2 * cs, 1 - alpha / a] - elif kind == "lowshelf": - sq = 2 * np.sqrt(a) * alpha - b = [ - a * ((a + 1) - (a - 1) * cs + sq), - 2 * a * ((a - 1) - (a + 1) * cs), - a * ((a + 1) - (a - 1) * cs - sq), - ] - aa = [ - (a + 1) + (a - 1) * cs + sq, - -2 * ((a - 1) + (a + 1) * cs), - (a + 1) + (a - 1) * cs - sq, - ] - elif kind == "highshelf": - sq = 2 * np.sqrt(a) * alpha - b = [ - a * ((a + 1) + (a - 1) * cs + sq), - -2 * a * ((a - 1) + (a + 1) * cs), - a * ((a + 1) + (a - 1) * cs - sq), - ] - aa = [ - (a + 1) - (a - 1) * cs + sq, - 2 * ((a - 1) - (a + 1) * cs), - (a + 1) - (a - 1) * cs - sq, - ] - else: - raise ValueError(f"unknown filter kind: {kind}") - arr = np.array(b + aa, dtype=np.float64) - return arr / arr[3] - - -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 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: - return biquad(x, rbj("highpass", sr, hz)) - - -def shelf(x: np.ndarray, sr: int, freq: float, gain_db: float, low: bool = True) -> np.ndarray: - if abs(gain_db) < 0.01: - return x - return biquad(x, rbj("lowshelf" if low else "highshelf", sr, freq, gain_db)) - - -def peak_eq(x: np.ndarray, sr: int, freq: float, gain_db: float, q: float = 1.0) -> np.ndarray: - if abs(gain_db) < 0.01: - return x - return biquad(x, rbj("peaking", sr, freq, gain_db, q)) - - -def compressor( - x: np.ndarray, - sr: int, - threshold_db: float, - ratio: float, - attack_ms: float, - release_ms: float, - knee_db: float = 6.0, -) -> np.ndarray: - 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(x32), a_att), 0.0, None)) - level_db = 20.0 * np.log10(env + _EPS) - over = level_db - np.float32(threshold_db) - k = np.float32(knee_db) - gr = np.where( - over <= -k / 2, - np.float32(0.0), - np.where( - over < k / 2, - (1.0 - 1.0 / ratio) * np.square(over + k / 2) / (2.0 * k), - (1.0 - 1.0 / ratio) * over, - ), - ) - gr = _onepole(np.clip(gr, 0.0, None), a_rel) - gain = 10.0 ** (-gr / 20.0) - return (x32 * gain).astype(np.float32, copy=False) - - -def deesser( - x: np.ndarray, sr: int, lo_hz: float, hi_hz: float, max_reduction_db: float -) -> np.ndarray: - if max_reduction_db < 0.05: - return x - 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").astype(np.float32), - high, - ) - win = max(3, int(0.005 * sr) | 1) - 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 = 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 * 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, copy=False) - - -def expander( - x: np.ndarray, - sr: int, - max_drop_db: float = 6.0, - ratio: float = 2.0, - attack_ms: float = 10.0, - release_ms: float = 120.0, -) -> np.ndarray: - if max_drop_db < 0.05: - return x - n = x.size - frame = max(1, int(0.020 * sr)) - nf = n // frame - if nf < 2: - return x - 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 - speech_ref = float(np.percentile(active, 90)) - floor_ref = float(np.percentile(active, 5)) - 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, 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, np.float32(max_drop_db)) / 20.0) - return (x.astype(np.float32, copy=False) * gain).astype(np.float32, copy=False) - - -def limit( - x: np.ndarray, - sr: int, - ceiling_db: float, - release_ms: float = 60.0, - lookahead_ms: float = 1.0, -) -> np.ndarray: - 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(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.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))) - held = np.empty(nb, dtype=np.float64) - prev = 0.0 - for i in range(nb): - prev = max(block_over[i], prev * decay) - held[i] = prev - 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 = x32 * gain - bad = np.abs(y) > lin - if np.any(bad): - 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: - if sr_in == sr_out or x.size == 0: - return x - from math import gcd - - g = gcd(int(sr_in), int(sr_out)) - 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 = 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 = 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: - """Dynamic reducer that digs into resonances only while they stick out. - - Targets boxy 200-450 Hz and harsh 2.5-6 kHz bands; static-only when turned - to zero. Deesser-style: band split -> smoothed envelope -> percentile - threshold -> gain reduce -> subtractive recombination. - """ - if amount < 0.02: - return x - 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))).astype(np.float32) - changed = False - for lo, hi, max_db in bands: - if max_db * float(s) < 0.1: - continue - band = signal.sosfiltfilt( - 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 = 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 * 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 - changed = True - if not changed or np.array_equal(y, x.astype(np.float32, copy=False)): - return x - return y.astype(np.float32, copy=False) diff --git a/lib/src/producer/engines/__init__.py b/lib/src/producer/engines/__init__.py deleted file mode 100644 index e69de29..0000000 --- a/lib/src/producer/engines/__init__.py +++ /dev/null diff --git a/lib/src/producer/engines/base.py b/lib/src/producer/engines/base.py deleted file mode 100644 index 9192a17..0000000 --- a/lib/src/producer/engines/base.py +++ /dev/null @@ -1,43 +0,0 @@ -from __future__ import annotations - -import numpy as np - - -def pick_device(pref: str = "auto") -> str: - if pref == "cpu": - return "cpu" - try: - import torch - except ImportError: - if pref == "cuda": - print("[producer] warning: torch not installed; using cpu", flush=True) - return "cpu" - if torch.cuda.is_available(): - return "cuda" - if pref == "cuda": - print("[producer] warning: CUDA requested but unavailable; using cpu", flush=True) - return "cpu" - - -def device_name(device: str) -> str: - if device == "cuda": - import torch - - return torch.cuda.get_device_name(0) - return "cpu" - - -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.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 deleted file mode 100644 index 627e3b8..0000000 --- a/lib/src/producer/engines/chunking.py +++ /dev/null @@ -1,166 +0,0 @@ -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 deleted file mode 100644 index 186ffd2..0000000 --- a/lib/src/producer/engines/denoise_dfn.py +++ /dev/null @@ -1,140 +0,0 @@ -from __future__ import annotations - -import sys -import types -import warnings -import zipfile -from collections.abc import Callable -from pathlib import Path - -import numpy as np - -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", -} -BASE_URL = f"https://raw.githubusercontent.com/Rikorose/DeepFilterNet/{TAG}" - - -def _shim_torchaudio_backend() -> None: - warnings.filterwarnings( - "ignore", - message=r".*AudioMetaData.*has been moved.*", - category=UserWarning, - module=r"df[./]io", - ) - try: - import torchaudio.backend # noqa: F401 - except Exception: - pkg = sys.modules.get("torchaudio") - if pkg is not None and "torchaudio.backend" not in sys.modules: - stub = types.ModuleType("torchaudio.backend") - stub.__path__ = [] - sys.modules["torchaudio.backend"] = stub - pkg.backend = stub - - -def _shim_df_git() -> None: - import df.io - import df.logger - import df.utils - - for mod in (df.utils, df.logger, df.io): - for name in ("get_git_root", "get_commit_hash", "get_branch_name"): - if hasattr(mod, name): - setattr(mod, name, lambda: None) - - -def ensure_model(model: str) -> Path: - target = MODELS_DIR / model - if (target / "config.ini").is_file(): - return target - url = f"{BASE_URL}/{MODEL_ZIPS[model]}" - MODELS_DIR.mkdir(parents=True, exist_ok=True) - zpath = MODELS_DIR / f"{model}.zip" - 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) - if not (target / "config.ini").is_file(): - inner = list(target.glob(f"**/{model}/config.ini")) - if inner: - src = inner[0].parent - for f in src.iterdir(): - f.rename(target / f.name) - if not (target / "config.ini").is_file(): - raise RuntimeError(f"{model} weights download failed") - return target - - -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, - post_filter: bool = False, -) -> tuple[np.ndarray, str, str]: - """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" - try: - model_dir = ensure_model(model_name) - except Exception: - model_name = "DeepFilterNet2" - model_dir = ensure_model(model_name) - from df.enhance import enhance as df_enhance - from df.enhance import init_df - - device = pick_device(device_pref) - model, df_state, _ = init_df( - model_base_dir=str(model_dir), - post_filter=post_filter, - log_level="error", - log_file=None, - ) - try: - model = model.to(device) - dev = device - except Exception: - dev = "cpu" - sr_df = int(df_state.sr()) - xin = dsp.resample(x, sr, sr_df) - import torch - - 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 deleted file mode 100644 index e1bb256..0000000 --- a/lib/src/producer/engines/denoise_spectral.py +++ /dev/null @@ -1,186 +0,0 @@ -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 deleted file mode 100644 index 3f37f61..0000000 --- a/lib/src/producer/engines/denoise_zip.py +++ /dev/null @@ -1,68 +0,0 @@ -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", - 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: 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) - - 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 deleted file mode 100644 index 3298c09..0000000 --- a/lib/src/producer/engines/enhance_mossformer.py +++ /dev/null @@ -1,70 +0,0 @@ -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", - 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)") - # 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") - 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) - if fs != sr_target: - y = dsp.resample(y, fs, sr_target) - y = dsp.resample(y, sr_target, sr) - y = blend(x, y, strength) - return y, "mossformer2 (MossFormer2_SE_48K)", device_name(device) diff --git a/lib/src/producer/engines/enhance_resemble.py b/lib/src/producer/engines/enhance_resemble.py deleted file mode 100644 index 8042215..0000000 --- a/lib/src/producer/engines/enhance_resemble.py +++ /dev/null @@ -1,95 +0,0 @@ -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, ui -from .base import blend - -VENV_DIR = lazy.DATA_DIR / "venvs" / "resemble" -WORKER = Path(__file__).resolve().parent / "resemble_worker.py" -REQ_HASH_FILE = VENV_DIR / ".req-hash" -REQS = ["resemble-enhance==0.0.1", "librosa", "soundfile"] - - -def _ensure_venv() -> Path: - marker = REQ_HASH_FILE.read_text() if REQ_HASH_FILE.exists() else None - want = "|".join(REQS) - py = VENV_DIR / "bin" / "python" - if py.is_file() and marker == want: - return py - uv = lazy.find_uv() - if not py.is_file(): - 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], - "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", - 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 - - py = _ensure_venv() - from .base import device_name as _dname - from .base import pick_device as _pick - - device = _pick(device_pref) - with tempfile.TemporaryDirectory(prefix="producer_res_") as td: - inp = Path(td) / "in.wav" - outp = Path(td) / "out.wav" - io.encode(x, sr, inp, "wav", 16) - env = dict(os.environ) - env["HF_HOME"] = str(lazy.DATA_DIR / "models" / "hf") - env["TORCH_HOME"] = str(lazy.DATA_DIR / "models" / "torch") - cmd = [ - str(py), - str(WORKER), - str(inp), - str(outp), - device, - str(float(np.clip(strength, 0.0, 1.0))), - str(float(chunk_s)), - str(float(overlap_s)), - ] - 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) - y = blend(x, y, strength) - return y, "resemble-enhance (generative, may alter timbre)", _dname(device) - - -if __name__ == "__main__": - sys.exit(0) diff --git a/lib/src/producer/engines/resemble_worker.py b/lib/src/producer/engines/resemble_worker.py deleted file mode 100644 index 493ceb3..0000000 --- a/lib/src/producer/engines/resemble_worker.py +++ /dev/null @@ -1,68 +0,0 @@ -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 = 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) - 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 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/lib/src/producer/io.py b/lib/src/producer/io.py deleted file mode 100644 index b6d38b6..0000000 --- a/lib/src/producer/io.py +++ /dev/null @@ -1,87 +0,0 @@ -from __future__ import annotations - -import shutil -import subprocess -import tempfile -from pathlib import Path - -import numpy as np -import soundfile as sf - -SUBTYPES = {16: "PCM_16", 24: "PCM_24", 32: "FLOAT"} - - -def ffmpeg_available() -> bool: - return shutil.which("ffmpeg") is not None - - -def decode(path: str | Path) -> tuple[np.ndarray, int]: - path = Path(path) - try: - data, sr = sf.read(str(path), dtype="float32", always_2d=True) - except RuntimeError: - if not ffmpeg_available(): - raise - with tempfile.TemporaryDirectory(prefix="producer_dec_") as tmpdir: - tmp = Path(tmpdir) / "dec.wav" - subprocess.run( - ["ffmpeg", "-v", "error", "-y", "-i", str(path), "-vn", "-ac", "1", str(tmp)], - check=True, - ) - data, sr = sf.read(str(tmp), dtype="float32", always_2d=True) - mono = data.mean(axis=1).astype(np.float32) - return mono, int(sr) - - -def encode( - x: np.ndarray, - sr: int, - path: str | Path, - out_format: str = "wav", - bit_depth: int = 32, -) -> Path: - path = Path(path) - fmt = out_format.lower() - if fmt in ("wav", "flac"): - subtype = SUBTYPES.get(bit_depth, "PCM_16") - sf.write(str(path), x, sr, subtype=subtype, format=fmt.upper()) - return path - if fmt == "mp3": - if not ffmpeg_available(): - raise RuntimeError("mp3 output requires ffmpeg on PATH") - with tempfile.TemporaryDirectory(prefix="producer_enc_") as td: - tmp_wav = Path(td) / "enc.wav" - sf.write(str(tmp_wav), x, sr, subtype="PCM_16") - subprocess.run( - [ - "ffmpeg", - "-v", - "error", - "-y", - "-i", - str(tmp_wav), - "-codec:a", - "libmp3lame", - "-b:a", - "192k", - str(path), - ], - check=True, - ) - return path - raise ValueError(f"unsupported output format: {out_format}") - - -def supported_globs() -> list[str]: - return ( - "*.wav", - "*.flac", - "*.mp3", - "*.m4a", - "*.aac", - "*.ogg", - "*.opus", - "*.aif", - "*.aiff", - "*.wma", - ) diff --git a/lib/src/producer/lazy.py b/lib/src/producer/lazy.py deleted file mode 100644 index 55bbf6c..0000000 --- a/lib/src/producer/lazy.py +++ /dev/null @@ -1,257 +0,0 @@ -from __future__ import annotations - -import importlib -import importlib.util -import json -import platform -import re -import shutil -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): - pass - - -def find_uv() -> str: - uv = shutil.which("uv") - if uv: - return uv - local = DATA_DIR / "bin" / "uv" - if local.is_file(): - return str(local) - raise EngineUnavailable("uv not found; run the producer script to bootstrap") - - -def has_module(name: str) -> bool: - return importlib.util.find_spec(name) is not None - - -def dist_name(spec: str) -> str: - return spec.split("==")[0].split(">=")[0].split("<")[0].split("[")[0] - - -def is_installed(spec: str) -> bool: - dist = dist_name(spec) - try: - metadata.distribution(dist) - return True - except metadata.PackageNotFoundError: - return has_module(DIST_ALIASES.get(dist, dist)) - - -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 - 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 = [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() - - -def ensure(packages: list[str], purpose: str) -> None: - missing = [p for p in packages if not is_installed(p)] - if not missing: - return - if any(m.startswith("torch") for m in missing): - ensure_torch() - missing = [m for m in missing if not m.startswith("torch") and not is_installed(m)] - if not missing: - return - 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/loudness.py b/lib/src/producer/loudness.py deleted file mode 100644 index 8ceae97..0000000 --- a/lib/src/producer/loudness.py +++ /dev/null @@ -1,30 +0,0 @@ -from __future__ import annotations - -import numpy as np - -from . import dsp -from .meters import meter, true_peak_db - - -def normalize( - x: np.ndarray, - sr: int, - mode: str, - target: float, - ceiling_db: float, - max_iters: int = 4, -) -> np.ndarray: - if not np.any(x): - return x - y = x - for _ in range(max_iters): - cur = meter(y, sr, mode) - gain = target - cur - if abs(gain) < 0.05: - break - y = y * (10.0 ** (gain / 20.0)) - y = dsp.limit(y, sr, ceiling_db) - tp = true_peak_db(y, sr) - if tp > ceiling_db + 0.05: - y = y * (10.0 ** ((ceiling_db - tp) / 20.0)) - return y.astype(np.float32) diff --git a/lib/src/producer/meters.py b/lib/src/producer/meters.py deleted file mode 100644 index dc4c566..0000000 --- a/lib/src/producer/meters.py +++ /dev/null @@ -1,112 +0,0 @@ -from __future__ import annotations - -import numpy as np -from scipy import signal - -SILENCE_DB = -120.0 -_METER_BLOCK = 1 << 20 - - -def _db(v: float) -> float: - if v <= 0.0 or not np.isfinite(v): - return SILENCE_DB - return max(20.0 * np.log10(v), SILENCE_DB) - - -def rms_db(x: np.ndarray) -> float: - 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: - 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) - 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 - 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 - - if not np.any(x): - return SILENCE_DB - if x.size < int(sr * 0.2): - return SILENCE_DB - meter = pyln.Meter(sr) - try: - val = meter.integrated_loudness(x) - except ValueError: - return SILENCE_DB - if not np.isfinite(val): - return SILENCE_DB - return float(val) - - -def meter(x: np.ndarray, sr: int, mode: str) -> float: - return lufs(x, sr) if mode == "lufs" else rms_db(x) - - -def all_meters(x: np.ndarray, sr: int) -> dict[str, float]: - return { - "rms_db": rms_db(x), - "sample_peak_db": sample_peak_db(x), - "true_peak_db": true_peak_db(x, sr), - "lufs": lufs(x, sr), - "noise_floor_db": noise_floor_db(x, sr), - "duration_s": round(x.size / sr, 3) if sr else 0.0, - } diff --git a/lib/src/producer/pipeline.py b/lib/src/producer/pipeline.py deleted file mode 100644 index e7aa4ae..0000000 --- a/lib/src/producer/pipeline.py +++ /dev/null @@ -1,269 +0,0 @@ -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, 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 -class Stage: - name: str - detail: str - enabled: bool = True - fn: Callable[[np.ndarray, int], np.ndarray] | None = None - - -@dataclass -class RunResult: - audio: np.ndarray - sr: int - stages: list[Stage] - timings: dict[str, float] = field(default_factory=dict) - notes: list[str] = field(default_factory=list) - before: dict[str, float] = field(default_factory=dict) - after: dict[str, float] = field(default_factory=dict) - - -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, - 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, - 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}") - return y - - return fn - - -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, - 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, - 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}") - return y - - return fn - - -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) - f, g, q = p.mud - x = dsp.peak_eq(x, sr, f, g * opts.eff("mud"), q) - wf, wg = p.warmth - x = dsp.shelf(x, sr, wf, wg * opts.eff("warmth"), low=True) - x = dsp.soothe(x, sr, opts.eff("soothe")) - c1 = p.comp1 - y = dsp.compressor(x, sr, c1[0], c1[1], c1[2], c1[3]) - x = _blend(x, y, opts.eff("compress") * 0.7) - c2 = p.comp2 - y = dsp.compressor(x, sr, c2[0], c2[1], c2[2], c2[3]) - x = _blend(x, y, opts.eff("compress") * 0.5) - x = dsp.tape(x, sr, opts.eff("tape")) - d = p.deess - x = dsp.deesser(x, sr, d[0], d[1], d[2] * opts.eff("deess")) - pf, pg = p.presence - 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) - # 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 - - -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.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]): - def fn(x: np.ndarray, sr: int) -> np.ndarray: - return loudness.normalize(x, sr, opts.loudness_mode(), opts.target_value(), opts.ceiling()) - - return fn - - -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" pf={'on' if opts.denoise_pf else 'off'}", - opts.denoise != "off", - _denoise_fn(opts, notes, reporter), - ) - ) - stages.append( - Stage( - "enhance", - f"engine={opts.enhance} strength={opts.enhance_strength:.2f}", - opts.enhance != "off", - _enhance_fn(opts, notes, reporter), - ) - ) - dsp_desc = "hpf→mud→warmth→soothe→comp2x→tape→deess→presence→air→breath" - stages.append(Stage("dsp", dsp_desc, opts.dsp, _dsp_fn(opts, notes) if opts.dsp else None)) - stages.append( - Stage( - "levelling", - f"mode={opts.loudness_mode()} target={opts.target_value()} ceiling={opts.ceiling()}", - opts.levelling, - _level_fn(opts, notes) if opts.levelling else None, - ) - ) - return stages - - -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) - y = x - timings: dict[str, float] = {} - 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) - 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/report.py b/lib/src/producer/report.py deleted file mode 100644 index 27fbb8f..0000000 --- a/lib/src/producer/report.py +++ /dev/null @@ -1,56 +0,0 @@ -from __future__ import annotations - -import json -from pathlib import Path - -from . import __version__ -from .config import Options - - -def build( - input_path: str, - out_path: str, - opts: Options, - before: dict[str, float], - after: dict[str, float], - timings: dict[str, float], - notes: list[str], -) -> dict: - return { - "tool": "producer", - "version": __version__, - "input": str(input_path), - "output": str(out_path), - "settings": opts.to_dict(), - "before": before, - "after": after, - "stage_seconds": timings, - "engine_notes": notes, - } - - -def print_human(rep: dict) -> None: - b = rep["before"] - a = rep["after"] - print(" before -> after:") - for key, label in ( - ("rms_db", "RMS"), - ("true_peak_db", "true peak"), - ("lufs", "LUFS"), - ("noise_floor_db", "noise floor"), - ): - print(f" {label:<12} {b[key]:>8.1f} dB -> {a[key]:>8.1f} dB") - if rep["engine_notes"]: - for note in rep["engine_notes"]: - print(f" {note}") - times = rep["stage_seconds"] - if times: - total = sum(times.values()) - detail = ", ".join(f"{k} {v:.1f}s" for k, v in times.items()) - print(f" stages: {detail} (total {total:.1f}s)") - - -def save(rep: dict, path: str | Path) -> Path: - path = Path(path) - path.write_text(json.dumps(rep, indent=2)) - return path diff --git a/lib/src/producer/ui.py b/lib/src/producer/ui.py deleted file mode 100644 index bc08756..0000000 --- a/lib/src/producer/ui.py +++ /dev/null @@ -1,319 +0,0 @@ -"""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 deleted file mode 100644 index 690e955..0000000 --- a/lib/src/producer/updates.py +++ /dev/null @@ -1,273 +0,0 @@ -"""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 |