220 lines
9.2 KiB
Python
220 lines
9.2 KiB
Python
"""百炼 CosyVoice / Qwen-Audio-TTS 双向流式适配。
|
|
|
|
每个 Pipecat 音频上下文对应一个任务和连接；文本增量发送，音频独立接收。
|
|
不复用连接，避免取消后没有 task_id 的二进制音频串入下一轮。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from collections.abc import AsyncGenerator
|
|
from dataclasses import dataclass
|
|
from urllib.parse import urlsplit
|
|
from uuid import uuid4
|
|
|
|
from pipecat.frames.frames import CancelFrame, EndFrame, ErrorFrame, Frame, TTSAudioRawFrame, TTSStoppedFrame
|
|
from pipecat.services.settings import TTSSettings
|
|
from pipecat.services.tts_service import TextAggregationMode, TTSService
|
|
from websockets.asyncio.client import connect
|
|
|
|
SOURCE_SAMPLE_RATE = 24000
|
|
TIMEOUT_SECONDS = 30
|
|
SUPPORTED_MODELS = (
|
|
"qwen-audio-3.1-tts-plus", "qwen-audio-3.1-tts-flash",
|
|
"qwen-audio-3.0-tts-plus", "qwen-audio-3.0-tts-flash",
|
|
"cosyvoice-v3.5-plus", "cosyvoice-v3.5-flash",
|
|
"cosyvoice-v3-plus", "cosyvoice-v3-flash", "cosyvoice-v2", "cosyvoice-v1",
|
|
)
|
|
|
|
|
|
def validate_config(url: str, model: str, voice: str, api_key: str) -> None:
|
|
parts = urlsplit(url)
|
|
if parts.scheme != "wss" or not parts.hostname or parts.path != "/api-ws/v1/inference":
|
|
raise ValueError("百炼 TTS API URL 必须为 wss://<域名>/api-ws/v1/inference，请更新旧 HTTP 配置")
|
|
if model not in SUPPORTED_MODELS:
|
|
raise ValueError(f"百炼流式 TTS 不支持模型: {model}")
|
|
if not voice or not api_key:
|
|
raise ValueError("百炼 TTS 需要 voice 和 apiKey")
|
|
|
|
|
|
class SynthesisSession:
|
|
"""一个百炼任务。连接测试和正式合成使用相同的协议实现。"""
|
|
|
|
def __init__(self, url: str, api_key: str, model: str, voice: str):
|
|
validate_config(url, model, voice, api_key)
|
|
self.url, self.api_key = url, api_key
|
|
self.model, self.voice = model, voice
|
|
self.task_id = uuid4().hex
|
|
self.socket = None
|
|
self.finished_input = False
|
|
|
|
async def send(self, action: str, input_: dict, **payload) -> None:
|
|
await self.socket.send(json.dumps({
|
|
"header": {"action": action, "task_id": self.task_id, "streaming": "duplex"},
|
|
"payload": {"input": input_, **payload},
|
|
}, ensure_ascii=False))
|
|
|
|
def event(self, raw: str) -> str:
|
|
header = json.loads(raw).get("header", {})
|
|
if header.get("task_id") != self.task_id:
|
|
raise RuntimeError("百炼 TTS 返回了不匹配的 task_id")
|
|
if header.get("event") == "task-failed":
|
|
raise RuntimeError(f"百炼 TTS: {header.get('error_code')}: {header.get('error_message')}")
|
|
return header.get("event", "")
|
|
|
|
async def start(self) -> None:
|
|
self.socket = await connect(
|
|
self.url, additional_headers={"Authorization": f"Bearer {self.api_key}"},
|
|
open_timeout=TIMEOUT_SECONDS, close_timeout=2, max_size=None,
|
|
)
|
|
await self.send("run-task", {}, task_group="audio", task="tts",
|
|
function="SpeechSynthesizer", model=self.model,
|
|
parameters={"text_type": "PlainText", "voice": self.voice,
|
|
"format": "pcm", "sample_rate": SOURCE_SAMPLE_RATE})
|
|
raw = await asyncio.wait_for(self.socket.recv(), TIMEOUT_SECONDS)
|
|
if not isinstance(raw, str) or self.event(raw) != "task-started":
|
|
raise RuntimeError("百炼 TTS 未返回 task-started")
|
|
|
|
async def text(self, text: str) -> None:
|
|
if self.finished_input:
|
|
raise RuntimeError("百炼 TTS 任务已结束文本输入")
|
|
# 单次最多 20000 字符；保留 token 间的空格。
|
|
for offset in range(0, len(text), 20000):
|
|
await self.send("continue-task", {"text": text[offset:offset + 20000]})
|
|
|
|
async def finish(self) -> None:
|
|
if not self.finished_input:
|
|
await self.send("finish-task", {})
|
|
self.finished_input = True
|
|
|
|
async def audio(self) -> AsyncGenerator[bytes, None]:
|
|
received_audio = False
|
|
while True:
|
|
raw = await asyncio.wait_for(self.socket.recv(), TIMEOUT_SECONDS)
|
|
if isinstance(raw, bytes):
|
|
if raw:
|
|
received_audio = True
|
|
yield raw
|
|
elif self.event(raw) == "task-finished":
|
|
if not received_audio:
|
|
raise RuntimeError("百炼 TTS 合成结束但没有返回音频")
|
|
return
|
|
|
|
async def close(self) -> None:
|
|
if self.socket:
|
|
await self.socket.close()
|
|
self.socket = None
|
|
|
|
|
|
@dataclass
|
|
class _Stream:
|
|
session: SynthesisSession
|
|
receiver: asyncio.Task | None = None
|
|
|
|
|
|
class DashScopeTTSService(TTSService):
|
|
"""核心参数只有模型、音色、地址和密钥；输出采样率由流水线决定。"""
|
|
|
|
def __init__(self, *, api_key: str, model: str, voice: str, url: str, **kwargs):
|
|
validate_config(url, model, voice, api_key)
|
|
super().__init__(
|
|
text_aggregation_mode=TextAggregationMode.TOKEN,
|
|
push_text_frames=True, push_start_frame=True, push_stop_frames=False,
|
|
pause_frame_processing=False,
|
|
settings=TTSSettings(model=model, voice=voice), **kwargs,
|
|
)
|
|
self._url, self._api_key = url, api_key
|
|
self._streams: dict[str, _Stream] = {}
|
|
# A failed turn must not silently restart and speak only its remaining tokens.
|
|
self._failed_contexts: set[str] = set()
|
|
|
|
@property
|
|
def supports_processing_metrics(self) -> bool:
|
|
return False # 音频在独立接收任务中生成。
|
|
|
|
async def stop(self, frame: EndFrame) -> None:
|
|
await super().stop(frame)
|
|
await self._disconnect()
|
|
|
|
async def cancel(self, frame: CancelFrame) -> None:
|
|
await super().cancel(frame)
|
|
await self._disconnect()
|
|
|
|
async def cleanup(self) -> None:
|
|
await self._disconnect()
|
|
await super().cleanup()
|
|
|
|
def _safe_error(self, exc: Exception) -> str:
|
|
return str(exc).replace(self._api_key, "***")[:300]
|
|
|
|
async def run_tts(self, text: str, context_id: str) -> AsyncGenerator[Frame | None, None]:
|
|
if not text or context_id in self._failed_contexts:
|
|
return
|
|
try:
|
|
stream = self._streams.get(context_id)
|
|
if stream is None:
|
|
session = SynthesisSession(self._url, self._api_key,
|
|
self._settings.model, self._settings.voice)
|
|
stream = _Stream(session)
|
|
self._streams[context_id] = stream
|
|
await session.start()
|
|
stream.receiver = self.create_task(self._receive(context_id, session))
|
|
await stream.session.text(text)
|
|
except Exception as exc:
|
|
self._failed_contexts.add(context_id)
|
|
await self._close_stream(context_id)
|
|
yield ErrorFrame(error=self._safe_error(exc))
|
|
await self._finish_context(context_id)
|
|
yield None
|
|
|
|
async def flush_audio(self, context_id: str | None = None) -> None:
|
|
stream = self._streams.get(context_id)
|
|
if stream:
|
|
try:
|
|
await stream.session.finish()
|
|
except Exception as exc:
|
|
self._failed_contexts.add(context_id)
|
|
await self.push_error(error_msg=self._safe_error(exc))
|
|
await self._close_stream(context_id)
|
|
await self._finish_context(context_id)
|
|
|
|
async def _receive(self, context_id: str, session: SynthesisSession) -> None:
|
|
try:
|
|
async for audio in session.audio():
|
|
if self.audio_context_available(context_id):
|
|
if self.sample_rate != SOURCE_SAMPLE_RATE:
|
|
audio = await self._resampler.resample(audio, SOURCE_SAMPLE_RATE, self.sample_rate)
|
|
await self.append_to_audio_context(
|
|
context_id, TTSAudioRawFrame(audio, self.sample_rate, 1, context_id=context_id))
|
|
except Exception as exc:
|
|
self._failed_contexts.add(context_id)
|
|
await self.push_error(error_msg=self._safe_error(exc))
|
|
finally:
|
|
await session.close()
|
|
self._streams.pop(context_id, None)
|
|
await self._finish_context(context_id)
|
|
|
|
async def _finish_context(self, context_id: str) -> None:
|
|
if self.audio_context_available(context_id):
|
|
await self.append_to_audio_context(context_id, TTSStoppedFrame(context_id=context_id))
|
|
await self.remove_audio_context(context_id)
|
|
|
|
async def _close_stream(self, context_id: str) -> None:
|
|
stream = self._streams.pop(context_id, None)
|
|
if stream:
|
|
if stream.receiver:
|
|
await self.cancel_task(stream.receiver)
|
|
await stream.session.close()
|
|
|
|
async def on_audio_context_interrupted(self, context_id: str) -> None:
|
|
await self._close_stream(context_id)
|
|
self._failed_contexts.discard(context_id)
|
|
|
|
async def on_audio_context_completed(self, context_id: str) -> None:
|
|
self._failed_contexts.discard(context_id)
|
|
|
|
async def _disconnect(self) -> None:
|
|
for context_id in list(self._streams):
|
|
await self._close_stream(context_id)
|
|
self._failed_contexts.clear()
|