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)}")