livekit-plugins-fishaudio: pooled websocket reuse + prebuffer to fix cold-start crackle

livekit-plugins/livekit-plugins-fishaudio/livekit/plugins/fishaudio/tts.py

@@ -31,6 +31,12 @@
 NUM_CHANNELS = 1
 USER_AGENT = f"livekit-plugins-fishaudio/{__version__}"
 
+# Hold the first N audio chunks before releasing any audio, so playout starts with
+# enough buffered to ride out the gap after Fish's small first chunk (else jitter
+# underruns it into a crackle). Internal stopgap for Fish's cold-start pacing; drop
+# to 1 once inference streams the opening chunks smoothly. Values <= 1 disable it.
+_PREBUFFER_CHUNKS = 2
+
 # Fish Audio's default sample rate per output format. Opus only supports 48 kHz;
 # the other formats default to 24 kHz, which matches the previous plugin default.
 _DEFAULT_SAMPLE_RATE: dict[OutputFormat, int] = {
@@ -148,6 +154,12 @@ def __init__(
         )
 
         self._session = http_session
+        self._pool = utils.ConnectionPool[aiohttp.ClientWebSocketResponse](
+            connect_cb=self._connect_ws,
+            close_cb=self._close_ws,
+            max_session_duration=300,
+            mark_refreshed_on_get=True,
+        )
         # min_sentence_len=1 emits each sentence as soon as the next one starts,
         # rather than batching short sentences together — minimizes TTFB on the
         # first sentence and keeps Fish synthesizing continuously.
@@ -183,6 +195,27 @@ def _ensure_session(self) -> aiohttp.ClientSession:
             self._session = utils.http_context.http_session()
         return self._session
 
+    async def _connect_ws(self, timeout: float) -> aiohttp.ClientWebSocketResponse:
+        session = self._ensure_session()
+        return await asyncio.wait_for(
+            session.ws_connect(
+                self._opts.get_ws_url("/v1/tts/live"),
+                headers={
+                    "Authorization": f"Bearer {self._opts.api_key}",
+                    "User-Agent": USER_AGENT,
+                    "model": self._opts.model,
+                },
+                heartbeat=30.0,
+            ),
+            timeout,
+        )
+
+    async def _close_ws(self, ws: aiohttp.ClientWebSocketResponse) -> None:
+        await ws.close()
+
+    def prewarm(self) -> None:
+        self._pool.prewarm()
+
     def update_options(
         self,
         *,
@@ -193,8 +226,13 @@ def update_options(
         speed: NotGivenOr[float] = NOT_GIVEN,
         volume: NotGivenOr[float] = NOT_GIVEN,
     ) -> None:
-        if is_given(model):
+        if is_given(model) and model != self._opts.model:
             self._opts.model = model
+            # The model is sent as a connection header at ws-handshake time, not in the
+            # per-request body, so a pooled socket keeps the old model. Drop pooled
+            # connections so the next stream reconnects with the new model. Other
+            # options ride in the per-request body and need no reconnect.
+            self._pool.invalidate()
         if is_given(voice_id):
             self._opts.voice_id = voice_id
         if is_given(latency_mode):
@@ -229,6 +267,7 @@ async def aclose(self) -> None:
         for stream in list(self._streams):
             await stream.aclose()
         self._streams.clear()
+        await self._pool.aclose()
 
 
 def _build_tts_request(opts: _TTSOptions, *, text: str = "") -> dict[str, Any]:
@@ -329,22 +368,9 @@ async def _run(self, output_emitter: tts.AudioEmitter) -> None:
         )
         output_emitter.start_segment(segment_id=request_id)
 
-        ws: aiohttp.ClientWebSocketResponse | None = None
         try:
-            ws = await asyncio.wait_for(
-                self._tts._ensure_session().ws_connect(
-                    self._opts.get_ws_url("/v1/tts/live"),
-                    headers={
-                        "Authorization": f"Bearer {self._opts.api_key}",
-                        "User-Agent": USER_AGENT,
-                        "model": self._opts.model,
-                    },
-                    heartbeat=30.0,
-                ),
-                self._conn_options.timeout,
-            )
-
-            await self._run_ws(ws, output_emitter)
+            async with self._tts._pool.connection(timeout=self._conn_options.timeout) as ws:
+                await self._run_ws(ws, output_emitter)
 
         except asyncio.TimeoutError:
             raise APITimeoutError() from None
@@ -357,8 +383,6 @@ async def _run(self, output_emitter: tts.AudioEmitter) -> None:
         except Exception as e:
             raise APIConnectionError() from e
         finally:
-            if ws is not None:
-                await ws.close()
             output_emitter.end_segment()
 
     async def _run_ws(
@@ -404,6 +428,33 @@ async def send_task() -> None:
 
             await ws.send_bytes(msgpack.packb({"event": "stop"}, use_bin_type=True))
 
+        prebuffer_chunks = _PREBUFFER_CHUNKS
+        prebuffer = bytearray()
+        chunks_seen = 0
+        prebuffering = prebuffer_chunks > 1
+
+        def push_audio(audio: bytes) -> None:
+            nonlocal prebuffering, chunks_seen
+            if prebuffering:
+                prebuffer.extend(audio)
+                chunks_seen += 1
+                if chunks_seen < prebuffer_chunks:
+                    return
+                output_emitter.push(bytes(prebuffer))
+                prebuffer.clear()
+                prebuffering = False
+            else:
+                output_emitter.push(audio)
+
+        def flush_prebuffer() -> None:
+            # Stream ended before we reached `prebuffer_chunks` (short utterance) —
+            # release whatever is held so nothing is dropped or left hanging.
+            nonlocal prebuffering
+            if prebuffering and prebuffer:
+                output_emitter.push(bytes(prebuffer))
+                prebuffer.clear()
+            prebuffering = False
+
         async def recv_task() -> None:
             # No per-receive timeout: Fish has natural inter-sentence gaps that
             # can exceed `_conn_options.timeout` when the LLM is slow. Dead
@@ -431,7 +482,7 @@ async def recv_task() -> None:
                 if event == "audio":
                     audio = data.get("audio")
                     if audio:
-                        output_emitter.push(audio)
+                        push_audio(audio)
                 elif event == "finish":
                     reason = data.get("reason")
                     if reason == "error":
@@ -441,6 +492,7 @@ async def recv_task() -> None:
                             request_id=None,
                             body=str(data),
                         )
+                    flush_prebuffer()
                     break
                 else:
                     logger.debug("unknown Fish Audio event: %s", data)