ace_step: support the ACE-Step 1.5 XL DiT variants - #235
Open
CaptainArni wants to merge 1 commit into
Open
Conversation
Adds acestep-v15-xl-turbo and acestep-v15-xl-sft as selectable DiT variants. They are the larger transformer — 32 layers of 2560 against turbo's 24 of 2048 — and they differ from the existing variants in three ways that the graph had to learn, all of them read off upstream's own modeling_acestep_v15_xl_*.py: 1. The encoder stack keeps turbo's width. AceStepConditionGenerationModel hands the condition encoder, audio tokenizer and detokenizer a copy of the config with the encoder_* values substituted, so those submodules stay at 2048 while the DiT runs at 2560, and decoder.condition_embedder bridges them. That copy is now AceStepConfig::encoder, and the encoder-side runtimes read it instead of the diffusion config. null_condition_emb is encoder-width for the same reason. 2. The attention width is no longer the model width. XL states head_dim outright: 32 heads x 128 = 4096 against a hidden size of 2560, so o_proj is rectangular. build_attention reshaped the context to hidden_size and built o_proj square, which happened to be right whenever head_dim was derived from hidden_size — every variant until this one. The condition encoder's own attention already did this correctly; this brings the DiT in line. 3. The XL timbre encoder prepends its CLS token. Earlier variants declare the same parameter but leave it out of the sequence and read the first audio frame back instead, so the tensor's presence says nothing about whether it is used and the config has to. TimbreEncoderGraph now concatenates it and extends positions and both masks by one. Packaging: the XL snapshots are ~19 GB each, so listing them in the spec's required `tensors` map would have forced the download on every install. Sources gain an `optional_tensors` map alongside the existing `optional_files`, and the XL entries live there; a package without them loads and behaves exactly as before, and selecting a variant that is not installed reports which directory is missing rather than a bare resource id. Their weights are four safetensors shards, which the spec points at through model.safetensors.index.json — already supported by open_tensor_source. Verified on an RTX 5090 (CUDA, safetensors): text2music renders real, prompt-responsive audio on xl-turbo, and turbo and base are unchanged, including base's CFG path. The XL snapshots are stored in float32, so passing ace_step.dit_weight_type=bf16 is worth it: 20 s of audio took 87 s at native against 24 s at bf16. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Owner
|
@CaptainArni Thanks! I will test after wrapping up the current minimax-music3 implementation. |
Owner
|
@CaptainArni Currently, the XL variants are not exposed as installable packages. Could you update Optional: Would you like to host the GGUFs on HF? We no longer actively maintain safetensors support, and the UI only supports GGUF. |
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
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.
Adds
acestep-v15-xl-turboandacestep-v15-xl-sftas selectable DiT variants forace_step.XL is the larger transformer — 32 layers of 2560 against turbo's 24 of 2048 — and it differs from the existing variants in three ways that the graph had to learn. All three are read off upstream's own
modeling_acestep_v15_xl_turbo.py, which is otherwise identical tomodeling_acestep_v15_turbo.pyapart from whitespace and a PyTorch autocast workaround that does not apply here.1. The encoder stack keeps turbo's width
AceStepConditionGenerationModel.__init__hands the condition encoder, audio tokenizer and detokenizer acopy.deepcopyof the config with theencoder_*values substituted, so those submodules stay at 2048 while the DiT runs at 2560, anddecoder.condition_embedder(2560×2048) bridges them.That copy is now
AceStepConfig::encoder, derived once inassets.cpp, and the encoder-side runtimes read it instead ofconfig.diffusion.null_condition_embis encoder-width for the same reason — it stands in for the encoder output under CFG, so the condition embedder projects it like any other conditioning. On a package that does not split the dimensions the two configs are identical, so nothing changes for turbo or base.2. The attention width is no longer the model width
XL states
head_dimoutright: 32 heads × 128 = 4096 against a hidden size of 2560, soo_projis rectangular.build_attentionindiffusion.cppreshaped the context toconfig.hidden_sizeand builto_projsquare — which was correct for every variant until now, becausehead_dimhad always been derived fromhidden_size. It now usesnum_attention_heads * head_dimon the input side. The condition encoder's ownencoder_layeralready did this correctly; this brings the DiT in line with it.3. The XL timbre encoder prepends its CLS token
AceStepTimbreEncoder.forwardconcatenatesself.special_tokenahead of the reference frames and reads position 0 back as the timbre embedding. The pre-XL class declares the same parameter but has that line commented out, so it reads the first audio frame instead — and the tensor ships in both checkpoints, which is why its presence says nothing and the config has to.TimbreEncoderGraphnow concatenates it and extends positions and both masks by one, matching upstream'scache_position, which is built after the prepend.The gate for 1 and 3 is the presence of
encoder_hidden_size, which is what the XL modeling class reads without a fallback.Packaging
The XL snapshots are ~19 GB each, so listing them in the spec's required
tensorsmap would have forced the download on every install. Sources gain anoptional_tensorsmap alongside the existingoptional_files, and the XL entries live there. A package without them loads and behaves exactly as before; selecting a variant that is not installed reports which directory is missing rather than a bare resource id.Their weights are four safetensors shards, which the spec points at through
model.safetensors.index.json—open_tensor_sourcealready follows that, so sharding needed no work.lyric_alignment_layers_config, the other key new to the XL config, is not referenced anywhere in the modeling file and is ignored.Testing
RTX 5090, CUDA, safetensors package:
text2musiconacestep-v15-xl-turborenders real, prompt-responsive audio — a lo-fi prompt gives a 166 Hz spectral centroid with 94% of energy below 500 Hz, a thrash-metal prompt on the same seed gives 587 Hz and a much stronger onset autocorrelation. No NaNs.acestep-v15-turboandacestep-v15-baseare unchanged, including base's CFG path, which exercisesnull_condition_emb.All three graph changes are load-bearing rather than cosmetic: a wrong encoder width fails the
[2048, 1024]text-projector shape check, a squareo_projfails against the[2560, 4096]weight, and a missing concat desynchronises the mask from the sequence — so a clean run is itself evidence that each path is taken.One practical note, documented in
docs/models/ace_step.md: the XL snapshots are stored in float32, not bf16, sonativeputs 19.9 GB of weights on the card. Passing--session-option ace_step.dit_weight_type=bf16took 20 s of audio from 87 s to 24 s here (turbo, for reference: 11 s).acestep-v15-xl-sftshares the graph withxl-turbo— upstream'smodeling_acestep_v15_xl_base.pydiffers only in its sampling loop, which audio.cpp implements itself and already keys onis_turbo— so it is registered too, but I have only runxl-turboend to end.I found this while adding music generation to a local Studio UI for audio.cpp, which is at https://github.com/CaptainArni/audiocpp-ui.
🤖 Generated with Claude Code