用 Rust 实现 MusicGen 音乐生成：candle-musicgen 示例源码级解析

导读

candle-musicgen 是 Candle 框架对 Meta MusicGen 音乐生成模型（源自论文 Simple and Controllable Music Generation）的纯 Rust 实现示例。它完整还原了"文本提示（Text Prompt）→ 文本语义编码 → 多 codebook 音频令牌解码"的完整链路，并内置了模型与分词器自动下载、CPU/GPU 自动选择等实用机制。读完本文，你将掌握该示例的运行方式、全部命令行参数的含义、三大子模型（T5 文本编码器、MusicGen 解码器、Encodec 音频编码器）的架构组成，以及如何基于仓库源码理解 MusicGen 的生成原理。

一、示例概览与文件布局

示例位于 candle-examples/examples/musicgen/ 目录，由两个 Rust 源文件和一个 README 组成：

文件	作用
README.md	示例说明与运行命令
main.rs	命令行入口：参数解析、权重加载、分词、文本编码
musicgen_model.rs	MusicGen 完整模型实现：配置、文本编码器、解码器、音频编码器的组装

从源码结构看，该示例目前完成了"文本编码"阶段（将提示词转成语义嵌入），音频解码与波形生成部分尚未接入完整生成循环，main.rs 顶部注释也明确标注了两个待办事项（TODO）：增加离线模式、增加 KV cache。

二、运行示例：一条命令跑通文本到音乐语义

README 中给出了最核心的运行方式，直接在 candle-examples 目录下执行：

$ cargo run --example musicgen -- --prompt "90s rock song with loud guitars and heavy drums"

首次运行时，程序会通过内置的 Hugging Face Hub 客户端自动下载 facebook/musicgen-small 仓库中的模型权重与分词器（详见后文），无需手动准备任何文件。正常运行后的输出分为三段：

> tokens: [2777, 7, 2480, 2324, 28, 8002, 5507, 7, 11, 2437, 5253, 7, 1]
> Tensor[dims 1, 13; u32]
> [[[ 0.0902,  0.1256, -0.0585, ...,  0.1057, -0.5141, -0.4675],
>   [ 0.1972, -0.0268, -0.3368, ..., -0.0495, -0.3597, -0.3940],
>   [-0.0855, -0.0007,  0.2225, ..., -0.2804, -0.5360, -0.2436],
>   ...
>   [ 0.0515,  0.0235, -0.3855, ..., -0.4728, -0.6858, -0.2923],
>   [-0.3728, -0.1442, -0.1179, ..., -0.4388, -0.0287, -0.3242],
>   [ 0.0163,  0.0012, -0.0020, ...,  0.0142,  0.0173, -0.0103]]]
> Tensor[[1, 13, 768], f32]

这三段输出分别对应：

1. tokens：提示词经分词器编码后的 token id 序列（[2777, 7, 2480, ...]）；
2. Tensor[dims 1, 13; u32]：token id 序列的张量形式，形状为 [1, 13]，即 1 个批次、13 个 token，元素类型为 u32；
3. Tensor[[1, 13, 768], f32]：T5 文本编码器输出的语义嵌入，形状为 [1, 13, 768]（批次 × 序列长度 × 隐藏维度），类型为 f32，这正是后续 MusicGen 解码器进行交叉注意力所需的文本条件表示。

三、命令行参数全解

示例通过 clap 解析命令行参数，定义位于 main.rs，共四个参数：

参数	类型	默认值	说明
--cpu	布尔开关	关闭（默认自动选设备）	强制在 CPU 上运行
--model	可选字符串	自动下载 model.safetensors	指定本地 safetensors 权重文件路径
--tokenizer	可选字符串	自动下载 tokenizer.json	指定本地分词器配置文件路径
--prompt	字符串	"90s rock song with loud guitars and heavy drums"	文本音乐描述

典型用法示例：

# 使用默认提示词、自动选择设备运行
$ cargo run --example musicgen

# 自定义提示词 + 强制 CPU
$ cargo run --example musicgen -- --cpu --prompt "lofi hip hop beats to relax"

# 使用本地权重与分词器，完全离线
$ cargo run --example musicgen -- \
    --model ./model.safetensors --tokenizer ./tokenizer.json

关于设备选择：main.rs 调用 candle_examples::device(args.cpu)（定义于 candle-examples/src/lib.rs）完成设备自动探测——当不指定 --cpu 时，优先使用 CUDA GPU（cuda_is_available()），其次尝试 Metal（metal_is_available()），都不满足则回退到 CPU，并在终端打印提示（例如在 macOS 上提示可用 --features metal 启用 GPU）。

四、模型权重与分词器：自动下载机制

示例默认从 Hugging Face 的 facebook/musicgen-small 仓库拉取资源，这一逻辑在 main.rs 中体现：

let tokenizer = match args.tokenizer {
    Some(tokenizer) => std::path::PathBuf::from(tokenizer),
    None => Api::new()?
        .model("facebook/musicgen-small")
        .get("tokenizer.json")?,
};
// ...
let model = match args.model {
    Some(model) => std::path::PathBuf::from(model),
    None => Api::new()?
        .model("facebook/musicgen-small")
        .with_revision("refs/pr/13")
        .get("model.safetensors")?,
};

关键细节：

• 默认权重源：facebook/musicgen-small 模型的 model.safetensors 文件，并使用 with_revision("refs/pr/13") 固定到指定修订版本，保证权重与源码实现的兼容性；
• 下载与缓存：底层走 candle-examples/src/hub.rs 封装 hf_hub 客户端的 Api/Repo 结构。Repo::get 会先检查本地缓存（local_files_only(true)），未命中才真正下载，并将下载进度（百分比、已下载/总大小，自动换算 KiB/MiB/GiB）打印到 stderr；同时支持通过 HF_TOKEN、HF_HUB_CACHE 等环境变量配置客户端；
• 本地模式：传入 --model 和 --tokenizer 即可完全跳过网络下载，用本地文件离线加载。

加载权重时使用内存映射方式读取 safetensors 文件，代码为 unsafe { VarBuilder::from_mmaped_safetensors(&<a href="https://link.gitcode.com/i/1bbedccfe1626b45c0ae9bb942da4df9" target="_blank">model], DTYPE, &device)? }，其中 DTYPE 在本示例中被固定为 DType::F32（见 [main.rs），即所有权重以 32 位浮点精度加载。

分词器通过 tokenizers crate 加载（Tokenizer::from_file），随后调用 .with_padding(None).with_truncation(None) 关闭自动 padding 与截断，即完全保留提示词原始长度。

五、文本编码流程：提示词如何变成语义条件

main.rs 中的推理主体非常简洁，核心步骤如下（main.rs）：

1. 分词：tokenizer.encode(args.prompt, true) 将文本提示词编码为 token id 序列，true 表示启用特殊 token；
2. 转张量：token id 通过 Tensor::new(...)?.unsqueeze(0)? 增加批次维度，得到形状 [1, seq_len] 的 u32 张量；
3. 文本编码：调用 model.text_encoder.forward(&tokens)?，得到形状 [1, seq_len, 768] 的 f32 语义嵌入；
4. 打印：将 token 序列、输入张量与嵌入结果依次打印到终端，供开发者观测中间状态。

其中 text_encoder 的类型是 candle_transformers::models::t5::T5EncoderModel（见 musicgen_model.rs）。在 MusicGen 论文的架构中，文本编码器的作用是把自然语言描述映射为语义条件向量，供解码器在生成音频令牌时通过交叉注意力对齐"文本描述的内容"与"音频令牌的内容"。T5 编码器的实现位于 candle-transformers/src/models/t5.rs，其 load 方法会自动适配权重命名（优先取 shared.weight，否则回退到 decoder.embed_tokens 或 encoder.embed_tokens），forward 则调用内部 T5 Stack 前向。

六、模型架构：三个子模型的完整组装

MusicgenForConditionalGeneration 是整个示例的顶层结构（musicgen_model.rs），它由三个子模型构成，与 MusicGen 论文架构一一对应：

pub struct MusicgenForConditionalGeneration {
    pub text_encoder: t5::T5EncoderModel,       // 文本 → 语义条件
    pub audio_encoder: encodec::Model,          // 音频 ↔ 离散令牌（Encodec）
    pub decoder: MusicgenForCausalLM,           // 语义条件 + 历史令牌 → 下一个令牌
    cfg: GenConfig,
}

三个子模型的加载职责划分如下（MusicgenForConditionalGeneration::load）：

子模型	权重前缀	类型	职责
text_encoder	text_encoder	T5 Encoder	把文本提示编码为语义嵌入
audio_encoder	audio_encoder	Encodec	音频编解码，将音频映射为离散 codebook 令牌（当前阶段未调用）
decoder	decoder	MusicGen Decoder + lm_heads	自回归式生成音频令牌 logits

其中：

• T5 文本编码器采用 t5::Config::musicgen_small()（定义于 t5.rs），关键配置为 d_model=768、num_heads=12、num_layers=12、d_ff=3072、d_kv=64、vocab_size=32128、relative_attention_num_buckets=32，激活函数为 ReLU；
• Encodec 音频编码器基于 candle-transformers/src/models/encodec.rs 的 Model，是官方 Encodec 的 Candle 移植，用于音频与离散令牌之间的双向映射。在完整的 MusicGen 生成流程中，解码器生成的令牌会交给 Encodec 解码器还原成 32 kHz 的音频波形，再经声码器输出为可听音乐。

七、解码器内部实现：多 codebook 条件生成

MusicGen 解码器是生成的核心，仓库中的实现分为三层：

7.1 MusicgenDecoder：多 codebook 嵌入与正弦位置编码

MusicgenDecoder 包含：

• 多 codebook 词嵌入：embed_tokens 是一个 Vec<Embedding>，数量等于 num_codebooks（small 配置为 4），每个 codebook 的嵌入维度是 vocab_size + 1（多出的 1 个位置对应 pad token）。前向时将每个 codebook 的嵌入累加求和得到 inputs_embeds；
• 正弦位置嵌入：MusicgenSinusoidalPositionalEmbedding 按 get_embedding 生成，使用标准正弦/余弦公式（cos/sin 拼接，奇数维度时补零），forward 时按当前序列长度 narrow 截取位置向量；
• 24 层解码器层（num_hidden_layers = 24）与最终 LayerNorm，LayerNorm 的 eps 为 1e-5。

7.2 MusicgenDecoderLayer：自注意力 + 交叉注意力 + FFN

每一层解码器（musicgen_model.rs）按标准 Transformer 解码器结构组织：

1. self_attn_layer_norm → 自注意力（self_attn）→ 残差连接；
2. 当存在 encoder_hidden_states 时：encoder_attn_layer_norm → 交叉注意力（encoder_attn，以文本编码器输出为 K/V）→ 残差连接；
3. final_layer_norm → fc1（升维到 ffn_dim=4096）→ GELU 激活 → fc2（降维回 hidden_size=1024）→ 残差连接。

自注意力与交叉注意力统一由 MusicgenAttention 实现（musicgen_model.rs）：四个无偏置线性层 q_proj/k_proj/v_proj/out_proj，注意力分数按 1 / sqrt(head_dim) 缩放（head_dim = hidden_size / num_attention_heads = 64），经 softmax 后加权求和，最后过 out_proj 投影。

7.3 MusicgenForCausalLM：多 codebook 语言模型头

MusicgenForCausalLM 是解码器之上的语言模型头：

• 每个 codebook 配一个 lm_heads 线性层（无偏置），将隐藏状态映射到 vocab_size=2048 的词汇表上；
• forward 输出形状为 (b_sz * num_codebooks, seq_len, vocab_size) 的 logits 张量，即每个 codebook 各自预测下一令牌的分布——这正是 MusicGen 并行（或逐 codebook）解码多音频流的结构基础。

当前阶段代码中对 prepare_decoder_attention_mask 使用 todo!() 占位，说明解码器前向的注意力掩码逻辑尚待补全，进一步印证该示例目前聚焦于"文本条件编码"部分。

八、GenConfig::small：三份配置的完整参数

GenConfig::small()（musicgen_model.rs）一次性配置了三个子模型，其中 MusicGen 解码器的 Config 默认值如下：

配置项	值	含义
vocab_size	2048	每个 codebook 的令牌词汇量
max_position_embeddings	2048	最大位置编码长度
num_hidden_layers	24	解码器层数
ffn_dim	4096	FFN 中间维度
num_attention_heads	16	注意力头数
hidden_size	1024	隐藏维度（d_model）
layerdrop	0.0	层丢弃率（未启用）
use_cache	true	是否启用 KV cache（当前示例尚未实现）
activation_function	Gelu	FFN 激活函数
dropout / attention_dropout / activation_dropout	0.1 / 0.0 / 0.0	三类 dropout 率
initializer_factor	0.02	初始化因子
scale_embedding	false	是否按 sqrt(d_model) 缩放嵌入
num_codebooks	4	音频 codebook 数量
pad_token_id / bos_token_id	2048 / 2048	填充与起始 token id
eos_token_id	None	结束 token id
tie_word_embeddings	false	是否绑定嵌入与输出投影

Encodec 部分的配置同样完整，其中与音频质量最相关的几项：

• sampling_rate: 32_000：目标音频采样率 32 kHz；
• codebook_size: 2048、codebook_dim: Some(128)：量化 codebook 规模；
• upsampling_ratios: [8, 5, 4, 4]：解码器逐级上采样比例；
• num_filters: 64、kernel_size: 7、num_lstm_layers: 2：卷积与 LSTM 结构参数；
• target_bandwidths: [2.2]：目标带宽（kbps）；
• 代码注释特别指出 pad_mode 应使用 Reflect 但当前因实现未支持而暂时使用 Replicate。

九、输出解读：T5 语义嵌入在生成流程中的位置

回到运行输出，Tensor[[1, 13, 768], f32] 就是"文本提示 → 语义条件"这一环节的产物。在完整的 MusicGen 生成管线中，它将被送入解码器的交叉注意力，作为每个生成步骤的 encoder_hidden_states（见 7.2 节中 MusicgenDecoderLayer::forward 对 encoder_hidden_states 的处理），引导 4 个 codebook 自回归生成音频令牌；令牌序列随后经 Encodec 解码器还原为 32 kHz 波形，最终得到可播放的音乐片段。

当前仓库示例运行到文本编码即告一段落，开发者可以在此基础上继续实现：

• 音频令牌解码：调用 audio_encoder（Encodec）把生成的令牌解码为波形；
• 自回归生成循环：实现 prepare_decoder_attention_mask，逐 token 生成 4 路 codebook 令牌；
• KV cache 加速：main.rs 顶部注释已明确将其列为待办事项，Config::use_cache 也已预留为 true；
• 离线模式：使用 --model 与 --tokenizer 指定本地文件即可脱离网络运行。

十、代码路径速查

本文涉及的所有关键实现均可直接在仓库中查看：

• 运行入口与参数解析：candle-examples/examples/musicgen/main.rs
• 模型组装与配置：candle-examples/examples/musicgen/musicgen_model.rs
• T5 文本编码器：candle-transformers/src/models/t5.rs
• Encodec 音频编解码器：candle-transformers/src/models/encodec.rs
• 设备自动选择：candle-examples/src/lib.rs
• Hub 下载与缓存封装：candle-examples/src/hub.rs