AWS Machine Learning来源发布时间

使用 Amazon SageMaker AI 上的 Kokoro TTS 流式传输来自 AI 代理的低延迟语音

一位客户向AI语音助手提问并等待。助手其实已经有了答案,它正在生成一段冗长的多段回复,但客户却整整五秒听不到任何声音。

一位客户向AI语音助手提问并等待。助手其实已经有了答案,它正在生成一段冗长的多段回复,但客户却整整五秒听不到任何声音。等答案终于传来时,客户已经感到沮丧和困惑,也不再专心听了。

文本转语音集成恰恰会造成这种差距。一句话的响应可能在600毫秒内生成,但完整响应需要五秒钟,而集成会等待全部内容完成后才开始合成。

本文将展示如何弥合这一差距。我们将每个句子流式传输到托管于服务器上的 Kokoro-82M 文本转语音(TTS)模型。 Amazon SageMaker AI 一旦智能体生成内容就开始处理,因此语音合成从第一句话就开始,而不是等完整回复结束。把音频传回浏览器也需要自己的解决方案,所以我们还介绍了实现这一点的两种传输方式以及每种方式在实际中的成本。

解决方案概述

该解决方案通过让语音合成与文本生成同步流式进行、而不是等待其完成来弥合这一差距。当智能体生成每一句话时,该句子会立即流式传输到托管的 Kokoro-82M 文本转语音模型,因此音频播放可以在智能体说完完整回复之前就开始。Kokoro-82M 仅有 8200 万参数,足够小,可以在单台中端 GPU 实例上低成本运行。这使得处理真实流量的语音智能体在延迟和推理成本上都保持在较低水平。

Kokoro-82M 运行在 SageMaker AI 实时端点上的自定义容器中,运行于 ml.g5.xlarge GPU 实例。该容器暴露了两个路由: POST /invocations,它会为一个完整输入流式传输响应,并且 WS /invocations-bidirectional-stream它可以在两个方向上通过多条消息保持一个连接处于打开状态。

下图展示了从客户端到模型的请求路径:

图 1:从客户端通过双向流到 SageMaker AI 实时端点的请求路径

前提条件

在开始之前,你需要:

  • 拥有创建 SageMaker AI 终端节点权限的 AWS 账户;AgentCore runtime(Amazon Bedrock AgentCore 的一项功能);AWS Identity and Access Management (IAM) 角色;以及 Amazon Elastic Container Registry (Amazon ECR) 存储库,至少包括 sagemaker:CreateEndpoint, sagemaker:CreateEndpointConfig, sagemaker:CreateModel, sagemaker:InvokeEndpointWithBidirectionalStream,以及 bedrock:InvokeModelWithResponseStream 行动。
  • 该示例容器固定了一个存在已知安全通告的旧版PyTorch(PYSEC-2026-2286,已在 torch 2.10.0 中修复)。在任何超出概念验证范围的用途之前,请先审查并升级该固定版本。
  • AWS Cloud Development Kit (AWS CDK) v2并在目标账户和 AWS 区域中完成安装和引导。
  • 一个由容器声明的 SageMaker AI 实时端点部署,该容器声明了 com.amazonaws.sagemaker.capabilities.bidirectional-streaming=true Docker 标签。
  • Python 3.12 或更高版本, uv,以及使用 Docker 来构建容器镜像(其底层的 HTTP/2 客户端仅为 Python 3.12 及更高版本提供预构建的 wheel 包)。
  • 支持以下功能的适用于 Python 的 AWS SDK (Boto3) invoke_endpoint_with_bidirectional_stream.
  • 熟悉 WebSockets 以及 Starlette 或类似的 ASGI 框架。
  • 本文中所引用的 GitHub 仓库 的部署代码。

演练

该演练首先部署端点和代理,然后逐步介绍使增量流式传输端到端工作的客户端更改。

部署端点和代理

Kokoro-82M 端点和代理运行时通过两个 CLI 工具部署,每个工具都封装了自己的 AWS CDK 堆栈:

uv run speech-infra deploy kokoro-82m
uv run agent-infra deploy --bedrock-model-id <your-bedrock-model-id>

speech-infra deploy 预置 Kokoro-82M 端点。这包括用于创建底层 IAM 角色的基础堆栈和模型缓存堆栈,以及将模型的权重同步到 Amazon Simple Storage Service (Amazon S3) 的堆栈。端点堆栈随后构建容器镜像,并创建 SageMaker AI 模型、端点配置和端点。 agent-infra deploy 预置托管 Strands 代理的 Amazon Bedrock AgentCore 运行时,该代理会调用你刚刚部署的端点,因此请先部署端点。

第 1 步:处理 SageMaker AI 实际发送的 WebSocket 帧类型

SageMaker AI 使用的双向传输将客户端消息的每个部分作为二进制 WebSocket 帧传递给你的容器。编写 Starlette WebSocket 处理器的一种常见方式是调用 receive_text(),它读取 message["text"] ,如果该键不存在则抛出 KeyError 。由于该客户端发送的每条消息都以二进制形式到达,该处理器在每个请求上都会失败。

这种故障常常被误诊。一个转发 str(e) 给客户的转次 KeyError('text') 转换为一个错误提示框,内容为 'text',这看起来像是缺少字段,而不是框架问题。该端点本身报告 InService ……始终如此,因为它接受连接并回复每条消息,每次都返回一个错误。

明确接受任一帧类型,而不是假设只会出现其中一种:

from starlette.websockets import WebSocket, WebSocketDisconnect

async def _receive_message(websocket: WebSocket) -> str:
    message = await websocket.receive()
    if message["type"] == "websocket.disconnect":
        raise WebSocketDisconnect(message.get("code", 1000))
    payload = message.get("text")
    if payload is None:
        payload = message.get("bytes", b"").decode("utf-8")
    return payload

当没有文本字段时,这会将字节解码为 UTF-8,而不是假定某一种帧类型。请通过发送一个二进制帧来进行测试,因为这种故障模式会产生一个格式良好的错误帧和正常的端点状态,而不是明显的崩溃。

第 2 步:在消息之间保持会话开启

当容器完成合成一条消息的音频后,直觉上是关闭会话,就像 HTTP 周期结束那样。但在这里这是错误的。在发送完一条消息的音频和一个 synthesis_complete 控制帧时,容器循环应返回继续等待 receive() 用于同一仍处于打开状态的会话上的下一条消息。仅当客户端发送明确的 {"type": "close"} 消息。

下图展示了一个会话依次承载两个句子的情形,只在最后才关闭:

图2:一个双向会话按顺序承载两个句子,仅在结束时关闭

这样,你可以在一个持久会话中逐条发送多个句子,句与句之间没有重连开销。如果在收到一个响应后就急切地关闭连接,就会导致每个句子都要新建连接,从而抵消大部分延迟优势。

步骤 3:在代理生成句子的同时进行流式输出

有了会话行为之后,针对延迟问题的客户端解决方案是一个类,它对以两种形式到达的文本进行规范化处理:一种是事先已知的完整字符串,另一种是实时的 token 流。无论哪种方式,它都产生相同的输出:完整的句子,每句话一完成就可以立即发送。

import re

_SENTENCE_RE = re.compile(r"(?<=[.!?])\s+")


class IncrementalSentenceChunker:
    def __init__(self) -> None:
        self._buffer = ""

    def feed(self, fragment: str) -> list[str]:
        self._buffer += fragment
        parts = _SENTENCE_RE.split(self._buffer)
        if len(parts) == 1:
            return []
        self._buffer = parts[-1]
        return [s.strip() for s in parts[:-1] if s.strip()]

    def flush(self) -> str | None:
        remainder = self._buffer.strip()
        self._buffer = ""
        return remainder or None

呼叫 feed() 每当文本片段到达时,对每个片段执行一次。它返回任何新完成的句子,并保留末尾未完成的文本。调用 flush() 在输入结束时执行一次,以释放剩余的任何部分文本,即使没有尾随标点符号。这可以防止在句号到来之前就结束的流丢失最后一个句子。

这不是并发合成。每个双向处理器都是一个严格顺序的循环:接收一条消息,合成其音频,发送该音频,然后发送 synthesis_complete,然后接收下一条消息。句子一的音频还在传输时,句子二的音频并不会生成。增量发送带来的收益更窄:句子一一存在,其合成便立刻开始,而不必等待完整响应。对于一个首句与末句之间有实际生成时间的多句响应来说,这就是全部的延迟优势。

我们通过一个实时端点进行了验证,在一个会话中发送了一段三句话的输入。三个分块全部到达,且每段的音频与所发送的文本相符:

片段 文本 音频字节 时长(秒)
1 “你好。” 73,200 1.525
2 “这是一项部署后检查。” 108,000 2.250
3 “它有三个句子。” 98,400 2.050

第4步:选择一种将音频传回客户端的传输方式

将文本以流式增量地送入模型解决了问题的一半。另一半是把生成的音频传输到负责渲染它的地方,通常是浏览器。大多数智能体和聊天前端已经通过 Server-Sent Events(SSE)流式传输文本:单向传输,并且大多数浏览器原生支持。SSE 是一种文本协议,并非为承载原始二进制数据而设计。要在现有 SSE 管道中添加语音功能的团队,自然会沿用已有的模式:将音频进行 base64 编码并封装在 JSON 事件中。

我们在24个音频片段上测量了该方法的开销,通过Amazon Transcribe确认不同编码之间的质量没有差异,并对比了MP3与AAC:

候选方案 原始字节数 Base64字节数 比特率
原始PCM 2,396,400 3,195,200 384 kbps
MP3 48k 308,304 411,072 49.5 kbps
AAC 32k fMP4 258,039 344,084 41.6 kbps

Base64 会使每个负载膨胀接近预期的 33%。更隐蔽的代价是帧对齐。SSE 事件由空行分隔,但内部的音频被划分为 MP3 帧。这两种划分方式未必与传输层对数据分块的方式一致。在一次实测响应中,四个逻辑 MP3 帧以六个传输分块到达,其中两个在帧中间结束。把每个传输分块当作一个完整单元的消费者在普通、健康的流量上也会失败。解决办法是缓冲传入的字节,只有在找到分隔符时才进行处理:

buffer = b""
for event in resp["Body"]:
    if "PayloadPart" not in event:
        continue
    buffer += event["PayloadPart"]["Bytes"]
    while b"\n\n" in buffer:
        frame, buffer = buffer.split(b"\n\n", 1)
        handle(frame)

如今用于面向模型一侧的替代方案完全跳过 base64 和 SSE 封装。它在已经在承载增量文本输入的同一条双向会话上多路复用原始二进制音频和 JSON 控制帧,并用一条规则将二者区分开来。

import json


def _decode_control_frame(chunk: bytes) -> dict | None:
    try:
        decoded = json.loads(chunk.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError):
        return None
    return decoded if isinstance(decoded, dict) else None

一个无法作为 JSON 对象解码的分块就是音频,按接收到的原样转发,路径中没有任何编码步骤。有一个边界情况很重要:原始 PCM 音频可能合法地以字节 0x7b开头,即 ASCII 字符 {,与 JSON 对象的起始字符相同。这条规则之所以能避免误分类该音频,是因为它要求整个分块都能解析为有效 JSON,而不仅仅是第一个字节——一个 PCM 帧基本上永远不会做到这一点。

注意: 如果你的前端已经有一条可用的文本 SSE 管道,那么在 SSE 事件中放入 base64 编码的音频是务实的选择。它带来大约 33% 的额外负载和一个缓冲重组步骤,但复用了你已经在运营的基础设施。如果你是从零构建语音通道,或者延迟是首要考虑,那么统一的二进制与 JSON 方式可以同时避免这两项成本,并复用已经在承载文本的连接。

第 5 步:在真实 agent 中整合起来

更完整的端到端案例是一个实时生成式 agent,一次一个 token 地组成响应,在第一个句子存在之前无法知道第四个句子。本示例作为一个 Strands agent 运行,由 Amazon Bedrock 支持,部署到 Amazon Bedrock AgentCore runtime,并使用第 4 步中的模式调用 Kokoro 端点。该设计连接了三条通信链路:面向浏览器、将 Server-Sent Events 送回调用方的链路,agent 为生成文本而对 Amazon Bedrock 发起的流式调用,以及 agent 与 Kokoro 端点之间用于合成的双向会话。

下图展示了这三条通信链路,以及生成与合成之间的重叠发生在何处:

图 3:连接调用方、Amazon Bedrock 和 Kokoro 端点的三条通信链路

请求负载带有一个 mode 字段,取值为 "bidi" 或 "batch",在同一入口函数内被分发到两条代码路径。对两者进行基准测试时只改变一个变量:合成何时开始。在 run_bidi模式下,两个任务并发运行。一个任务排空 Bedrock 文本流并将完成的句子送入分块器。另一个任务排空 Kokoro 会话并在音频分块到达时将其发出。两者都推入同一个共享队列,因此消费者按真实的到达顺序产生事件:较早句子的音频可以在 Bedrock 仍在生成较晚句子时到达调用方。 run_batch 模式则是有意的反面:它等待完整响应之后才进行一次合成调用,作为公平的基准。

反方向上,从 agent 容器返回调用方的链路确实使用了 base64,原因与第 4 步所述不同。 BedrockAgentCoreApp的入口必须产生普通的、可 JSON 序列化的字典,其 SDK 会自动将其转换为 data: {json}\n\n 帧。没有办法向它传递原始二进制帧。这是 JSON 原生托管运行时的结果,而不是第 4 步所测量并弃用的选择。两种模式 run_bidi 和 run_batch 在这条链路上付出完全相同的代价,因此它们之间测得的差异纯粹是时序问题:第一个音频分块何时被发出。

测试与验证

我们通过用两种方式发送同一段 267 字符、四个句子的文本,来隔离 TTS 传输本身对首个音频时间的贡献:作为一条消息(batch),以及只发送其第一句(incremental)。我们针对在线端点各运行了五次试验。

正如之前所测量的,只发送第一句使平均首个音频字节时间从 329.3 ms 降到 156.7 ms,相对改进 52.4%,且完全归因于 TTS 传输本身。

随后我们测量了由真实语言模型产生文本的更完整端到端案例,将相同的提示词分别通过 run_bidi 和 run_batch 针对同一个在线 Bedrock 模型和端点运行,因此唯一的变量就是合成何时开始。

使用相同的方法,bidi 模式对本地容器的首个音频时间比 batch 模式快 2.4 倍,对已部署的 AgentCore runtime 快 3.5 倍。

下图展示了两种环境中这一差距:

图 4:batch 与 bidi 模式在本地和已部署 runtime 上的相对首个音频时间

一旦加入真实的网络跳转,在部署场景下这一差距不是缩小而是扩大。 run_bidi 将该网络跳转吸收进它在生成与合成之间已经在利用的重叠之中。 run_batch 没有等效的重叠可以掩盖这一成本,因为它一次只做一件事,因此每一毫秒新增的网络开销都直接落在它的串行等待上。

要在你自己的部署环境上复现这一比较,请运行 agent 容器提供的相同的 bidi-vs-batch 测试:

uv run agent-infra demo --endpoint-url <your-endpoint-url> --prompt "<your prompt>"

这会将提示词通过两种模式发送,并报告每者的第一个 audio_chunk 事件相对于 agent_start 的时间戳。如果你看到的是错误,请检查步骤 1 中的 KeyError: 'text' 来自步骤1的症状。一个报告 InService 但每个请求都失败的端点表明存在该 bug,而不是部署问题。

清理

为避免持续产生费用,请按照与部署相反的顺序删除本演练创建的资源,先删除 agent,再删除端点:

uv run agent-infra destroy --bedrock-model-id <your-bedrock-model-id>
uv run speech-infra destroy kokoro-82m

agent-infra destroy 删除 Amazon Bedrock AgentCore 运行时资源。 speech-infra destroy 删除 SageMaker AI 模型、端点配置和端点。这两条命令都不会清空缓存模型权重的 Amazon S3 存储桶,也不会删除存放容器镜像的 Amazon ECR 仓库,因为两者可以在多个模型之间共享。如果您不再需要它们,请手动清空并删除 Amazon S3 存储桶以及删除 Amazon ECR 仓库,因为存储的模型检查点和容器镜像在删除之前会持续产生少量存储费用。

结论

我们展示了如何消除语音代理中一种可避免的延迟来源,即在开始语音合成之前等待代理的完整响应,以及如何量化第二种更隐蔽的成本:SSE 中的 base64 编码音频会使负载膨胀约三分之一,并且如果不进行缓冲,可能会与 MP3 帧边界错位。在单独测试中,代理一生成句子就将其流式传输到 SageMaker AI 双向端点,将首个音频出现的时间缩短了一半以上;在已部署的代理中则缩短了 2.4 到 3.5 倍。在该会话上多路复用二进制音频和 JSON 控制帧,而不是将音频编码为文本,消除了朴素的 SSE 中继所需的大小膨胀和帧重新对齐工作。

如果您正在为一个已经通过 SSE 流式传输文本的代理添加语音功能,请从第 4 步中的 base64 方法开始,并在延迟成为优先考虑因素时再重新审视统一传输方案。要了解更多信息,请参阅 Amazon SageMaker AI real-time inference documentation 和 Amazon Bedrock AgentCore documentation。如果您希望获得将此模式应用于您自己的架构的帮助,请联系您的 AWS 客户团队。

 


关于作者

原始出处

AWS Machine Learning

内容说明

原始发布及相关权利归来源方。

机器翻译 · 请以原文为准