ACE-Step 1.5 Training Pipeline

Overview

What Is This Pipeline?

This pipeline fine-tunes ACE-Step 1.5 — a state-of-the-art Diffusion Transformer (DiT) for text-to-music generation with 815 million parameters — using LoRA (Low-Rank Adaptation). Unlike autoregressive models (MusicGen, AudioLM), ACE-Step generates audio through flow matching: iteratively denoising a latent representation to produce high-fidelity music.

Key Capabilities

• Multi-genre training — Train on 50-200 tracks across any genres
• Text + lyrics conditioning — Separate cross-attention for style tags and lyrics
• Serverless execution — Fully containerized with Docker, runs on Lambda/RunPod/Modal
• Hot-swap adapters — Switch between adapters without reloading the base model
• WTA attribution — Integrated Wasserstein Trajectory Attribution for IP tracking
• S3 storage — Centralized checkpoint and dataset management on AWS S3
• Validated configs — Tested presets with 11/11 validation tests passing

Architecture

Pipeline Flow

The complete training pipeline follows this flow:

Model Architecture

LoRA adapters target self-attention and cross-attention layers: linear_q, linear_k, linear_v, to_q, to_k, to_v, to_out.0 across all 48 DiT layers.

Quick Start

# Clone the training-pipelines repo
git clone https://github.com/aenfr/training-pipelines.git
cd training-pipelines/ace-step-1-5

# Build the Docker image
cd docker
docker build -f Dockerfile.trainer -t lumina-trainer:v1 .

# Build HuggingFace dataset from your audio files
python scripts/build_multi_style_dataset.py \
    --data-dir ~/my_audio \
    --output-dir ~/my_hf_dataset \
    --validate-audio

MODE=smoke ./docker/train.sh

# Runs with validated Preset C defaults
./docker/train.sh

Training Presets

Preset C — Multi-Style (Production)

LOO Baseline — Research Validation

Used for Leave-One-Out causal validation experiments on GTZAN (90 tracks, 10 genres).

Docker Setup

Container Volumes

Entrypoint Auto-Patching

The container entrypoint automatically handles these known issues:

1. TorchCodec incompatibility — Patches torchaudio.load() → librosa.load()
2. Audio save failure — Patches torchaudio.save() → soundfile.write()
3. Step-0 plot crash — Skips inference at step 0 to prevent gradient corruption

Custom Presets

Override any parameter via environment variables:

# Example: High-rank, low-LR for single artist
EPOCHS=200 \
LEARNING_RATE=2e-5 \
GRAD_ACCUM=8 \
EXP_NAME="my_custom_preset" \
./docker/train.sh

For LoRA architecture changes, create a JSON config:

{
    "r": 128,
    "lora_alpha": 256,
    "lora_dropout": 0.1,
    "target_modules": [
        "linear_q", "linear_k", "linear_v",
        "to_q", "to_k", "to_v", "to_out.0"
    ],
    "use_rslora": false
}

LORA_CONFIG=/path/to/my_config.json ./docker/train.sh

Dataset Preparation

Required Files Per Track

Build the Dataset

python scripts/build_multi_style_dataset.py \
    --data-dir ~/my_audio \
    --output-dir ~/my_hf_dataset \
    --validate-audio

S3 Storage

Bucket Layout

s3://lumina-data-foldartists/
├── models/ace-step-1.5/         # Base model weights
├── lora/                        # Fine-tuned adapters
│   ├── multi-style-gen-c/       # ✅ Production
│   └── loo-subsets/             # Validation models
├── datasets/                    # HF datasets
├── trajectories/                # WTA trajectory data
└── results/                     # Experiment outputs

Common Operations

# Download production adapter
aws s3 sync s3://lumina-data-foldartists/lora/multi-style-gen-c/ \
    ~/my_adapter/

# Upload new checkpoint after training
aws s3 sync /output/my_experiment/ \
    s3://lumina-data-foldartists/lora/my-new-adapter/

See the full S3 Setup Guide for bucket creation, IAM roles, and cost estimates.

Monitoring Training

Key Metrics

Duration Estimates

Troubleshooting

Related Documentation