From 1147bdd7a746745bfc3568f4eaef537b64488c3d Mon Sep 17 00:00:00 2001 From: CaptainArni Date: Fri, 14 Aug 2026 15:07:07 +0200 Subject: [PATCH] ace_step: support the ACE-Step 1.5 XL DiT variants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- docs/models/ace_step.md | 27 +++++- include/engine/models/ace_step/assets.h | 17 ++++ .../models/ace_step/dit_weights_runtime.h | 2 + model_specs/ace_step.json | 32 ++++++++ model_specs_v1/ace_step.json | 38 ++++++++- src/framework/model_spec/package.cpp | 29 ++++++- src/framework/model_spec/schema.cpp | 2 +- src/models/ace_step/assets.cpp | 82 +++++++++++++++++-- src/models/ace_step/condition_encoder.cpp | 80 ++++++++++++------ src/models/ace_step/cover_tokenizer.cpp | 10 +-- src/models/ace_step/detokenizer.cpp | 8 +- src/models/ace_step/diffusion.cpp | 31 ++++--- src/models/ace_step/dit_weights_runtime.cpp | 22 ++++- src/models/ace_step/loader.cpp | 11 ++- 14 files changed, 326 insertions(+), 65 deletions(-) diff --git a/docs/models/ace_step.md b/docs/models/ace_step.md index 2ff01e60..ae64b083 100644 --- a/docs/models/ace_step.md +++ b/docs/models/ace_step.md @@ -16,6 +16,7 @@ audiocpp_cli --task gen --family ace_step --model models/Ace-Step1.5 --backend c | Model directory | `models/Ace-Step1.5` | | Task | `gen` | | Default DiT | `acestep-v15-turbo` | +| Optional DiT | `acestep-v15-xl-turbo`, `acestep-v15-xl-sft` | | Default LM | `acestep-5Hz-lm-1.7B` | | Prompt input | `--text` | | Lyrics input | `--lyrics` | @@ -183,7 +184,7 @@ audiocpp_cli --task gen --family ace_step --model models/Ace-Step1.5 --backend c | Option | Values | Default | Meaning | |---|---|---:|---| -| `--load-option ace_step.dit_model_path=` | `acestep-v15-turbo`, `acestep-v15-base` | `acestep-v15-turbo` | Select DiT variant inside the model root. | +| `--load-option ace_step.dit_model_path=` | `acestep-v15-turbo`, `acestep-v15-base`, `acestep-v15-xl-turbo`, `acestep-v15-xl-sft` | `acestep-v15-turbo` | Select DiT variant inside the model root. | | `--session-option ace_step.dit_weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0` | `native` | DiT weight type. | | `--session-option ace_step.planner_weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0` | `native` | Planner LM weight type. | | `--session-option ace_step.mem_saver=true\|false` | bool | `false` | Release staged graph/cache state after request phases to reduce resident VRAM. Later requests may rebuild released graphs. | @@ -191,3 +192,27 @@ audiocpp_cli --task gen --family ace_step --model models/Ace-Step1.5 --backend c ACE-Step GGUF packages are variant-specific. Use the Turbo GGUF for the default `acestep-v15-turbo` path, and pass `--load-option ace_step.dit_model_path=acestep-v15-base` when loading a Base GGUF package. + +### XL variants + +`acestep-v15-xl-turbo` and `acestep-v15-xl-sft` are the larger DiT: 32 layers of +2560 against turbo's 24 of 2048, with 32 attention heads of 128 (so the attention +width is 4096, wider than the model). The condition encoder, audio tokenizer and +detokenizer stay at 2048 — the `encoder_hidden_size` group in the XL config — and +the DiT's condition embedder bridges the two. The XL timbre encoder also prepends +a CLS token to the reference frames and reads that position back, where earlier +variants read the first audio frame. + +Both are **optional package resources**: they are only loadable when their +directory is present, and a package without them loads and runs exactly as +before. Selecting one that is not installed reports which directory is missing. +The upstream snapshots ship four safetensors shards plus a +`model.safetensors.index.json`, which the package spec points at directly. + +```bash +audiocpp_cli --task gen --family ace_step --model models/Ace-Step1.5 --backend cuda --task-route text2music --text "warm lo-fi hip hop with a soft rhodes piano" --duration-seconds 60 --load-option ace_step.dit_model_path=acestep-v15-xl-turbo --session-option ace_step.dit_weight_type=bf16 --out song.wav +``` + +`dit_weight_type=bf16` is worth passing. The XL snapshots are stored in float32, +so `native` puts 19.9 GB of weights on the card: measured on an RTX 5090, 20 s of +audio took 87 s at `native` against 24 s at `bf16` (turbo, for reference: 11 s). diff --git a/include/engine/models/ace_step/assets.h b/include/engine/models/ace_step/assets.h index c4d2adb2..b4df6072 100644 --- a/include/engine/models/ace_step/assets.h +++ b/include/engine/models/ace_step/assets.h @@ -75,6 +75,18 @@ struct AceStepDiffusionConfig { int64_t sliding_window = 0; bool use_sliding_window = false; bool is_turbo = true; + // XL packages size the condition encoder, audio tokenizer and detokenizer + // independently of the DiT (2048 against 2560), which upstream expresses by + // handing those submodules a copy of the config with the encoder_* values + // substituted. The same copy lives in AceStepConfig::encoder, and this flag + // marks the two configs as genuinely different so the code that has to + // bridge them — the condition embedder, the cross-attention KV — can say so. + bool has_separate_encoder = false; + // XL's timbre encoder prepends a CLS token to the reference frames and reads + // position 0 back as the timbre embedding; the pre-XL class carries the same + // parameter but leaves that line commented out, so the tensor's presence says + // nothing and the config has to. + bool timbre_special_token = false; float rms_norm_eps = 1.0e-6F; float rope_theta = 1000000.0F; std::vector fsq_input_levels; @@ -94,7 +106,12 @@ struct AceStepVAEConfig { struct AceStepConfig { AceStepPlannerConfig planner; AceStepTextEncoderConfig text_encoder; + // The DiT itself. AceStepDiffusionConfig diffusion; + // Everything that feeds it: the condition encoder, the audio tokenizer and + // the detokenizer. Identical to `diffusion` except for the four attention and + // MLP dimensions, and identical outright on packages that do not split them. + AceStepDiffusionConfig encoder; AceStepVAEConfig vae; }; diff --git a/include/engine/models/ace_step/dit_weights_runtime.h b/include/engine/models/ace_step/dit_weights_runtime.h index 57df79e4..c2d0952f 100644 --- a/include/engine/models/ace_step/dit_weights_runtime.h +++ b/include/engine/models/ace_step/dit_weights_runtime.h @@ -59,6 +59,8 @@ struct AceStepConditionEncoderWeights { core::TensorValue timbre_embed_bias; std::vector timbre_layers; core::TensorValue timbre_norm; + // Empty unless the variant prepends a CLS token to the timbre sequence. + std::vector timbre_special_token_host; }; struct AceStepTimeEmbeddingWeights { diff --git a/model_specs/ace_step.json b/model_specs/ace_step.json index b0762800..cc0d80f2 100644 --- a/model_specs/ace_step.json +++ b/model_specs/ace_step.json @@ -119,6 +119,28 @@ "text_encoder_chat_template": "model:Qwen3-Embedding-0.6B/chat_template.jinja", "vae_config": "model:vae/config.json" }, + "optional_files": { + "dit_xl_turbo_config": "model:acestep-v15-xl-turbo/config.json", + "dit_xl_sft_config": "model:acestep-v15-xl-sft/config.json" + }, + "optional_tensors": { + "dit_xl_turbo_weights": { + "source": "weights:", + "prefix": "dit_xl_turbo_weights" + }, + "dit_xl_turbo_silence_latent": { + "source": "weights:", + "prefix": "dit_xl_turbo_silence_latent" + }, + "dit_xl_sft_weights": { + "source": "weights:", + "prefix": "dit_xl_sft_weights" + }, + "dit_xl_sft_silence_latent": { + "source": "weights:", + "prefix": "dit_xl_sft_silence_latent" + } + }, "tensors": { "dit_turbo_weights": { "source": "weights:", @@ -172,6 +194,16 @@ "text_encoder_chat_template": "model:Qwen3-Embedding-0.6B/chat_template.jinja", "vae_config": "model:vae/config.json" }, + "optional_files": { + "dit_xl_turbo_config": "model:acestep-v15-xl-turbo/config.json", + "dit_xl_sft_config": "model:acestep-v15-xl-sft/config.json" + }, + "optional_tensors": { + "dit_xl_turbo_weights": "model:acestep-v15-xl-turbo/model.safetensors.index.json", + "dit_xl_turbo_silence_latent": "model:acestep-v15-xl-turbo/silence_latent.safetensors", + "dit_xl_sft_weights": "model:acestep-v15-xl-sft/model.safetensors.index.json", + "dit_xl_sft_silence_latent": "model:acestep-v15-xl-sft/silence_latent.safetensors" + }, "tensors": { "dit_turbo_weights": "model:acestep-v15-turbo/model.safetensors", "dit_turbo_silence_latent": "model:acestep-v15-turbo/silence_latent.safetensors", diff --git a/model_specs_v1/ace_step.json b/model_specs_v1/ace_step.json index 2cb74096..2b018921 100644 --- a/model_specs_v1/ace_step.json +++ b/model_specs_v1/ace_step.json @@ -305,10 +305,12 @@ { "name": "dit_variant", "type": "enum", - "description": "DiT variant inside the model package; default acestep-v15-turbo.", + "description": "DiT variant inside the model package; default acestep-v15-turbo. The XL variants are optional package resources and are only selectable when installed.", "values": [ "acestep-v15-turbo", - "acestep-v15-base" + "acestep-v15-base", + "acestep-v15-xl-turbo", + "acestep-v15-xl-sft" ], "required": false, "default": "acestep-v15-turbo" @@ -344,6 +346,28 @@ "text_encoder_chat_template": "model:Qwen3-Embedding-0.6B/chat_template.jinja", "vae_config": "model:vae/config.json" }, + "optional_files": { + "dit_xl_turbo_config": "model:acestep-v15-xl-turbo/config.json", + "dit_xl_sft_config": "model:acestep-v15-xl-sft/config.json" + }, + "optional_tensors": { + "dit_xl_turbo_weights": { + "source": "weights:", + "prefix": "dit_xl_turbo_weights" + }, + "dit_xl_turbo_silence_latent": { + "source": "weights:", + "prefix": "dit_xl_turbo_silence_latent" + }, + "dit_xl_sft_weights": { + "source": "weights:", + "prefix": "dit_xl_sft_weights" + }, + "dit_xl_sft_silence_latent": { + "source": "weights:", + "prefix": "dit_xl_sft_silence_latent" + } + }, "tensors": { "dit_turbo_weights": { "source": "weights:", @@ -397,6 +421,16 @@ "text_encoder_chat_template": "model:Qwen3-Embedding-0.6B/chat_template.jinja", "vae_config": "model:vae/config.json" }, + "optional_files": { + "dit_xl_turbo_config": "model:acestep-v15-xl-turbo/config.json", + "dit_xl_sft_config": "model:acestep-v15-xl-sft/config.json" + }, + "optional_tensors": { + "dit_xl_turbo_weights": "model:acestep-v15-xl-turbo/model.safetensors.index.json", + "dit_xl_turbo_silence_latent": "model:acestep-v15-xl-turbo/silence_latent.safetensors", + "dit_xl_sft_weights": "model:acestep-v15-xl-sft/model.safetensors.index.json", + "dit_xl_sft_silence_latent": "model:acestep-v15-xl-sft/silence_latent.safetensors" + }, "tensors": { "dit_turbo_weights": "model:acestep-v15-turbo/model.safetensors", "dit_turbo_silence_latent": "model:acestep-v15-turbo/silence_latent.safetensors", diff --git a/src/framework/model_spec/package.cpp b/src/framework/model_spec/package.cpp index 6bbac497..5433c414 100644 --- a/src/framework/model_spec/package.cpp +++ b/src/framework/model_spec/package.cpp @@ -316,6 +316,27 @@ void add_tensor_map(assets::ResourceBundle & bundle, } } +// Tensor sources a package may or may not ship. The required `tensors` map is +// checked eagerly, which is what a package wants for the weights it cannot run +// without; a family whose variants are separate multi-gigabyte downloads needs +// the other answer, or installing one variant means downloading all of them. +// A model that selects a missing variant reports it itself, where it can name +// the variant instead of a resource id. +void add_optional_tensor_map(assets::ResourceBundle & bundle, + const ResourceRoots & roots, + const engine::io::json::Value * map_value) { + if (map_value == nullptr || map_value->is_null()) { + return; + } + for (const auto & [id, ref] : map_value->as_object()) { + std::string prefix; + const auto path = resolve_tensor_source_ref(roots, ref, prefix); + if (engine::io::is_existing_file(path)) { + bundle.add_tensor_source(id, path, std::move(prefix)); + } + } +} + void add_optional_resource_map(assets::ResourceBundle & bundle, const ResourceRoots & roots, const engine::io::json::Value * map_value) { if (map_value == nullptr || map_value->is_null()) { @@ -354,6 +375,7 @@ assets::ResourceBundle load_source(const std::filesystem::path & model_root, con add_resource_map(bundle, roots, source.find("files")); add_optional_resource_map(bundle, roots, source.find("optional_files")); add_tensor_map(bundle, roots, source.find("tensors")); + add_optional_tensor_map(bundle, roots, source.find("optional_tensors")); return bundle; } @@ -362,10 +384,9 @@ std::vector discover_safetensors_source_resources(const en const ResourceRoots & roots) { auto resources = resources_from_resource_map( roots, source.find(kind == ResourceKind::Files ? "files" : "tensors"), true); - if (kind == ResourceKind::Files) { - auto optional = resources_from_resource_map(roots, source.find("optional_files"), false); - resources.insert(resources.end(), optional.begin(), optional.end()); - } + auto optional = resources_from_resource_map( + roots, source.find(kind == ResourceKind::Files ? "optional_files" : "optional_tensors"), false); + resources.insert(resources.end(), optional.begin(), optional.end()); return resources; } diff --git a/src/framework/model_spec/schema.cpp b/src/framework/model_spec/schema.cpp index 86fdfc11..149e5a89 100644 --- a/src/framework/model_spec/schema.cpp +++ b/src/framework/model_spec/schema.cpp @@ -485,7 +485,7 @@ void validate_layout(const json::Value & value, std::string_view path) { } (void) require_spec_string(root_value, std::string(path) + ".roots." + root_id); } - for (const std::string map_name : {"files", "optional_files", "tensors"}) { + for (const std::string map_name : {"files", "optional_files", "tensors", "optional_tensors"}) { const auto * map_value = value.find(map_name); if (map_value == nullptr) { continue; diff --git a/src/models/ace_step/assets.cpp b/src/models/ace_step/assets.cpp index e324f588..bbeedbe0 100644 --- a/src/models/ace_step/assets.cpp +++ b/src/models/ace_step/assets.cpp @@ -101,11 +101,38 @@ AceStepDiffusionConfig parse_diffusion_config(const engine::io::json::Value & va config.use_sliding_window = json::optional_bool(value, "use_sliding_window", false); config.layer_types = json::optional_string_array(value, "layer_types"); config.is_turbo = json::optional_bool(value, "is_turbo", config.is_turbo); + // `encoder_hidden_size` is what upstream's XL modeling class reads without a + // fallback, so a config that declares it is an XL-class package: the encoder + // stack is a different width from the DiT, and the timbre encoder is the + // variant that actually uses its CLS token. + config.has_separate_encoder = value.find("encoder_hidden_size") != nullptr; + config.timbre_special_token = config.has_separate_encoder; config.rms_norm_eps = json::optional_f32(value, "rms_norm_eps", config.rms_norm_eps); config.rope_theta = json::optional_f32(value, "rope_theta", config.rope_theta); return config; } +// Upstream builds the condition encoder, audio tokenizer and detokenizer from a +// copy of the config with four values substituted (copy.deepcopy in +// AceStepConditionGenerationModel.__init__). Doing the same here keeps every +// encoder-side shape derived from one place instead of spreading `is this XL?` +// across the weight loaders. +AceStepDiffusionConfig derive_encoder_config( + const AceStepDiffusionConfig & diffusion, + const engine::io::json::Value & value) { + AceStepDiffusionConfig config = diffusion; + if (!diffusion.has_separate_encoder) { + return config; + } + config.hidden_size = json::require_i64(value, "encoder_hidden_size"); + config.intermediate_size = json::optional_i64(value, "encoder_intermediate_size", diffusion.intermediate_size); + config.num_attention_heads = + json::optional_i64(value, "encoder_num_attention_heads", diffusion.num_attention_heads); + config.num_key_value_heads = + json::optional_i64(value, "encoder_num_key_value_heads", diffusion.num_key_value_heads); + return config; +} + AceStepVAEConfig parse_vae_config(const engine::io::json::Value & value) { AceStepVAEConfig config; config.sample_rate = static_cast(json::optional_i64(value, "sampling_rate", config.sample_rate)); @@ -118,14 +145,38 @@ AceStepVAEConfig parse_vae_config(const engine::io::json::Value & value) { return config; } -std::string dit_resource_id(const AceStepModelSelection & selection, std::string_view suffix) { - if (selection.dit_model_path == "acestep-v15-turbo") { - return "dit_turbo_" + std::string(suffix); +// Directory inside the package -> prefix of the resource ids the model spec +// registers for it. The XL variants are optional resources, so a package may +// name a variant here that it does not ship; that is caught when the resources +// are opened, with a message that says which directory is missing. +constexpr std::pair kDitVariants[] = { + {"acestep-v15-turbo", "dit_turbo_"}, + {"acestep-v15-base", "dit_base_"}, + {"acestep-v15-xl-turbo", "dit_xl_turbo_"}, + {"acestep-v15-xl-sft", "dit_xl_sft_"}, +}; + +std::string known_dit_variants() { + std::string names; + for (const auto & [directory, prefix] : kDitVariants) { + (void)prefix; + if (!names.empty()) { + names += ", "; + } + names += directory; } - if (selection.dit_model_path == "acestep-v15-base") { - return "dit_base_" + std::string(suffix); + return names; +} + +std::string dit_resource_id(const AceStepModelSelection & selection, std::string_view suffix) { + for (const auto & [directory, prefix] : kDitVariants) { + if (selection.dit_model_path == directory) { + return std::string(prefix) + std::string(suffix); + } } - throw std::runtime_error("ACE-Step package spec supports only acestep-v15-turbo and acestep-v15-base DiT variants"); + throw std::runtime_error( + "unknown ACE-Step DiT variant '" + selection.dit_model_path + "'; the package spec knows " + + known_dit_variants()); } void validate_selection(const AceStepModelSelection & selection) { @@ -134,13 +185,29 @@ void validate_selection(const AceStepModelSelection & selection) { AceStepConfig parse_config(const assets::ResourceBundle & resources, const AceStepModelSelection & selection) { AceStepConfig config; - config.diffusion = parse_diffusion_config(resources.parse_json(dit_resource_id(selection, "config"))); + const auto diffusion_json = resources.parse_json(dit_resource_id(selection, "config")); + config.diffusion = parse_diffusion_config(diffusion_json); + config.encoder = derive_encoder_config(config.diffusion, diffusion_json); config.planner = parse_planner_config(resources.parse_json("lm_config")); config.text_encoder = parse_text_encoder_config(resources.parse_json("text_encoder_config")); config.vae = parse_vae_config(resources.parse_json("vae_config")); return config; } +// The XL variants are registered as optional package resources so that a +// package holding only turbo does not have to ship 20 GB it will never load. +// The cost is that "not installed" surfaces here rather than at spec-load time, +// where a bare "missing asset resource: dit_xl_turbo_config" would not say why. +void require_installed_variant(const assets::ResourceBundle & resources, const AceStepModelSelection & selection) { + if (resources.has_file(dit_resource_id(selection, "config"))) { + return; + } + throw std::runtime_error( + "ACE-Step DiT variant '" + selection.dit_model_path + "' is not installed in this package: " + + (resources.model_root() / selection.dit_model_path).string() + + " is missing. Download that variant, or load one of the variants the package ships."); +} + void validate_config(const AceStepConfig & config) { if (config.diffusion.model_type != "acestep") { throw std::runtime_error("ACE-Step diffusion config must have model_type=acestep"); @@ -167,6 +234,7 @@ std::shared_ptr load_ace_step_assets( assets->resources = engine::model_spec::load_resource_bundle( model_path, engine::model_spec::default_spec_path("ace_step")); + require_installed_variant(assets->resources, assets->selection); assets->config = parse_config(assets->resources, assets->selection); validate_config(assets->config); assets->dit_weights = assets->resources.open_tensor_source(dit_resource_id(assets->selection, "weights")); diff --git a/src/models/ace_step/condition_encoder.cpp b/src/models/ace_step/condition_encoder.cpp index 36843ee8..1cc4857b 100644 --- a/src/models/ace_step/condition_encoder.cpp +++ b/src/models/ace_step/condition_encoder.cpp @@ -313,11 +313,11 @@ class AceStepConditionEncoderRuntime::Impl { AceStepTextConditioning run(const AceStepTextConditioning & input) const { if (input.tokens <= 0 || input.tokens > tokens_ || - input.hidden_size != assets_->config.diffusion.text_hidden_dim) { + input.hidden_size != assets_->config.encoder.text_hidden_dim) { throw std::runtime_error("ACE-Step text projector input shape mismatch"); } std::vector padded( - static_cast(tokens_ * assets_->config.diffusion.text_hidden_dim), + static_cast(tokens_ * assets_->config.encoder.text_hidden_dim), 0.0F); std::copy(input.values.begin(), input.values.end(), padded.begin()); core::write_tensor_f32(input_value_, padded); @@ -328,7 +328,7 @@ class AceStepConditionEncoderRuntime::Impl { } AceStepTextConditioning out; out.tokens = input.tokens; - out.hidden_size = assets_->config.diffusion.hidden_size; + out.hidden_size = assets_->config.encoder.hidden_size; std::vector full; core::read_tensor_f32_into(output_, full); out.values.assign( @@ -339,7 +339,7 @@ class AceStepConditionEncoderRuntime::Impl { private: void build() { - const auto & config = assets_->config.diffusion; + const auto & config = assets_->config.encoder; ggml_init_params params{8ull * 1024ull * 1024ull, nullptr, true}; ctx_.reset(ggml_init(params)); if (ctx_ == nullptr) { @@ -410,11 +410,11 @@ class AceStepConditionEncoderRuntime::Impl { AceStepTextConditioning run(const AceStepTextConditioning & input) const { if (input.tokens <= 0 || input.tokens > tokens_ || - input.hidden_size != assets_->config.diffusion.text_hidden_dim) { + input.hidden_size != assets_->config.encoder.text_hidden_dim) { throw std::runtime_error("ACE-Step lyric encoder input shape mismatch"); } std::vector padded( - static_cast(tokens_ * assets_->config.diffusion.text_hidden_dim), + static_cast(tokens_ * assets_->config.encoder.text_hidden_dim), 0.0F); std::copy(input.values.begin(), input.values.end(), padded.begin()); core::write_tensor_f32(input_value_, padded); @@ -428,7 +428,7 @@ class AceStepConditionEncoderRuntime::Impl { } AceStepTextConditioning out; out.tokens = input.tokens; - out.hidden_size = assets_->config.diffusion.hidden_size; + out.hidden_size = assets_->config.encoder.hidden_size; std::vector full; core::read_tensor_f32_into(output_, full); out.values.assign( @@ -439,7 +439,7 @@ class AceStepConditionEncoderRuntime::Impl { private: void build() { - const auto & config = assets_->config.diffusion; + const auto & config = assets_->config.encoder; ggml_init_params params{128ull * 1024ull * 1024ull, nullptr, true}; ctx_.reset(ggml_init(params)); if (ctx_ == nullptr) { @@ -563,7 +563,7 @@ class AceStepConditionEncoderRuntime::Impl { } AceStepTextConditioning run_one(const std::vector & packed_reference, int64_t frames) const { - const auto & config = assets_->config.diffusion; + const auto & config = assets_->config.encoder; if (frames <= 0 || frames > frames_ || static_cast(packed_reference.size()) != frames * config.timbre_hidden_dim) { throw std::runtime_error("ACE-Step timbre encoder input shape mismatch"); @@ -575,7 +575,7 @@ class AceStepConditionEncoderRuntime::Impl { core::write_tensor_f32(input_value_, padded); core::write_tensor_f16( padding_mask_value_, - build_padding_attention_mask_values(frames_, frames)); + build_padding_attention_mask_values(tokens_, frames + cls_tokens())); core::set_backend_threads(backend_, threads_); const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); if (status != GGML_STATUS_SUCCESS) { @@ -589,8 +589,17 @@ class AceStepConditionEncoderRuntime::Impl { } private: + // XL prepends a CLS token to the reference frames and reads it back as the + // timbre embedding. Earlier variants declare the same parameter but leave + // it out of the sequence, so position 0 there is the first audio frame — + // which is why this is a config question and not a tensor lookup. + int64_t cls_tokens() const noexcept { + return assets_->config.encoder.timbre_special_token ? 1 : 0; + } + void build() { - const auto & config = assets_->config.diffusion; + const auto & config = assets_->config.encoder; + tokens_ = frames_ + cls_tokens(); ggml_init_params params{128ull * 1024ull * 1024ull, nullptr, true}; ctx_.reset(ggml_init(params)); if (ctx_ == nullptr) { @@ -603,28 +612,35 @@ class AceStepConditionEncoderRuntime::Impl { GGML_TYPE_F32, core::TensorShape::from_dims({1, frames_, config.timbre_hidden_dim})); input_value_ = input_; - positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, frames_); - auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({frames_}), GGML_TYPE_I32); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, tokens_); + auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({tokens_}), GGML_TYPE_I32); sliding_mask_value_ = core::make_tensor( build_ctx, GGML_TYPE_F16, - core::TensorShape::from_dims({1, 1, frames_, frames_})); + core::TensorShape::from_dims({1, 1, tokens_, tokens_})); padding_mask_value_ = core::make_tensor( build_ctx, GGML_TYPE_F16, - core::TensorShape::from_dims({1, 1, frames_, frames_})); + core::TensorShape::from_dims({1, 1, tokens_, tokens_})); auto hidden = modules::LinearModule({config.timbre_hidden_dim, config.hidden_size, true}) .build( build_ctx, input_, {weights_->timbre_embed_weight, weights_->timbre_embed_bias}); + if (cls_tokens() > 0) { + special_token_value_ = core::make_tensor( + build_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, 1, config.hidden_size})); + hidden = modules::ConcatModule({1}).build(build_ctx, special_token_value_, hidden); + } const auto sliding_mask = core::wrap_tensor( sliding_mask_value_.tensor, - core::TensorShape::from_dims({1, 1, frames_, frames_}), + core::TensorShape::from_dims({1, 1, tokens_, tokens_}), GGML_TYPE_F16); const auto padding_mask = core::wrap_tensor( padding_mask_value_.tensor, - core::TensorShape::from_dims({1, 1, frames_, frames_}), + core::TensorShape::from_dims({1, 1, tokens_, tokens_}), GGML_TYPE_F16); for (int64_t i = 0; i < config.num_timbre_encoder_hidden_layers; ++i) { const std::optional mask = @@ -648,14 +664,18 @@ class AceStepConditionEncoderRuntime::Impl { throw std::runtime_error("ACE-Step timbre encoder backend buffer allocation failed"); } - std::vector position_values(static_cast(frames_), 0); - for (int64_t i = 0; i < frames_; ++i) { + // Positions and both masks run over the extended sequence: upstream + // builds its cache_position from the embeddings *after* the CLS token + // is prepended, so the token sits at position 0 and everything else + // shifts by one. + std::vector position_values(static_cast(tokens_), 0); + for (int64_t i = 0; i < tokens_; ++i) { position_values[static_cast(i)] = static_cast(i); } ggml_backend_tensor_set(positions_, position_values.data(), 0, position_values.size() * sizeof(int32_t)); const std::vector sliding_mask_values = ace_step_bidirectional_sliding_mask_values( - frames_, + tokens_, config.sliding_window, "ACE-Step"); std::vector sliding_mask_f16(sliding_mask_values.size()); @@ -667,6 +687,17 @@ class AceStepConditionEncoderRuntime::Impl { sliding_mask_f16.data(), 0, sliding_mask_f16.size() * sizeof(ggml_fp16_t)); + if (cls_tokens() > 0) { + const auto & token = weights_->timbre_special_token_host; + if (static_cast(token.size()) != config.hidden_size) { + throw std::runtime_error("ACE-Step timbre encoder CLS token shape mismatch"); + } + ggml_backend_tensor_set( + special_token_value_.tensor, + token.data(), + 0, + token.size() * sizeof(float)); + } } ggml_backend_t backend_ = nullptr; @@ -674,9 +705,12 @@ class AceStepConditionEncoderRuntime::Impl { std::shared_ptr assets_; std::shared_ptr weights_; int64_t frames_ = 0; + // frames_ plus the CLS token, when the variant has one. + int64_t tokens_ = 0; std::unique_ptr ctx_; core::TensorValue input_; core::TensorValue input_value_; + core::TensorValue special_token_value_; ggml_tensor * positions_ = nullptr; core::TensorValue sliding_mask_value_; core::TensorValue padding_mask_value_; @@ -715,10 +749,10 @@ class AceStepConditionEncoderRuntime::Impl { int64_t refer_audio_frames, const std::vector & refer_audio_order_mask) const { const auto total_start = Clock::now(); - if (text_hidden_states.hidden_size != assets_->config.diffusion.text_hidden_dim) { + if (text_hidden_states.hidden_size != assets_->config.encoder.text_hidden_dim) { throw std::runtime_error("ACE-Step condition encoder text hidden size mismatch"); } - if (lyric_token_embeddings.hidden_size != assets_->config.diffusion.text_hidden_dim) { + if (lyric_token_embeddings.hidden_size != assets_->config.encoder.text_hidden_dim) { throw std::runtime_error("ACE-Step condition encoder lyric hidden size mismatch"); } const auto project_text_start = Clock::now(); @@ -795,7 +829,7 @@ class AceStepConditionEncoderRuntime::Impl { int64_t refer_audio_count, int64_t refer_audio_frames, const std::vector & refer_audio_order_mask) const { - const auto & config = assets_->config.diffusion; + const auto & config = assets_->config.encoder; if (refer_audio_count <= 0 || refer_audio_frames <= 0) { throw std::runtime_error("ACE-Step timbre encoder requires positive reference count and frame count"); } diff --git a/src/models/ace_step/cover_tokenizer.cpp b/src/models/ace_step/cover_tokenizer.cpp index 4c1424ba..5ba10366 100644 --- a/src/models/ace_step/cover_tokenizer.cpp +++ b/src/models/ace_step/cover_tokenizer.cpp @@ -94,7 +94,7 @@ std::shared_ptr load_cover_tokenizer_weights backend_type, "ace_step.cover_tokenizer.weights", 256ull * 1024ull * 1024ull); - const auto & config = assets.config.diffusion; + const auto & config = assets.config.encoder; const auto & source = *assets.dit_weights; auto weights = std::make_shared(); weights->store = store; @@ -215,7 +215,7 @@ class AceStepCoverTokenizerRuntime::Impl { int64_t silence_channels, std::vector & input_buffer) const { const auto total_start = Clock::now(); - const auto & config = assets_->config.diffusion; + const auto & config = assets_->config.encoder; if (code_count <= 0 || code_count > code_capacity_) { throw std::runtime_error("ACE-Step cover tokenizer chunk exceeds graph capacity"); } @@ -272,7 +272,7 @@ class AceStepCoverTokenizerRuntime::Impl { private: void build() { - const auto & config = assets_->config.diffusion; + const auto & config = assets_->config.encoder; const int64_t patch_tokens = config.pool_window_size + 1; ggml_init_params params{128ull * 1024ull * 1024ull, nullptr, true}; ctx_.reset(ggml_init(params)); @@ -423,7 +423,7 @@ class AceStepCoverTokenizerRuntime::Impl { threads_(std::max(1, execution.config().threads)), storage_type_(storage_type), quantizer_(ace_step_build_fsq_quantizer_table( - assets_->config.diffusion, + assets_->config.encoder, "ACE-Step native cover tokenizer")) { if (assets_ == nullptr) { throw std::runtime_error("ACE-Step cover tokenizer requires assets"); @@ -439,7 +439,7 @@ class AceStepCoverTokenizerRuntime::Impl { int64_t silence_frames, int64_t silence_channels) const { const auto total_start = Clock::now(); - const auto & config = assets_->config.diffusion; + const auto & config = assets_->config.encoder; if (latents.frames <= 0 || latents.channels != config.latent_channels) { throw std::runtime_error("ACE-Step cover tokenizer requires positive latent frames and matching channels"); } diff --git a/src/models/ace_step/detokenizer.cpp b/src/models/ace_step/detokenizer.cpp index c890be84..f4ebada0 100644 --- a/src/models/ace_step/detokenizer.cpp +++ b/src/models/ace_step/detokenizer.cpp @@ -123,7 +123,7 @@ class AceStepAudioDetokenizerRuntime::Impl { AceStepLatents decode_audio_codes(const int32_t * audio_code_ids, int64_t code_count, std::vector & expanded) const { const auto total_start = Clock::now(); - const auto & config = assets_->config.diffusion; + const auto & config = assets_->config.encoder; if (code_count <= 0) { return {}; } @@ -186,7 +186,7 @@ class AceStepAudioDetokenizerRuntime::Impl { private: void build() { - const auto & config = assets_->config.diffusion; + const auto & config = assets_->config.encoder; ggml_init_params params{96ull * 1024ull * 1024ull, nullptr, true}; ctx_.reset(ggml_init(params)); if (ctx_ == nullptr) { @@ -350,7 +350,7 @@ class AceStepAudioDetokenizerRuntime::Impl { assets_(std::move(assets)), weights_(dit_weights_runtime->detokenizer_weights()), quantizer_(ace_step_build_fsq_quantizer_table( - assets_->config.diffusion, + assets_->config.encoder, "ACE-Step native detokenizer")) { if (backend_ == nullptr) { throw std::runtime_error("ACE-Step detokenizer backend initialization failed"); @@ -378,7 +378,7 @@ class AceStepAudioDetokenizerRuntime::Impl { engine::debug::timing_log_scalar( "ace_step.detokenizer.graph.ensure_ms", engine::debug::elapsed_ms(ensure_start, Clock::now())); - const auto & config = assets_->config.diffusion; + const auto & config = assets_->config.encoder; AceStepLatents out; out.frames = code_count * config.pool_window_size; out.channels = config.latent_channels; diff --git a/src/models/ace_step/diffusion.cpp b/src/models/ace_step/diffusion.cpp index 65a16125..103ab916 100644 --- a/src/models/ace_step/diffusion.cpp +++ b/src/models/ace_step/diffusion.cpp @@ -161,11 +161,15 @@ core::TensorValue build_attention( v_heads = ensure_contiguous(ctx, v_heads); auto context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask, backend_type); context = ensure_contiguous(ctx, context); + // The attention width is heads x head_dim, which only equals hidden_size when + // head_dim was derived from it. XL states head_dim outright (32 x 128 against + // a hidden size of 2560), so o_proj is the one rectangular projection here. + const int64_t attention_size = config.num_attention_heads * dim; context = core::reshape_tensor( ctx, context, - core::TensorShape::from_dims({hidden_states.shape.dims[0], hidden_states.shape.dims[1], config.hidden_size})); - return modules::LinearModule({config.hidden_size, config.hidden_size, false, GGML_PREC_F32}) + core::TensorShape::from_dims({hidden_states.shape.dims[0], hidden_states.shape.dims[1], attention_size})); + return modules::LinearModule({attention_size, config.hidden_size, false, GGML_PREC_F32}) .build(ctx, context, {weights.out_weight, std::nullopt}); } @@ -977,7 +981,10 @@ class AceStepDiffusionRuntime::Impl { Output run(const std::vector & encoder_hidden_states) const { const auto & config = assets_->config.diffusion; - if (static_cast(encoder_hidden_states.size()) != encoder_tokens_ * config.hidden_size) { + // Still in the encoder's width here: the condition embedder inside the + // graph is what lifts it to the DiT's. + const int64_t encoder_hidden_size = assets_->config.encoder.hidden_size; + if (static_cast(encoder_hidden_states.size()) != encoder_tokens_ * encoder_hidden_size) { throw std::runtime_error("ACE-Step diffusion cross-attention cache encoder hidden state shape mismatch"); } core::write_tensor_f32(encoder_value_, encoder_hidden_states); @@ -1012,12 +1019,14 @@ class AceStepDiffusionRuntime::Impl { } core::ModuleBuildContext ctx{ctx_.get(), "ace_step.diffusion.cross_cache", backend_type_}; + const int64_t encoder_hidden_size = assets_->config.encoder.hidden_size; encoder_value_ = core::make_tensor( ctx, GGML_TYPE_F32, - core::TensorShape::from_dims({1, encoder_tokens_, config.hidden_size})); - auto encoder_hidden_states = modules::LinearModule({config.hidden_size, config.hidden_size, true, GGML_PREC_F32}) - .build(ctx, encoder_value_, weights_->condition_embedder); + core::TensorShape::from_dims({1, encoder_tokens_, encoder_hidden_size})); + auto encoder_hidden_states = + modules::LinearModule({encoder_hidden_size, config.hidden_size, true, GGML_PREC_F32}) + .build(ctx, encoder_value_, weights_->condition_embedder); key_outputs_.reserve(weights_->layers.size()); value_outputs_.reserve(weights_->layers.size()); for (const auto & layer : weights_->layers) { @@ -1467,10 +1476,12 @@ class AceStepDiffusionRuntime::Impl { const auto total_start = Clock::now(); const auto & pre = conditioning.pre_dit; const auto & config = assets_->config.diffusion; + // Everything the condition encoder produced is still in its own width. + const int64_t encoder_hidden_size = assets_->config.encoder.hidden_size; if (pre.context_latents.frames <= 0 || pre.context_latents.channels != config.latent_channels * 2) { throw std::runtime_error("ACE-Step diffusion requires valid context latents"); } - if (pre.encoder_hidden_states.tokens <= 0 || pre.encoder_hidden_states.hidden_size != config.hidden_size) { + if (pre.encoder_hidden_states.tokens <= 0 || pre.encoder_hidden_states.hidden_size != encoder_hidden_size) { throw std::runtime_error("ACE-Step diffusion requires valid encoder hidden states"); } const int64_t encoder_token_capacity = std::max( @@ -1512,7 +1523,7 @@ class AceStepDiffusionRuntime::Impl { cfg_context_padded = duplicate_batch_values(context_padded, diffusion_batch_size); } const auto encoder_hidden_padded = - pad_encoder_hidden_values(pre.encoder_hidden_states, graph_->encoder_token_capacity(), config.hidden_size); + pad_encoder_hidden_values(pre.encoder_hidden_states, graph_->encoder_token_capacity(), encoder_hidden_size); const auto encoder_attention_mask_padded = pad_encoder_attention_mask(pre.encoder_hidden_states, graph_->encoder_token_capacity()); engine::debug::timing_log_scalar("ace_step.diffusion.padding_ms", engine::debug::elapsed_ms(padding_start, Clock::now())); @@ -1537,7 +1548,7 @@ class AceStepDiffusionRuntime::Impl { const auto null_encoder_hidden_padded = null_encoder_hidden_values( weights_->null_condition_emb_host, graph_->encoder_token_capacity(), - config.hidden_size); + encoder_hidden_size); null_cross_attention_cache = cross_cache_graph_->run(null_encoder_hidden_padded); cross_attention_cache = combine_cross_attention_cache(cross_attention_cache, *null_cross_attention_cache); } @@ -1563,7 +1574,7 @@ class AceStepDiffusionRuntime::Impl { non_cover_encoder_hidden_padded = pad_encoder_hidden_values( pre.encoder_hidden_states_non_cover, graph_->encoder_token_capacity(), - config.hidden_size); + encoder_hidden_size); non_cover_encoder_attention_mask_padded = pad_encoder_attention_mask( pre.encoder_hidden_states_non_cover, graph_->encoder_token_capacity()); diff --git a/src/models/ace_step/dit_weights_runtime.cpp b/src/models/ace_step/dit_weights_runtime.cpp index e955e29a..a16db12d 100644 --- a/src/models/ace_step/dit_weights_runtime.cpp +++ b/src/models/ace_step/dit_weights_runtime.cpp @@ -64,7 +64,7 @@ std::shared_ptr load_condition_encoder_wei const std::shared_ptr & store, const AceStepAssets & assets, assets::TensorStorageType storage_type) { - const auto & config = assets.config.diffusion; + const auto & config = assets.config.encoder; const auto & source = *assets.dit_weights; auto weights = std::make_shared(); weights->store = store; @@ -116,6 +116,14 @@ std::shared_ptr load_condition_encoder_wei config)); } weights->timbre_norm = store->load_f32_tensor(source, "encoder.timbre_encoder.norm.weight", {config.hidden_size}); + // Every ACE-Step package carries this parameter; only the XL class prepends + // it to the reference frames, so loading it anywhere else would pin a tensor + // the graph never reads. + if (config.timbre_special_token) { + weights->timbre_special_token_host = source.require_f32( + "encoder.timbre_encoder.special_token", + {1, 1, config.hidden_size}); + } return weights; } @@ -123,7 +131,7 @@ std::shared_ptr load_detokenizer_weights( const std::shared_ptr & store, const AceStepAssets & assets, assets::TensorStorageType storage_type) { - const auto & config = assets.config.diffusion; + const auto & config = assets.config.encoder; const auto & source = *assets.dit_weights; const int64_t dim = config.head_dim; @@ -230,6 +238,9 @@ std::shared_ptr load_diffusion_weights( const AceStepAssets & assets, assets::TensorStorageType storage_type) { const auto & config = assets.config.diffusion; + // The two places the DiT meets the condition encoder, and the only shapes in + // this function that are not square in the DiT's own width. + const int64_t encoder_hidden_size = assets.config.encoder.hidden_size; const auto & source = *assets.dit_weights; auto weights = std::make_shared(); weights->store = store; @@ -241,11 +252,14 @@ std::shared_ptr load_diffusion_weights( weights->time_embed = load_time_embedding_weights(*store, source, "decoder.time_embed", storage_type, config.hidden_size); weights->time_embed_r = load_time_embedding_weights(*store, source, "decoder.time_embed_r", storage_type, config.hidden_size); weights->condition_embedder = { - store->load_tensor(source, "decoder.condition_embedder.weight", storage_type, {config.hidden_size, config.hidden_size}), + store->load_tensor(source, "decoder.condition_embedder.weight", storage_type, {config.hidden_size, encoder_hidden_size}), store->load_tensor(source, "decoder.condition_embedder.bias", assets::TensorStorageType::F32, {config.hidden_size}), }; if (!config.is_turbo) { - weights->null_condition_emb_host = source.require_f32("null_condition_emb", {1, 1, config.hidden_size}); + // Stands in for the encoder output under classifier-free guidance, so it + // is sized in the encoder's width and the condition embedder projects it + // like any other conditioning. + weights->null_condition_emb_host = source.require_f32("null_condition_emb", {1, 1, encoder_hidden_size}); } weights->layers.reserve(static_cast(config.num_hidden_layers)); for (int64_t i = 0; i < config.num_hidden_layers; ++i) { diff --git a/src/models/ace_step/loader.cpp b/src/models/ace_step/loader.cpp index b2b5e3a0..90ccc44c 100644 --- a/src/models/ace_step/loader.cpp +++ b/src/models/ace_step/loader.cpp @@ -22,10 +22,11 @@ AceStepModelSelection selection_from_request(const runtime::ModelLoadRequest & r std::transform(model_name.begin(), model_name.end(), model_name.begin(), [](unsigned char ch) { return static_cast(std::tolower(ch)); }); - if (model_name.find("base") != std::string::npos) { - selection.dit_model_path = "acestep-v15-base"; + const bool xl = model_name.find("xl") != std::string::npos; + if (model_name.find("base") != std::string::npos || model_name.find("sft") != std::string::npos) { + selection.dit_model_path = xl ? "acestep-v15-xl-sft" : "acestep-v15-base"; } else if (model_name.find("turbo") != std::string::npos) { - selection.dit_model_path = "acestep-v15-turbo"; + selection.dit_model_path = xl ? "acestep-v15-xl-turbo" : "acestep-v15-turbo"; } return selection; } @@ -89,7 +90,9 @@ runtime::ModelCliInterface cli(const AceStepAssets &) { {"ace_step.mem_saver", "true|false", "Release staged runtime graphs after each request; default false."}, }; out.load_options = { - {"ace_step.dit_model_path", "acestep-v15-turbo|acestep-v15-base", "DiT variant inside the model package."}, + {"ace_step.dit_model_path", + "acestep-v15-turbo|acestep-v15-base|acestep-v15-xl-turbo|acestep-v15-xl-sft", + "DiT variant inside the model package; the XL variants are optional and must be installed."}, }; return out; }