QVAC-23944 audiogen: add ACE-Step Multi-Track (lego) task with base-model CFG/APG sampling

README.md

@@ -140,7 +140,8 @@ the [TTS capability table](engines/tts/README.md#capabilities).
 | Model | Engine | Task | Rate | Quantization | Backends | Notes |
 |---|---|---|---|---|---|---|
 | ACE-Step v15 turbo | audiogen | text-to-music | 48 kHz stereo | `f32`, `f16`, `bf16`, `q8_0`, `q4_k_m` | CPU, Vulkan, Metal, OpenCL (Adreno 700+) | 8 diffusion steps by default |
-| ACE-Step v15 sft | audiogen | text-to-music | 48 kHz stereo | `f32`, `f16`, `bf16`, `q8_0` | CPU, Vulkan, Metal, OpenCL (Adreno 700+) | 50 diffusion steps by default |
+| ACE-Step v15 sft | audiogen | text-to-music | 48 kHz stereo | `f32`, `f16`, `bf16`, `q8_0` | CPU, Vulkan, Metal, OpenCL (Adreno 700+) | 50 diffusion steps by default, CFG via APG |
+| ACE-Step v15 base | audiogen | text-to-music, multi-track (lego) stems | 48 kHz stereo | `f32`, `f16`, `bf16`, `q8_0` | CPU, Vulkan, Metal, OpenCL (Adreno 700+) | 50 diffusion steps by default, CFG via APG, `--task lego --track <layer>` |
 | MiniMax-Music3 | audiogen | text-to-music | 44.1 kHz stereo | `f16`, `q8_0` | desktop CPU + GPU (CUDA, Vulkan, Metal via `EngineOptions::device`) | 25 fps, 30 flow steps, two GGUF files |
 
 ## Build

engines/audiogen/.gitignore

@@ -13,3 +13,4 @@ src/acestep/vae_decode_bin.cpp
 
 # macOS
 .DS_Store
+output/

engines/audiogen/README.md

@@ -103,12 +103,30 @@ The most specific stems (`embedding`, `vae`) are tested first so no short token
 
 ### DiT variants
 
-Detected from the `acestep.is_turbo` GGUF key; absent means base or sft.
+Detected from the `acestep.is_turbo` GGUF key (absent means base or sft) and
+`general.name` (an "sft" substring marks the sft fine-tune).
+
+| Variant | Default steps | Default shift | Default guidance | Notes |
+|---|--:|--:|--:|---|
+| turbo | 8 | 3.0 | 1.0 | fastest; guidance-distilled, CFG overrides are clamped to 1.0 |
+| base | 50 | 1.0 | 7.0 | CFG via APG; the only variant supporting the lego stem task |
+| sft | 50 | 1.0 | 7.0 | CFG via APG; no stem tasks |
+
+### Multi-Track (lego)
+
+`task_type: "lego"` generates a new instrument layer that follows
+`source_audio` and returns only that stem, trimmed to the source length for
+sample-for-sample mixing. It requires a base DiT (turbo and sft are rejected)
+and a `track` name — one of vocals, backing_vocals, drums, bass, guitar,
+keyboard, percussion, strings, synth, fx, brass, woodwinds:
+
+```bash
+music-cli --models <dir> --task lego --track guitar \
+  --caption "clean electric guitar with syncopated fills" \
+  --lyrics "[Instrumental]" --src-audio song.wav --out stem.wav
+```
 
-| Variant | Default steps | Default shift | Notes |
-|---|--:|--:|---|
-| turbo | 8 | 3.0 | fastest, no CFG on the DiT |
-| base / sft | 50 | 1.0 | DiT CFG and APG (`guidance > 1`) are deferred |
+`--guidance F` overrides the DiT guidance scale on base/sft (0 = auto).
 
 Weights load quantized. `f32`, `f16`, and `bf16` are handled for norms and biases, and the detokenizer's `special_tokens` may be `q8_0`. This tree ships no converter: the stage GGUFs come from `convert.py` in [acestep.cpp](https://github.com/ServeurpersoCom/acestep.cpp), the upstream C++/ggml implementation this port follows, which also publishes pre-quantized GGUFs at [Serveurperso/ACE-Step-1.5-GGUF](https://huggingface.co/Serveurperso/ACE-Step-1.5-GGUF).
 

engines/audiogen/include/audiogen-cpp/acestep/engine.h

@@ -129,6 +129,7 @@ struct GenerateParams {
     float       duration = 20.0f;        // target seconds (drives LM code count)
     int         inference_steps = 0;     // 0 = auto (turbo: 8, base/sft: 50)
     float       shift = 0.0f;            // 0 = auto (turbo: 3.0, base/sft: 1.0)
+    float       guidance_scale = 0.0f;   // 0 = auto (turbo: 1.0, base/sft: 7.0); >1 runs CFG via APG
     std::string vocal_language;          // optional hint, e.g. "en"
     int         bpm = 0;                 // optional; 0 => N/A (LM/DiT infer)
     std::string keyscale;                // optional, e.g. "C major"
@@ -141,7 +142,8 @@ struct GenerateParams {
     int         lm_top_k       = 0;      // 0 = disabled (top_p only)
     float       lm_cfg_scale   = 2.0f;   // classifier-free guidance for codes
     bool        lm_phase1      = true;   // auto-fill missing metadata (FSM CoT)
-    // Official sampler-side Haar DCW "double" correction.
+    // Official sampler-side Haar DCW "double" correction. Applied on turbo
+    // DiTs only: the official preset disables DCW for base/sft models.
     bool        dcw_enabled     = true;
     float       dcw_scaler      = 0.05f;  // low band coefficient: t * scaler
     float       dcw_high_scaler = 0.02f;  // high band coefficient: (1-t) * scaler
@@ -159,10 +161,17 @@ struct GenerateParams {
     std::vector<float> source_audio;
 
     // Task discriminator (mirrors acestep.cpp AceRequest::task_type).
-    // Supported today: "text2music" | "cover-nofsq".
+    // Supported today: "text2music" | "cover-nofsq" | "lego".
     // "cover" (FSQ roundtrip) is accepted at the API but not implemented yet.
+    // "lego" generates a new instrument layer that follows source_audio and
+    // returns only that layer; it requires a base/sft DiT (turbo is rejected).
     std::string task_type = "text2music";
 
+    // Lego target layer. Required when task_type is "lego"; one of:
+    // vocals|backing_vocals|drums|bass|guitar|keyboard|percussion|strings|
+    // synth|fx|brass|woodwinds.
+    std::string track;
+
     // Fraction of DiT steps that keep the source context (0..1). Default 1.0
     // keeps source context for every step. Values < 1.0 need DiT context
     // switching and are rejected until that path is ported.

engines/audiogen/src/acestep/dit_ggml.cpp

@@ -830,6 +830,123 @@ static void preserve_repaint_latent(const DitSampleParams & params, size_t laten
         (size_t) params.T, params.repaint_crossfade_frames, channels);
 }
 
+static constexpr double DIT_APG_MOMENTUM       = -0.75;
+static constexpr double DIT_APG_NORM_THRESHOLD = 2.5;
+
+static void apg_accumulate_momentum(std::vector<double> & running, std::vector<double> & diff) {
+    for (size_t i = 0; i < diff.size(); i++) {
+        running[i] = diff[i] + DIT_APG_MOMENTUM * running[i];
+        diff[i]    = running[i];
+    }
+}
+
+static double apg_channel_norm(const double * values, int T, int Oc, int channel) {
+    double sum = 0.0;
+    for (int t = 0; t < T; t++) {
+        const double v = values[(size_t) t * Oc + channel];
+        sum += v * v;
+    }
+    return std::sqrt(sum);
+}
+
+static void apg_scale_channel(double * values, int T, int Oc, int channel, double scale) {
+    for (int t = 0; t < T; t++) {
+        values[(size_t) t * Oc + channel] *= scale;
+    }
+}
+
+static void apg_clip_channel_norms(double * diff, int T, int Oc) {
+    for (int c = 0; c < Oc; c++) {
+        const double norm = apg_channel_norm(diff, T, Oc, c);
+        if (norm > DIT_APG_NORM_THRESHOLD) {
+            apg_scale_channel(diff, T, Oc, c, DIT_APG_NORM_THRESHOLD / norm);
+        }
+    }
+}
+
+static double apg_channel_norm_f32(const float * values, int T, int Oc, int channel) {
+    double sum = 0.0;
+    for (int t = 0; t < T; t++) {
+        const double v = values[(size_t) t * Oc + channel];
+        sum += v * v;
+    }
+    return std::sqrt(sum);
+}
+
+static void apg_remove_parallel_component(double * diff, const float * cond, int T, int Oc, int channel) {
+    const double norm = apg_channel_norm_f32(cond, T, Oc, channel);
+    if (norm <= 0.0) return;
+    const double inv_norm = 1.0 / norm;
+    double dot = 0.0;
+    for (int t = 0; t < T; t++) {
+        const size_t idx = (size_t) t * Oc + channel;
+        dot += diff[idx] * (double) cond[idx] * inv_norm;
+    }
+    for (int t = 0; t < T; t++) {
+        const size_t idx = (size_t) t * Oc + channel;
+        diff[idx] -= dot * (double) cond[idx] * inv_norm;
+    }
+}
+
+static void apg_project_orthogonal(double * diff, const float * cond, int T, int Oc) {
+    for (int c = 0; c < Oc; c++) {
+        apg_remove_parallel_component(diff, cond, T, Oc, c);
+    }
+}
+
+static std::vector<double> apg_velocity_difference(const std::vector<float> & velocity,
+                                                   const std::vector<float> & velocity_uncond) {
+    std::vector<double> diff(velocity.size());
+    for (size_t i = 0; i < diff.size(); i++) {
+        diff[i] = (double) velocity[i] - (double) velocity_uncond[i];
+    }
+    return diff;
+}
+
+static void apg_shape_batch_updates(std::vector<double> & diff, const std::vector<float> & velocity,
+                                    int T, int Oc, int N) {
+    const size_t n_per = (size_t) T * Oc;
+    for (int b = 0; b < N; b++) {
+        apg_clip_channel_norms(diff.data() + (size_t) b * n_per, T, Oc);
+        apg_project_orthogonal(diff.data() + (size_t) b * n_per, velocity.data() + (size_t) b * n_per, T, Oc);
+    }
+}
+
+static void apg_apply_guided_update(std::vector<float> & velocity, const std::vector<double> & diff,
+                                    float guidance_scale) {
+    const double weight = (double) guidance_scale - 1.0;
+    for (size_t i = 0; i < diff.size(); i++) {
+        velocity[i] = (float) ((double) velocity[i] + weight * diff[i]);
+    }
+}
+
+void dit_apg_guide(std::vector<float> &       velocity,
+                   const std::vector<float> & velocity_uncond,
+                   std::vector<double> &      momentum,
+                   float                      guidance_scale,
+                   int                        T,
+                   int                        Oc,
+                   int                        N) {
+    std::vector<double> diff = apg_velocity_difference(velocity, velocity_uncond);
+    apg_accumulate_momentum(momentum, diff);
+    apg_shape_batch_updates(diff, velocity, T, Oc, N);
+    apg_apply_guided_update(velocity, diff, guidance_scale);
+}
+
+static std::vector<float> make_null_enc_hidden(const float * null_emb, int H_enc, int enc_S, int N) {
+    std::vector<float> hidden((size_t) H_enc * enc_S * N);
+    for (int b = 0; b < N; b++) {
+        for (int s = 0; s < enc_S; s++) {
+            memcpy(&hidden[((size_t) b * enc_S + s) * H_enc], null_emb, (size_t) H_enc * sizeof(float));
+        }
+    }
+    return hidden;
+}
+
+static std::vector<uint16_t> make_visible_ca_mask(int enc_S, int S, int N) {
+    return std::vector<uint16_t>((size_t) enc_S * S * N, ggml_fp32_to_fp16(0.0f));
+}
+
 bool dit_sample(DitModel * m, const DitSampleParams & p, std::vector<float> & latent_out) {
     const DitConfig & c      = m->cfg;
     const int         Oc     = c.out_channels;                 // 64 (noisy latent channels)
@@ -895,6 +1012,18 @@ bool dit_sample(DitModel * m, const DitSampleParams & p, std::vector<float> & la
         xt_before.resize(n_per * N);
         denoised.resize(n_per * N);
     }
+
+    const bool use_cfg = p.guidance_scale > 1.0f && p.null_cond_emb != nullptr && p.H_enc > 0;
+    std::vector<float>    vt_uncond;
+    std::vector<float>    null_enc_hidden;
+    std::vector<uint16_t> null_ca_mask;
+    std::vector<double>   apg_momentum;
+    if (use_cfg) {
+        null_enc_hidden = make_null_enc_hidden(p.null_cond_emb, p.H_enc, enc_S, N);
+        null_ca_mask    = make_visible_ca_mask(enc_S, S, N);
+        apg_momentum.assign(n_per * N, 0.0);
+    }
+
     for (int step = 0; step < p.num_steps; step++) {
         if (p.on_step && !p.on_step(step, p.num_steps)) return false;
         const float t_curr = p.schedule[step];
@@ -916,10 +1045,8 @@ bool dit_sample(DitModel * m, const DitSampleParams & p, std::vector<float> & la
         fin.enc_S         = enc_S;
         fin.H_enc         = p.H_enc;
         fin.t             = t_curr;
-        // t_r == t (t_diff == 0, so time_embed_r sees 0). Holds for turbo
-        // text2music, which is also why the sampler runs a single conditional
-        // pass (N == 1, no CFG). base/sft (50-step, CFG) parity is not yet
-        // verified against the reference and would need t_r / uncond wiring.
+        // t_r == t (t_diff == 0, so time_embed_r sees 0) for both turbo and
+        // base/sft: the reference passes timestep_r = timestep unconditionally.
         fin.t_r           = t_curr;
         fin.sa_mask_sw    = sa_mask.data();
         fin.ca_mask       = ca_mask.data();
@@ -929,6 +1056,17 @@ bool dit_sample(DitModel * m, const DitSampleParams & p, std::vector<float> & la
             return false;
         }
 
+        if (use_cfg) {
+            DitForwardInputs fin_uncond = fin;
+            fin_uncond.enc_hidden       = null_enc_hidden.data();
+            fin_uncond.ca_mask          = null_ca_mask.data();
+            if (!dit_model_forward(m, fin_uncond, vt_uncond)) {
+                fprintf(stderr, "[acestep-dit] sample: uncond forward failed at step %d\n", step);
+                return false;
+            }
+            dit_apg_guide(vt, vt_uncond, apg_momentum, p.guidance_scale, T, Oc, N);
+        }
+
         // Euler ODE step. Final step integrates all the way to x0 (t_next = 0).
         const float t_next = (step == p.num_steps - 1) ? 0.0f : p.schedule[step + 1];
         const float dt     = t_curr - t_next;

engines/audiogen/src/acestep/dit_ggml.h

@@ -45,8 +45,16 @@ struct DitConfig {
     float rms_norm_eps      = 0.0f;
     int   enc_hidden_size   = 0;     // condition_embedder input dim (from weight shape)
     bool  is_turbo          = true;  // acestep.is_turbo: turbo=8 steps/shift 3; base/sft=50/shift 1
+    bool  is_sft            = false; // general.name contains "sft": CFG-capable, no stem tasks
+    std::string model_name;          // general.name from the GGUF (checkpoint dir name)
 };
 
+// convert.py stamps general.name with the checkpoint directory name
+// (acestep-v15-base / acestep-v15-sft / acestep-v15-xl-sft / ...).
+inline bool is_sft_model_name(const std::string & model_name) {
+    return model_name.find("sft") != std::string::npos;
+}
+
 struct DitModel;  // opaque: fused weight tensors + backend weight buffer
 
 // Load DiT weights from `path` onto `backend` (borrowed). Returns nullptr on
@@ -94,7 +102,11 @@ void dit_apply_haar_dcw(std::vector<float> &       x_next,
                         float                      low_scale,
                         float                      high_scale);
 
-// One full flow-matching denoise (Euler, no CFG — turbo runs guidance=1.0).
+// One full flow-matching denoise (Euler). guidance_scale > 1 with a
+// null-condition embedding runs classifier-free guidance via APG (Adaptive
+// Projected Guidance, official ACE-Step base/sft path): a second unconditional
+// forward per step whose encoder states are the null embedding broadcast over
+// enc_S, combined as pred_cond + (scale-1) * orthogonal(momentum(diff)).
 struct DitSampleParams {
     const float * noise           = nullptr;  // [out_channels, T, N] initial x_T
     const float * context_latents = nullptr;  // [in_channels-out_channels, T, N] conditioning
@@ -105,6 +117,8 @@ struct DitSampleParams {
     int           N               = 1;
     const float * schedule        = nullptr;  // [num_steps] descending timesteps
     int           num_steps       = 0;
+    float         guidance_scale  = 1.0f;     // <= 1 disables CFG
+    const float * null_cond_emb   = nullptr;  // [H_enc] cond-model null embedding
     const int *   real_enc_S      = nullptr;  // [N] valid encoder lengths; null = all enc_S
     bool          dcw_enabled     = true;     // official ACE-Step Haar "double" mode
     float         dcw_scaler      = 0.05f;    // low band: t_curr * scaler
@@ -127,6 +141,17 @@ struct DitSampleParams {
 // DiT graph per step (bring-up simplicity); correctness first, fusion later.
 bool dit_sample(DitModel * m, const DitSampleParams & p, std::vector<float> & latent_out);
 
+// APG combine for one step: velocity holds the conditional prediction on entry
+// and the guided result on exit. momentum is the caller-held running average
+// ([out_channels * T * N] doubles, zero-initialized before the first step).
+void dit_apg_guide(std::vector<float> &       velocity,
+                   const std::vector<float> & velocity_uncond,
+                   std::vector<double> &      momentum,
+                   float                      guidance_scale,
+                   int                        T,
+                   int                        Oc,
+                   int                        N);
+
 struct DitFlowEditCondition {
     const float * context_latents = nullptr;
     const float * enc_hidden      = nullptr;

engines/audiogen/src/acestep/dit_gguf.cpp

@@ -101,6 +101,12 @@ bool dit_gguf_bool(const DitGGUF & g, const std::string & key, bool def) {
     return gguf_get_val_bool(g.ctx, id);
 }
 
+std::string dit_gguf_str(const DitGGUF & g, const std::string & key, const std::string & def) {
+    int64_t id = gguf_find_key(g.ctx, key.c_str());
+    if (id < 0) return def;
+    return gguf_get_val_str(g.ctx, id);
+}
+
 bool dit_gguf_read_config(const DitGGUF & g, DitConfig & cfg) {
     try {
         cfg.n_layers          = (int) dit_gguf_u32(g, "acestep-dit.block_count");
@@ -117,6 +123,8 @@ bool dit_gguf_read_config(const DitGGUF & g, DitConfig & cfg) {
         cfg.rms_norm_eps      = dit_gguf_f32(g, "acestep-dit.attention.layer_norm_rms_epsilon");
         // convert.py only writes acestep.is_turbo when true; absent => base/sft.
         cfg.is_turbo          = dit_gguf_bool(g, "acestep.is_turbo", false);
+        cfg.model_name        = dit_gguf_str(g, "general.name", "");
+        cfg.is_sft            = is_sft_model_name(cfg.model_name);
     } catch (const std::exception & e) {
         fprintf(stderr, "[acestep-dit] %s\n", e.what());
         return false;

engines/audiogen/src/acestep/dit_gguf.h

@@ -95,6 +95,7 @@ inline size_t dit_gguf_mapped_bytes(ggml_context * ctx, const DitGGUF & g) {
 uint32_t dit_gguf_u32(const DitGGUF & g, const std::string & key);
 float    dit_gguf_f32(const DitGGUF & g, const std::string & key);
 bool     dit_gguf_bool(const DitGGUF & g, const std::string & key, bool def);  // def if key missing
+std::string dit_gguf_str(const DitGGUF & g, const std::string & key, const std::string & def);
 
 // Read the full DiT config from metadata. Returns false if any key is missing.
 bool dit_gguf_read_config(const DitGGUF & g, DitConfig & cfg);