livekit-plugins-speakerbeamss 0.1.0__tar.gz

livekit_plugins_speakerbeamss-0.1.0/.gitignore

@@ -0,0 +1,27 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info/
+
+# Virtual environments
+.venv
+
+# ONNX / model artifacts (downloaded separately via export_onnx.py)
+exports/*.onnx
+*.onnx
+
+# Pytest
+.pytest_cache/
+.coverage
+htmlcov/
+
+# Editor
+.vscode/
+.idea/
+*.swp
+
+# OS
+.DS_Store

livekit_plugins_speakerbeamss-0.1.0/PKG-INFO

@@ -0,0 +1,169 @@
+Metadata-Version: 2.5
+Name: livekit-plugins-speakerbeamss
+Version: 0.1.0
+Summary: LiveKit Agents STT plugin that runs OpenSpeakerBeam-SS target-speaker extraction in front of any inner STT.
+Author-email: linsan <bin.zaq@foxmail.com>
+License: Apache-2.0
+Keywords: asr,audio,livekit,realtime,speaker-extraction,target-speaker,webrtc
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Multimedia :: Sound/Audio
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: livekit-agents>=1.5.13
+Requires-Dist: numpy>=1.24
+Requires-Dist: onnxruntime>=1.17.0
+Requires-Dist: resemblyzer>=0.1.1
+Requires-Dist: soundfile>=0.12
+Provides-Extra: examples
+Requires-Dist: livekit-plugins-volcengine>=1.3.26; extra == 'examples'
+Provides-Extra: test
+Requires-Dist: numpy>=1.24; extra == 'test'
+Requires-Dist: onnx>=1.14; extra == 'test'
+Requires-Dist: pytest-asyncio>=0.21; extra == 'test'
+Requires-Dist: pytest>=7.0; extra == 'test'
+Requires-Dist: setuptools<81; extra == 'test'
+Description-Content-Type: text/markdown
+
+# livekit-plugins-speakerbeamss
+
+> A [LiveKit Agents](https://github.com/livekit/agents) STT plugin that
+> runs [OpenSpeakerBeam-SS](../OpenSpeakerBeam-SS) target-speaker
+> extraction in front of any inner STT instance. The agent hears only
+> the user — even with competing speakers, music, or noise in the room.
+
+The plugin is **ONNX-only**. The PyTorch separator never ships here.
+The **trained INT8 model is bundled inside the wheel** at
+`livekit/plugins/speakerbeamss/models/speakerbeam_ss_trained_int8.onnx`
+(≈13 MB), so a fresh `pip install livekit-plugins-speakerbeamss` is
+enough — no separate download. Follow step 0 below only if you want to
+swap in a different export (FP32, the base release, or your own
+fine-tune).
+
+The reference audio for the target speaker is taken from the
+*agent's first user turn*: the first sentence the user speaks is
+forwarded to the inner STT **as-is** so the agent can still
+understand it, and is *also* kept in a buffer. When the inner STT
+emits `END_OF_SPEECH`, the buffer is collapsed to a Resemblyzer
+d-vector and the plugin flips to extraction mode for every
+subsequent utterance.
+
+You can also pass a pre-recorded enrollment audio via
+`reference_audio="path.wav"` to skip the capture phase entirely.
+
+## Installation
+
+```bash
+# 1. Install this plugin (from the plugin source root):
+cd livekit/plugins/speakerbeamss
+pip install -e .
+```
+
+That's it — `pip install` wires up `livekit.plugins.speakerbeamss`
+*and* ships the trained ONNX bundle, which the loader resolves via
+`importlib.resources`. No env vars, no separate download.
+
+### Using a different exported variant
+
+If you have your own export (FP32, base release, a re-trained
+checkpoint, …) you'd rather use, point the loader at it explicitly:
+
+```python
+from livekit.plugins.speakerbeamss import STT
+
+sb_stt = STT(
+    inner=inner,
+    onnx_dir="/Users/linsan/Public/projects/OpenSpeakerBeam-SS/exports",
+    onnx_filename="speakerbeam_ss_int8.onnx",   # base INT8 release
+)
+```
+
+The loader's lookup order is:
+
+1. **Bundled** — `livekit.plugins.speakerbeamss.models/<​filename>`
+   (the trained INT8 lives here by default).
+2. **Disk** — `{onnx_dir}/{filename}` if set, otherwise the
+   `~/Public/projects/OpenSpeakerBeam-SS/exports/` fallback.
+
+To export a different variant yourself:
+
+```bash
+cd ~/Public/projects/OpenSpeakerBeam-SS
+pip install -r requirements.txt
+python export_onnx.py --ckpt checkpoints/best_model.pth
+ls exports/   # -> speakerbeam_ss_{fp32,int8,trained_fp32,trained_int8}.onnx
+```
+
+## Usage — wrap any inner STT
+
+```python
+from livekit.plugins.speakerbeamss import STT
+from livekit.plugins.volcengine import STT as VolcSTT  # any streaming inner STT
+
+inner = VolcSTT()                          # read keys from env, etc.
+stt = STT(inner=inner)
+# ... pass `stt` to your VoicePipelineAgent as usual.
+```
+
+That's it. The agent will see exactly the transcripts it would have
+seen without the wrapper, but with the first turn's enrollment
+frozen into a d-vector and applied to every chunk that follows.
+
+## Usage — offline
+
+```python
+import numpy as np
+from livekit.plugins.speakerbeamss import OnnxSpeakerBeamExtractor
+
+extractor = OnnxSpeakerBeamExtractor()                     # INT8 by default
+extractor.enroll(np.random.randn(16000).astype(np.float32))  # reference audio
+
+chunk = np.random.randn(8000).astype(np.float32)             # < 1 s
+print(extractor.append(chunk))    # []  (no full chunk yet)
+print(extractor.append(np.zeros(8000, dtype=np.float32)))    # [ndarray of len 16000]
+```
+
+## How the state machine works
+
+| Phase | Inner STT input | Plugin behavior |
+|------|-----------------|------------------|
+| `CAPTURING` (first user turn) | Raw audio, downmixed to mono 16 kHz | Buffers audio; on `END_OF_SPEECH`, computes d-vector and locks |
+| `LOCKED` (after first turn) | ONNX-extracted mono 16 kHz | Embedding reused for all subsequent user turns until stream ends |
+| `EXTERNAL` (when `reference_audio` was provided) | Raw audio, downmixed to mono 16 kHz | No enrollment phase; extraction applied from the first frame |
+
+## Configuration knobs
+
+`STT(...)` arguments:
+
+| Name | Default | Purpose |
+|------|---------|---------|
+| `inner` | *(required)* | Any `stt.STT` with `streaming=True`. |
+| `reference_audio` | `None` | Path to an enrollment WAV/PCM file. Bypasses the capture phase. |
+| `onnx_dir` | `~/Public/projects/OpenSpeakerBeam-SS/exports` | Fallback directory for `.onnx` files (used only when the bundled model is missing). |
+| `onnx_filename` | `speakerbeam_ss_trained_int8.onnx` (bundled) | Override the exact model filename. Looked up in the bundled `models/` directory first, then on disk. |
+| `chunk_seconds` | `1.0` | ONNX window length. Must match the time axis baked into the export. |
+| `sample_rate` | `16_000` | Sample rate for both extractor and the framework resampler. |
+| `min_enrollment_seconds` | `0.5` | Refuse to lock enrollment from < this much captured audio. |
+| `pass_through_first_utterance` | `True` | When `True`, the first turn reaches the inner STT untouched so the agent's first reply still makes sense. |
+| `interim_results` | *(inherit from inner)* | Capability flag for downstream `VoicePipelineAgent`. |
+
+Environment variables:
+
+| Var | Purpose |
+|-----|---------|
+| `SPEAKERBEAMSS_ONNX_DIR` | Override `onnx_dir` for one-off runs. |
+
+## Verified against
+
+- LiveKit Agents `>= 1.5.13`
+- OpenSpeakerBeam-SS commit exporting
+  `speakerbeam_ss_trained_int8.onnx` (~13 MB) at opset 20, chunk length
+  16 000 samples / 16 kHz — the bundle that ships inside the wheel.
+
+## License
+
+Apache-2.0 (same as LiveKit Agents).

livekit_plugins_speakerbeamss-0.1.0/README.md

@@ -0,0 +1,138 @@
+# livekit-plugins-speakerbeamss
+
+> A [LiveKit Agents](https://github.com/livekit/agents) STT plugin that
+> runs [OpenSpeakerBeam-SS](../OpenSpeakerBeam-SS) target-speaker
+> extraction in front of any inner STT instance. The agent hears only
+> the user — even with competing speakers, music, or noise in the room.
+
+The plugin is **ONNX-only**. The PyTorch separator never ships here.
+The **trained INT8 model is bundled inside the wheel** at
+`livekit/plugins/speakerbeamss/models/speakerbeam_ss_trained_int8.onnx`
+(≈13 MB), so a fresh `pip install livekit-plugins-speakerbeamss` is
+enough — no separate download. Follow step 0 below only if you want to
+swap in a different export (FP32, the base release, or your own
+fine-tune).
+
+The reference audio for the target speaker is taken from the
+*agent's first user turn*: the first sentence the user speaks is
+forwarded to the inner STT **as-is** so the agent can still
+understand it, and is *also* kept in a buffer. When the inner STT
+emits `END_OF_SPEECH`, the buffer is collapsed to a Resemblyzer
+d-vector and the plugin flips to extraction mode for every
+subsequent utterance.
+
+You can also pass a pre-recorded enrollment audio via
+`reference_audio="path.wav"` to skip the capture phase entirely.
+
+## Installation
+
+```bash
+# 1. Install this plugin (from the plugin source root):
+cd livekit/plugins/speakerbeamss
+pip install -e .
+```
+
+That's it — `pip install` wires up `livekit.plugins.speakerbeamss`
+*and* ships the trained ONNX bundle, which the loader resolves via
+`importlib.resources`. No env vars, no separate download.
+
+### Using a different exported variant
+
+If you have your own export (FP32, base release, a re-trained
+checkpoint, …) you'd rather use, point the loader at it explicitly:
+
+```python
+from livekit.plugins.speakerbeamss import STT
+
+sb_stt = STT(
+    inner=inner,
+    onnx_dir="/Users/linsan/Public/projects/OpenSpeakerBeam-SS/exports",
+    onnx_filename="speakerbeam_ss_int8.onnx",   # base INT8 release
+)
+```
+
+The loader's lookup order is:
+
+1. **Bundled** — `livekit.plugins.speakerbeamss.models/<​filename>`
+   (the trained INT8 lives here by default).
+2. **Disk** — `{onnx_dir}/{filename}` if set, otherwise the
+   `~/Public/projects/OpenSpeakerBeam-SS/exports/` fallback.
+
+To export a different variant yourself:
+
+```bash
+cd ~/Public/projects/OpenSpeakerBeam-SS
+pip install -r requirements.txt
+python export_onnx.py --ckpt checkpoints/best_model.pth
+ls exports/   # -> speakerbeam_ss_{fp32,int8,trained_fp32,trained_int8}.onnx
+```
+
+## Usage — wrap any inner STT
+
+```python
+from livekit.plugins.speakerbeamss import STT
+from livekit.plugins.volcengine import STT as VolcSTT  # any streaming inner STT
+
+inner = VolcSTT()                          # read keys from env, etc.
+stt = STT(inner=inner)
+# ... pass `stt` to your VoicePipelineAgent as usual.
+```
+
+That's it. The agent will see exactly the transcripts it would have
+seen without the wrapper, but with the first turn's enrollment
+frozen into a d-vector and applied to every chunk that follows.
+
+## Usage — offline
+
+```python
+import numpy as np
+from livekit.plugins.speakerbeamss import OnnxSpeakerBeamExtractor
+
+extractor = OnnxSpeakerBeamExtractor()                     # INT8 by default
+extractor.enroll(np.random.randn(16000).astype(np.float32))  # reference audio
+
+chunk = np.random.randn(8000).astype(np.float32)             # < 1 s
+print(extractor.append(chunk))    # []  (no full chunk yet)
+print(extractor.append(np.zeros(8000, dtype=np.float32)))    # [ndarray of len 16000]
+```
+
+## How the state machine works
+
+| Phase | Inner STT input | Plugin behavior |
+|------|-----------------|------------------|
+| `CAPTURING` (first user turn) | Raw audio, downmixed to mono 16 kHz | Buffers audio; on `END_OF_SPEECH`, computes d-vector and locks |
+| `LOCKED` (after first turn) | ONNX-extracted mono 16 kHz | Embedding reused for all subsequent user turns until stream ends |
+| `EXTERNAL` (when `reference_audio` was provided) | Raw audio, downmixed to mono 16 kHz | No enrollment phase; extraction applied from the first frame |
+
+## Configuration knobs
+
+`STT(...)` arguments:
+
+| Name | Default | Purpose |
+|------|---------|---------|
+| `inner` | *(required)* | Any `stt.STT` with `streaming=True`. |
+| `reference_audio` | `None` | Path to an enrollment WAV/PCM file. Bypasses the capture phase. |
+| `onnx_dir` | `~/Public/projects/OpenSpeakerBeam-SS/exports` | Fallback directory for `.onnx` files (used only when the bundled model is missing). |
+| `onnx_filename` | `speakerbeam_ss_trained_int8.onnx` (bundled) | Override the exact model filename. Looked up in the bundled `models/` directory first, then on disk. |
+| `chunk_seconds` | `1.0` | ONNX window length. Must match the time axis baked into the export. |
+| `sample_rate` | `16_000` | Sample rate for both extractor and the framework resampler. |
+| `min_enrollment_seconds` | `0.5` | Refuse to lock enrollment from < this much captured audio. |
+| `pass_through_first_utterance` | `True` | When `True`, the first turn reaches the inner STT untouched so the agent's first reply still makes sense. |
+| `interim_results` | *(inherit from inner)* | Capability flag for downstream `VoicePipelineAgent`. |
+
+Environment variables:
+
+| Var | Purpose |
+|-----|---------|
+| `SPEAKERBEAMSS_ONNX_DIR` | Override `onnx_dir` for one-off runs. |
+
+## Verified against
+
+- LiveKit Agents `>= 1.5.13`
+- OpenSpeakerBeam-SS commit exporting
+  `speakerbeam_ss_trained_int8.onnx` (~13 MB) at opset 20, chunk length
+  16 000 samples / 16 kHz — the bundle that ships inside the wheel.
+
+## License
+
+Apache-2.0 (same as LiveKit Agents).

livekit_plugins_speakerbeamss-0.1.0/pyproject.toml

@@ -0,0 +1,84 @@
+[project]
+name = "livekit-plugins-speakerbeamss"
+dynamic = ["version"]
+description = "LiveKit Agents STT plugin that runs OpenSpeakerBeam-SS target-speaker extraction in front of any inner STT."
+readme = "README.md"
+authors = [{ name = "linsan", email = "bin.zaq@foxmail.com" }]
+keywords = ["webrtc", "realtime", "audio", "livekit", "speaker-extraction", "target-speaker", "asr"]
+requires-python = ">=3.10"
+license = { text = "Apache-2.0" }
+dependencies = [
+    # LiveKit Agents framework (provides STT base classes, AudioByteStream, AudioBuffer, Plugin)
+    "livekit-agents>=1.5.13",
+    # ONNX runtime — the *only* inference backend the plugin ships with.
+    "onnxruntime>=1.17.0",
+    # Resemblyzer provides the d-vector speaker encoder used for enrollment.
+    # Resemblyzer itself pulls in torch (used only by the VoiceEncoder, not for the
+    # separator — the plugin model is ONNX-only).
+    "resemblyzer>=0.1.1",
+    # numpy for audio numpy interop; soundfile for offline example I/O.
+    "numpy>=1.24",
+    "soundfile>=0.12",
+]
+
+classifiers = [
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: Apache Software License",
+    "Topic :: Multimedia :: Sound/Audio",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+]
+
+[project.optional-dependencies]
+test = [
+    "pytest>=7.0",
+    "pytest-asyncio>=0.21",
+    # for the SpeechStream-pattern fake STT used in tests / examples
+    "numpy>=1.24",
+    # for building the no-op ONNX fixture used in tests
+    "onnx>=1.14",
+    # resemblyzer needs setuptools < 81 to expose pkg_resources
+    "setuptools<81",
+]
+examples = [
+    "livekit-plugins-volcengine>=1.3.26",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.version]
+path = "src/livekit/plugins/speakerbeamss/version.py"
+
+# Hatch bundles everything under `packages = ["src/livekit"]`, including
+# the trained ONNX file at
+# `src/livekit/plugins/speakerbeamss/models/*.onnx`. The explicit
+# `artifacts` declarations below make that intent obvious and guarantee
+# the binary file lands in both wheel and sdist even if a future revision
+# moves it around.
+[tool.hatch.build.targets.wheel]
+packages = ["src/livekit"]
+artifacts = [
+    "src/livekit/plugins/speakerbeamss/models/*.onnx",
+]
+
+[tool.hatch.build.targets.sdist]
+include = ["/src"]
+artifacts = [
+    "src/livekit/plugins/speakerbeamss/models/*.onnx",
+]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+testpaths = ["tests"]
+log_cli = true
+log_cli_level = "INFO"
+log_cli_format = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+log_cli_date_format = "%H:%M:%S"
+markers = [
+    "integration: marks tests that require a real exported ONNX bundle under OpenSpeakerBeam-SS/exports/",
+]

livekit_plugins_speakerbeamss-0.1.0/src/livekit/plugins/speakerbeamss/__init__.py

@@ -0,0 +1,72 @@
+"""LiveKit Agents plugin: target-speaker extraction in front of any STT.
+
+Public surface:
+
+    >>> from livekit.plugins.speakerbeamss import STT
+    >>> from livekit.plugins.volcengine import STT as VolcSTT
+    >>> stt = STT(inner=VolcSTT())          # generic wrapper
+    >>> s = stt.stream(language="zh-CN")
+
+Or use the standalone preprocessor without LiveKit:
+
+    >>> from livekit.plugins.speakerbeamss import (
+    ...     OnnxSpeakerBeamExtractor,
+    ...     EnrollmentState,
+    ...     ExtractorConfig,
+    ... )
+    >>> ext = OnnxSpeakerBeamExtractor()
+    >>> ext.enroll(reference_clip)         # 1-D numpy float32 at 16 kHz
+    >>> out = ext.process(chunk)            # chunk of any size
+
+The model is ONNX-only: ensure you ran ``python export_onnx.py`` in the
+OpenSpeakerBeam-SS repository (producing ``exports/speakerbeam_ss_*.onnx``)
+and pointed ``ExtractorConfig.onnx_dir`` (or the env var
+``SPEAKERBEAMSS_ONNX_DIR``) at the right directory before importing
+this module.
+"""
+
+from __future__ import annotations
+
+from .plugin import SpeakerBeamSSPlugin
+from .processor import (
+    DEFAULT_SAMPLE_RATE,
+    EMBEDDING_DIM,
+    EnrollmentState,
+    ExtractorConfig,
+    OnnxSpeakerBeamExtractor,
+)
+from .stt import SpeakerBeamSTT
+from .version import __version__
+
+# Alias `STT` to the wrapper class so callers can ``from livekit.plugins
+# .speakerbeamss import STT`` — matches every other LiveKit plugin's
+# primary export shape.
+STT = SpeakerBeamSTT
+
+__all__ = [
+    "STT",
+    "SpeakerBeamSTT",
+    "OnnxSpeakerBeamExtractor",
+    "ExtractorConfig",
+    "EnrollmentState",
+    "SpeakerBeamSSPlugin",
+    "DEFAULT_SAMPLE_RATE",
+    "EMBEDDING_DIM",
+    "__version__",
+]
+
+
+# ---------------------------------------------------------------------------
+# Plugin registration
+# ---------------------------------------------------------------------------
+from livekit.agents import Plugin as _Plugin
+
+_Plugin.register_plugin(SpeakerBeamSSPlugin())
+
+
+# Cleanup docs of unexported modules — mirrors livekit-plugins-volcengine
+_module = dir()
+NOT_IN_ALL = [m for m in _module if m not in __all__]
+__pdoc__ = {}
+for n in NOT_IN_ALL:
+    __pdoc__[n] = False

livekit_plugins_speakerbeamss-0.1.0/src/livekit/plugins/speakerbeamss/log.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+import logging
+
+logger = logging.getLogger("livekit.plugins.speakerbeamss")

livekit_plugins_speakerbeamss-0.1.0/src/livekit/plugins/speakerbeamss/models/__init__.py

@@ -0,0 +1,7 @@
+"""Bundled SpeakerBeam-SS ONNX models.
+
+The trained INT8 model ships inside the wheel so a fresh
+``pip install livekit-plugins-speakerbeamss`` is enough — no separate
+download is needed. ``processor._bundled_model_path`` resolves files
+relative to this package via :mod:`importlib.resources`.
+"""

livekit_plugins_speakerbeamss-0.1.0/src/livekit/plugins/speakerbeamss/plugin.py

@@ -0,0 +1,32 @@
+"""LiveKit plugin entry-point.
+
+Mirrors the pattern in
+``livekit-plugins-volcengine/livekit/plugins/volcengine/__init__.py``.
+Registering an instance via :py:meth:`Plugin.register_plugin` lets the
+framework discover this package at import time (e.g. when LiveKit's
+plugin loader scans for ``livekit.plugins.*`` modules).
+"""
+
+from __future__ import annotations
+
+from livekit.agents import Plugin
+
+from .log import logger
+from .version import __version__
+
+
+class SpeakerBeamSSPlugin(Plugin):
+    """Plugin descriptor for ``livekit.plugins.speakerbeamss``.
+
+    Held by LiveKit's global plugin registry. The plugin does not run
+    any background work; it just exposes metadata so the framework can
+    see this package is installed.
+    """
+
+    def __init__(self) -> None:
+        super().__init__(__name__, __version__, __package__, logger)
+        logger.info(
+            "speakerbeam-ss plugin registered (version=%s, package=%s)",
+            __version__,
+            __package__,
+        )