Pocket-TTS ExecuTorch 模型

导出为 ExecuTorch PTE 格式的 GPU 加速 Pocket-TTS 模型，用于移动端推理。

模型流程

Text → [Tokenizer] → Tokens
         ↓
    text_conditioner.pte → 文本嵌入 (B, T, 1024)
         ↓
    flow_lm_main_bundled.pte → 条件向量 (B, 1024) + EOS logit
         ↓
    flow_net.pte (ODE 步骤) → 音频编码 (B, 32)
         ↓
    mimi_decoder.pte → 音频样本 (Float32, 24kHz)

文件

使用方法 (Python)

from executorch.runtime import Runtime
import torch

runtime = Runtime.get()

# 加载模型
tc = runtime.load_program("text_conditioner.pte").load_method("forward")
backbone = runtime.load_program("flow_lm_main_bundled.pte")
bb_0 = backbone.load_method("forward_0")
bb_32 = backbone.load_method("forward_32")
flow = runtime.load_program("flow_net.pte").load_method("forward")
decoder = runtime.load_program("mimi_decoder.pte").load_method("forward")

# 推理
tokens = torch.randint(0, 100, (1, 20), dtype=torch.int64)
text_emb = tc.execute([tokens])[0]  # (1, 20, 1024)

# 主干网络步骤 0
seq = torch.randn(1, 1, 32)  # 初始音频潜变量
k_cache = torch.zeros(6, 1, 512, 16, 64)
v_cache = torch.zeros(6, 1, 512, 16, 64)
conditioning, eos, k_new, v_new = bb_0.execute([seq, k_cache, v_cache])

# Flow 步骤 (ODE)
c = conditioning
s, t, x = torch.tensor([[0.0]]), torch.tensor([[1.0]]), torch.randn(1, 32)
flow_dir = flow.execute([c, s, t, x])[0]

# 解码为音频
audio = decoder.execute([final_codes])[0]  # (samples,)

Android 集成

// 在 PocketTtsVulkanEngine.kt 中
val module = Module.load(context.filesDir.resolve("pocket/pte/flow_lm_main_bundled.pte"))
val output = module.forward(EValue.from(inputTensor))

依赖要求

• ExecuTorch 0.6.0+
• Android: org.pytorch:executorch-android:0.6.0

致谢

基于 Kyutai 的 Pocket-TTS。 导出模式参考 Kokoro ExecuTorch。

许可证

Apache 2.0