bolo-torch 0.1.0__py3-none-any.whl

bolo_torch/__init__.py

@@ -0,0 +1,20 @@
+"""bolo-torch: the PyTorch engine plugin of the bolo speech server.
+
+bolo never imports PyTorch. It starts this package as a separate process (``python -m bolo_torch``) and talks to
+it over the plugin protocol (bolo's docs/plugins.md). This module is what bolo's entry-point discovery imports in
+the *server* process, so it must stay free of heavy imports: no torch here.
+"""
+
+from __future__ import annotations
+
+import sys
+
+__version__ = "0.1.0"
+
+
+def plugin() -> list[str]:
+    """The command line bolo runs to start the plugin (entry point ``bolo.engines: torch``)."""
+    return [sys.executable, "-m", "bolo_torch"]
+
+
+__all__ = ["__version__", "plugin"]

bolo_torch/__main__.py

@@ -0,0 +1,8 @@
+"""``python -m bolo_torch``: serve the plugin protocol on stdin/stdout (what bolo starts)."""
+
+from __future__ import annotations
+
+from bolo_torch.cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main(["serve"]))

bolo_torch/cli.py

@@ -0,0 +1,163 @@
+"""``bolo-torch``: serve the plugin (what bolo runs), check the environment, or separate a recording directly.
+
+bolo-torch serve                       # the plugin protocol on stdin/stdout (bolo starts this)
+bolo-torch doctor                      # torch build, devices, which model families are installed
+bolo-torch separate mix.wav -o out/    # speech separation without bolo (SepFormer / PixIT)
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import logging
+import os
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+from bolo_torch import __version__
+
+#: separation models for ``bolo-torch separate`` (the same files as bolo's catalog entries)
+SEPARATION_MODELS: dict[str, dict[str, Any]] = {
+    "sepformer-whamr": {
+        "kind": "sepformer", "repo": "speechbrain/sepformer-whamr",
+        "revision": "9950fd8055b38679d2af5bef2366b5dba600ec66",
+        "allow": ["hyperparams.yaml", "masknet.ckpt", "encoder.ckpt", "decoder.ckpt"],
+        "about": "SpeechBrain SepFormer, WHAMR! (noisy, reverberant), 8 kHz, 2 speakers. Apache-2.0; WSJ0-derived data",
+    },
+    "sepformer-whamr16k": {
+        "kind": "sepformer", "repo": "speechbrain/sepformer-whamr16k",
+        "revision": "21a5b500c6f52fddc387c5d9e5fb13ffd6f039c5",
+        "allow": ["hyperparams.yaml", "masknet.ckpt", "encoder.ckpt", "decoder.ckpt"],
+        "about": "SpeechBrain SepFormer, WHAMR! at 16 kHz, 2 speakers. Apache-2.0; WSJ0-derived data",
+    },
+    "pixit-ami": {
+        "kind": "pixit", "repo": "pyannote/speech-separation-ami-1.0",
+        "revision": "9486b106945ae0cc0784041a08bfcdba5edadfb9", "allow": None,
+        "about": "pyannote PixIT, joint diarization and separation, AMI meetings, 16 kHz. MIT, gated",
+    },
+}  # fmt: skip
+
+FAMILIES = {
+    "parler": ["parler_tts", "transformers"],
+    "f5": ["transformers", "vocos", "x_transformers", "torchdiffeq", "pydub", "jieba", "pypinyin"],
+    "orpheus": ["transformers", "snac"],
+    "sepformer": ["speechbrain", "torchaudio"],
+    "pyannote": ["pyannote.audio"],
+}
+
+
+def _installed(module: str) -> bool:
+    try:
+        return importlib.util.find_spec(module) is not None
+    except (ImportError, ValueError):
+        return False
+
+
+def doctor() -> dict[str, Any]:
+    import torch
+
+    from bolo_torch.runtime import hf_token, pick_device
+
+    mps = getattr(torch.backends, "mps", None)
+    return {
+        "bolo_torch": __version__,
+        "python": sys.version.split()[0],
+        "torch": torch.__version__,
+        "cuda_build": torch.version.cuda,
+        "cuda_available": torch.cuda.is_available(),
+        "cuda_devices": [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())],
+        "mps_available": bool(mps is not None and mps.is_available()),
+        "auto_device": pick_device("auto"),
+        "threads": torch.get_num_threads(),
+        "hf_token": bool(hf_token()),
+        "families": {name: all(_installed(m) for m in mods) for name, mods in FAMILIES.items()},
+    }
+
+
+def separate(args: argparse.Namespace) -> dict[str, Any]:
+    import numpy as np
+    import soundfile as sf
+    from bolo.engines.base import LoadOptions, SeparationModel
+
+    from bolo_torch.engine import TorchEngine
+    from bolo_torch.models.separation import SR, resample
+    from bolo_torch.runtime import hub_snapshot
+
+    entry = SEPARATION_MODELS[args.model]
+    t0 = time.perf_counter()
+    root = hub_snapshot(entry["repo"], entry["revision"], entry["allow"])
+    t_download = time.perf_counter() - t0
+    t0 = time.perf_counter()
+    model = TorchEngine().load(
+        entry["kind"],
+        "separation",
+        root,
+        {"model_id": args.model},
+        LoadOptions(threads=args.threads, device=args.device),
+    )
+    t_load = time.perf_counter() - t0
+    assert isinstance(model, SeparationModel)
+    samples, rate = sf.read(args.input, dtype="float32", always_2d=True)
+    mono = resample(samples.mean(axis=1).astype(np.float32), int(rate), SR)
+    t0 = time.perf_counter()
+    tracks = model.separate(mono)
+    t_run = time.perf_counter() - t0
+    out = Path(args.output)
+    out.mkdir(parents=True, exist_ok=True)
+    files = []
+    for i, track in enumerate(tracks):
+        path = out / f"SPEAKER_{i:02d}.wav"
+        sf.write(str(path), track, SR)
+        files.append(str(path))
+    duration = len(mono) / SR
+    turns = getattr(model, "last_turns", [])
+    return {
+        "model": args.model, "device": args.device, "threads": args.threads, "audio_s": round(duration, 2),
+        "download_s": round(t_download, 2), "load_s": round(t_load, 2), "separate_s": round(t_run, 2),
+        "rtf": round(t_run / duration, 3) if duration else None, "tracks": files,
+        "turns": [[round(s, 2), round(e, 2), k] for s, e, k in turns],
+    }  # fmt: skip
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        prog="bolo-torch", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument("--version", action="version", version=f"bolo-torch {__version__}")
+    sub = parser.add_subparsers(dest="command")
+    sub.add_parser("serve", help="serve the plugin protocol on stdin/stdout")
+    sub.add_parser("doctor", help="report torch, devices and installed model families")
+    sep = sub.add_parser("separate", help="split a recording into one track per speaker")
+    sep.add_argument("input")
+    sep.add_argument("-o", "--output", default="separated")
+    sep.add_argument("-m", "--model", default="sepformer-whamr16k", choices=sorted(SEPARATION_MODELS))
+    sep.add_argument("--device", default=os.environ.get("BOLO_TORCH_DEVICE", "auto"))
+    sep.add_argument("--threads", type=int, default=0)
+    args = parser.parse_args(argv)
+    if args.command in (None, "serve"):
+        logging.basicConfig(level=logging.INFO, stream=sys.stderr, format="%(name)s: %(message)s")
+        from bolo_torch.engine import TorchEngine
+        from bolo_torch.server import serve
+
+        return serve(TorchEngine())
+    if args.command == "doctor":
+        print(json.dumps(doctor(), indent=2))
+        return 0
+    if args.command == "separate":
+        from bolo_torch.runtime import PluginError
+
+        try:
+            print(json.dumps(separate(args), indent=2, ensure_ascii=False))
+        except PluginError as exc:
+            print(f"error ({exc.code}): {exc}", file=sys.stderr)
+            return 2
+        return 0
+    parser.print_help()
+    return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

bolo_torch/engine.py

@@ -0,0 +1,82 @@
+"""The plugin's engine: model kinds, and the device/dtype/thread choices made at load time."""
+
+from __future__ import annotations
+
+import importlib
+import logging
+import time
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+from bolo.engines.base import Engine, EngineStatus, LoadOptions, Model
+
+from bolo_torch import __version__
+from bolo_torch.runtime import InvalidRequest, LoadContext, pick_device, pick_dtype, set_threads
+
+log = logging.getLogger("bolo_torch")
+
+Loader = Callable[[LoadContext], Model]
+
+#: kind -> (module, function); modules are imported only when a model of that kind is loaded
+LOADERS: dict[str, tuple[str, str]] = {
+    "parler-tts": ("bolo_torch.models.parler", "load"),
+    "f5-tts": ("bolo_torch.models.f5", "load"),
+    "orpheus-snac": ("bolo_torch.models.orpheus", "load"),
+    "indic-speak": ("bolo_torch.models.indic_speak", "load"),
+    "sepformer": ("bolo_torch.models.separation", "load_sepformer"),
+    "pixit": ("bolo_torch.models.separation", "load_pixit"),
+    "dicow": ("bolo_torch.models.dicow", "load"),
+}
+#: kinds the catalog lists for this plugin that are not implemented yet
+NOT_IMPLEMENTED = {
+    "sravaani": "SraVaani (TorchScript STT)",
+    "indicconformer-official": "IndicConformer official",
+    "diarizen": "DiariZen diarization (the plugin protocol has no diarization request yet)",
+    "spark-tts": "Spark-TTS (no packaged inference code)",
+    "maskgct": "MaskGCT (no packaged inference code)",
+}
+
+
+class TorchEngine(Engine):
+    id = "torch"
+    name = "torch"
+    license = "Apache-2.0"
+    url = "https://github.com/cb-Sigma/bolo/tree/main/plugins/bolo-torch"
+
+    def __init__(self, loaders: dict[str, Loader | tuple[str, str]] | None = None) -> None:
+        self.loaders: dict[str, Loader | tuple[str, str]] = dict(LOADERS if loaders is None else loaders)
+        self.kinds = tuple(self.loaders)
+
+    def status(self) -> EngineStatus:
+        return EngineStatus(self.id, True, __version__, None)
+
+    def _loader(self, kind: str) -> Loader:
+        entry = self.loaders.get(kind)
+        if entry is None:
+            if kind in NOT_IMPLEMENTED:
+                raise InvalidRequest(f"bolo-torch {__version__} cannot load {NOT_IMPLEMENTED[kind]} models yet")
+            raise InvalidRequest(
+                f"bolo-torch {__version__} has no model kind {kind!r} (it has {', '.join(self.kinds)})"
+            )
+        if isinstance(entry, tuple):
+            module, func = entry
+            return getattr(importlib.import_module(module), func)  # type: ignore[no-any-return]
+        return entry
+
+    def load(self, kind: str, task: str, path: Path, spec: dict[str, Any], opts: LoadOptions) -> Model:
+        loader = self._loader(kind)
+        device = pick_device(opts.device, spec)
+        dtype = pick_dtype(device, spec)
+        threads = set_threads(opts.threads)
+        ctx = LoadContext(
+            kind, task, path, dict(spec), device=device, dtype=dtype, threads=threads,
+            model_id=str(spec.get("model_id", path.name)),
+        )  # fmt: skip
+        t0 = time.perf_counter()
+        model = loader(ctx)
+        log.info(
+            "loaded %s (%s) on %s, %s, %d threads in %.1f s",
+            ctx.model_id, kind, device, str(dtype).replace("torch.", ""), threads, time.perf_counter() - t0,
+        )  # fmt: skip
+        return model

bolo_torch/models/__init__.py

@@ -0,0 +1,30 @@
+"""Model families served by the plugin. Each module imports its libraries only when a model is loaded."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+import numpy as np
+from bolo.engines.base import FloatArray, Voice
+
+
+def voices_from_options(options: dict[str, Any]) -> list[dict[str, Any]]:
+    raw = options.get("voices") or []
+    return [dict(v) for v in raw if isinstance(v, dict) and v.get("id")]
+
+
+def to_voices(entries: Iterable[dict[str, Any]]) -> tuple[Voice, ...]:
+    return tuple(
+        Voice(str(v["id"]), str(v.get("name", v["id"])), str(v.get("language", "")), str(v.get("gender", "")), i)
+        for i, v in enumerate(entries)
+    )
+
+
+def peak_safe(audio: Any) -> FloatArray:
+    """float32 mono, scaled down only if it would clip."""
+    out = np.asarray(audio, dtype=np.float32).reshape(-1)
+    peak = float(np.max(np.abs(out))) if out.size else 0.0
+    if peak > 1.0:
+        out = out / peak
+    return out