Merge pull request #68 from xcc-zach/refactor/api

feature: add PyWebRTC speech enhancer far support

Author: xcc-zach

.gitignore

@@ -128,7 +128,7 @@ celerybeat.pid
 *.sage.py
 
 # Environments
-.env
+*.env
 .venv
 env/
 venv/
@@ -184,4 +184,4 @@ examples/sample_server/node_modules/
 # User defined
 /logs/
 server_configs/
-/data/
+/data/

docs/docs/speech_enhancer_design.md

@@ -0,0 +1,140 @@
+# Speech Enhancer Design
+
+```python
+class SpeechEnhancer(ABC):
+    """Abstract base class for speech enhancement engines."""
+
+    @abstractmethod
+    def enhance(self, audio: bytes, far: bytes) -> bytes:
+        ...
+
+    def flush(self) -> bytes:
+        return b""
+
+    async def async_enhance(
+        self,
+        audio: bytes,
+        far: bytes,
+    ) -> bytes:
+        ...
+
+    async def async_flush(self) -> bytes:
+        ...
+
+    @abstractmethod
+    def reset(self) -> None:
+        ...
+
+    @abstractmethod
+    def clone(self) -> "SpeechEnhancer":
+        ...
+```
+
+## Audio Format
+
+`SpeechEnhancer` input and output use:
+
+- PCM 16-bit
+- Mono
+- 16000 Hz
+- Raw PCM bytes without a WAV header
+
+`audio` is the near-end microphone signal. `far` is the far-end reference signal, usually derived from the TTS audio being played to the user, and is used for acoustic echo cancellation.
+
+The upstream audio pipeline always provides `far` and guarantees that:
+
+- `far` uses the same audio format as `audio`
+- `len(far) == len(audio)`
+- both buffers describe the near-end input and far-end reference for the same time window
+
+## `enhance` and `async_enhance`
+
+The service layer primarily calls `async_enhance`. If the underlying implementation only has a synchronous API, implement `enhance` and reuse the base class thread-pool wrapper for `async_enhance`.
+
+If the underlying implementation is asynchronous or remote, prefer implementing `async_enhance` directly and wrapping it for `enhance`.
+
+```python
+import asyncio
+
+def enhance(self, audio: bytes, far: bytes) -> bytes:
+    return self._run_coro(self.async_enhance(audio, far=far))
+
+def _run_coro(self, coro: "asyncio.Future[bytes]") -> bytes:
+    loop = asyncio.new_event_loop()
+    try:
+        return loop.run_until_complete(coro)
+    finally:
+        loop.close()
+```
+
+## Meaning of `far`
+
+`far` is a required interface argument and enables enhancement engines that need a far-end reference, such as acoustic echo cancellation.
+
+- When nothing is being played, the service layer passes a silent `far` buffer with the same length as `audio`.
+- When TTS is being played, the service layer reads a same-length slice from the TTS reference buffer.
+- If the reference buffer is short, the service layer pads silence on the right and still preserves equal length.
+- FastEnhancer accepts the interface argument but ignores `far` in both local and remote modes.
+- A concrete enhancer implementation that targets an echo-cancellation service may use `far`.
+
+Therefore, implementations may assume that `far` has the same byte length as `audio`. The upstream service pipeline owns this invariant, so concrete enhancers do not need to check it again.
+
+## How the Service Layer Builds the Far-End Reference
+
+`EnhancerManager` consumes `AudioFrameReceived` and publishes `EnhancedAudioFrameReceived`. When a `SpeechEnhancer` model is configured, each user audio frame is processed like this:
+
+```python
+far = far_reference.take(len(audio))
+enhanced = await enhancer.async_enhance(audio, far=far)
+```
+
+The far-end reference comes from `TTSChunkReady`. The service layer copies the TTS chunk that is about to be sent to the client, converts it to the speech-enhancer format, and writes it into a reference buffer.
+
+Current conversion rules:
+
+- Treat input TTS chunks as PCM 16-bit mono
+- Resample from `TTSChunkReady.sample_rate` to 16000 Hz
+- Store the converted audio in a bounded FIFO buffer
+- Keep up to 5 seconds by default, configurable with `far_reference_buffer_seconds`
+
+When user audio arrives, the service layer reads a far-end reference buffer with the same byte length as the current `audio`. If no reference is available, it returns same-length silence.
+
+## Role of `TTSChunkPlayed`
+
+`TTSChunkPlayed` means the frontend has confirmed that a TTS chunk finished playback. The current event does not include a `chunk_id`, client playback timestamp, or sample offset, so it cannot provide sample-level alignment.
+
+The service layer uses it to discard already played reference audio that microphone frames did not consume. This prevents stale TTS audio from being used as `far` when the user waits until TTS playback completes before speaking.
+
+More precise alignment should come from future client-side playback reference feedback or timestamped playback events.
+
+## FastEnhancer Behavior
+
+`FastEnhancer` implements the `SpeechEnhancer` interface but does not use `far`.
+This is true for both local ONNX mode and remote FastEnhancer WebSocket mode.
+Remote FastEnhancer continues to use its legacy binary PCM protocol.
+
+`PyWebRTCAudio` is the concrete adapter for the pywebrtc-audio service. Its
+constructor only accepts `base_url`, and it sends the upstream-provided `audio`
+and `far` to the service's `/v1/stream` JSON WebSocket endpoint.
+
+## `flush`
+
+`flush` / `async_flush` drains tail audio buffered inside the enhancer. The service layer inserts a flush barrier after `VADSpeechEnd` so all earlier audio frames are enhanced and dispatched downstream before the flush runs.
+
+For remote FastEnhancer mode, `flush` sends any locally pending padded audio and then uses the remote FastEnhancer flush command.
+
+## `reset` and `clone`
+
+`reset` should clear the current session's streaming state, such as model caches, input/output buffers, remote connections, and pending audio. It should not reload model weights.
+
+`clone` should create an independent runtime instance for a new session. Clones may share weights, configuration, or read-only resources, but must not share streaming buffers, remote connections, pending audio, or session state.
+
+See [Semantics of `clone()` and `reset()` on Model Objects](model_clone_reset.md).
+
+## Implementation Suggestions
+
+- Keep output length aligned with input `audio` whenever possible, so downstream VAD and ASR timelines do not drift.
+- Do not try to infer TTS alignment inside the model interface; length padding and far selection belong to the service layer.
+- If the implementation requires `far`, it may assume the upstream pipeline already matched its length with `audio`.
+- If the implementation does not support `far`, ignore the argument but do not fail on a silent reference.
+- Remote implementations should clear pending audio on `reset` and connection rebuilds.

docs/docs/speech_enhancer_design.zh.md

@@ -0,0 +1,136 @@
+# 语音增强设计
+
+```python
+class SpeechEnhancer(ABC):
+    """语音增强引擎抽象基类。"""
+
+    @abstractmethod
+    def enhance(self, audio: bytes, far: bytes) -> bytes:
+        ...
+
+    def flush(self) -> bytes:
+        return b""
+
+    async def async_enhance(
+        self,
+        audio: bytes,
+        far: bytes,
+    ) -> bytes:
+        ...
+
+    async def async_flush(self) -> bytes:
+        ...
+
+    @abstractmethod
+    def reset(self) -> None:
+        ...
+
+    @abstractmethod
+    def clone(self) -> "SpeechEnhancer":
+        ...
+```
+
+## 音频格式约定
+
+`SpeechEnhancer` 的输入和输出均使用：
+
+- PCM 16-bit
+- 单声道
+- 16000 Hz
+- 裸 PCM 字节流，不包含 WAV 头
+
+`audio` 表示近端麦克风音频。`far` 表示远端参考音频，通常来自系统正在播放给用户的 TTS 音频，用于声学回声消除。
+
+上游音频管线总是提供 `far`，并保证：
+
+- `far` 与 `audio` 使用相同音频格式
+- `len(far) == len(audio)`
+- 二者描述同一段时间窗口内的近端输入与远端参考
+
+## `enhance` 与 `async_enhance`
+
+服务层实际优先调用 `async_enhance`。如果底层实现只有同步 API，可以实现 `enhance`，并复用基类提供的 `async_enhance` 线程池包装。
+
+如果底层实现本身是异步或远程服务，最佳实践是直接实现 `async_enhance`，再用同步包装实现 `enhance`。
+
+```python
+import asyncio
+
+def enhance(self, audio: bytes, far: bytes) -> bytes:
+    return self._run_coro(self.async_enhance(audio, far=far))
+
+def _run_coro(self, coro: "asyncio.Future[bytes]") -> bytes:
+    loop = asyncio.new_event_loop()
+    try:
+        return loop.run_until_complete(coro)
+    finally:
+        loop.close()
+```
+
+## `far` 的语义
+
+`far` 是必传接口参数，用于支持带远端参考的语音增强或回声消除。
+
+- 没有远端播放时，服务层会传入与 `audio` 等长的静音 `far`。
+- 有 TTS 播放时，服务层会从 TTS 参考缓冲中取出与当前 `audio` 等长的片段。
+- 参考缓冲不足时，服务层会在右侧补静音，仍然保证长度一致。
+- FastEnhancer 会接收接口参数，但本地模式和远程模式都会忽略 `far`。
+- 面向回声消除服务的具体 enhancer 实现可以使用 `far`。
+
+因此，语音增强实现可以安全地假设：`far` 与 `audio` 等长。该约束由上游服务管线负责，具体 enhancer 实现不需要重复检查。
+
+## 服务层如何生成远端参考
+
+`EnhancerManager` 消费 `AudioFrameReceived`，并发布 `EnhancedAudioFrameReceived`。当存在 `SpeechEnhancer` 模型时，它会在处理每个用户音频帧时调用：
+
+```python
+far = far_reference.take(len(audio))
+enhanced = await enhancer.async_enhance(audio, far=far)
+```
+
+远端参考来自 `TTSChunkReady` 事件。服务层会把将要发送给客户端的 TTS chunk 复制一份，转换为语音增强接口要求的格式后写入参考缓冲。
+
+当前转换规则为：
+
+- 输入 TTS chunk 视为 PCM 16-bit 单声道
+- 根据 `TTSChunkReady.sample_rate` 重采样到 16000 Hz
+- 写入一个有限长度的 FIFO 缓冲
+- 默认缓冲上限为 5 秒，可通过 `far_reference_buffer_seconds` 配置
+
+当用户音频到达时，服务层从该缓冲取出与当前 `audio` 等长的远端参考。没有可用参考时返回等长静音。
+
+## `TTSChunkPlayed` 的作用
+
+`TTSChunkPlayed` 表示前端确认某个 TTS chunk 已播放完成。当前事件不携带 `chunk_id`、客户端播放时间戳或样本偏移，因此它不能提供样本级对齐。
+
+服务层将其用于清理已经播放但尚未被麦克风帧消耗的旧参考音频，避免用户等 TTS 播完后再说话时仍然拿旧 TTS 作为 `far`。
+
+更精确的对齐应由未来的客户端播放参考回传或带时间戳的播放事件实现。
+
+## FastEnhancer 行为
+
+`FastEnhancer` 实现 `SpeechEnhancer` 接口，但不使用 `far`。本地 ONNX 模式和远程 FastEnhancer WebSocket 模式都是如此。远程 FastEnhancer 仍然使用旧的二进制 PCM 协议。
+
+`PyWebRTCAudio` 是面向 pywebrtc-audio 服务的具体适配器。它的构造参数仅包含 `base_url`，并把上游提供的 `audio` 和 `far` 发送到服务的 `/v1/stream` JSON WebSocket 接口。
+
+## `flush`
+
+`flush` / `async_flush` 用于排出增强器内部缓冲的尾部音频。服务层会在 `VADSpeechEnd` 后插入 flush barrier，确保所有更早的音频帧先完成增强和下游分发。
+
+对远程 FastEnhancer 模式，`flush` 会发送本地 pending audio 的补齐帧，然后使用远程 FastEnhancer 的 flush 命令。
+
+## `reset` 与 `clone`
+
+`reset` 应清空当前会话的流式状态，例如模型 cache、输入输出缓冲、远程连接和 pending audio，但不应重新加载权重。
+
+`clone` 应为新会话创建独立运行时实例。多个克隆可以共享权重、配置或只读资源，但不能共享流式缓冲、远程连接、pending audio 或会话状态。
+
+请参阅[模型对象的 `clone()` 与 `reset()` 语义](model_clone_reset.zh.md)。
+
+## 实现建议
+
+- 输出长度应尽量与输入 `audio` 长度一致，避免下游 VAD/ASR 时间轴漂移。
+- 不要在模型接口层自行猜测 TTS 对齐；长度补齐和 far 选择由服务层负责。
+- 如果实现需要 `far`，可以假设上游管线已经将其长度与 `audio` 对齐。
+- 如果实现不支持 `far`，可以忽略该参数，但不应因为收到静音 `far` 而失败。
+- 远程实现应在 `reset` 和连接重建时清空 pending audio。

docs/docs/supported_models.md

@@ -117,3 +117,18 @@ Turn detector is used to determine whether the user has finished speaking and de
 [Original Repository](https://github.com/aask1357/fastenhancer)
 
 </details>
+
+<details markdown="1">
+<summary>PyWebRTCAudio</summary>
+
+**Dependency:** `pip install "xtalk[pywebrtc-audio] @ git+https://github.com/xcc-zach/xtalk.git@main"`
+
+**Path:** [`src/xtalk/models/speech_enhancer/pywebrtc_audio.py`](https://github.com/xcc-zach/xtalk/blob/main/src/xtalk/models/speech_enhancer/pywebrtc_audio.py)
+
+**Config params:** only `base_url`.
+
+[Quick Start](https://github.com/xcc-zach/xtalk-pywebrtc-audio)
+
+[Original Repository](https://github.com/strands-labs/pywebrtc-audio)
+
+</details>

docs/docs/supported_models.zh.md

@@ -116,4 +116,19 @@ Turn detector 用于判断用户是否已经说完，并决定系统何时开始
 
 [原始仓库](https://github.com/aask1357/fastenhancer)
 
-</details>
+</details>
+
+<details markdown="1">
+<summary>PyWebRTCAudio</summary>
+
+**依赖：** `pip install "xtalk[pywebrtc-audio] @ git+https://github.com/xcc-zach/xtalk.git@main"`
+
+**路径：** [`src/xtalk/models/speech_enhancer/pywebrtc_audio.py`](https://github.com/xcc-zach/xtalk/blob/main/src/xtalk/models/speech_enhancer/pywebrtc_audio.py)
+
+**配置参数：** 仅 `base_url`。
+
+[快速开始](https://github.com/xcc-zach/xtalk-pywebrtc-audio)
+
+[原始仓库](https://github.com/strands-labs/pywebrtc-audio)
+
+</details>

docs/mkdocs.yml

@@ -57,6 +57,7 @@ plugins:
             'Service Configuration': 服务配置项
             'ASR Design': ASR设计
             'TTS Design': TTS设计
+            'Speech Enhancer Design': 语音增强设计
             'VAD Design': VAD设计
             'Turn Detector Design': Turn Detector设计
             'Model clone() and reset()': 模型对象的clone()与reset()语义
@@ -83,6 +84,7 @@ nav:
       - 'Service Configuration': docs/service_config.md
       - 'ASR Design': docs/asr_design.md
       - 'TTS Design': docs/tts_design.md
+      - 'Speech Enhancer Design': docs/speech_enhancer_design.md
       - 'VAD Design': docs/vad_design.md
       - 'Turn Detector Design': docs/turn_detector_design.md
       - 'Model clone() and reset()': docs/model_clone_reset.md

pyproject.toml

@@ -135,6 +135,9 @@ fast-enhancer = [
   "onnxruntime",
   "websockets",
 ]
+pywebrtc-audio = [
+  "websockets",
+]
 rubberband = [
   "pyrubberband"
 ]