NanoCodec 解码器 - ONNX

NeMo NanoCodec 音频编解码器的 ONNX 优化解码器。

与 PyTorch 版本相比，该模型为 KaniTTS 及类似的 TTS 系统提供了 2.5 倍更快的推理速度。

模型详情

• 模型类型： 音频编解码器解码器
• 格式： ONNX (Opset 14)
• 输入： Token 索引 [batch, 4, num_frames]
• 输出： 音频波形 [batch, samples] @ 22050 Hz
• 大小： 122 MB
• 参数量： ~31.5M（仅解码器，占完整模型的 15.8%）

性能

实时因子 (RTF)： GPU 上为 0.44x（生成音频速度快于播放速度！）

快速开始

安装

pip install onnxruntime-gpu numpy

仅 CPU 版本：

pip install onnxruntime numpy

使用方法

import numpy as np
import onnxruntime as ort

# Load model
session = ort.InferenceSession(
    \"nano_codec_decoder.onnx\",
    providers=[\"CUDAExecutionProvider\", \"CPUExecutionProvider\"]
)

# Prepare input
tokens = np.random.randint(0, 500, (1, 4, 10), dtype=np.int64)  # [batch, codebooks, frames]
tokens_len = np.array([10], dtype=np.int64)

# Run inference
outputs = session.run(
    None,
    {\"tokens\": tokens, \"tokens_len\": tokens_len}
)

audio, audio_len = outputs
print(f\"Generated audio: {audio.shape}\")  # [1, 17640] samples

与 KaniTTS 集成

from onnx_decoder_optimized import ONNXKaniTTSDecoderOptimized

# Initialize decoder
decoder = ONNXKaniTTSDecoderOptimized(
    onnx_model_path=\"nano_codec_decoder.onnx\",
    device=\"cuda\"
)

# Decode frame (4 codec tokens)
codes = [100, 200, 300, 400]
audio = decoder.decode_frame(codes)  # Returns int16 numpy array

模型架构

解码器包含两个阶段：

1. 反量化 (FSQ)： 将 token 索引转换为潜在表示

   • 输入：[batch, 4, frames] → 输出：[batch, 16, frames]

2. 音频解码器 (HiFiGAN)： 从潜在表示生成音频

   • 输入：[batch, 16, frames] → 输出：[batch, samples]
   • 上采样因子：~1764x（22050 Hz 下每帧 80ms）

导出详情

• 源模型： nvidia/nemo-nano-codec-22khz-0.6kbps-12.5fps
• 导出方法： PyTorch → ONNX (legacy exporter)
• Opset 版本： 14
• 动态轴： 帧维度和音频样本数
• 优化： 启用图优化、常量折叠

使用场景

• 文本转语音系统： 快速神经编解码器解码
• 实时音频生成： GPU 上实现亚实时性能
• 流式 TTS： 低延迟逐帧解码
• KaniTTS 集成： PyTorch 解码器的直接替代

系统要求

GPU（推荐）

• CUDA 11.8+ 或 12.x
• cuDNN 8.x 或 9.x
• ONNX Runtime GPU：pip install onnxruntime-gpu

CPU

• 任何现代 CPU
• ONNX Runtime：pip install onnxruntime

输入

• tokens (int64)：编解码器 token 索引

  • 形状：[batch_size, 4, num_frames]
  • 范围：[0, 499]（FSQ 码本索引）

• tokens_len (int64)：帧数

  • 形状：[batch_size]
  • 值：序列中的帧数

输出

• audio (float32)：生成的音频波形

  • 形状：[batch_size, num_samples]
  • 范围：[-1.0, 1.0]
  • 采样率：22050 Hz

• audio_len (int64)：音频长度

  • 形状：[batch_size]
  • 值：音频样本数

精度

与 PyTorch 参考实现相比：

• 平均绝对误差： 0.0087
• 相关性： 1.000000（完美）
• 相对误差： 0.0006%

音频质量与 PyTorch 版本几乎完全相同。

限制

• 固定采样率 (22050 Hz)
• 仅支持单声道音频
• 需要有效的 FSQ token 索引（0-499 范围）
• 在支持 CUDA 的 NVIDIA GPU 上性能最佳

许可证

Apache 2.0（与源模型相同）

链接

• 原始模型： nvidia/nemo-nano-codec-22khz-0.6kbps-12.5fps
• KaniTTS： nineninesix/kani-tts-400m-en
• ONNX Runtime： onnxruntime.ai

致谢

• NVIDIA NeMo 团队提供的原始 NanoCodec
• ONNX Runtime 团队提供的推理引擎
• KaniTTS 团队提供的 TTS 系统