ace_step: support the ACE-Step 1.5 XL DiT variants

[CaptainArni]

Adds acestep-v15-xl-turbo and acestep-v15-xl-sft as selectable DiT variants for ace_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 to modeling_acestep_v15_turbo.py apart 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 a copy.deepcopy of the config with the encoder_* values substituted, so those submodules stay at 2048 while the DiT runs at 2560, and decoder.condition_embedder (2560×2048) bridges them.

That copy is now AceStepConfig::encoder, derived once in assets.cpp, and the encoder-side runtimes read it instead of config.diffusion. null_condition_emb is 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_dim outright: 32 heads × 128 = 4096 against a hidden size of 2560, so o_proj is rectangular. build_attention in diffusion.cpp reshaped the context to config.hidden_size and built o_proj square — which was correct for every variant until now, because head_dim had always been derived from hidden_size. It now uses num_attention_heads * head_dim on the input side. The condition encoder's own encoder_layer already did this correctly; this brings the DiT in line with it.

3. The XL timbre encoder prepends its CLS token

AceStepTimbreEncoder.forward concatenates self.special_token ahead 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. TimbreEncoderGraph now concatenates it and extends positions and both masks by one, matching upstream's cache_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 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; 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_source already 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:

• text2music on acestep-v15-xl-turbo renders 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-turbo and acestep-v15-base are unchanged, including base's CFG path, which exercises null_condition_emb.
• A package with the XL directory removed still loads and runs; selecting the missing variant produces the intended message.

All three graph changes are load-bearing rather than cosmetic: a wrong encoder width fails the [2048, 1024] text-projector shape check, a square o_proj fails 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, so native puts 19.9 GB of weights on the card. Passing --session-option ace_step.dit_weight_type=bf16 took 20 s of audio from 87 s to 24 s here (turbo, for reference: 11 s).

acestep-v15-xl-sft shares the graph with xl-turbo — upstream's modeling_acestep_v15_xl_base.py differs only in its sampling loop, which audio.cpp implements itself and already keys on is_turbo — so it is registered too, but I have only run xl-turbo end 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

[0xShug0]

@CaptainArni Thanks! I will test after wrapping up the current minimax-music3 implementation.

[0xShug0]

@CaptainArni Currently, the XL variants are not exposed as installable packages. Could you update model_specs/ace_step.json? Note that model_specs_v1/ mostly serves as examples, so you do not need to update it.

Optional: Would you like to host the GGUFs on HF? We no longer actively maintain safetensors support, and the UI only supports GGUF.