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; }