services(elevenlabs): implement word-by-word support through websockets

This commit is contained in:
Aleix Conchillo Flaqué
2024-09-13 09:31:35 -07:00
parent f08b25dbb2
commit 434493b8aa
6 changed files with 216 additions and 134 deletions

View File

@@ -6,7 +6,6 @@
import asyncio
import io
import time
import wave
from abc import abstractmethod
@@ -171,7 +170,7 @@ class TTSService(AIService):
# if True, TTSService will push TTSStoppedFrames, otherwise subclass must do it
push_stop_frames: bool = False,
# if push_stop_frames is True, wait for this idle period before pushing TTSStoppedFrame
stop_frame_timeout_s: float = 0.8,
stop_frame_timeout_s: float = 1.0,
**kwargs):
super().__init__(**kwargs)
self._aggregate_sentences: bool = aggregate_sentences
@@ -319,16 +318,16 @@ class AsyncTTSService(TTSService):
class AsyncWordTTSService(AsyncTTSService):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self._start_word_timestamp = None
self._initial_word_timestamp = -1
self._words_queue = asyncio.Queue()
self._words_task = self.get_event_loop().create_task(self._words_task_handler())
def init_word_timestamps(self):
if not self._start_word_timestamp:
self._start_word_timestamp = self.get_clock().get_time()
def start_word_timestamps(self):
if self._initial_word_timestamp == -1:
self._initial_word_timestamp = self.get_clock().get_time()
def reset_word_timestamps(self):
self._start_word_timestamp = None
self._initial_word_timestamp = -1
self._word_timestamps = []
async def add_word_timestamps(self, word_times: List[Tuple[str, float]]):
@@ -366,7 +365,7 @@ class AsyncWordTTSService(AsyncTTSService):
await self.push_frame(LLMFullResponseEndFrame())
else:
frame = TextFrame(word)
frame.pts = self._start_word_timestamp + timestamp
frame.pts = self._initial_word_timestamp + timestamp
await self.push_frame(frame)
self._words_queue.task_done()
except asyncio.CancelledError:

View File

@@ -181,7 +181,7 @@ class CartesiaTTSService(AsyncWordTTSService):
)
elif msg["type"] == "chunk":
await self.stop_ttfb_metrics()
self.init_word_timestamps()
self.start_word_timestamps()
frame = AudioRawFrame(
audio=base64.b64decode(msg["data"]),
sample_rate=self._output_format["sample_rate"],

View File

@@ -4,17 +4,30 @@
# SPDX-License-Identifier: BSD 2-Clause License
#
from typing import AsyncGenerator, Literal
import asyncio
import base64
import json
from typing import Any, AsyncGenerator, List, Literal, Mapping, Tuple
from pydantic import BaseModel
from pipecat.frames.frames import AudioRawFrame, Frame, TTSStartedFrame, TTSStoppedFrame
from pipecat.services.ai_services import TTSService
from pipecat.frames.frames import (
AudioRawFrame,
CancelFrame,
EndFrame,
Frame,
StartFrame,
StartInterruptionFrame,
TTSStartedFrame,
TTSStoppedFrame)
from pipecat.processors.frame_processor import FrameDirection
from pipecat.services.ai_services import AsyncWordTTSService
from loguru import logger
# See .env.example for ElevenLabs configuration needed
try:
from elevenlabs.client import AsyncElevenLabs
import websockets
except ModuleNotFoundError as e:
logger.error(f"Exception: {e}")
logger.error(
@@ -35,7 +48,35 @@ def sample_rate_from_output_format(output_format: str) -> int:
return 16000
class ElevenLabsTTSService(TTSService):
def calculate_word_times(
alignment_info: Mapping[str, Any], cumulative_time: float
) -> List[Tuple[str, float]]:
end_times = [
s + d for s,
d in zip(
alignment_info["charStartTimesMs"],
alignment_info["charDurationsMs"])]
zipped_end_times = list(zip(alignment_info["chars"], end_times))
# Get the start time of every character that appears after a space and
# match this to the word
words = "".join(alignment_info["chars"]).split(" ")
# Calculate end time for each word. We do this by finding a space character
# and using the previous word time, also taking into account there might not
# be a space at the end.
times = []
for (i, (a, b)) in enumerate(zipped_end_times):
if a == " " or i == len(zipped_end_times) - 1:
t = cumulative_time + (zipped_end_times[i - 1][1] / 1000.0)
times.append(t)
word_times = list(zip(words, times))
return word_times
class ElevenLabsTTSService(AsyncWordTTSService):
class InputParams(BaseModel):
output_format: Literal["pcm_16000", "pcm_22050", "pcm_24000", "pcm_44100"] = "pcm_16000"
@@ -45,49 +86,186 @@ class ElevenLabsTTSService(TTSService):
api_key: str,
voice_id: str,
model: str = "eleven_turbo_v2_5",
url: str = "wss://api.elevenlabs.io",
params: InputParams = InputParams(),
**kwargs):
super().__init__(**kwargs)
# Aggregating sentences still gives cleaner-sounding results and fewer
# artifacts than streaming one word at a time. On average, waiting for a
# full sentence should only "cost" us 15ms or so with GPT-4o or a Llama
# 3 model, and it's worth it for the better audio quality.
#
# We also don't want to automatically push LLM response text frames,
# because the context aggregators will add them to the LLM context even
# if we're interrupted. ElevenLabs gives us word-by-word timestamps. We
# can use those to generate text frames ourselves aligned with the
# playout timing of the audio!
#
# Finally, ElevenLabs doesn't provide information on when the bot stops
# speaking for a while, so we want the parent class to send TTSStopFrame
# after a short period not receiving any audio.
super().__init__(
aggregate_sentences=True,
push_text_frames=False,
push_stop_frames=True,
stop_frame_timeout_s=2.0,
**kwargs
)
self._api_key = api_key
self._voice_id = voice_id
self._model = model
self._url = url
self._params = params
self._client = AsyncElevenLabs(api_key=api_key)
self._sample_rate = sample_rate_from_output_format(params.output_format)
# Websocket connection to ElevenLabs.
self._websocket = None
# Indicates if we have sent TTSStartedFrame. It will reset to False when
# there's an interruption or TTSStoppedFrame.
self._started = False
self._cumulative_time = 0
def can_generate_metrics(self) -> bool:
return True
async def set_model(self, model: str):
logger.debug(f"Switching TTS model to: [{model}]")
self._model = model
await self._disconnect()
await self._connect()
async def set_voice(self, voice: str):
logger.debug(f"Switching TTS voice to: [{voice}]")
self._voice_id = voice
await self._disconnect()
await self._connect()
async def start(self, frame: StartFrame):
await super().start(frame)
await self._connect()
async def stop(self, frame: EndFrame):
await super().stop(frame)
await self._disconnect()
async def cancel(self, frame: CancelFrame):
await super().cancel(frame)
await self._disconnect()
async def flush_audio(self):
if self._websocket:
msg = {"text": " ", "flush": True}
await self._websocket.send(json.dumps(msg))
async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM):
await super().push_frame(frame, direction)
if isinstance(frame, (TTSStoppedFrame, StartInterruptionFrame)):
self._started = False
if isinstance(frame, TTSStoppedFrame):
await self.add_word_timestamps([("LLMFullResponseEndFrame", 0)])
async def _connect(self):
try:
voice_id = self._voice_id
model = self._model
output_format = self._params.output_format
url = f"{
self._url}/v1/text-to-speech/{voice_id}/stream-input?model_id={model}&output_format={output_format}"
self._websocket = await websockets.connect(url)
self._receive_task = self.get_event_loop().create_task(self._receive_task_handler())
self._keepalive_task = self.get_event_loop().create_task(self._keepalive_task_handler())
# According to ElevenLabs, we should always start with a single space.
msg = {
"text": " ",
"xi_api_key": self._api_key,
}
await self._websocket.send(json.dumps(msg))
except Exception as e:
logger.exception(f"{self} initialization error: {e}")
self._websocket = None
async def _disconnect(self):
try:
await self.stop_all_metrics()
if self._receive_task:
self._receive_task.cancel()
await self._receive_task
self._receive_task = None
if self._keepalive_task:
self._keepalive_task.cancel()
await self._keepalive_task
self._keepalive_task = None
if self._websocket:
await self._websocket.close()
self._websocket = None
self._started = False
except Exception as e:
logger.exception(f"{self} error closing websocket: {e}")
async def _receive_task_handler(self):
try:
async for message in self._websocket:
msg = json.loads(message)
if msg.get("audio"):
await self.stop_ttfb_metrics()
self.start_word_timestamps()
audio = base64.b64decode(msg["audio"])
frame = AudioRawFrame(audio, self._sample_rate, 1)
await self.push_frame(frame)
if msg.get("alignment"):
word_times = calculate_word_times(msg["alignment"], self._cumulative_time)
await self.add_word_timestamps(word_times)
self._cumulative_time = word_times[-1][1]
except asyncio.CancelledError:
pass
except Exception as e:
logger.exception(f"{self} exception: {e}")
async def _keepalive_task_handler(self):
while True:
try:
await asyncio.sleep(10)
await self._send_text("")
except asyncio.CancelledError:
break
except Exception as e:
logger.exception(f"{self} exception: {e}")
async def _send_text(self, text: str):
if self._websocket:
msg = {"text": text + " "}
await self._websocket.send(json.dumps(msg))
async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
logger.debug(f"Generating TTS: [{text}]")
await self.start_tts_usage_metrics(text)
await self.start_ttfb_metrics()
try:
if not self._websocket:
await self._connect()
results = await self._client.generate(
text=text,
voice=self._voice_id,
model=self._model,
output_format=self._params.output_format
)
try:
if not self._started:
await self.push_frame(TTSStartedFrame())
await self.start_ttfb_metrics()
self._started = True
self._cumulative_time = 0
tts_started = False
async for audio in results:
# This is so we send TTSStartedFrame when we have the first audio
# bytes.
if not tts_started:
await self.push_frame(TTSStartedFrame())
tts_started = True
await self.stop_ttfb_metrics()
frame = AudioRawFrame(audio, self._sample_rate, 1)
yield frame
await self.push_frame(TTSStoppedFrame())
await self._send_text(text)
await self.start_tts_usage_metrics(text)
except Exception as e:
logger.error(f"{self} error sending message: {e}")
await self.push_frame(TTSStoppedFrame())
await self._disconnect()
await self._connect()
return
yield None
except Exception as e:
logger.exception(f"{self} exception: {e}")

View File

@@ -60,6 +60,8 @@ class LmntTTSService(AsyncTTSService):
self._speech = None
self._connection = None
self._receive_task = None
# Indicates if we have sent TTSStartedFrame. It will reset to False when
# there's an interruption or TTSStoppedFrame.
self._started = False
def can_generate_metrics(self) -> bool: