316 lines
11 KiB
Python
316 lines
11 KiB
Python
import asyncio
|
||
import json
|
||
import websockets
|
||
import numpy as np
|
||
import sounddevice as sd
|
||
from typing import Optional, Callable, Dict, Any, Coroutine
|
||
from dataclasses import dataclass, field
|
||
import uuid
|
||
|
||
|
||
@dataclass
|
||
class TTSRequest:
|
||
"""TTS请求对象（带唯一标识）"""
|
||
tts_text: str
|
||
mode: str = "预训练音色"
|
||
sft_spk: str = ""
|
||
seed: int = field(default_factory=lambda: np.random.randint(1, 100000000))
|
||
stream: bool = True
|
||
speed: float = 1.0
|
||
prompt_wav: str = ""
|
||
instruct_text: str = ""
|
||
request_id: str = field(default_factory=lambda: str(uuid.uuid4())) # 唯一请求ID
|
||
|
||
|
||
class CosyVoiceTTSSocketClient:
|
||
"""CosyVoice TTS WebSocket 客户端（异步/流式/带任务队列）"""
|
||
|
||
def __init__(self, ws_url: str = "ws://localhost:50000/ws/tts", max_queue_size: int = 100):
|
||
"""
|
||
初始化客户端
|
||
:param ws_url: WebSocket 服务端地址
|
||
:param max_queue_size: 最大队列长度（防止内存溢出）
|
||
"""
|
||
self.ws_url = ws_url
|
||
self.websocket: Optional[websockets.WebSocketClientProtocol] = None
|
||
self.is_connected = False
|
||
self.is_processing = False # 是否正在处理请求
|
||
|
||
# 异步任务队列（FIFO）
|
||
self.request_queue: asyncio.Queue[TTSRequest] = asyncio.Queue(maxsize=max_queue_size)
|
||
|
||
# 回调函数定义（所有回调都带request_id，方便关联请求）
|
||
self.on_task_enqueue: Callable[[str], None] = lambda req_id: None # 任务入队回调
|
||
self.on_start: Callable[[str, Dict[str, Any]], None] = lambda req_id, data: None # 合成开始回调
|
||
self.on_audio_chunk: Callable[[str, np.ndarray], None] = lambda req_id, chunk: None # 音频块回调
|
||
self.on_end: Callable[[str, Dict[str, Any]], None] = lambda req_id, data: None # 合成结束回调
|
||
self.on_error: Callable[[str, str], None] = lambda req_id, msg: None # 错误回调
|
||
self.on_queue_full: Callable[[str], None] = lambda req_id: None # 队列满回调
|
||
|
||
async def connect(self):
|
||
"""建立 WebSocket 连接（初始化一次）"""
|
||
if not self.is_connected:
|
||
try:
|
||
self.websocket = await websockets.connect(self.ws_url)
|
||
self.is_connected = True
|
||
print(f"成功连接到 TTS 服务端: {self.ws_url}")
|
||
|
||
# 启动队列消费协程（后台运行）
|
||
asyncio.create_task(self._consume_queue())
|
||
except Exception as e:
|
||
raise ConnectionError(f"连接失败: {str(e)}")
|
||
|
||
async def disconnect(self):
|
||
"""关闭 WebSocket 连接"""
|
||
if self.is_connected and self.websocket:
|
||
await self.websocket.close()
|
||
self.is_connected = False
|
||
self.websocket = None
|
||
print("已断开与 TTS 服务端的连接")
|
||
|
||
def add_tts_request(self, **kwargs) -> str:
|
||
"""
|
||
添加TTS请求到队列（非阻塞，立即返回请求ID）
|
||
:param kwargs: TTS参数（同TTSRequest）
|
||
:return: 唯一请求ID
|
||
"""
|
||
# 创建请求对象
|
||
request = TTSRequest(**kwargs)
|
||
req_id = request.request_id
|
||
|
||
# 尝试入队（非阻塞）
|
||
try:
|
||
self.request_queue.put_nowait(request)
|
||
self.on_task_enqueue(req_id)
|
||
print(f"请求 {req_id} 已加入队列，当前队列长度: {self.request_queue.qsize()}")
|
||
except asyncio.QueueFull:
|
||
self.on_queue_full(req_id)
|
||
print(f"队列已满，请求 {req_id} 入队失败")
|
||
|
||
return req_id
|
||
|
||
async def _consume_queue(self):
|
||
"""消费队列（后台协程，自动处理排队请求）"""
|
||
print("队列消费协程已启动")
|
||
while True:
|
||
try:
|
||
# 等待队列中有请求（阻塞）
|
||
request = await self.request_queue.get()
|
||
req_id = request.request_id
|
||
|
||
# 标记为处理中
|
||
self.is_processing = True
|
||
print(f"开始处理请求 {req_id}，剩余队列长度: {self.request_queue.qsize()}")
|
||
|
||
# 处理当前请求
|
||
await self._process_single_request(request)
|
||
|
||
# 标记任务完成
|
||
self.request_queue.task_done()
|
||
self.is_processing = False
|
||
|
||
except Exception as e:
|
||
print(f"队列消费异常: {str(e)}")
|
||
self.is_processing = False
|
||
# 短暂等待后继续消费，避免死循环
|
||
await asyncio.sleep(0.1)
|
||
|
||
async def _process_single_request(self, request: TTSRequest):
|
||
"""处理单个TTS请求"""
|
||
req_id = request.request_id
|
||
|
||
# 参数校验
|
||
if not request.tts_text:
|
||
self.on_error(req_id, "合成文本不能为空")
|
||
return
|
||
|
||
# 确保已连接
|
||
if not self.is_connected:
|
||
try:
|
||
await self.connect()
|
||
except Exception as e:
|
||
self.on_error(req_id, f"连接失败: {str(e)}")
|
||
return
|
||
|
||
# 构造请求数据
|
||
request_data = {
|
||
"tts_text": request.tts_text,
|
||
"mode": request.mode,
|
||
"sft_spk": request.sft_spk,
|
||
"seed": request.seed,
|
||
"stream": request.stream,
|
||
"speed": request.speed,
|
||
"prompt_wav": request.prompt_wav,
|
||
"instruct_text": request.instruct_text
|
||
}
|
||
|
||
try:
|
||
# 发送请求
|
||
await self.websocket.send(json.dumps(request_data))
|
||
|
||
# 处理返回的流式数据
|
||
await self._handle_response(req_id)
|
||
|
||
except websockets.exceptions.ConnectionClosed:
|
||
self.is_connected = False
|
||
self.on_error(req_id, "连接已关闭")
|
||
except Exception as e:
|
||
self.on_error(req_id, f"处理请求失败: {str(e)}")
|
||
|
||
async def _handle_response(self, req_id: str):
|
||
"""处理单个请求的服务端响应"""
|
||
if not self.websocket:
|
||
return
|
||
|
||
try:
|
||
while True:
|
||
# 接收服务端消息（异步）
|
||
response = await self.websocket.recv()
|
||
data = json.loads(response)
|
||
|
||
# 根据状态分发到不同回调（都带req_id）
|
||
status = data.get("status")
|
||
if status == "start":
|
||
# 合成开始 - 返回采样率等信息
|
||
self.on_start(req_id, data)
|
||
|
||
elif status == "stream":
|
||
# 流式音频块 - 转为 numpy 数组
|
||
audio_chunk = np.array(data["audio_chunk"], dtype=np.float32)
|
||
self.on_audio_chunk(req_id, audio_chunk)
|
||
|
||
elif status == "end":
|
||
# 合成结束
|
||
self.on_end(req_id, data)
|
||
break
|
||
|
||
elif status == "error":
|
||
# 错误处理
|
||
self.on_error(req_id, data["msg"])
|
||
break
|
||
|
||
except Exception as e:
|
||
self.on_error(req_id, f"接收数据异常: {str(e)}")
|
||
|
||
async def wait_queue_empty(self):
|
||
"""等待队列所有任务处理完成（阻塞）"""
|
||
await self.request_queue.join()
|
||
print("所有队列任务已处理完成")
|
||
|
||
|
||
# ------------------------------
|
||
# 示例：带队列的客户端使用演示
|
||
# ------------------------------
|
||
async def demo():
|
||
# 1. 创建客户端实例（初始化一次）
|
||
client = CosyVoiceTTSSocketClient(
|
||
ws_url="ws://10.10.10.202:50000/ws/tts",
|
||
max_queue_size=50 # 最大队列长度
|
||
)
|
||
|
||
# 2. 初始化音频播放流（全局）
|
||
play_stream = None
|
||
current_req_id = None
|
||
|
||
# 3. 定义各类回调函数
|
||
# 任务入队回调
|
||
def on_task_enqueue(req_id):
|
||
print(f"✅ 任务 {req_id[:8]} 已入队")
|
||
|
||
# 队列满回调
|
||
def on_queue_full(req_id):
|
||
print(f"❌ 队列已满，任务 {req_id[:8]} 入队失败")
|
||
|
||
# 合成开始回调
|
||
def on_tts_start(req_id, data):
|
||
nonlocal play_stream, current_req_id
|
||
current_req_id = req_id
|
||
print(f"\n🎤 开始合成 [{req_id[:8]}] - 采样率: {data['sample_rate']}")
|
||
|
||
# 初始化播放流
|
||
if play_stream:
|
||
play_stream.stop()
|
||
play_stream.close()
|
||
play_stream = sd.OutputStream(
|
||
samplerate=data["sample_rate"],
|
||
channels=1,
|
||
dtype=np.float32
|
||
)
|
||
play_stream.start()
|
||
|
||
# 音频块回调（实时播放）
|
||
def on_audio_chunk(req_id, chunk):
|
||
if chunk.size > 0 and req_id == current_req_id:
|
||
# 只播放当前正在处理的请求音频
|
||
play_stream.write(chunk)
|
||
# 可选：保存音频块（按req_id区分文件）
|
||
# with open(f"output_{req_id[:8]}.raw", "ab") as f:
|
||
# f.write(chunk.tobytes())
|
||
|
||
# 合成结束回调
|
||
def on_tts_end(req_id, data):
|
||
nonlocal play_stream
|
||
print(f"🏁 合成完成 [{req_id[:8]}] - {data['msg']}")
|
||
if play_stream:
|
||
play_stream.stop()
|
||
play_stream.close()
|
||
play_stream = None
|
||
|
||
# 错误回调
|
||
def on_tts_error(req_id, msg):
|
||
nonlocal play_stream
|
||
print(f"❌ 合成失败 [{req_id[:8]}] - {msg}")
|
||
if play_stream:
|
||
play_stream.stop()
|
||
play_stream.close()
|
||
play_stream = None
|
||
|
||
# 注册回调
|
||
client.on_task_enqueue = on_task_enqueue
|
||
client.on_queue_full = on_queue_full
|
||
client.on_start = on_tts_start
|
||
client.on_audio_chunk = on_audio_chunk
|
||
client.on_end = on_tts_end
|
||
client.on_error = on_tts_error
|
||
|
||
# 4. 建立初始连接
|
||
await client.connect()
|
||
|
||
# 5. 模拟持续输入请求（批量添加到队列）
|
||
test_texts = [
|
||
"你好，这是第一个排队的TTS请求。",
|
||
"我是第二个请求，会等第一个处理完再执行。",
|
||
"第三个请求，支持流式播放和队列管理。",
|
||
"第四个请求，测试队列的自动消费功能。",
|
||
"最后一个请求，处理完成后会自动结束。"
|
||
]
|
||
|
||
# 批量添加请求到队列（非阻塞）
|
||
req_ids = []
|
||
for i, text in enumerate(test_texts):
|
||
req_id = client.add_tts_request(
|
||
tts_text=text,
|
||
mode="预训练音色",
|
||
sft_spk="中文女",
|
||
speed=1.0
|
||
)
|
||
req_ids.append(req_id)
|
||
# 模拟间隔输入
|
||
await asyncio.sleep(0.1)
|
||
|
||
# 6. 等待所有队列任务完成
|
||
await client.wait_queue_empty()
|
||
|
||
# 7. 断开连接
|
||
await client.disconnect()
|
||
print("\n所有请求处理完毕，客户端已关闭")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 安装依赖：pip install websockets numpy sounddevice
|
||
try:
|
||
asyncio.run(demo())
|
||
except KeyboardInterrupt:
|
||
print("\n程序被用户中断")
|
||
except Exception as e:
|
||
print(f"程序异常: {str(e)}") |