Add ace step lora support#14193
Merged
sayakpaul merged 16 commits intoJul 18, 2026
Merged
Conversation
Add standard diffusers LoRA API (load_lora_weights, set_adapters, fuse_lora, save_lora_weights) for AceStepPipeline, targeting the transformer (AceStepTransformer1DModel) only. Includes automatic conversion of ACE-Step-1.5 original PEFT format LoRA weights (q_proj/k_proj/v_proj/o_proj naming) to diffusers format (to_q/to_k/to_v/to_out.0 naming) during loading. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…ply ruff formatting - Replace hardcoded local path with "Qwen/Qwen3-Embedding-0.6B" so tests run on CI - Apply ruff formatting fix (slice spacing) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
github-actions
Bot
added
fixes-issue
size/L
Member
|
Could you also share some outputs? |
Contributor
Author
Friendly ping @yiyixuxu @dg845 @asomoza @sayakpaul @DN6 .Please check the outputs. Thank you for your help! baseline_no_lora.wav Environment
Generation Parametersprompt = "chinese traditional music, erhu solo, peaceful and elegant"
lyrics = "[verse]\n春风又绿江南岸\n明月何时照我还"
audio_duration = 10.0 # seconds
num_inference_steps = 8
seed = 42How to ReproduceStep 1: Install the modified diffuserscd /path/to/diffusers
git checkout add-ace-step-lora-support
pip install -e .
pip install peft soundfileStep 2: Run the generation scriptimport torch
import soundfile as sf
from diffusers import AceStepPipeline
model_path = "Runware/acestep-v15-turbo-diffusers"
lora_path = "ACE-Step/ACE-Step-v1.5-chinese-new-year-LoRA"
pipe = AceStepPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
pipe = pipe.to("cuda")
kwargs = dict(
prompt="chinese traditional music, erhu solo, peaceful and elegant",
lyrics="[verse]\n春风又绿江南岸\n明月何时照我还",
audio_duration=10.0,
num_inference_steps=8,
max_text_length=256,
output_type="np",
)
# Baseline (no LoRA)
generator = torch.Generator(device="cuda").manual_seed(42)
audio_baseline = pipe(**kwargs, generator=generator).audios
sf.write("baseline_no_lora.wav", audio_baseline[0].T, 48000)
# Load LoRA (auto-converts from original ACE-Step-1.5 PEFT format)
pipe.load_lora_weights(lora_path, adapter_name="chinese_new_year")
# LoRA scale=1.0
generator = torch.Generator(device="cuda").manual_seed(42)
audio_lora_1_0 = pipe(**kwargs, generator=generator).audios
sf.write("lora_scale_1.0.wav", audio_lora_1_0[0].T, 48000)
# LoRA scale=0.5
pipe.set_adapters(["chinese_new_year"], [0.5])
generator = torch.Generator(device="cuda").manual_seed(42)
audio_lora_0_5 = pipe(**kwargs, generator=generator).audios
sf.write("lora_scale_0.5.wav", audio_lora_0_5[0].T, 48000)
# LoRA scale=0.3
pipe.set_adapters(["chinese_new_year"], [0.3])
generator = torch.Generator(device="cuda").manual_seed(42)
audio_lora_0_3 = pipe(**kwargs, generator=generator).audios
sf.write("lora_scale_0.3.wav", audio_lora_0_3[0].T, 48000)Results
Audio files are saved under Conclusions
|
sayakpaul
reviewed
Jul 16, 2026
Comment on lines
+196
to
+198
| @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") | ||
| def test_lora_scale_kwargs_match_fusion(self): | ||
| pass |
Contributor
Author
There was a problem hiding this comment.
Why is this the case?
These tests were skipped because AceStepPipeline.call does not have parameters cross_attention_kwargs / joint_attention_kwargs / attention_kwargs. The base class tests look for these parameter names in the signature to pass per-step LoRA scale, and if they are not found, it crashes directly.
All core LoRA functionalities (load, set_adapters, fuse, save) are working properly, but the optional feature of dynamically passing scale via pipe(..., joint_attention_kwargs={"scale": 0.5}) during inference is missing.
Member
There was a problem hiding this comment.
We should add attention_kwargs to the pipeline call then like other LoRA-supported pipelines do.
Contributor
Author
There was a problem hiding this comment.
We should add
attention_kwargsto the pipeline call then like other LoRA-supported pipelines do.
Fixed in 2c49f32
sayakpaul
reviewed
Jul 16, 2026
Comment on lines
+212
to
+214
| @unittest.skip("AceStep uses 'layers' not 'transformer_blocks'/'blocks'; base test hardcodes block names.") | ||
| def test_lora_fuse_nan(self): | ||
| pass |
Member
There was a problem hiding this comment.
We can modify the test to include the block name of AceStep.
sayakpaul
reviewed
Jul 16, 2026
sayakpaul
left a comment
sayakpaul
left a comment
Member
There was a problem hiding this comment.
Left some comments around attention_kawargs we should support it.
The base test hardcodes block attribute names (transformer_blocks, blocks, etc.) but AceStep uses 'layers'. Override with AceStep-specific block access path instead of skipping. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
… scale
- Add @apply_lora_scale("joint_attention_kwargs") decorator and
joint_attention_kwargs param to AceStepTransformer1DModel.forward
- Add joint_attention_kwargs param to AceStepPipeline.__call__ and wire
it through to all three transformer calls
- Un-skip 4 LoRA scale tests that now pass (39 passed, 13 skipped)
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The base test hardcodes block attribute names (transformer_blocks, blocks, etc.) but AceStep uses 'layers'. Override with AceStep-specific block access path instead of skipping. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
… scale
- Add @apply_lora_scale("joint_attention_kwargs") decorator and
joint_attention_kwargs param to AceStepTransformer1DModel.forward
- Add joint_attention_kwargs param to AceStepPipeline.__call__ and wire
it through to all three transformer calls
- Un-skip 4 LoRA scale tests that now pass (39 passed, 13 skipped)
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
sayakpaul
reviewed
Jul 16, 2026
| timestep_r: torch.Tensor, | ||
| encoder_hidden_states: torch.Tensor, | ||
| context_latents: torch.Tensor, | ||
| joint_attention_kwargs: Optional[dict] = None, |
Member
There was a problem hiding this comment.
Let's use attention_kwargs, instead and update the docstrings.
Contributor
Author
There was a problem hiding this comment.
Let's use
attention_kwargs, instead and update the docstrings.
Fixed in 2cd2039
|
|
||
| components, _, denoiser_lora_config = self.get_dummy_components() | ||
| pipe = self.pipeline_class(**components) | ||
| pipe = pipe.to("cpu") |
… torch_device in test - Rename joint_attention_kwargs to attention_kwargs in transformer forward and pipeline __call__ for consistency with the diffusers naming convention for new pipelines - Add docstring for the attention_kwargs parameter - Use torch_device instead of hardcoded "cpu" in test_lora_fuse_nan Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
sayakpaul
reviewed
Jul 16, 2026
| cfg_interval_start: float = 0.0, | ||
| cfg_interval_end: float = 1.0, | ||
| timesteps: Optional[List[float]] = None, | ||
| attention_kwargs: Optional[dict] = None, |
Member
|
@bot /style |
Contributor
|
Style bot fixed some files and pushed the changes. |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
sayakpaul
approved these changes
Jul 16, 2026
Contributor
Author
|
Friendly ping @yiyixuxu @dg845 @asomoza @sayakpaul @DN6. Please take a look at this PR. Thank you for your help! |
Member
|
We need to fix the CI failures. |
github-actions
Bot
added
the
documentation
Contributor
Author
Fixed in 86f9afd |
Member
|
Failing tests are unrelated. |
|
Hi @chenyangzhu1 , thanks for working on the LoRa support for ACE-Step Turbo! |
Contributor
Author
You can follow this tutorial. Summarized by Claude agent. ACE-Step LoRA Training & Inference with DiffusersThis guide walks through the complete workflow: training a LoRA adapter on the ACE-Step base model using the original ACE-Step-1.5 repo, then loading it for inference in diffusers. Prerequisites
# Clone the training repo
git clone https://github.com/AceStepAI/ACE-Step-1.5
cd ACE-Step-1.5
pip install -e .
# Install diffusers with LoRA support
pip install diffusers peftOverviewStep 1: Prepare Your DatasetAudio Requirements
Create Dataset JSONCreate a [
{
"audio_path": "./songs/track1.wav",
"caption": "An energetic Chinese New Year celebration with traditional erhu and drums",
"lyrics": "[Verse]\n春风又绿江南岸\n明月何时照我还",
"genre": "Chinese Traditional",
"bpm": 120,
"keyscale": "D major",
"timesignature": "4",
"language": "zh",
"is_instrumental": false
},
{
"audio_path": "./songs/track2.wav",
"caption": "Festive instrumental with guzheng and bamboo flute",
"genre": "Chinese Traditional",
"is_instrumental": true
}
]Key fields:
Quick start — if you just have audio files and no metadata, point directly at a folder: python train.py fixed \
--checkpoint-dir ./checkpoints/acestep-v15-turbo \
--model-variant turbo \
--preprocess \
--audio-dir ./my_audio \
--tensor-output ./my_tensorsFilenames are auto-converted to captions (e.g., Step 2: Download the Base Model# Download the original ACE-Step checkpoint (used for training)
# The model will be auto-downloaded, or manually:
huggingface-cli download ACE-Step/Ace-Step1.5 --local-dir ./checkpointsThe checkpoint directory should contain: Step 3: Preprocess DataPreprocessing encodes audio to latents and text to embeddings. This runs in two passes to reduce VRAM usage: python train.py fixed \
--checkpoint-dir ./checkpoints \
--model-variant turbo \
--preprocess \
--audio-dir ./my_audio \
--dataset-json ./dataset.json \
--tensor-output ./preprocessed_tensorsPass 1 (~3 GB VRAM): VAE encodes audio → latents, text encoder encodes prompts/lyrics → embeddings Output: one Step 4: Train LoRAUsing the CLI (Recommended)python train.py fixed \
--checkpoint-dir ./checkpoints \
--model-variant turbo \
--dataset-dir ./preprocessed_tensors \
--output-dir ./lora_output \
--rank 64 \
--alpha 128 \
--learning-rate 1e-4 \
--epochs 100 \
--batch-size 1 \
--gradient-accumulation 4 \
--warmup-steps 100 \
--save-every 10 \
--cfg-ratio 0.15 \
--gradient-checkpointingUsing the Interactive Wizardpython train.py
# Follow the prompts to select model, dataset, and hyperparametersKey Hyperparameters
VRAM PresetsFor low-VRAM GPUs, use built-in presets: # 8 GB VRAM
python train.py fixed --preset vram_8gb ...
# 12 GB VRAM
python train.py fixed --preset vram_12gb ...Training Details
Monitor Trainingtensorboard --logdir ./lora_output/runsStep 5: Verify the OutputAfter training completes, the output directory contains: The key files are {
"peft_type": "LORA",
"r": 64,
"lora_alpha": 128,
"lora_dropout": 0.1,
"target_modules": ["k_proj", "v_proj", "o_proj", "q_proj"],
"bias": "none",
"task_type": "FEATURE_EXTRACTION"
}Step 6: Inference with DiffusersLoad and Generateimport torch
import soundfile as sf
from diffusers import AceStepPipeline
# Load the base model (diffusers format)
pipe = AceStepPipeline.from_pretrained(
"Runware/acestep-v15-turbo-diffusers",
torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")
# Load your trained LoRA weights
# Diffusers auto-converts the original ACE-Step key names
# (q_proj → to_q, k_proj → to_k, etc.)
pipe.load_lora_weights("./lora_output/final", adapter_name="my_style")
# Generate music
audio = pipe(
prompt="Chinese traditional music, erhu solo, peaceful and elegant",
lyrics="[verse]\n春风又绿江南岸\n明月何时照我还",
audio_duration=30.0,
num_inference_steps=8,
).audios
# Save as WAV
sf.write("output.wav", audio[0].T.cpu().numpy(), 48000)Control LoRA Strength# Scale down LoRA effect (0.0 = no LoRA, 1.0 = full LoRA)
pipe.set_adapters(["my_style"], [0.5])
audio = pipe(prompt="...", lyrics="...", audio_duration=30.0).audios
# Or pass scale per inference call via attention_kwargs
audio = pipe(
prompt="...",
lyrics="...",
audio_duration=30.0,
attention_kwargs={"scale": 0.7},
).audiosFuse LoRA for Faster Inference# Merge LoRA weights into base model (no per-step overhead)
pipe.fuse_lora()
audio = pipe(prompt="...", lyrics="...", audio_duration=30.0).audios
# Undo fusion
pipe.unfuse_lora()Multiple LoRA Adapters# Load multiple adapters
pipe.load_lora_weights("./lora_chinese_new_year/", adapter_name="chinese_new_year")
pipe.load_lora_weights("./lora_jazz/", adapter_name="jazz")
# Use one at a time
pipe.set_adapters(["chinese_new_year"], [1.0])
audio_cny = pipe(prompt="festive music", lyrics="...", audio_duration=30.0).audios
# Switch to another
pipe.set_adapters(["jazz"], [1.0])
audio_jazz = pipe(prompt="jazz ballad", lyrics="...", audio_duration=30.0).audios
# Blend two adapters
pipe.set_adapters(["chinese_new_year", "jazz"], [0.6, 0.4])
audio_blend = pipe(prompt="jazz with chinese elements", lyrics="...", audio_duration=30.0).audiosSave in Diffusers Formatfrom peft.utils import get_peft_model_state_dict
lora_state_dict = get_peft_model_state_dict(pipe.transformer, adapter_name="my_style")
pipe.save_lora_weights("./my_style_diffusers/", transformer_lora_layers=lora_state_dict)Unload LoRApipe.unload_lora_weights()
# Pipeline is back to the base modelKey Name MappingWhen loading LoRA weights trained with the original ACE-Step-1.5 repo into diffusers, the key names are automatically converted:
This conversion happens transparently inside |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds standard diffusers LoRA support to
AceStepPipeline, enabling users to load, apply, fuse, and save LoRA adapters for theAceStepTransformer1DModelusing the familiar diffusers API:Key features
load_lora_weights,set_adapters,fuse_lora,unfuse_lora,save_lora_weights,delete_adapters,enable_lora,disable_loraq_proj/k_proj/v_proj/o_projnaming, while diffusers usesto_q/to_k/to_v/to_out.0. The loader detects the original format automatically and converts on-the-fly — no manual conversion step needed.set_adapters.Changes
src/diffusers/models/transformers/ace_step_transformer.pyPeftAdapterMixintoAceStepTransformer1DModelinheritancesrc/diffusers/loaders/lora_conversion_utils.py_convert_non_diffusers_ace_step_lora_to_diffusers()conversion functionsrc/diffusers/loaders/lora_pipeline.pyAceStepLoraLoaderMixin(LoraBaseMixin)classsrc/diffusers/loaders/__init__.pyAceStepLoraLoaderMixinsrc/diffusers/pipelines/ace_step/pipeline_ace_step.pyAceStepLoraLoaderMixintoAceStepPipelineinheritancetests/lora/test_lora_layers_ace_step.pyConversion details
The conversion function handles the key name mapping between the original ACE-Step-1.5 PEFT format and diffusers format:
Mapping rules:
base_model.model.prefix.q_proj.→.to_q.,.k_proj.→.to_k.,.v_proj.→.to_v.,.o_proj.→.to_out.0.transformer.prefixDesign decisions
AceStepTransformer1DModelreceives LoRA adapters, matching the original ACE-Step-1.5 training setup where only the DiT decoder is fine-tuned.Mochi1LoraLoaderMixin/CogVideoXLoraLoaderMixinpattern (simplest transformer-only LoRA mixin).AceStepPipelinedoes not exposeattention_kwargsfor per-step LoRA scale control, the transformer useslayersinstead oftransformer_blocks/blocks, and attention layers have no bias.Self-Review Report
Reviewed against
.ai/review-rules.md,.ai/AGENTS.md,.ai/pipelines.md,.ai/models.md.Blocking Issues
All resolved.
1. Test file hardcoded local absolute path for tokenizerFIXED in commit
de90a1e7a. Changed from/nas/zcy/github_issue/models/Qwen3-Embedding-0.6B-tokenizerto"Qwen/Qwen3-Embedding-0.6B".2.make styleandmake fix-copiesnot runFIXED in commit
de90a1e7a.ruffapplied one spacing fix inlora_conversion_utils.py.make fix-copiesproduced no changes.Non-blocking Issues (for reviewer)
1.
save_lora_weightserror message mentionstext_encoder_lora_layersbut AceStep doesn't support itsrc/diffusers/loaders/lora_pipeline.py—AceStepLoraLoaderMixin.save_lora_weights:This is copied verbatim from
CogVideoXLoraLoaderMixin(annotated with# Copied from). AceStep only supports transformer LoRA, so thetext_encoder_lora_layersreference is misleading. However, modifying this line would break the# Copied fromsync mechanism. Leaving for reviewer to decide whether to break the copy link or keep as-is.2.
lora_state_dictmethod has no# Copied fromannotationThe method body is nearly identical to
SD3LoraLoaderMixin.lora_state_dict/Mochi1LoraLoaderMixin.lora_state_dict, but has 3 extra lines for ACE-Step format detection (is_original_ace_stepblock). This custom logic prevents using a standard# Copied fromannotation. The deviation is intentional and minimal.3. Documentation not updated
docs/source/en/api/pipelines/ace_step.mddoes not mention LoRA support. A short usage example could be added. Not blocking since the API is standard and discoverable via the existing LoRA documentation.Dead Code Analysis
tests/lora/test_lora_layers_ace_step.py—noiseandinput_idsinget_dummy_inputs(noise, input_ids, pipeline_inputs)but AceStepPipeline's__call__does not accept them as arguments. Kept for base class compatibility — some inherited tests may reference these values.Skipped Tests Rationale
test_simple_inference_with_text_denoiser_block_scaletest_simple_inference_with_text_denoiser_block_scale_for_all_dict_optionstest_modify_padding_modetest_simple_inference_with_text_denoiser_multi_adapter_block_loratest_correct_lora_configs_with_different_rankstest_simple_inference_with_text_denoiser_lora_and_scaleAceStepPipeline.__call__has noattention_kwargs/cross_attention_kwargs/joint_attention_kwargsparametertest_lora_scale_kwargs_match_fusiontest_set_adapters_match_attention_kwargstest_simple_inference_with_text_lora_and_scaletest_lora_B_biasbias=False;lora_biasis not applicabletest_lora_fuse_nantransformer_blocks,blocks, etc.); AceStep useslayersSummary
Verdict: READY
ACE-Step-v1.5-chinese-new-year-LoRA, r=64, alpha=128)make styleandmake fix-copiescleanFixes #14060
Before submitting
self-reviewskill on the diff?documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@yiyixuxu @dg845 @asomoza @sayakpaul @DN6