tianyan/aistream-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

xx

Browse Source
This commit is contained in:
Home
2025-12-04 22:45:28 +08:00
parent 795fdcf2ad
commit c2758373b4
16 changed files with 297 additions and 276 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/output_pretrained.wav
BIN
View File
Binary file not shown.
python/test.py
+292 -276
View File
@@ -1,316 +1,332 @@
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 base64
import uuid
import numpy as np
import soundfile as sf
from typing import List, Dict, Optional, Any, Set
from websockets import connect
from websockets.exceptions import ConnectionClosed, ConnectionClosedError
import warnings
# 忽略websockets的废弃警告
warnings.filterwarnings("ignore", category=DeprecationWarning)
# 连接池配置
POOL_SIZE = 3
WS_URL = "ws://10.10.10.202:8765"
@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 AsyncConnectionPool:
"""异步WebSocket连接池（完整实现）"""
def __init__(self, create_fn, destroy_fn, max_size: int = 5):
self.create_fn = create_fn # 创建连接的函数
self.destroy_fn = destroy_fn # 销毁连接的函数
self.max_size = max_size # 最大连接数
self.pool = asyncio.Queue(maxsize=max_size) # 空闲连接队列
self._in_use: Set[Any] = set() # 正在使用的连接
self._closed = False # 连接池是否已关闭
class CosyVoiceTTSSocketClient:
"""CosyVoice TTS WebSocket 客户端（异步/流式/带任务队列）"""
async def acquire(self) -> Any:
"""获取连接（从池子里取，没有则创建）"""
if self._closed:
raise RuntimeError("连接池已关闭，无法获取新连接")
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} 入队失败")
conn = self.pool.get_nowait()
# 检查连接是否还活着
if hasattr(conn, 'open') and not conn.open:
# 连接已关闭，创建新的
await self.destroy_fn(conn)
conn = await self.create_fn()
except asyncio.QueueEmpty:
# 队列空，检查是否可以创建新连接
if len(self._in_use) < self.max_size:
conn = await self.create_fn()
else:
# 等待有空闲连接
conn = await self.pool.get()
return req_id
self._in_use.add(conn)
return conn
async def _consume_queue(self):
"""消费队列（后台协程，自动处理排队请求）"""
print("队列消费协程已启动")
while True:
async def release(self, conn: Any):
"""释放连接（放回池子里）"""
if conn not in self._in_use:
raise ValueError("连接不在使用中，无法释放")
self._in_use.remove(conn)
if self._closed:
# 连接池已关闭，直接销毁连接
await self.destroy_fn(conn)
else:
try:
# 等待队列中有请求（阻塞）
request = await self.request_queue.get()
req_id = request.request_id
# 检查连接是否还可用
if hasattr(conn, 'open') and conn.open:
# 放回队列
self.pool.put_nowait(conn)
else:
# 连接已关闭，销毁
await self.destroy_fn(conn)
except asyncio.QueueFull:
# 队列满了，销毁连接
await self.destroy_fn(conn)
# 标记为处理中
self.is_processing = True
print(f"开始处理请求 {req_id}，剩余队列长度: {self.request_queue.qsize()}")
async def close(self):
"""关闭连接池，清理所有连接"""
self._closed = True
# 处理当前请求
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:
# 关闭队列中的空闲连接
while not self.pool.empty():
try:
await self.connect()
except Exception as e:
self.on_error(req_id, f"连接失败: {str(e)}")
return
conn = self.pool.get_nowait()
await self.destroy_fn(conn)
except asyncio.QueueEmpty:
break
# 构造请求数据
# 关闭正在使用的连接
for conn in list(self._in_use):
await self.destroy_fn(conn)
self._in_use.remove(conn)
# ========== CosyVoice WebSocket客户端 ==========
class CosyVoiceWSClient:
"""CosyVoice WebSocket客户端（修复版）"""
def __init__(self, ws_url: str = WS_URL, pool_size: int = POOL_SIZE):
self.ws_url = ws_url
self.pool = AsyncConnectionPool(
create_fn=self._create_connection,
destroy_fn=self._destroy_connection,
max_size=pool_size
)
async def _create_connection(self):
"""创建新的WebSocket连接（修复废弃API）"""
# 新版本websockets不再需要显式引用WebSocketClientProtocol
conn = await connect(self.ws_url)
print(f"✅ 创建新连接: {id(conn)}")
return conn
async def _destroy_connection(self, conn):
"""销毁WebSocket连接"""
try:
if hasattr(conn, 'open') and conn.open:
await conn.close()
print(f"❌ 销毁连接: {id(conn)}")
except Exception as e:
print(f"销毁连接出错: {e}")
async def list_speakers(self) -> List[str]:
"""获取预训练音色列表"""
conn = await self.pool.acquire()
try:
# 发送请求
await conn.send(json.dumps({
"type": "list_spks"
}))
# 接收响应
response = await conn.recv()
data = json.loads(response)
if data.get("status") == "success":
return data.get("data", [])
else:
raise Exception(f"获取音色列表失败: {data.get('message', '未知错误')}")
finally:
await self.pool.release(conn)
async def tts(
self,
tts_text: str,
mode: str = "预训练音色",
sft_spk: str = "",
prompt_text: str = "",
prompt_audio_path: Optional[str] = None,
instruct_text: str = "",
seed: Optional[int] = None,
stream: bool = False,
speed: float = 1.0,
output_path: Optional[str] = None
) -> np.ndarray:
"""
执行TTS合成（修复版）
Args:
tts_text: 合成文本
mode: 推理模式（预训练音色/3s极速复刻/跨语种复刻/自然语言控制）
sft_spk: 预训练音色名称
prompt_text: prompt文本（3s极速复刻用）
prompt_audio_path: prompt音频路径（3s极速复刻/跨语种复刻用）
instruct_text: 自然语言控制文本
seed: 随机种子
stream: 是否流式输出
speed: 语速（0.5-2.0）
output_path: 音频保存路径
Returns:
合成的音频数据（numpy数组）
"""
# 准备请求数据
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
"type": "tts",
"task_id": str(uuid.uuid4()),
"tts_text": tts_text,
"mode": mode,
"sft_spk": sft_spk,
"prompt_text": prompt_text,
"instruct_text": instruct_text,
"seed": seed if seed else np.random.randint(1, 100000000),
"stream": stream,
"speed": speed
}
# 处理prompt音频
if prompt_audio_path:
try:
with open(prompt_audio_path, "rb") as f:
audio_data = f.read()
request_data["prompt_audio"] = base64.b64encode(audio_data).decode("utf-8")
except Exception as e:
raise Exception(f"读取prompt音频失败: {e}")
# 获取连接
conn = await self.pool.acquire()
audio_chunks = []
sample_rate = 24000 # 默认采样率
try:
# 发送请求
await self.websocket.send(json.dumps(request_data))
await conn.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)
try:
response = await conn.recv()
data = json.loads(response)
# 根据状态分发到不同回调（都带req_id）
status = data.get("status")
if status == "start":
# 合成开始 - 返回采样率等信息
self.on_start(req_id, data)
status = data.get("status")
elif status == "stream":
# 流式音频块 - 转为 numpy 数组
audio_chunk = np.array(data["audio_chunk"], dtype=np.float32)
self.on_audio_chunk(req_id, audio_chunk)
if status == "start":
sample_rate = data.get("sample_rate", 24000)
print(f"📢 开始合成 - Task ID: {data.get('task_id')}")
elif status == "end":
# 合成结束
self.on_end(req_id, data)
break
elif status == "stream":
# 解码音频数据
audio_base64 = data.get("audio")
if audio_base64:
audio_bytes = base64.b64decode(audio_base64)
audio_chunk = np.frombuffer(audio_bytes, dtype=np.float32)
audio_chunks.append(audio_chunk)
print(f"🔊 接收音频块 {data.get('chunk_index', 0)} - 长度: {len(audio_chunk)}")
elif status == "error":
# 错误处理
self.on_error(req_id, data["msg"])
break
elif status == "complete":
print(f"✅ 合成完成 - {data.get('message')}")
break
except Exception as e:
self.on_error(req_id, f"接收数据异常: {str(e)}")
elif status == "error":
raise Exception(f"TTS合成失败: {data.get('message', '未知错误')}")
async def wait_queue_empty(self):
"""等待队列所有任务处理完成（阻塞）"""
await self.request_queue.join()
print("所有队列任务已处理完成")
except ConnectionClosedError:
raise Exception("连接意外关闭")
except ConnectionClosed:
raise Exception("连接已关闭")
# 合并音频
if audio_chunks:
audio_data = np.concatenate(audio_chunks)
# 保存音频
if output_path:
sf.write(output_path, audio_data, sample_rate)
print(f"💾 音频已保存到: {output_path}")
return audio_data
else:
raise Exception("未接收到音频数据")
finally:
await self.pool.release(conn)
async def close(self):
"""关闭连接池"""
await self.pool.close()
print("🔌 连接池已关闭")
# ------------------------------
# 示例：带队列的客户端使用演示
# ------------------------------
# ========== 使用示例 ==========
async def demo():
# 1. 创建客户端实例（初始化一次）
client = CosyVoiceTTSSocketClient(
ws_url="ws://10.10.10.202:50000/ws/tts",
max_queue_size=50 # 最大队列长度
)
"""客户端使用示例"""
# 创建客户端
client = CosyVoiceWSClient(ws_url="ws://10.10.10.202:8765", pool_size=3)
# 2. 初始化音频播放流（全局）
play_stream = None
current_req_id = None
try:
print("=== CosyVoice WebSocket客户端测试 ===")
# 3. 定义各类回调函数
# 任务入队回调
def on_task_enqueue(req_id):
print(f"✅ 任务 {req_id[:8]} 已入队")
# 1. 获取预训练音色列表
print("\n1. 获取预训练音色列表...")
try:
speakers = await client.list_speakers()
print(f"可用预训练音色: {speakers}")
except Exception as e:
print(f"获取音色列表失败: {e}")
speakers = []
# 队列满回调
def on_queue_full(req_id):
print(f"❌ 队列已满，任务 {req_id[:8]} 入队失败")
# 2. 预训练音色模式合成
if speakers:
print("\n2. 预训练音色模式合成...")
try:
audio = await client.tts(
tts_text="我来为你添加实时流式播放音频的示例，客户端会在接收音频块的同时后再播放",
mode="预训练音色",
sft_spk=speakers[0] if speakers else "",
output_path="output_pretrained.wav"
)
print(f"预训练音色合成完成，音频长度: {len(audio)}")
except Exception as e:
print(f"合成失败: {e}")
else:
# 测试3s极速复刻（需要先准备prompt.wav文件）
# print("\n2. 3s极速复刻模式合成...")
# try:
# audio = await client.tts(
# tts_text="这是极速复刻的测试文本",
# mode="3s极速复刻",
# prompt_text="这是prompt音频对应的文本",
# prompt_audio_path="prompt.wav",
# output_path="output_zeroshot.wav"
# )
# print(f"极速复刻合成完成，音频长度: {len(audio)}")
# except Exception as e:
# print(f"合成失败: {e}")
# 合成开始回调
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']}")
# 测试基础合成（无音色）
print("\n2. 基础模式合成...")
try:
audio = await client.tts(
tts_text="你好，这是CosyVoice的WebSocket服务测试",
mode="预训练音色",
sft_spk="",
output_path="output_basic.wav"
)
print(f"基础合成完成，音频长度: {len(audio)}")
except Exception as e:
print(f"合成失败: {e}")
# 初始化播放流
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所有请求处理完毕，客户端已关闭")
except Exception as e:
print(f"❌ 测试出错: {e}")
finally:
# 关闭连接池
await client.close()
if __name__ == "__main__":
# 安装依赖：pip install websockets numpy sounddevice
try:
asyncio.run(demo())
except KeyboardInterrupt:
print("\n程序被用户中断")
except Exception as e:
print(f"程序异常: {str(e)}")
# 运行测试
asyncio.run(demo())
python/tts_output_05c42041.pcm
BIN
View File
Binary file not shown.
python/tts_output_17c5db5c.pcm
BIN
View File
Binary file not shown.
python/tts_output_3428cc4d.pcm
BIN
View File
Binary file not shown.
python/tts_output_38aeb17d.pcm
BIN
View File
Binary file not shown.
python/tts_output_5a7fcf2d.pcm
BIN
View File
Binary file not shown.
python/tts_output_623cd970.pcm
BIN
View File
Binary file not shown.
python/tts_output_6fc93368.pcm
BIN
View File
Binary file not shown.
python/tts_output_76b4600f.pcm
BIN
View File
Binary file not shown.
python/tts_output_7e80df55.pcm
BIN
View File
Binary file not shown.
python/tts_output_85ba099e.pcm
BIN
View File
Binary file not shown.
python/tts_output_96aff098.pcm
BIN
View File
Binary file not shown.
python/tts_output_b0cfd3e4.pcm
BIN
View File
Binary file not shown.
python/tts_output_d8f880b4.pcm
BIN
View File
Binary file not shown.