wehub-resource-sync
532 lines
18 KiB
Python
532 lines
18 KiB
Python
"""
|
|
Uplift TTS Plugin for LiveKit, this will soon be available as a python lib
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import os
|
|
import time
|
|
import uuid
|
|
import weakref
|
|
from dataclasses import dataclass
|
|
from typing import Any, Literal
|
|
|
|
import socketio # type: ignore[import-not-found]
|
|
|
|
from livekit.agents import (
|
|
APIConnectionError,
|
|
APIConnectOptions,
|
|
APIError,
|
|
APITimeoutError,
|
|
tokenize,
|
|
tts,
|
|
utils,
|
|
)
|
|
from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS, NOT_GIVEN, NotGivenOr
|
|
from livekit.agents.utils import is_given
|
|
|
|
from .log import logger
|
|
|
|
# Output format options
|
|
OutputFormat = Literal[
|
|
"PCM_22050_16",
|
|
"WAV_22050_16",
|
|
"WAV_22050_32",
|
|
"MP3_22050_32",
|
|
"MP3_22050_64",
|
|
"MP3_22050_128",
|
|
"OGG_22050_16",
|
|
"ULAW_8000_8",
|
|
]
|
|
|
|
# Default configuration
|
|
DEFAULT_BASE_URL = "wss://api.upliftai.org"
|
|
DEFAULT_SAMPLE_RATE = 22050
|
|
DEFAULT_NUM_CHANNELS = 1
|
|
DEFAULT_VOICE_ID = "v_meklc281"
|
|
DEFAULT_OUTPUT_FORMAT: OutputFormat = "MP3_22050_32"
|
|
WEBSOCKET_NAMESPACE = "/text-to-speech/multi-stream"
|
|
|
|
|
|
def get_content_type_from_output_format(output_format: OutputFormat) -> str:
|
|
"""Get MIME type based on output format"""
|
|
if output_format == "PCM_22050_16":
|
|
return "audio/pcm"
|
|
elif output_format == "WAV_22050_16":
|
|
return "audio/wav"
|
|
elif output_format == "WAV_22050_32":
|
|
return "audio/wav"
|
|
elif output_format.startswith("MP3"):
|
|
return "audio/mpeg"
|
|
elif output_format.startswith("OGG"):
|
|
return "audio/ogg"
|
|
elif output_format == "ULAW_8000_8":
|
|
return "audio/x-mulaw"
|
|
else:
|
|
raise ValueError(f"Unsupported output format: {output_format}")
|
|
|
|
|
|
@dataclass
|
|
class VoiceSettings:
|
|
"""Voice configuration settings"""
|
|
|
|
voice_id: str = DEFAULT_VOICE_ID
|
|
output_format: OutputFormat = DEFAULT_OUTPUT_FORMAT
|
|
|
|
|
|
@dataclass
|
|
class _TTSOptions:
|
|
"""Internal TTS options"""
|
|
|
|
base_url: str
|
|
api_key: str
|
|
voice_settings: VoiceSettings
|
|
word_tokenizer: tokenize.WordTokenizer | tokenize.SentenceTokenizer
|
|
sample_rate: int
|
|
num_channels: int
|
|
phrase_replacement_config_id: str | None
|
|
|
|
|
|
class TTS(tts.TTS):
|
|
"""Uplift TTS implementation for LiveKit"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
base_url: NotGivenOr[str] = NOT_GIVEN,
|
|
api_key: NotGivenOr[str] = NOT_GIVEN,
|
|
voice_id: str = DEFAULT_VOICE_ID,
|
|
output_format: OutputFormat = DEFAULT_OUTPUT_FORMAT,
|
|
num_channels: int = DEFAULT_NUM_CHANNELS,
|
|
phrase_replacement_config_id: NotGivenOr[str] = NOT_GIVEN,
|
|
word_tokenizer: NotGivenOr[tokenize.WordTokenizer | tokenize.SentenceTokenizer] = NOT_GIVEN,
|
|
) -> None:
|
|
"""
|
|
Create a new instance of Uplift TTS.
|
|
|
|
Args:
|
|
base_url: Base URL for TTS service. Defaults to wss://api.upliftai.org
|
|
api_key: API key for authentication
|
|
voice_id: Voice ID to use. Defaults to "17"
|
|
output_format: Audio output format. Options:
|
|
- 'PCM_22050_16': PCM format, 22.05kHz, 16-bit
|
|
- 'WAV_22050_16': WAV format, 22.05kHz, 16-bit
|
|
- 'WAV_22050_32': WAV format, 22.05kHz, 32-bit
|
|
- 'MP3_22050_32': MP3 format, 22.05kHz, 32kbps (default)
|
|
- 'MP3_22050_64': MP3 format, 22.05kHz, 64kbps
|
|
- 'MP3_22050_128': MP3 format, 22.05kHz, 128kbps
|
|
- 'OGG_22050_16': OGG format, 22.05kHz, 16-bit
|
|
- 'ULAW_8000_8': μ-law format, 8kHz, 8-bit
|
|
sample_rate: Sample rate for audio output. Defaults to 22050
|
|
num_channels: Number of audio channels. Defaults to 1 (mono)
|
|
phrase_replacement_config_id: Optional ID for phrase replacement configuration
|
|
word_tokenizer: Tokenizer for processing text. Defaults to `livekit.agents.tokenize.basic.WordTokenizer`.
|
|
"""
|
|
super().__init__(
|
|
capabilities=tts.TTSCapabilities(
|
|
streaming=True,
|
|
aligned_transcript=False,
|
|
),
|
|
sample_rate=DEFAULT_SAMPLE_RATE,
|
|
num_channels=num_channels,
|
|
)
|
|
|
|
# Get configuration from environment if not provided
|
|
resolved_base_url: str = (
|
|
base_url
|
|
if is_given(base_url)
|
|
else os.environ.get("UPLIFTAI_BASE_URL", DEFAULT_BASE_URL)
|
|
)
|
|
resolved_api_key: str | None = (
|
|
api_key if is_given(api_key) else os.environ.get("UPLIFTAI_API_KEY")
|
|
)
|
|
|
|
if not resolved_api_key:
|
|
raise ValueError(
|
|
"API key is required, either as argument or set UPLIFTAI_API_KEY environment variable"
|
|
)
|
|
|
|
# Use provided tokenizer or create default
|
|
resolved_word_tokenizer: tokenize.WordTokenizer | tokenize.SentenceTokenizer
|
|
if is_given(word_tokenizer):
|
|
resolved_word_tokenizer = word_tokenizer
|
|
else:
|
|
resolved_word_tokenizer = tokenize.basic.WordTokenizer(ignore_punctuation=False)
|
|
|
|
self._opts = _TTSOptions(
|
|
base_url=resolved_base_url,
|
|
api_key=resolved_api_key,
|
|
voice_settings=VoiceSettings(voice_id=voice_id, output_format=output_format),
|
|
word_tokenizer=resolved_word_tokenizer,
|
|
sample_rate=DEFAULT_SAMPLE_RATE,
|
|
num_channels=num_channels,
|
|
phrase_replacement_config_id=phrase_replacement_config_id
|
|
if is_given(phrase_replacement_config_id)
|
|
else None,
|
|
)
|
|
|
|
self._client: WebSocketClient | None = None
|
|
self._streams = weakref.WeakSet[SynthesizeStream]()
|
|
|
|
def update_options(
|
|
self,
|
|
*,
|
|
voice_id: NotGivenOr[str] = NOT_GIVEN,
|
|
output_format: NotGivenOr[OutputFormat] = NOT_GIVEN,
|
|
) -> None:
|
|
"""
|
|
Update TTS configuration options.
|
|
|
|
Args:
|
|
voice_id: New voice ID
|
|
output_format: New output format (see __init__ for options)
|
|
"""
|
|
if is_given(voice_id):
|
|
self._opts.voice_settings.voice_id = voice_id
|
|
if is_given(output_format):
|
|
self._opts.voice_settings.output_format = output_format
|
|
|
|
def synthesize(
|
|
self, text: str, *, conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS
|
|
) -> ChunkedStream:
|
|
"""Synthesize text to speech using chunked stream."""
|
|
return ChunkedStream(tts=self, input_text=text, conn_options=conn_options)
|
|
|
|
def stream(
|
|
self, *, conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS
|
|
) -> SynthesizeStream:
|
|
"""Create a streaming synthesis session."""
|
|
stream = SynthesizeStream(tts=self, conn_options=conn_options)
|
|
self._streams.add(stream)
|
|
return stream
|
|
|
|
async def aclose(self) -> None:
|
|
"""Clean up resources"""
|
|
for stream in list(self._streams):
|
|
await stream.aclose()
|
|
|
|
self._streams.clear()
|
|
|
|
if self._client:
|
|
await self._client.disconnect()
|
|
self._client = None
|
|
|
|
|
|
class WebSocketClient:
|
|
"""Manages WebSocket connection to TTS service"""
|
|
|
|
def __init__(self, opts: _TTSOptions):
|
|
self.opts = opts
|
|
self.sio: socketio.AsyncClient | None = None
|
|
self.connected = False
|
|
self.audio_callbacks: dict[str, asyncio.Queue[bytes | None]] = {}
|
|
self.active_requests: dict[str, bool] = {}
|
|
|
|
async def connect(self) -> bool:
|
|
"""Establish WebSocket connection"""
|
|
if self.connected:
|
|
return True
|
|
|
|
try:
|
|
self.sio = socketio.AsyncClient(
|
|
reconnection=True,
|
|
reconnection_attempts=3,
|
|
reconnection_delay=1,
|
|
logger=False,
|
|
engineio_logger=False,
|
|
)
|
|
|
|
# Register handlers
|
|
self.sio.on("message", self._on_message, namespace=WEBSOCKET_NAMESPACE)
|
|
self.sio.on("connect", self._on_connect, namespace=WEBSOCKET_NAMESPACE)
|
|
self.sio.on("disconnect", self._on_disconnect, namespace=WEBSOCKET_NAMESPACE)
|
|
|
|
# Prepare auth
|
|
auth_data = {"token": self.opts.api_key}
|
|
|
|
# Connect
|
|
await self.sio.connect(
|
|
self.opts.base_url,
|
|
auth=auth_data,
|
|
namespaces=[WEBSOCKET_NAMESPACE],
|
|
transports=["websocket"],
|
|
wait_timeout=10,
|
|
)
|
|
|
|
# Wait for connection
|
|
max_wait = 5.0
|
|
start_time = time.time()
|
|
while not self.connected and (time.time() - start_time) < max_wait:
|
|
await asyncio.sleep(0.1)
|
|
|
|
if not self.connected and self.sio.connected:
|
|
self.connected = True
|
|
|
|
return self.connected
|
|
|
|
except Exception as e:
|
|
logger.error(f"Connection failed: {e}")
|
|
return False
|
|
|
|
async def synthesize(
|
|
self, text: str, request_id: str | None = None
|
|
) -> asyncio.Queue[bytes | None]:
|
|
"""Send synthesis request and return audio queue"""
|
|
if not self.sio or not self.connected:
|
|
if not await self.connect():
|
|
raise ConnectionError("Failed to connect to TTS service")
|
|
|
|
if not request_id:
|
|
request_id = str(uuid.uuid4())
|
|
|
|
# Create audio queue
|
|
audio_queue: asyncio.Queue[bytes | None] = asyncio.Queue()
|
|
self.audio_callbacks[request_id] = audio_queue
|
|
self.active_requests[request_id] = True
|
|
|
|
# Build message
|
|
message: dict[str, Any] = {
|
|
"type": "synthesize",
|
|
"requestId": request_id,
|
|
"text": text,
|
|
"voiceId": self.opts.voice_settings.voice_id,
|
|
"outputFormat": self.opts.voice_settings.output_format,
|
|
}
|
|
|
|
if self.opts.phrase_replacement_config_id:
|
|
message["phraseReplacementConfigId"] = self.opts.phrase_replacement_config_id
|
|
|
|
logger.debug(f"Sending synthesis request {request_id[:8]} for text: '{text[:50]}...'")
|
|
|
|
try:
|
|
if self.sio is not None:
|
|
await self.sio.emit("synthesize", message, namespace=WEBSOCKET_NAMESPACE)
|
|
except Exception as e:
|
|
logger.error(f"Failed to emit synthesis: {e}")
|
|
del self.audio_callbacks[request_id]
|
|
del self.active_requests[request_id]
|
|
raise
|
|
|
|
return audio_queue
|
|
|
|
async def disconnect(self) -> None:
|
|
"""Disconnect from service"""
|
|
if self.sio and self.connected:
|
|
await self.sio.disconnect()
|
|
self.connected = False
|
|
|
|
async def _on_connect(self) -> None:
|
|
"""Handle connection"""
|
|
logger.debug("WebSocket connected")
|
|
|
|
async def _on_message(self, data: Any) -> None:
|
|
"""Handle messages"""
|
|
message_type = data.get("type")
|
|
|
|
if message_type == "ready":
|
|
self.connected = True
|
|
logger.debug(f"Ready with session: {data.get('sessionId')}")
|
|
|
|
elif message_type == "audio":
|
|
request_id = data.get("requestId")
|
|
audio_b64 = data.get("audio")
|
|
|
|
if audio_b64 and request_id in self.audio_callbacks:
|
|
audio_bytes = base64.b64decode(audio_b64)
|
|
if self.active_requests.get(request_id, False):
|
|
await self.audio_callbacks[request_id].put(audio_bytes)
|
|
|
|
elif message_type == "audio_end":
|
|
request_id = data.get("requestId")
|
|
if request_id in self.audio_callbacks:
|
|
await self.audio_callbacks[request_id].put(None)
|
|
del self.audio_callbacks[request_id]
|
|
if request_id in self.active_requests:
|
|
del self.active_requests[request_id]
|
|
|
|
elif message_type == "error":
|
|
request_id = data.get("requestId", "unknown")
|
|
error_msg = data.get("message", str(data))
|
|
logger.error(f"Error for {request_id}: {error_msg}")
|
|
|
|
if request_id in self.audio_callbacks:
|
|
await self.audio_callbacks[request_id].put(None)
|
|
del self.audio_callbacks[request_id]
|
|
if request_id in self.active_requests:
|
|
del self.active_requests[request_id]
|
|
|
|
async def _on_disconnect(self) -> None:
|
|
"""Handle disconnection"""
|
|
self.connected = False
|
|
for queue in self.audio_callbacks.values():
|
|
await queue.put(None)
|
|
self.audio_callbacks.clear()
|
|
self.active_requests.clear()
|
|
|
|
|
|
class ChunkedStream(tts.ChunkedStream):
|
|
"""Chunked synthesis implementation"""
|
|
|
|
def __init__(self, *, tts: TTS, input_text: str, conn_options: APIConnectOptions) -> None:
|
|
super().__init__(tts=tts, input_text=input_text, conn_options=conn_options)
|
|
self._tts: TTS = tts
|
|
|
|
async def _run(self, output_emitter: tts.AudioEmitter) -> None:
|
|
"""Execute synthesis"""
|
|
request_id = utils.shortuuid()
|
|
|
|
try:
|
|
# Initialize emitter
|
|
output_emitter.initialize(
|
|
request_id=request_id,
|
|
sample_rate=self._tts._opts.sample_rate,
|
|
num_channels=self._tts._opts.num_channels,
|
|
mime_type=get_content_type_from_output_format(
|
|
self._tts._opts.voice_settings.output_format
|
|
),
|
|
)
|
|
|
|
# Create client if needed
|
|
if not self._tts._client:
|
|
self._tts._client = WebSocketClient(self._tts._opts)
|
|
|
|
# Get audio queue
|
|
audio_queue = await self._tts._client.synthesize(self._input_text, request_id)
|
|
|
|
# Stream audio
|
|
while True:
|
|
try:
|
|
audio_data = await asyncio.wait_for(audio_queue.get(), timeout=30.0)
|
|
|
|
if audio_data is None:
|
|
break
|
|
|
|
output_emitter.push(audio_data)
|
|
|
|
except asyncio.TimeoutError:
|
|
logger.warning("Audio timeout")
|
|
break
|
|
|
|
output_emitter.flush()
|
|
|
|
except asyncio.TimeoutError as e:
|
|
raise APITimeoutError() from e
|
|
except Exception as e:
|
|
raise APIConnectionError(f"TTS synthesis failed: {str(e)}") from e
|
|
|
|
|
|
class SynthesizeStream(tts.SynthesizeStream):
|
|
"""Streaming synthesis implementation"""
|
|
|
|
def __init__(self, *, tts: TTS, conn_options: APIConnectOptions):
|
|
super().__init__(tts=tts, conn_options=conn_options)
|
|
self._tts: TTS = tts
|
|
|
|
async def _run(self, output_emitter: tts.AudioEmitter) -> None:
|
|
"""Execute streaming synthesis"""
|
|
request_id = utils.shortuuid()
|
|
segments_ch = utils.aio.Chan[tokenize.WordStream | tokenize.SentenceStream]()
|
|
|
|
output_emitter.initialize(
|
|
request_id=request_id,
|
|
sample_rate=self._tts._opts.sample_rate,
|
|
num_channels=self._tts._opts.num_channels,
|
|
stream=True,
|
|
mime_type=get_content_type_from_output_format(
|
|
self._tts._opts.voice_settings.output_format
|
|
),
|
|
)
|
|
|
|
async def _tokenize_input() -> None:
|
|
"""Tokenize input text"""
|
|
word_stream = None
|
|
async for input in self._input_ch:
|
|
if isinstance(input, str):
|
|
if word_stream is None:
|
|
word_stream = self._tts._opts.word_tokenizer.stream()
|
|
segments_ch.send_nowait(word_stream)
|
|
|
|
word_stream.push_text(input)
|
|
elif isinstance(input, self._FlushSentinel):
|
|
if word_stream is not None:
|
|
word_stream.end_input()
|
|
word_stream = None
|
|
|
|
if word_stream is not None:
|
|
word_stream.end_input()
|
|
|
|
segments_ch.close()
|
|
|
|
async def _process_segments() -> None:
|
|
"""Process segments"""
|
|
async for word_stream in segments_ch:
|
|
await self._run_segment(word_stream, output_emitter)
|
|
|
|
tasks = [
|
|
asyncio.create_task(_tokenize_input()),
|
|
asyncio.create_task(_process_segments()),
|
|
]
|
|
|
|
try:
|
|
await asyncio.gather(*tasks)
|
|
except asyncio.TimeoutError:
|
|
raise APITimeoutError() from None
|
|
except Exception as e:
|
|
raise APIConnectionError() from e
|
|
finally:
|
|
await utils.aio.gracefully_cancel(*tasks)
|
|
|
|
async def _run_segment(
|
|
self,
|
|
word_stream: tokenize.WordStream | tokenize.SentenceStream,
|
|
output_emitter: tts.AudioEmitter,
|
|
) -> None:
|
|
"""Process a single segment"""
|
|
segment_id = utils.shortuuid()
|
|
output_emitter.start_segment(segment_id=segment_id)
|
|
|
|
try:
|
|
# Create client if needed
|
|
if not self._tts._client:
|
|
self._tts._client = WebSocketClient(self._tts._opts)
|
|
|
|
# Collect text
|
|
text_parts = []
|
|
async for data in word_stream:
|
|
text_parts.append(data.token)
|
|
|
|
if not text_parts:
|
|
return
|
|
|
|
# Format text
|
|
if isinstance(self._tts._opts.word_tokenizer, tokenize.WordTokenizer):
|
|
full_text = self._tts._opts.word_tokenizer.format_words(text_parts)
|
|
else:
|
|
full_text = " ".join(text_parts)
|
|
|
|
self._mark_started()
|
|
|
|
# Synthesize
|
|
audio_queue = await self._tts._client.synthesize(full_text, segment_id)
|
|
|
|
# Stream audio
|
|
while True:
|
|
try:
|
|
audio_data = await asyncio.wait_for(audio_queue.get(), timeout=30.0)
|
|
|
|
if audio_data is None:
|
|
break
|
|
|
|
output_emitter.push(audio_data)
|
|
|
|
except asyncio.TimeoutError:
|
|
break
|
|
|
|
output_emitter.end_input()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Segment synthesis error: {e}")
|
|
raise APIError(f"Segment synthesis failed: {str(e)}") from e
|