ASR-demo/scripts/auxiliary_server.py

615 lines
28 KiB
Python
Raw Normal View History

2026-09-10 05:47:09 +00:00
#!/usr/bin/env python3
"""常驻加载 VAD、说话人聚类和声纹识别模型的独立服务。"""
from __future__ import annotations
import asyncio
import math
import os
import tempfile
import time
from collections.abc import Mapping
from pathlib import Path
from typing import Any
from aiohttp import web
from aiohttp.web_request import FileField
from dotenv import load_dotenv
try:
from .model_manifest import auxiliary_models, load_manifest, model_directory
except ImportError:
from model_manifest import auxiliary_models, load_manifest, model_directory
# 将辅助服务端口固定在代码变量中,服务器启动时只需执行脚本,便于部署和排查。
AUXILIARY_HOST = "0.0.0.0"
AUXILIARY_PORT = 8010
PROJECT_ROOT = Path(__file__).resolve().parents[1]
# 独立启动辅助服务也必须读取部署配置,不能只在启动 vLLM 时才加载 .env。
load_dotenv(PROJECT_ROOT / ".env")
MODELS_DIR = Path(os.getenv("MODEL_DIR", str(PROJECT_ROOT / "models"))).resolve()
AUXILIARY_DEVICE = os.getenv("AUXILIARY_DEVICE", "cuda:0")
ONLINE_SPEAKER_MATCH_THRESHOLD = 0.68
MIN_ONLINE_SPEAKER_AUDIO_MS = 800
# 实时 WebSocket 必须使用 VAD + CAM++ 声纹模型进行在线聚类。完整的
# speech_campplus_speaker-diarization_common 是整段离线 diarization 接口,
# 与实时每个 turn 的 CAM++ embedding 不是同一加载路径;它仍可按需加载。
DEFAULT_PRELOAD_KINDS = {"vad", "speaker_verification"}
def _coerce_finite_float(value: object) -> float | None:
"""把表单或模型返回的数值安全转换为有限浮点数。"""
if isinstance(value, bool):
return None
if isinstance(value, (int, float)):
parsed = float(value)
elif isinstance(value, (str, bytes)):
try:
parsed = float(value.strip())
except (TypeError, ValueError):
return None
else:
return None
return parsed if math.isfinite(parsed) else None
def _parse_form_float(value: object, field_name: str, default: float | None = None) -> float:
"""解析 multipart 数值字段,避免直接把 FileField/bytes 传给 float。"""
parsed = _coerce_finite_float(value)
if parsed is not None:
return parsed
if default is not None:
return default
raise ValueError(f"{field_name} must be a number")
def _asset_ready(path: Path, config: dict[str, Any]) -> bool:
"""在导入或加载模型前,先检查清单声明的文件和大小要求。"""
if not path.is_dir():
return False
for relative_path in config.get("required_files", []):
if not (path / str(relative_path)).is_file():
return False
any_files = config.get("any_files", [])
if any_files and not any(
file_path.is_file()
for pattern in any_files
for file_path in path.rglob(str(pattern))
):
return False
minimum_size = int(config.get("min_total_size_bytes", 0) or 0)
return not minimum_size or sum(
file_path.stat().st_size for file_path in path.rglob("*") if file_path.is_file()
) >= minimum_size
class AuxiliaryRuntime:
"""管理常驻辅助模型,并串行化 GPU 推理调用以避免显存竞争。"""
def __init__(self) -> None:
self.manifest = load_manifest()
self.assets = auxiliary_models(self.manifest)
self.models: dict[str, Any] = {}
self.status: dict[str, dict[str, Any]] = {}
self.inference_lock = asyncio.Lock()
# 每个 WebSocket session 独立维护聚类中心,避免不同浏览器会话互相污染。
self.speaker_clusters: dict[str, list[dict[str, Any]]] = {}
self.speaker_last_seen: dict[str, float] = {}
def _preload_kinds(self) -> set[str]:
"""读取需要在启动时加载的模型类型,默认不加载完整 diarization。"""
raw = os.getenv("AUXILIARY_PRELOAD_KINDS", "")
if not raw.strip():
return set(DEFAULT_PRELOAD_KINDS)
# 无论环境变量如何设置VAD 和 CAM++ speaker_verification 都是核心
# 依赖;额外类型只会增加预加载项,不能绕过核心模型校验。
return DEFAULT_PRELOAD_KINDS | {item.strip() for item in raw.split(",") if item.strip()}
def _load_asset(self, model_id: str, config: dict[str, Any], path: Path) -> Any | None:
"""只加载当前运行接口需要的模型;依赖模型和对齐模型先保持本地资产就绪。"""
kind = str(config.get("kind") or "")
if kind == "vad":
from funasr import AutoModel
return AutoModel(
model=str(path),
device=AUXILIARY_DEVICE,
disable_update=True,
disable_pbar=True,
disable_log=True,
local_files_only=True,
)
if kind in {"diarization", "speaker_verification", "realtime_speaker_verification"}:
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
task = Tasks.speaker_diarization if kind == "diarization" else Tasks.speaker_verification
return pipeline(task=task, model=str(path), device=AUXILIARY_DEVICE)
# CAM++ 依赖模型和 ForcedAligner 会先确认文件已落盘,后续由各自的专用
# 推理路径使用;这里不猜测它们的通用加载方式,避免错误占用显存。
return None
def preload(self) -> None:
"""预加载核心模型;可选模型失败只记录状态,避免服务整体退出。"""
failures: list[str] = []
preload_kinds = self._preload_kinds()
loaded_kinds: set[str] = set()
for model_id, config in self.assets.items():
path = model_directory(model_id, self.manifest, MODELS_DIR)
record: dict[str, Any] = {"path": str(path), "asset_ready": _asset_ready(path, config)}
kind = str(config.get("kind") or "")
if kind not in preload_kinds:
record["state"] = "optional_not_preloaded" if record["asset_ready"] else "optional_missing"
record["preload"] = False
self.status[model_id] = record
continue
# 清单中可能同时存在 iic/damo 两个同类型 CAM++ 资产;实时路径
# 只需一份,按清单顺序选第一个成功加载的模型,避免重复占显存。
if kind == "speaker_verification" and kind in loaded_kinds:
record["state"] = "duplicate_not_preloaded"
record["preload"] = False
self.status[model_id] = record
continue
record["preload"] = True
if not record["asset_ready"]:
record["state"] = "missing"
failures.append(model_id)
self.status[model_id] = record
continue
try:
loaded = self._load_asset(model_id, config, path)
if loaded is not None:
self.models[model_id] = loaded
record["state"] = "loaded"
loaded_kinds.add(kind)
else:
record["state"] = "load_error"
record["error"] = "model loader returned no model"
if kind in {"vad", "speaker_verification"}:
failures.append(model_id)
except Exception as exc:
record["state"] = "load_error"
record["error"] = str(exc)
failures.append(model_id)
self.status[model_id] = record
# 启动日志必须包含每个资产的路径、是否完整和底层异常;不能只打印
# 一个笼统的“startup failed”否则远程部署时无法判断缺文件还是版本错误。
for model_id, record in self.status.items():
print(
f"[model] {model_id}: state={record.get('state')}, "
f"asset_ready={record.get('asset_ready')}, path={record.get('path')}"
+ (f", error={record['error']}" if record.get("error") else ""),
flush=True,
)
# VAD 和 CAM++ speaker_verification 都是核心依赖,缺失/加载异常时
# 立即失败并列出路径与底层错误,避免页面一直显示“未确认”。
required_failures = [
model_id for model_id in failures
if self.assets[model_id].get("kind") == "vad"
or (
self.assets[model_id].get("kind") == "speaker_verification"
and self._speaker_embedding_model_id() is None
)
]
if required_failures:
details = "; ".join(
f"{model_id} -> {self.status[model_id]['path']}"
f" [{self.status[model_id].get('state')}: {self.status[model_id].get('error', 'asset missing')}]"
for model_id in required_failures
)
raise RuntimeError(
"Auxiliary core model is missing or failed to load: " + details
+ ". Run `python scripts/download_models.py --auxiliary-only` "
"or set MODEL_DIR to the directory containing the downloaded assets."
)
def _find_model(self, kind: str) -> Any:
"""按模型清单中的 kind 查找一个已经加载完成的模型。"""
for model_id, config in self.assets.items():
if config.get("kind") == kind and model_id in self.models:
return self.models[model_id]
raise RuntimeError(f"Auxiliary model is not loaded: {kind}")
def _load_optional_kind(self, kind: str) -> Any:
"""按需加载可选模型,例如完整 diarization 接口首次被调用时。"""
for model_id, config in self.assets.items():
if config.get("kind") != kind:
continue
path = model_directory(model_id, self.manifest, MODELS_DIR)
if not _asset_ready(path, config):
raise RuntimeError(f"Auxiliary model asset is missing: {model_id} ({path})")
try:
loaded = self._load_asset(model_id, config, path)
except Exception as exc:
self.status.setdefault(model_id, {})["state"] = "load_error"
self.status[model_id]["error"] = str(exc)
raise RuntimeError(f"Auxiliary {kind} model failed to load: {exc}") from exc
if loaded is None:
raise RuntimeError(f"Auxiliary model has no loader for kind: {kind}")
self.models[model_id] = loaded
self.status.setdefault(model_id, {})["state"] = "loaded_on_demand"
return loaded
raise RuntimeError(f"Auxiliary model asset is not configured: {kind}")
def _speaker_embedding_model_id(self) -> str | None:
"""选择实时声纹模型,并在实时模型不可用时回退到普通声纹模型。"""
# CAM++ 是实时聚类的主模型;其它声纹模型不能静默替代它。
for preferred_kind in ("speaker_verification",):
for model_id, config in self.assets.items():
if config.get("kind") == preferred_kind and model_id in self.models:
return model_id
return None
async def vad(self, audio_path: str) -> Any:
"""使用临时音频文件执行一次串行化的 VAD 推理。"""
model = self._find_model("vad")
async with self.inference_lock:
return await asyncio.to_thread(model.generate, input=audio_path, cache={})
async def diarization(self, audio_path: str) -> Any:
"""使用 CAM++ 对完整会话执行说话人聚类,保持跨片段的标签一致性。"""
try:
model = self._find_model("diarization")
except RuntimeError:
# 完整 diarization 不在核心启动路径,第一次调用接口时才加载。
model = self._load_optional_kind("diarization")
async with self.inference_lock:
# ModelScope 的 CAM++ pipeline 以位置参数接收音频路径;使用关键字
# input 在不同版本中可能被忽略或直接报参数错误。
return await asyncio.to_thread(model, audio_path)
@staticmethod
def _normalize_embedding(embedding: Any) -> Any:
"""将声纹模型输出转成单位向量,并拒绝 NaN、无穷值和零向量。"""
import numpy as np
vector = np.asarray(embedding, dtype=np.float32)
# 单次请求只能对应一条新声纹,不能把多个样本矩阵拼接成伪造特征。
if vector.ndim > 2 or (vector.ndim == 2 and vector.shape[0] != 1):
raise RuntimeError("speaker embedding must contain exactly one vector")
vector = vector.reshape(-1)
if vector.size == 0 or not np.isfinite(vector).all():
raise RuntimeError("speaker embedding is empty or non-finite")
norm = float(np.linalg.norm(vector))
if not np.isfinite(norm) or norm < 1e-8:
raise RuntimeError("speaker embedding has zero norm")
return vector / norm
@staticmethod
def _extract_embedding_value(result: Any) -> Any | None:
"""从不同 ModelScope 版本的 pipeline 返回值中提取 embedding。"""
if result is None:
return None
# ERes2Net pipeline 在 output_emb=True 时返回 {'embs': numpy.ndarray,
# 'outputs': ...};部分版本或其它声纹 pipeline 使用 embedding 变体字段。
if isinstance(result, Mapping):
for key in ("embs", "embedding", "embeddings", "speaker_embedding", "output_embedding"):
if key in result:
return AuxiliaryRuntime._extract_embedding_value(result[key])
return None
# 某些 ModelScope 版本把结果包装成带 embs/embedding 属性的对象。
for key in ("embs", "embedding", "embeddings", "speaker_embedding", "output_embedding"):
value = getattr(result, key, None)
if value is not None:
return AuxiliaryRuntime._extract_embedding_value(value)
# torch.Tensor 不能直接依赖 numpy.asarray 的 object 转换;先显式移到 CPU。
detach = getattr(result, "detach", None)
if callable(detach):
detached = detach()
cpu = getattr(detached, "cpu", None)
if callable(cpu):
detached = cpu()
numpy_method = getattr(detached, "numpy", None)
if callable(numpy_method):
return numpy_method()
# 单条音频通常返回 [embedding],递归拆开这一层;数值列表则保留为向量。
if isinstance(result, (list, tuple)) and len(result) == 1:
return AuxiliaryRuntime._extract_embedding_value(result[0])
return result
@staticmethod
def _run_embedding_pipeline(model_pipeline: Any, audio_path: str) -> Any:
"""调用声纹 pipeline 的公开预处理和 embedding 输出接口。"""
# ModelScope 的 ERes2Net pipeline 要求输入为音频路径列表,并通过
# output_emb=True 返回 embedding不能直接把原始 waveform Tensor 喂给
# pipeline.model因为那会跳过采样率、声道和 waveform 预处理。
try:
result = model_pipeline([audio_path], output_emb=True)
except TypeError:
# 兼容不支持 output_emb 参数的旧 pipeline仍然使用 pipeline 自带
# preprocess/forward而不是直接调用内部 model确保输入格式一致。
preprocess = getattr(model_pipeline, "preprocess", None)
forward = getattr(model_pipeline, "forward", None)
if not callable(preprocess) or not callable(forward):
raise RuntimeError("speaker pipeline does not expose embedding inference")
result = forward(preprocess([audio_path]))
embedding = AuxiliaryRuntime._extract_embedding_value(result)
if embedding is None:
raise RuntimeError(
"speaker pipeline returned no embedding "
f"(result_type={type(result).__name__})"
)
return embedding
def _extract_embedding_sync(self, audio_path: str) -> Any:
"""在工作线程中读取当前 turn并使用已加载的声纹模型提取特征。"""
import librosa
audio, _ = librosa.load(audio_path, sr=16000, mono=True)
audio_array = audio.reshape(-1)
if audio_array.size < int(16000 * MIN_ONLINE_SPEAKER_AUDIO_MS / 1000):
return None
model_id = self._speaker_embedding_model_id()
if model_id is None:
raise RuntimeError("no loaded speaker verification model is available")
model_pipeline = self.models[model_id]
output = self._run_embedding_pipeline(model_pipeline, audio_path)
return self._normalize_embedding(output)
async def resolve_speaker(
self,
audio_path: str,
session_id: str,
start_time_ms: float,
end_time_ms: float,
) -> dict[str, Any] | None:
"""对一个实时 turn 提取声纹,并更新该 session 的在线聚类中心。"""
async with self.inference_lock:
# 异常断网时客户端可能来不及 reset过期状态在下一次请求时回收。
now = time.monotonic()
for stale_id, seen in list(self.speaker_last_seen.items()):
if now - seen > 1800:
self.reset_speaker_session(stale_id)
self.speaker_last_seen[session_id] = now
embedding = await asyncio.to_thread(self._extract_embedding_sync, audio_path)
if embedding is None:
return {"speaker_id": -1, "speaker_evidence": "pending", "speaker_confidence": 0.0,
"speaker_status": "insufficient_audio", "speaker_reason": "音频不足 800ms未提取声纹"}
import numpy as np
# reset 可能在模型线程运行期间到达;结束的会话不能被晚到结果重新创建。
if session_id not in self.speaker_last_seen:
return None
embedding = self._normalize_embedding(embedding)
clusters = self.speaker_clusters.setdefault(session_id, [])
best_cluster: dict[str, Any] | None = None
best_score = -1.0
for cluster in clusters:
if embedding.shape != cluster["embedding"].shape:
raise RuntimeError("speaker embedding dimension changed within the session")
score = float(np.dot(embedding, cluster["embedding"]))
if score > best_score:
best_score = score
best_cluster = cluster
if best_cluster is not None and best_score >= ONLINE_SPEAKER_MATCH_THRESHOLD:
count = int(best_cluster["count"])
best_cluster["embedding"] = self._normalize_embedding(
(best_cluster["embedding"] * count) + embedding
)
best_cluster["count"] = count + 1
speaker_id = int(best_cluster["speaker_id"])
confidence = best_score
strategy = "online_embedding_cluster_match"
else:
speaker_id = len(clusters)
clusters.append({"speaker_id": speaker_id, "embedding": embedding, "count": 1})
confidence = 0.75
strategy = "online_embedding_cluster_new"
return {
"speaker_id": speaker_id,
"speaker_name": f"说话人 {speaker_id + 1}",
"speaker_evidence": "fresh",
"speaker_confidence": round(max(0.6, min(1.0, confidence)), 3),
"speaker_strategy": strategy,
"speaker_status": "confirmed",
"speaker_reason": "当前片段独立声纹已完成在线聚类",
"start_time": start_time_ms,
"end_time": end_time_ms,
}
def reset_speaker_session(self, session_id: str) -> None:
"""释放已结束 WebSocket 的聚类中心,防止长时间运行时内存增长。"""
self.speaker_clusters.pop(session_id, None)
self.speaker_last_seen.pop(session_id, None)
MODEL_SERVICE_KEY = web.AppKey("auxiliary_runtime", AuxiliaryRuntime)
async def health_handler(request: web.Request) -> web.Response:
"""返回模型资产完整性和预加载状态,供 WebSocket 编排服务检查。"""
runtime: AuxiliaryRuntime = request.app[MODEL_SERVICE_KEY]
speaker_model_id = runtime._speaker_embedding_model_id()
vad_model_id = next(
(model_id for model_id, config in runtime.assets.items()
if config.get("kind") == "vad" and model_id in runtime.models),
None,
)
# ready 表示实时链路的两个核心模型都可用;完整 diarization 是否
# 预加载不影响这里的结果。
ready = vad_model_id is not None and speaker_model_id is not None
return web.json_response(
{
"ready": ready,
"speaker_protocol_version": 2,
"vad_model": vad_model_id,
"vad_ready": vad_model_id is not None,
"device": AUXILIARY_DEVICE,
"speaker_embedding_model": speaker_model_id,
"speaker_embedding_ready": speaker_model_id is not None,
"models": runtime.status,
}
)
async def vad_handler(request: web.Request) -> web.Response:
"""接收 WAV 文件上传,并返回 FunASR 生成的语音活动区间。"""
runtime: AuxiliaryRuntime = request.app[MODEL_SERVICE_KEY]
form = await request.post()
upload = form.get("file")
if not isinstance(upload, FileField):
return web.json_response({"error": "multipart field 'file' is required"}, status=400)
temp_path: str | None = None
try:
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
temp_file.write(upload.file.read())
temp_path = temp_file.name
result = await runtime.vad(temp_path)
return web.json_response({"segments": result})
finally:
if temp_path:
Path(temp_path).unlink(missing_ok=True)
def _normalize_diarization_segments(result: Any) -> list[dict[str, Any]]:
"""将不同 ModelScope 版本的聚类输出统一为 start/end/speaker 字段。"""
# ModelScope 通常返回 {'text': [[start_sec, end_sec, speaker_id], ...]}
# 某些版本返回对象而不是字典,因此这里同时读取属性形式。
if isinstance(result, dict):
for key in ("segments", "output", "text", "result"):
candidate = result.get(key)
if isinstance(candidate, list):
result = candidate
break
else:
for key in ("segments", "output", "text", "result"):
candidate = getattr(result, key, None)
if isinstance(candidate, list):
result = candidate
break
if not isinstance(result, list):
return []
normalized: list[dict[str, Any]] = []
for item in result:
values_are_milliseconds = False
if isinstance(item, dict):
values_are_milliseconds = "start_time" in item or "end_time" in item
start = item.get("start", item.get("start_time", item.get("begin")))
end = item.get("end", item.get("end_time", item.get("stop")))
speaker = item.get("speaker", item.get("speaker_id", item.get("label")))
elif isinstance(item, (list, tuple)) and len(item) >= 3:
start, end, speaker = item[0], item[1], item[2]
else:
continue
start_value = _coerce_finite_float(start)
end_value = _coerce_finite_float(end)
if start_value is None or end_value is None:
continue
# 列表形式是 CAM++ 的秒单位;明确命名为 start_time/end_time 的
# 字段按毫秒处理,避免用“超过多少数值”猜单位导致长录音误判。
if values_are_milliseconds:
start_value /= 1000
end_value /= 1000
if end_value > start_value:
normalized.append({"start_time": round(start_value * 1000, 1), "end_time": round(end_value * 1000, 1), "speaker": str(speaker)})
return normalized
async def diarization_handler(request: web.Request) -> web.Response:
"""接收完整 WAV返回 CAM++ 说话人聚类时间段。"""
runtime: AuxiliaryRuntime = request.app[MODEL_SERVICE_KEY]
form = await request.post()
upload = form.get("file")
if not isinstance(upload, FileField):
return web.json_response({"error": "multipart field 'file' is required"}, status=400)
temp_path: str | None = None
try:
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
temp_file.write(upload.file.read())
temp_path = temp_file.name
result = await runtime.diarization(temp_path)
return web.json_response({"segments": _normalize_diarization_segments(result)})
finally:
if temp_path:
Path(temp_path).unlink(missing_ok=True)
async def speaker_resolve_handler(request: web.Request) -> web.Response:
"""接收一个实时 turn提取声纹并返回当前会话的在线聚类标签。"""
runtime: AuxiliaryRuntime = request.app[MODEL_SERVICE_KEY]
form = await request.post()
upload = form.get("file")
if not isinstance(upload, FileField):
return web.json_response({"error": "multipart field 'file' is required"}, status=400)
session_id = str(form.get("session_id") or "").strip()
if not session_id:
return web.json_response({"error": "multipart field 'session_id' is required"}, status=400)
try:
# aiohttp 的 MultiDictProxy 值可能是 str、bytes 或 FileField先收窄
# 为有限浮点数,避免静态检查告警和异常类型值进入声纹服务。
start_time_ms = _parse_form_float(form.get("start_time_ms"), "start_time_ms", default=0.0)
end_time_ms = _parse_form_float(form.get("end_time_ms"), "end_time_ms", default=start_time_ms)
except ValueError:
return web.json_response({"error": "turn time fields must be numbers"}, status=400)
temp_path: str | None = None
try:
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
temp_file.write(upload.file.read())
temp_path = temp_file.name
result = await runtime.resolve_speaker(
temp_path,
session_id,
start_time_ms,
end_time_ms,
)
return web.json_response(result or {
"speaker_id": -1, "speaker_evidence": "pending", "speaker_confidence": 0.0,
"speaker_status": "no_embedding", "speaker_reason": "当前片段未生成可用声纹",
})
except Exception as exc:
# 将模型推理异常返回给 WebSocket 客户端,避免客户端只能看到笼统的 500。
print(
f"[speaker] resolve failed: session_id={session_id}, "
f"start={start_time_ms}, end={end_time_ms}, error={exc}",
flush=True,
)
return web.json_response({"error": str(exc)}, status=500)
finally:
if temp_path:
Path(temp_path).unlink(missing_ok=True)
async def speaker_reset_handler(request: web.Request) -> web.Response:
"""释放已经结束的实时会话聚类中心。"""
runtime: AuxiliaryRuntime = request.app[MODEL_SERVICE_KEY]
payload = await request.json()
session_id = str(payload.get("session_id") or "").strip() if isinstance(payload, dict) else ""
if session_id:
runtime.reset_speaker_session(session_id)
return web.json_response({"ok": True})
async def create_app() -> web.Application:
"""创建辅助 HTTP 服务,并在服务启动前完成模型预加载。"""
runtime = AuxiliaryRuntime()
runtime.preload()
app = web.Application(client_max_size=64 * 1024 * 1024)
app[MODEL_SERVICE_KEY] = runtime
app.router.add_get("/health", health_handler)
app.router.add_post("/v1/vad", vad_handler)
app.router.add_post("/v1/diarization", diarization_handler)
app.router.add_post("/v1/speaker/resolve", speaker_resolve_handler)
app.router.add_post("/v1/speaker/reset", speaker_reset_handler)
return app
def main() -> None:
"""启动宿主机上的常驻辅助模型服务。"""
print(f"Auxiliary model service: http://127.0.0.1:{AUXILIARY_PORT}", flush=True)
print(f"Device: {AUXILIARY_DEVICE}", flush=True)
web.run_app(create_app(), host=AUXILIARY_HOST, port=AUXILIARY_PORT)
if __name__ == "__main__":
main()