stable-audio-tools

stable-audio-tools
Training and inference code for audio generation models

Install

Requires PyTorch 2.5 or later for Flash Attention and Flex Attention support. Development for the repo is done in Python 3.10.

We use uv for fast, reproducible dependency management. Install it with pip install uv or see the uv docs.

Clone the repo and install dependencies:
bash
$ git clone https://github.com/Stability-AI/stable-audio-tools.git
$ cd stable-audio-tools

Inference only
$ uv sync

Training
$ uv sync --extra train

Everything (training + Gradio UI)
$ uv sync --extra train --extra ui

Flash Attention

Flash Attention is recommended for performance. Follow the installation instructions in the Flash Attention repository to install it after running uv sync.

Run scripts with uv run:
bash
$ uv run python train.py --dataset-config /path/to/config ...
$ uv run python run_gradio.py --pretrained-name stabilityai/stable-audio-open-1.0

Alternatively, you can use pip directly:
bash
$ pip install "stable-audio-tools[train]"

Interface

A basic Gradio interface is provided to test out trained models. 

For example, to create an interface for the stable-audio-open-1.0 model, once you've accepted the terms for the model on Hugging Face, you can run:
bash
$ python3 ./run_gradio.py --pretrained-name stabilityai/stable-audio-open-1.0

The run_gradio.py script accepts the following command line arguments:

- --pretrained-name
  - Hugging Face repository name for a Stable Audio Tools model
  - Will prioritize model.safetensors over model.ckpt in the repo
  - Optional, used in place of model-config and ckpt-path when using pre-trained model checkpoints on Hugging Face
- --model-config
  - Path to the model config file for a local model
- --ckpt-path
  - Path to unwrapped model checkpoint file for a local model
- --pretransform-ckpt-path 
  - Path to an unwrapped pretransform checkpoint, replaces the pretransform in the model, useful for testing out fine-tuned decoders
  - Optional
- --share
  - If true, a publicly shareable link will be created for the Gradio demo
  - Optional
- --username and --password
  - Used together to set a login for the Gradio demo
  - Optional
- --model-half
  - If true, the model weights to half-precision
  - Optional

Training

Prerequisites
Before starting your training run, you'll need a model config file, as well as a dataset config file. For more information about those, refer to the Configurations section below

The training code also requires a Weights & Biases account to log the training outputs and demos. Create an account and log in with:
bash
$ wandb login

Start training
To start a training run, run the train.py script in the repo root with:
bash
$ python3 ./train.py --dataset-config /path/to/dataset/config --model-config /path/to/model/config --name harmonai_train

The --name parameter will set the project name for your Weights and Biases run.

Training wrappers and model unwrapping
stable-audio-tools uses PyTorch Lightning to facilitate multi-GPU and multi-node training. 

When a model is being trained, it is wrapped in a "training wrapper", which is a pl.LightningModule that contains all of the relevant objects needed only

……（后略，见原文）