Open MusicLM

这是 Google 发布的 SOTA 文本到音乐模型 MusicLM 的 PyTorch 实现，并做了一些修改。我们使用 CLAP 替代 MuLan，Encodec 替代 SoundStream，以及 MERT 替代 w2v-BERT。

为什么选择 CLAP？

CLAP 是一个联合音频-文本模型，基于 LAION-Audio-630K 数据集训练而成。与 MuLan 类似，它由音频塔和文本塔组成，分别将各自模态映射到一个共享的潜在空间（CLAP 为 512 维，而 MuLan 为 128 维）。

MuLan 是在 5000 万对文本-音乐数据上训练的。遗憾的是，我并没有这些数据来复现它，因此我依赖 CLAP 的预训练检查点来尽可能接近。CLAP 在 LAION-630k（约 63.3 万对文本-音频）和 AudioSet（200 万个带有关键词转写字幕的样本）的数据上进行了训练，总共使用了 260 万对文本-音频数据。尽管这远少于 MuLan 训练时使用的数据量，但我们已经成功利用 CLAP 生成了多样化的音乐样本，你可以在 这里 收听这些早期结果。如果 CLAP 的潜在空间对于音乐生成来说不够丰富，我们可以用音乐数据继续训练 CLAP，或者等 @lucidrain 的 MuLan 实现 训练完成后将其作为替代方案。

为什么选择 Encodec？

SoundStream 和 Encodec 都是神经网络音频编解码器，能够将任意波形编码为一系列声学标记，随后再解码回与原始波形相似的信号。这些中间标记可以被建模为序列到序列任务。Encodec 由 Facebook 开源，并提供了公开可用的预训练检查点，而 SoundStream 则没有。

与 @lucidrains 实现的区别

• 自回归地将 CLAP/MuLan 的条件信号作为离散标记传递给 Transformer，正如论文第 3.1 节所述。而在 musiclm-pytorch 中，则通过交叉注意力对其进行条件化。
• TokenConditionedTransformer 可以支持可变长度的标记序列，这使得进一步的实验更加便捷（例如组合多个条件信号、生成立体声波形等）。
• 使用现有的开源模型，而不是重新训练 MuLan 和 SoundStream。
• 对代码进行了一些修改，以提高模型成功训练的可能性。

最终目标

该项目的目标是在不完全拘泥于论文中架构的前提下，尽快复现 MusicLM 的效果。如果你希望获得更贴近原论文实现的版本，请查看 musiclm-pytorch。

我们还希望更好地理解 CLAP 的潜在空间。

如果你想参与其中，欢迎加入我们的 Discord！

使用方法

安装

conda env create -f environment.yaml
conda activate open-musiclm

配置文件

“模型配置”包含有关模型架构的信息，比如层数、量化器数量、各阶段的目标音频长度等，用于在训练和推理时实例化模型。

“训练配置”则包含训练模型的超参数，用于在训练过程中实例化训练器类。

示例配置文件请参见 ./configs 目录。

训练

CLAP RVQ

第一步是训练残差向量量化器，将连续的 CLAP 嵌入映射为离散的标记序列。

python ./scripts/train_clap_rvq.py \
    --results_folder ./results/clap_rvq \ # 保存结果和检查点的路径
    --model_config ./configs/model/musiclm_small.json \ # 模型配置文件路径
    --training_config ./configs/training/train_musiclm_fma.json # 训练配置文件路径

Hubert K-means

接下来，我们需要学习一个 K-means 层，用于将 MERT 嵌入量化为语义标记。

python ./scripts/train_hubert_kmeans.py \
    --results_folder ./results/hubert_kmeans \ # 保存结果和检查点的路径
    --model_config ./configs/model/musiclm_small.json \
    --training_config ./configs/training/train_musiclm_fma.json

语义阶段 + 粗粒度阶段 + 细粒度阶段

一旦 K-means 和 RVQ 准备就绪，我们就可以开始训练语义、粗粒度和细粒度阶段。这些阶段可以同时进行训练。

python ./scripts/train_semantic_stage.py \
    --results_folder ./results/semantic \ # 保存结果和检查点的路径
    --model_config ./configs/model/musiclm_small.json \
    --training_config ./configs/training/train_musiclm_fma.json \
    --rvq_path PATH_TO_RVQ_CHECKPOINT \ # 之前训练好的 RVQ 检查点路径
    --kmeans_path PATH_TO_KMEANS_CHECKPOINT # 之前训练好的 K-means 检查点路径

python ./scripts/train_coarse_stage.py \
    --results_folder ./results/coarse \ # 保存结果和检查点的路径
    --model_config ./configs/model/musiclm_small.json \
    --training_config ./configs/training/train_musiclm_fma.json \
    --rvq_path PATH_TO_RVQ_CHECKPOINT \ # 之前训练好的 RVQ 检查点路径
    --kmeans_path PATH_TO_KMEANS_CHECKPOINT # 之前训练好的 K-means 检查点路径

python ./scripts/train_fine_stage.py \
    --results_folder ./results/fine \ # 保存结果和检查点的路径
    --model_config ./configs/model/musiclm_small.json \
    --training_config ./configs/training/train_musiclm_fma.json \
    --rvq_path PATH_TO_RVQ_CHECKPOINT \ # 之前训练好的 RVQ 检查点路径
    --kmeans_path PATH_TO_KMEANS_CHECKPOINT # 之前训练好的 K-means 检查点路径

预处理

在上述情况下，我们在训练过程中实时使用 CLAP、Hubert 和 Encodec 生成 CLAP 令牌、语义令牌和声学令牌。然而，这些模型会占用 GPU 的显存，如果对同一数据集进行多次运行，每次都重新计算这些令牌效率较低。我们可以预先计算这些令牌，并在训练过程中迭代使用它们。

为此，在配置文件中填写 data_preprocessor_cfg 字段，并在训练器配置中将 use_preprocessed_data 设置为 True（可参考 train_fma_preprocess.json 获取灵感）。然后运行以下命令预处理数据集，再执行你的训练脚本：

python ./scripts/preprocess_data.py \
    --model_config ./configs/model/musiclm_small.json \
    --training_config ./configs/training/train_fma_preprocess.json \
    --rvq_path RVQ检查点路径 \ # 之前训练好的RVQ检查点路径
    --kmeans_path K-means检查点路径 # 之前训练好的K-means检查点路径

推理

生成多个样本，并使用 CLAP 选择最佳的几个：

python scripts/infer_top_match.py \
    "你的文本提示"
    --num_samples 4                                 # 生成的样本数量
    --num_top_matches 1                             # 返回的最佳匹配数量
    --semantic_path 语义阶段检查点路径 \   # 之前训练好的语义阶段检查点
    --coarse_path 粗略阶段检查点路径 \       # 之前训练好的粗略阶段检查点
    --fine_path 精细阶段检查点路径 \           # 之前训练好的精细阶段检查点
    --rvq_path RVQ检查点路径 \             # 之前训练好的RVQ检查点
    --kmeans_path K-means检查点路径         # 之前训练好的K-means检查点
    --model_config ./configs/model/musiclm_small.json \
    --duration 4

为不同的测试提示生成样本：

python scripts/infer.py \
    --semantic_path 语义阶段检查点路径 \   # 之前训练好的语义阶段检查点
    --coarse_path 粗略阶段检查点路径 \       # 之前训练好的粗略阶段检查点
    --fine_path 精细阶段检查点路径 \           # 之前训练好的精细阶段检查点
    --rvq_path RVQ检查点路径 \             # 之前训练好的RVQ检查点
    --kmeans_path K-means检查点路径         # 之前训练好的K-means检查点
    --model_config ./configs/model/musiclm_small.json \
    --duration 4

你可以使用 --return_coarse_wave 标志跳过精细阶段，仅从粗略令牌重建音频。

检查点

你可以从 这里 下载 musiclm_large_small_context 模型的实验性检查点。要对该模型进行微调，请在训练脚本中使用 --fine_tune_from 标志。

感谢

• Okio 提供了训练模型所需的算力！Okio 是一家正在开发 Nendo——一个开源的生成式音乐工具套件——的初创公司，旨在重新构想音乐。如果你感兴趣，可以访问 okio.ai 了解更多信息。
• @lucidrains 提供了 audiolm-pytorch 的实现。该仓库包含了对 audiolm-pytorch 中许多代码的重构版本。
• LAION 提供了 CLAP。
• Music Audio Pretrain 团队 提供了 MERT。

引用

@inproceedings{Agostinelli2023MusicLMGM,
    title     = {MusicLM: 从文本生成音乐},
    author    = {Andrea Agostinelli、Timo I. Denk、Zalán Borsos、Jesse Engel、Mauro Verzetti、Antoine Caillon、Qingqing Huang、Aren Jansen、Adam Roberts、Marco Tagliasacchi、Matthew Sharifi、Neil Zeghidour、C. Frank},
    year      = {2023}
}

@article{wu2022large,
  title     = {基于特征融合和关键词到字幕增强的大规模对比语言-音频预训练},
  author    = {Wu, Yusong、Chen, Ke、Zhang, Tianyu、Hui, Yuchen、Berg-Kirkpatrick, Taylor、Dubnov, Shlomo},
  journal   = {arXiv 预印本 arXiv:2211:06687},
  year      = {2022},
}

@article{defossez2022highfi,
  title     = {高保真神经网络音频压缩},
  author    = {Défossez, Alexandre、Copet, Jade、Synnaeve, Gabriel、Adi, Yossi},
  journal   = {arXiv 预印本 arXiv:2210.13438},
  year      = {2022},
}

@misc{li2023mert,
  title     = {MERT：具有大规模自监督训练的声学音乐理解模型}, 
  author    = {Yizhi Li、Ruibin Yuan、Ge Zhang、Yinghao Ma、Xingran Chen、Hanzhi Yin、Chenghua Lin、Anton Ragni、Emmanouil Benetos、Norbert Gyenge、Roger Dannenberg、Ruibo Liu、Wenhu Chen、Gus Xia、Yemin Shi、Wenhao Huang、Yike Guo、Jie Fu},
  year      = {2023},
  eprint    = {2306.00107},
  archivePrefix = {arXiv},
  primaryClass  = {cs.SD}
}

Open MusicLM 快速上手指南

Open MusicLM 是 Google SOTA 文本生成音乐模型 MusicLM 的 PyTorch 实现。该项目使用 CLAP 替代 MuLan，Encodec 替代 SoundStream，MERT 替代 w2v-BERT，旨在利用现有的开源预训练模型快速复现音乐生成效果。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.8+
• 包管理器: Conda (强烈推荐用于管理依赖)
• 硬件: 支持 CUDA 的 NVIDIA GPU (训练和推理均需要较大的显存)
• 前置知识: 熟悉基本的命令行操作和深度学习训练流程

安装步骤

本项目通过 Conda 环境进行依赖管理。请依次执行以下命令：

1. 克隆仓库 (如果尚未下载):

   git clone https://github.com/your-repo/open-musiclm.git
   cd open-musiclm
   

2. 创建并激活 Conda 环境:

   conda env create -f environment.yaml
   conda activate open-musiclm
   

   提示: 如果 environment.yaml 中的源下载缓慢，国内用户可在创建环境后，手动将 pip 源配置为清华源或阿里源以加速后续可能的 pip 安装： pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

本项目的核心流程分为：训练量化器 -> 训练生成阶段 -> 推理生成。若仅需测试效果，可直接下载官方提供的实验性检查点（Checkpoints）进行推理。

1. 准备检查点 (可选)

如果您不想从头训练，可以从 Google Drive 下载预训练的 musiclm_large_small_context 模型文件。记下以下路径以备后用：

• PATH_TO_SEMANTIC_CHECKPOINT
• PATH_TO_COARSE_CHECKPOINT
• PATH_TO_FINE_CHECKPOINT
• PATH_TO_RVQ_CHECKPOINT
• PATH_TO_KMEANS_CHECKPOINT

2. 推理生成 (Inference)

方案 A：生成并自动筛选最佳样本

此脚本会生成多个样本，并利用 CLAP 模型计算与文本提示的相似度，自动返回最匹配的一个。

python scripts/infer_top_match.py \
    "your text prompt" \
    --num_samples 4 \
    --num_top_matches 1 \
    --semantic_path PATH_TO_SEMANTIC_CHECKPOINT \
    --coarse_path PATH_TO_COARSE_CHECKPOINT \
    --fine_path PATH_TO_FINE_CHECKPOINT \
    --rvq_path PATH_TO_RVQ_CHECKPOINT \
    --kmeans_path PATH_TO_KMEANS_CHECKPOINT \
    --model_config ./configs/model/musiclm_small.json \
    --duration 4

• 将 "your text prompt" 替换为您想要的音乐描述（英文效果更佳）。
• 将 PATH_TO_... 替换为实际的模型文件路径。

方案 B：直接生成样本

此脚本直接根据提示词生成音频，不进行自动筛选。

python scripts/infer.py \
    --semantic_path PATH_TO_SEMANTIC_CHECKPOINT \
    --coarse_path PATH_TO_COARSE_CHECKPOINT \
    --fine_path PATH_TO_FINE_CHECKPOINT \
    --rvq_path PATH_TO_RVQ_CHECKPOINT \
    --kmeans_path PATH_TO_KMEANS_CHECKPOINT \
    --model_config ./configs/model/musiclm_small.json \
    --duration 4

技巧: 添加 --return_coarse_wave 标志可以跳过精细阶段（Fine Stage），仅使用粗粒度令牌重建音频，从而加快生成速度。

3. 训练流程简述 (如需从头训练)

若需使用自定义数据训练，请按顺序执行以下步骤：

1. 训练 CLAP RVQ (将连续嵌入映射为离散令牌):

   python ./scripts/train_clap_rvq.py \
       --results_folder ./results/clap_rvq \
       --model_config ./configs/model/musiclm_small.json \
       --training_config ./configs/training/train_musiclm_fma.json
   

2. 训练 Hubert K-means (量化 MERT 嵌入为语义令牌):

   python ./scripts/train_hubert_kmeans.py \
       --results_folder ./results/hubert_kmeans \
       --model_config ./configs/model/musiclm_small.json \
       --training_config ./configs/training/train_musiclm_fma.json
   

3. 训练生成阶段 (语义、粗粒度、细粒度可并发训练):

   python ./scripts/train_semantic_stage.py \
       --results_folder ./results/semantic \
       --model_config ./configs/model/musiclm_small.json \
       --training_config ./configs/training/train_musiclm_fma.json \
       --rvq_path PATH_TO_RVQ_CHECKPOINT \
       --kmeans_path PATH_TO_KMEANS_CHECKPOINT
   

   (同理运行 train_coarse_stage.py 和 train_fine_stage.py)

4. 数据预处理 (优化选项): 为避免每次训练重复计算令牌，可预先处理数据集。修改配置文件中的 data_preprocessor_cfg 并设置 use_preprocessed_data: True，然后运行：

   python ./scripts/preprocess_data.py \
       --model_config ./configs/model/musiclm_small.json \
       --training_config ./configs/training/train_fma_preprocess.json \
       --rvq_path PATH_TO_RVQ_CHECKPOINT \
       --kmeans_path PATH_TO_KMEANS_CHECKPOINT