ace_step: support the ACE-Step 1.5 XL DiT variants

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,11 +184,35 @@ 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=<dir>` | `acestep-v15-turbo`, `acestep-v15-base` | `acestep-v15-turbo` | Select DiT variant inside the model root. |
+| `--load-option ace_step.dit_model_path=<dir>` | `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=<type>` | `native`, `f32`, `f16`, `bf16`, `q8_0` | `native` | DiT weight type. |
 | `--session-option ace_step.planner_weight_type=<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. |
 
 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).

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

include/engine/models/ace_step/dit_weights_runtime.h

@@ -59,6 +59,8 @@ struct AceStepConditionEncoderWeights {
     core::TensorValue timbre_embed_bias;
     std::vector<AceStepConditionEncoderLayerWeights> timbre_layers;
     core::TensorValue timbre_norm;
+    // Empty unless the variant prepends a CLS token to the timbre sequence.
+    std::vector<float> timbre_special_token_host;
 };
 
 struct AceStepTimeEmbeddingWeights {

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",

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",

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<assets::ResourceFile> 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;
 }
 

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;

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<int>(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<std::string_view, std::string_view> 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<const AceStepAssets> 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"));