ComfyUI-FL-AceStep-Training

FL AceStep Training

LoRA training nodes for ComfyUI powered by ACE-Step 1.5, the open-source music generation foundation model. Train custom LoRAs to personalize music generation with your own style, voice, or genre — entirely within ComfyUI's node graph.

Features

• End-to-End Training — Full LoRA training pipeline inside ComfyUI's node graph
• Dataset Management — Scan audio directories, auto-label with LLM, load sidecar metadata
• Tiled VAE Encoding — Handles long audio via 30-second chunks with 2-second overlap
• Real-Time Training UI — Live loss chart, progress bar, and stats via WebSocket widget
• Auto Model Download — LLM models download automatically from HuggingFace on first use
• Native ComfyUI Types — Uses MODEL, VAE, and CLIP from ComfyUI's built-in checkpoint loader

Nodes

| Node | Category | Description | |------|----------|-------------| | FL AceStep LLM Loader | Loaders | Load 5Hz causal LM (0.6B / 1.7B / 4B) for auto-labeling | | FL AceStep Scan Audio Directory | Dataset | Recursively scan folders for audio files with sidecar metadata | | FL AceStep Auto-Label Samples | Dataset | Generate captions, BPM, key, genre, and lyrics via LLM | | FL AceStep Preprocess Dataset | Dataset | VAE-encode audio and CLIP-encode text, save as .pt tensors | | FL AceStep Training Configuration | Training | Configure LoRA rank/alpha/dropout and training hyperparameters | | FL AceStep Train LoRA | Training | Run flow matching training loop with real-time progress widget |

Installation

Manual

cd ComfyUI/custom_nodes
git clone https://github.com/filliptm/ComfyUI-FL-AceStep-Training.git
cd ComfyUI-FL-AceStep-Training
pip install -r requirements.txt

Frontend (optional rebuild)

npm install
npm run build

The pre-built JS is included in js/, so rebuilding is only needed if modifying the training widget UI.

Quick Start

Training Pipeline

1. Load Checkpoint — Use ComfyUI's native Load Checkpoint node with an ACE-Step model to get MODEL, VAE, and CLIP
2. Load LLM (optional) — Add FL AceStep LLM Loader if you want auto-labeling
3. Scan Dataset — Use FL AceStep Scan Audio Directory to find your audio files
4. Label — Connect MODEL, VAE, and LLM to FL AceStep Auto-Label Samples for LLM-generated metadata
5. Preprocess — Run FL AceStep Preprocess Dataset with MODEL, VAE, and CLIP to encode audio/text to tensors
6. Configure — Set LoRA rank, learning rate, epochs in FL AceStep Training Configuration
7. Train — Connect MODEL and config to FL AceStep Train LoRA and execute

Using Trained LoRAs

Use ComfyUI's native LoRA loading nodes to apply your trained LoRA for inference with the built-in ACE-Step nodes.

Node Details

FL AceStep LLM Loader

Loads one of three 5Hz causal language models for auto-labeling audio samples.

| Input | Type | Default | Notes | |-------|------|---------|-------| | model_name | Dropdown | acestep-5Hz-lm-1.7B | Also: 0.6B, 4B | | device | Dropdown | auto | auto / cuda / cpu | | backend | Dropdown | pt | pt / vllm | | checkpoint_path | STRING | (empty) | Optional, leave empty for auto-download |

Output: ACESTEP_LLM

FL AceStep Scan Audio Directory

Recursively scans a directory for audio files and loads accompanying metadata.

| Input | Type | Default | Notes | |-------|------|---------|-------| | directory | STRING | — | Path to audio folder | | all_instrumental | BOOLEAN | True | Mark all samples as instrumental | | custom_tag | STRING | (empty) | LoRA activation tag (e.g., my_style) | | tag_position | Dropdown | prepend | prepend / append / replace |

Outputs: ACESTEP_DATASET, sample count, status

FL AceStep Auto-Label Samples

Uses the loaded LLM to generate metadata for each audio sample.

| Input | Type | Default | Notes | |-------|------|---------|-------| | dataset | ACESTEP_DATASET | — | From Scan Directory | | model | MODEL | — | ACE-Step model (purple) | | vae | VAE | — | ACE-Step VAE (red) — used for audio-to-codes | | llm | ACESTEP_LLM | — | From LLM Loader | | skip_metas | BOOLEAN | False | Skip BPM/key/time signature | | only_unlabeled | BOOLEAN | False | Process only unlabeled samples | | format_lyrics | BOOLEAN | False | Format user-provided lyrics with LLM | | transcribe_lyrics | BOOLEAN | False | Transcribe lyrics from audio |

Outputs: ACESTEP_DATASET, labeled count, status

FL AceStep Preprocess Dataset

VAE-encodes audio and CLIP-encodes text to .pt tensor files for training.

| Input | Type | Default | Notes | |-------|------|---------|-------| | dataset | ACESTEP_DATASET | — | From label or scan node | | model | MODEL | — | ACE-Step model (purple) | | vae | VAE | — | ACE-Step VAE (red) | | clip | CLIP | — | ACE-Step CLIP (yellow) | | output_dir | STRING | ./output/acestep/datasets | — | | max_duration | FLOAT | 240.0 | 10–600 seconds | | genre_ratio | INT | 0 | 0–100% chance to use genre instead of caption |

Outputs: output path, sample count, status

FL AceStep Training Configuration

| Parameter | Default | Range | |-----------|---------|-------| | LoRA Rank | 8 | 4–256 (step 4) | | LoRA Alpha | 16 | 4–512 (step 4) | | LoRA Dropout | 0.1 | 0–0.5 | | Learning Rate | 1e-4 | 1e-6 – 1e-2 | | Max Epochs | 100 | 10–10000 | | Batch Size | 1 | 1–8 | | Gradient Accumulation | 4 | 1–16 | | Save Every N Epochs | 10 | 5–1000 | | Seed | 42 | — | | Optional | | | | Warmup Steps | 100 | 0–1000 | | Weight Decay | 0.01 | 0–0.1 | | Max Grad Norm | 1.0 | 0.1–10.0 | | Target Modules | q_proj,k_proj,v_proj,o_proj | Comma-separated |

Output: ACESTEP_TRAINING_CONFIG

Mixed precision is fixed to bf16. The turbo model uses 8-step discrete timesteps with shift=3.0.

FL AceStep Train LoRA

Runs the training loop with flow matching loss: MSE(predicted_v, x1 - x0).

| Input | Type | Default | Notes | |-------|------|---------|-------| | model | MODEL | — | ACE-Step model (purple) | | config | ACESTEP_TRAINING_CONFIG | — | From config node | | tensor_dir | STRING | ./output/acestep/datasets | Directory of .pt files | | lora_name | STRING | my_lora | Name for the trained LoRA (used as subfolder) | | resume_from | STRING | (empty) | Path to checkpoint to resume from |

Outputs: MODEL (with LoRA), final LoRA path, status

The training widget displays a live loss chart, progress bar, and per-epoch stats via WebSocket (acestep.training.progress).

Supported Audio Formats

.wav, .mp3, .flac, .ogg, .opus, .m4a

Sidecar files for metadata:

• .txt files alongside audio for lyrics
• key_bpm.csv or metadata.csv for BPM, key, and caption data

Requirements

• Python 3.10+
• NVIDIA GPU with 8GB+ VRAM (bf16 training)
• PyTorch 2.0+
• PEFT (Parameter-Efficient Fine-Tuning)
• Transformers, Diffusers, Accelerate

License

MIT