$ show ace-step/ACE-Step
ACE-Step is a production-grade music generation system that implements a flow-matching diffusion model (ACEStepTransformer2D) with a custom DCAE autoencoder and vocoder. It provides multiple deployment options: a Gradio UI, CLI inference, and FastAPI REST endpoints. The system supports text-to-music, lyrics-conditioned generation, audio-to-audio editing, LoRA fine-tuning, and multilingual lyrics (17+ languages). Architecturally, it extends Hugging Face diffusers with custom attention processors (LiteLA linear attention), schedulers (Euler/Heun/PingPong flow matching), and a Conformer-based lyric encoder. It includes CPU offloading, overlapped decode, and memory optimizations for GPU-constrained environments. While primarily an application, the repository also contains training scripts (PyTorch Lightning) and dataset conversion utilities, reflecting its hybrid research-production nature.
The LLM identified 6 files worth highlighting for reuse:
Contains two self-contained attention processors: CustomLiteLAProcessor2_0 (ReLU-based linear attention with O(n) complexity) and CustomerAttnProcessor2_0 (standard scaled dot-product with optional shared KV). These are implemented as diffusers AttentionProcessor subclasses and can be extracted and used in any transformer model with minimal adaptation.
A custom Euler scheduler for flow-matching diffusion, built as a diffusers SchedulerMixin subclass. Implements logistic omega rescaling and configurable timestep spacing. Self-contained and reusable in any flow-matching pipeline.
A Conformer encoder (ConformerEncoder) with relative position encoding, causal convolution, and Macaron-style feed-forward. This is a self-contained sequence encoder that can be reused for any speech or text encoding task. Demonstrates careful implementation of the Conformer architecture with proper masking and multi-head attention.
A standalone language segmentation module that detects and segments multilingual text using py3langid and a custom priority-based filter. Supports 96+ languages with configurable thresholds. Can be extracted and used in any NLP pipeline requiring language-aware text splitting.
A context manager (CpuOffloader) that offloads model weights to CPU during idle periods and reloads them when needed. This is a clean, reusable pattern for running large models on limited GPU memory. The implementation is self-contained and can be adapted for any PyTorch model.
ADaMoSHiFiGANV1 vocoder: a neural vocoder using a ConvNeXt backbone with HiFiGAN-style upsampling. Self-contained as a PyTorch module (nn.Module) with configurable parameters. Can be extracted and used as a standalone vocoder in other audio generation systems.