332 lines
12 KiB
Python
332 lines
12 KiB
Python
import asyncio
|
||
import json
|
||
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"
|
||
|
||
|
||
# ========== 修复：内置连接池实现 ==========
|
||
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 # 连接池是否已关闭
|
||
|
||
async def acquire(self) -> Any:
|
||
"""获取连接（从池子里取，没有则创建）"""
|
||
if self._closed:
|
||
raise RuntimeError("连接池已关闭，无法获取新连接")
|
||
|
||
# 尝试从空闲队列获取
|
||
try:
|
||
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()
|
||
|
||
self._in_use.add(conn)
|
||
return conn
|
||
|
||
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:
|
||
# 检查连接是否还可用
|
||
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)
|
||
|
||
async def close(self):
|
||
"""关闭连接池，清理所有连接"""
|
||
self._closed = True
|
||
|
||
# 关闭队列中的空闲连接
|
||
while not self.pool.empty():
|
||
try:
|
||
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 = {
|
||
"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 conn.send(json.dumps(request_data))
|
||
|
||
# 接收响应
|
||
while True:
|
||
try:
|
||
response = await conn.recv()
|
||
data = json.loads(response)
|
||
|
||
status = data.get("status")
|
||
|
||
if status == "start":
|
||
sample_rate = data.get("sample_rate", 24000)
|
||
print(f"📢 开始合成 - Task ID: {data.get('task_id')}")
|
||
|
||
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 == "complete":
|
||
print(f"✅ 合成完成 - {data.get('message')}")
|
||
break
|
||
|
||
elif status == "error":
|
||
raise Exception(f"TTS合成失败: {data.get('message', '未知错误')}")
|
||
|
||
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():
|
||
"""客户端使用示例"""
|
||
# 创建客户端
|
||
client = CosyVoiceWSClient(ws_url="ws://10.10.10.202:8765", pool_size=3)
|
||
|
||
try:
|
||
print("=== CosyVoice WebSocket客户端测试 ===")
|
||
|
||
# 1. 获取预训练音色列表
|
||
print("\n1. 获取预训练音色列表...")
|
||
try:
|
||
speakers = await client.list_speakers()
|
||
print(f"可用预训练音色: {speakers}")
|
||
except Exception as e:
|
||
print(f"获取音色列表失败: {e}")
|
||
speakers = []
|
||
|
||
# 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}")
|
||
|
||
# 测试基础合成（无音色）
|
||
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}")
|
||
|
||
except Exception as e:
|
||
print(f"❌ 测试出错: {e}")
|
||
finally:
|
||
# 关闭连接池
|
||
await client.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 运行测试
|
||
asyncio.run(demo()) |