livekit-cameras/enhance_daemon.py

"""Shared MetricGAN daemon: one model, many camera clients."""
from __future__ import annotations

import logging
import os
import socket
import threading
import time
from typing import Callable, Optional

import numpy as np

from enhance_ipc import read_frame, write_frame

log = logging.getLogger("audio.enhance_daemon")

Handler = Callable[[str, np.ndarray], np.ndarray]


class SharedEnhanceEngine:
    """One AudioCleaner model; per-identity context and TDOA."""

    def __init__(self, make_cleaner: Callable):
        self._template = make_cleaner()
        self._sessions: dict[str, object] = {}
        self._lock = threading.Lock()

    def identities(self) -> list[str]:
        return list(self._sessions)

    def _session(self, ident: str):
        sess = self._sessions.get(ident)
        if sess is not None:
            return sess
        sess = self._clone()
        self._sessions[ident] = sess
        return sess

    def _clone(self):
        import copy
        sess = copy.copy(self._template)
        sess._ctx = np.zeros(0, dtype=np.float32)
        sess._last_tdoas = None
        sess._tdoa_age = 0
        sess._client = None
        sess.last_dt_ms = -1.0
        compiled = getattr(sess, "_ov_compiled", None)
        if compiled is not None:
            sess._ov_request = compiled.create_infer_request()
        return sess

    def process(self, ident: str, pcm: np.ndarray) -> np.ndarray:
        with self._lock:
            sess = self._session(ident)
        return sess.process(pcm)


def serve_unix(path: str, handler: Handler, stop: Optional[threading.Event] = None) -> None:
    """Accept unix-socket clients; handler(ident, pcm) -> int16 mono."""
    if os.path.exists(path):
        os.unlink(path)
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    srv.bind(path)
    srv.listen(64)
    srv.settimeout(0.2)

    def _client(conn: socket.socket) -> None:
        try:
            while stop is None or not stop.is_set():
                ident, pcm = read_frame(conn)
                out = handler(ident, pcm)
                out = np.asarray(out, dtype=np.int16)
                write_frame(conn, ident, out)
        except (EOFError, ConnectionError, BrokenPipeError, TimeoutError, OSError):
            pass
        finally:
            conn.close()

    threads: list[threading.Thread] = []
    try:
        while stop is None or not stop.is_set():
            try:
                conn, _ = srv.accept()
            except TimeoutError:
                continue
            conn.settimeout(15.0)
            t = threading.Thread(target=_client, args=(conn,), daemon=True)
            t.start()
            threads.append(t)
    finally:
        srv.close()
        if os.path.exists(path):
            os.unlink(path)


def wait_for_socket(path: str, timeout: float = 90.0) -> bool:
    """Return True once a unix socket at path accepts a connection."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        if os.path.exists(path):
            try:
                s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
                s.settimeout(0.3)
                s.connect(path)
                s.close()
                return True
            except OSError:
                pass
        time.sleep(0.05)
    return False


def make_production_cleaner(**kw):
    from audio_cleanup import AudioCleaner, apply_torch_thread_limits

    threads = int(kw.pop("torch_num_threads", 2))
    apply_torch_thread_limits(threads)
    return AudioCleaner(torch_num_threads=threads, **kw)


def main(argv: Optional[list[str]] = None) -> int:
    import argparse

    ap = argparse.ArgumentParser(description="Shared MetricGAN enhance daemon")
    ap.add_argument("--socket", default="/tmp/livekit-enhance.sock")
    ap.add_argument("--enhance-mode", default="force")
    ap.add_argument("--enhance-model",
                    default="speechbrain/metricgan-plus-voicebank")
    ap.add_argument("--enhance-chunk-s", type=float, default=1.0)
    ap.add_argument("--enhance-hop-s", type=float, default=0.1)
    ap.add_argument("--enhance-infer", default="auto")
    ap.add_argument("--beamform-mode", default="auto")
    ap.add_argument("--torch-num-threads", type=int, default=2)
    ap.add_argument("--tdoa-every", type=int, default=5)
    ap.add_argument("--audio-rate", type=int, default=16000)
    args = ap.parse_args(argv)

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [enhance-daemon] %(levelname)s %(message)s")
    log.info("loading shared MetricGAN (one copy for all cameras)")

    def make():
        return make_production_cleaner(
            mode=args.enhance_mode,
            model_source=args.enhance_model,
            sample_rate=args.audio_rate,
            chunk_s=args.enhance_chunk_s,
            hop_s=args.enhance_hop_s,
            beamform=args.beamform_mode,
            torch_num_threads=args.torch_num_threads,
            tdoa_every=args.tdoa_every,
            enhance_infer=args.enhance_infer,
        )

    engine = SharedEnhanceEngine(make)
    log.info("model ready backend=%s socket=%s",
             getattr(engine._template, "backend", "?"), args.socket)
    serve_unix(args.socket, engine.process)
    return 0


if __name__ == "__main__":
    import sys
    sys.exit(main())