Friendly Stable Audio Tools: 重构和增强的音频生成工具集

介绍 Friendly Stable Audio Tools

Friendly Stable Audio Tools 是一个基于 Stability AI 开源的 stable-audio-tools 项目进行重构和增强的音频生成工具集。这个项目旨在为音频和音乐生成模型提供更加友好和实用的开发环境。

Stable Audio 2.0

主要特性

重构的代码结构,提高了可读性和可用性
新增了评估和使用自训练模型的实用脚本
提供了如何训练 Stable Audio 2.0 等模型的详细指导
为 Stable Audio Open 1.0 提供了详细文档和便捷使用脚本

Stable Audio Open 的开源发布

值得注意的是,Stability AI 最近开源了 Stable Audio 的预训练模型。这是音频生成领域的一个重要里程碑。

官方发布公告:Introducing Stable Audio Open

对于想要尝试 Stable Audio Open 的用户,本项目提供了详细的使用说明文档:

Stable Audio Open 1.0 使用文档

此外,本项目还支持通过 YAML 文件输入提示词,实现多 GPU/节点的批量生成功能。示例配置文件可以在这里找到:generate_conditions.yaml

环境要求与安装

环境要求

PyTorch 2.0 或更高版本(用于支持 Flash Attention)
Python 3.8.10 或更高版本

安装步骤

克隆仓库:

git clone https://github.com/yukara-ikemiya/friendly-stable-audio-tools.git

进入项目目录:
```
cd friendly-stable-audio-tools
```
安装依赖:
```
pip install .
```
解决可能的 Accelerate 导入错误:
```
pip uninstall -y transformer-engine
```

构建训练环境

为了简化训练环境的搭建,推荐使用容器系统如 Docker 或 Singularity,而不是在每台 GPU 机器上单独安装依赖。以下是创建 Docker 和 Singularity 容器的步骤:

创建 Docker 镜像:

NAME=friendly-stable-audio-tools
docker build -t ${NAME} -f ./container/${NAME}.Dockerfile .

将 Docker 镜像转换为 Singularity 容器:

singularity build friendly-stable-audio-tools.sif docker-daemon://friendly-stable-audio-tools

执行上述脚本后,将在工作目录中创建 friendly-stable-audio-tools.sif 文件。

Gradio 界面

本项目提供了一个基于 Gradio 的简单界面,用于测试训练好的模型。例如,要为 stable-audio-open-1.0 模型创建界面,可以运行:

python3 ./run_gradio.py --pretrained-name stabilityai/stable-audio-open-1.0

Gradio interface

run_gradio.py 脚本支持多个命令行参数,包括:

--pretrained-name: Hugging Face 上的模型名称
--model-config: 本地模型配置文件路径
--ckpt-path: 本地模型检查点文件路径
--pretransform-ckpt-path: 预转换检查点路径
--username 和 --password: 设置 Gradio 演示的登录信息
--model-half: 是否使用半精度
--tmp-dir: 临时输出文件保存目录

训练流程

配置文件准备

在开始训练之前,需要准备两个配置文件:

模型配置文件
数据集配置文件

从头开始训练

要开始训练,在仓库根目录运行 train.py 脚本:

python3 train.py --dataset-config /path/to/dataset/config --model-config /path/to/model/config --name my_experiment

微调模型

微调涉及从预训练检查点恢复训练:

使用 --ckpt-path 参数从包装的检查点恢复训练
使用 --pretrained-ckpt-path 参数从预训练的解包模型开始新的训练

模型解包

stable-audio-tools 使用 PyTorch Lightning 来实现多 GPU 和多节点训练。训练过程中,模型被包装在一个"训练包装器"中,这增加了检查点文件的大小。unwrap_model.py 脚本用于解包模型检查点:

python3 unwrap_model.py --model-config /path/to/model/config --ckpt-path /path/to/wrapped/ckpt.ckpt --name /new/path/to/new_ckpt_name

训练 Stable Audio 2.0

准备工作

准备 CLAP 编码器检查点:
- 从 LAION CLAP 仓库下载预训练的 CLAP 检查点
- 将检查点文件存储在选定的目录中
- 编辑 Stable Audio 2.0 的模型配置文件,设置 clap_ckpt_path
准备音频和元数据:
- 使用 JSON 格式的元数据,包含 prompt 字段
- 将元数据文件与对应的音频文件放在同一目录下,文件名相同

第一阶段: VAE-GAN (压缩模型)

训练 VAE-GAN:

singularity exec --nv --pwd $ROOT_DIR -B $ROOT_DIR -B $DATASET_DIR \
  --env WANDB_API_KEY=$WANDB_API_KEY \
  ${CONTAINER_PATH} \
  torchrun --nproc_per_node gpu --master_port ${PORT} \
  ${ROOT_DIR}/train.py \
    --dataset-config ${DATASET_CONFIG} \
    --model-config ${MODEL_CONFIG} \
    --name "vae_training" \
    --num-gpus 8 \
    --batch-size ${BATCH_SIZE} \
    --num-workers 8 \
    --save-dir ${OUTPUT_DIR}

模型解包:

singularity exec --nv --pwd $ROOT_DIR -B $ROOT_DIR \
  --env WANDB_API_KEY=$WANDB_API_KEY \
  ${CONTAINER_PATH} \
  torchrun --nproc_per_node gpu --master_port ${PORT} \
    ${ROOT_DIR}/unwrap_model.py \
    --model-config ${MODEL_CONFIG} \
    --ckpt-path ${CKPT_PATH} \
    --name ${OUTPUT_PATH}

重构测试:

singularity exec --nv --pwd $ROOT_DIR -B $ROOT_DIR -B $DATASET_DIR \
  --env WANDB_API_KEY=$WANDB_API_KEY \
  ${CONTAINER_PATH} \
  torchrun --nproc_per_node gpu --master_port ${PORT} \
    ${ROOT_DIR}/reconstruct_audios.py \
    --model-config ${MODEL_CONFIG} \
    --ckpt-path ${UNWRAP_CKPT_PATH} \
    --audio-dir ${AUDIO_DIR} \
    --output-dir ${OUTPUT_DIR} \
    --frame-duration ${FRAME_DURATION} \
    --overlap-rate ${OVERLAP_RATE} \
    --batch-size ${BATCH_SIZE}

第二阶段: 扩散变换器 (DiT)

确保完成所有先决条件并训练了 VAE 模型后,可以开始训练 DiT 模型:

singularity exec --nv --pwd $ROOT_DIR -B $ROOT_DIR -B $DATASET_DIR \
  --env WANDB_API_KEY=$WANDB_API_KEY \
  ${CONTAINER_PATH} \
  torchrun --nproc_per_node gpu --master_port ${PORT} \
    ${ROOT_DIR}/train.py \
    --dataset-config ${DATASET_CONFIG} \
    --model-config ${MODEL_CONFIG} \
    --pretransform-ckpt-path ${PRETRANSFORM_CKPT} \
    --name "dit_training" \
    --num-gpus ${NUM_GPUS} \
    --batch-size ${BATCH_SIZE} \
    --save-dir ${OUTPUT_DIR}