Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
freddy311082 wants to merge 5 commits into
base: master
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
3dfdf9a
QVAC-23944 audiogen: add ACE-Step Multi-Track (lego) task with base-m…
freddy311082 Aug 21, 2026
1aabecc
QVAC-23944 fix[audiogen]: address lego review findings
freddy311082 Aug 21, 2026
babc013
QVAC-23944 test[audiogen]: unit-cover sft detection from the GGUF name
freddy311082 Aug 21, 2026
3ec3007
Merge origin/master into QVAC-23944/multi-track-lego
freddy311082 Aug 21, 2026
95758aa
QVAC-23944 chore[audiogen]: drop committed test output artifacts
freddy311082 Aug 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions engines/audiogen/include/audiogen-cpp/acestep/engine.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand All @@ -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.
Expand Down
133 changes: 129 additions & 4 deletions engines/audiogen/src/acestep/dit_ggml.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -830,6 +830,110 @@ static void preserve_repaint_latent(const DitSampleParams & params, size_t laten
(size_t) params.T, params.repaint_crossfade_frames, channels);
}

// APG (Adaptive Projected Guidance), official ACE-Step base/sft CFG combine.
// Math in double precision; norms and projection run per channel over the
// temporal axis (upstream apg_forward with dims=[1] on [B, T, C]).
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);
}
}

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) {
const size_t n_per = (size_t) T * Oc;
std::vector<double> diff(n_per * N);
for (size_t i = 0; i < diff.size(); i++) {
diff[i] = (double) velocity[i] - (double) velocity_uncond[i];
}
apg_accumulate_momentum(momentum, diff);
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);
}
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]);
}
}

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)
Expand Down Expand Up @@ -895,6 +999,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];
Expand All @@ -916,10 +1032,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();
Expand All @@ -929,6 +1043,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;
Expand Down
19 changes: 18 additions & 1 deletion engines/audiogen/src/acestep/dit_ggml.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,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
Expand All @@ -105,6 +109,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
Expand All @@ -127,6 +133,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;
Expand Down
53 changes: 47 additions & 6 deletions engines/audiogen/src/acestep/engine.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,19 @@ static const char * DIT_INSTR_TEXT2MUSIC = "Fill the audio semantic mask based o
static const char * DIT_INSTR_COVER = "Generate audio semantic tokens based on the given conditions:";
static const char * DIT_INSTR_REPAINT = "Repaint the mask area based on the given conditions:";

static std::string uppercase_track_name(const std::string & track) {
std::string upper = track;
std::transform(upper.begin(), upper.end(), upper.begin(),
[](unsigned char ch) { return (char) std::toupper(ch); });
return upper;
}

// Lego instruction, uppercase track per the reference implementation
// (task_utils.py formats TASK_INSTRUCTIONS["lego"] with track_name.upper()).
static std::string make_lego_instruction(const std::string & track) {
return "Generate the " + uppercase_track_name(track) + " track based on the audio context:";
}

namespace fs = std::filesystem;

struct Engine::Impl {
Expand Down Expand Up @@ -484,6 +497,10 @@ static constexpr int TURBO_STEPS = 8;
static constexpr int STANDARD_STEPS = 50;
static constexpr float TURBO_SHIFT = 3.0f;
static constexpr float STANDARD_SHIFT = 1.0f;
static constexpr float TURBO_GUIDANCE = 1.0f;
static constexpr float STANDARD_GUIDANCE = 7.0f;
static constexpr const char * LEGO_ERROR_TURBO =
"acestep engine: task 'lego' requires a base/sft DiT (turbo does not support stem tasks)";
static constexpr int DIT_BATCH_SIZE = 1;
static constexpr int EDIT_CONTEXT_PLANES = 2;
static constexpr int EDIT_NO_SOURCE_FRAMES = 0;
Expand Down Expand Up @@ -684,6 +701,7 @@ struct PromptEncoding {
struct EncoderConditioning {
std::vector<float> context;
std::vector<float> hidden;
std::vector<float> null_emb;
int frames = 0;
int context_channels = 0;
int sequence = 0;
Expand All @@ -704,9 +722,9 @@ static GenerationState make_generation_state(const GenerateParams & params, bool
}
state.plan = make_generation_plan(params, state.task);
state.seed = resolve_seed(params.seed);
const bool language_neutral = !params.edit_plan.empty() || is_lego_task(state.task.type);
state.language = params.vocal_language.empty()
? (params.edit_plan.empty() ? DEFAULT_VOCAL_LANGUAGE
: EDIT_VOCAL_LANGUAGE)
? (language_neutral ? EDIT_VOCAL_LANGUAGE : DEFAULT_VOCAL_LANGUAGE)
: params.vocal_language;
if (params.augment_caption_with_metadata) {
state.original_caption = params.caption;
Expand Down Expand Up @@ -941,11 +959,23 @@ static void encode_cross_attention(EngineImpl & engine, const PromptEncoding & p
dump.write("06_enc_hidden", output.hidden, output.sequence, output.hidden_size);
}

static void validate_lego_model(const DitConfig & config, const GenerateTask & task) {
if (is_lego_task(task.type) && config.is_turbo) {
throw std::invalid_argument(LEGO_ERROR_TURBO);
}
}

static std::string resolve_dit_instruction(const GenerateTask & task) {
if (is_lego_task(task.type)) return make_lego_instruction(task.track);
return DIT_INSTR_COVER;
}

template <typename EngineImpl>
static EncoderConditioning prepare_encoder_conditioning(EngineImpl & engine,
GenerationState & state,
StageDump & dump, StageTimes & timing) {
const DitConfig & config = engine.dit_cfg;
validate_lego_model(config, state.task);
const int patch = config.patch_size;
EncoderConditioning output;
output.frames = ((state.latent_frames + patch - 1) / patch) * patch;
Expand All @@ -955,8 +985,9 @@ static EncoderConditioning prepare_encoder_conditioning(EngineImpl & engine,
output.context = make_dit_context(
state.context_latents, cond_model_silence_latent(engine.cond), output.frames,
state.latent_frames, output.context_channels, config.out_channels);
output.null_emb = cond_model_null_emb(engine.cond);
const PromptTokens tokens = tokenize_prompt(
engine.bpe_text, state.prompt, state.language, DIT_INSTR_COVER);
engine.bpe_text, state.prompt, state.language, resolve_dit_instruction(state.task).c_str());
const PromptEncoding prompt = encode_prompt(engine, tokens, state, dump, timing);
encode_cross_attention(engine, prompt, state, output, dump, timing);
return output;
Expand Down Expand Up @@ -1013,15 +1044,21 @@ static void dump_parity_inputs(const std::vector<float> & latent, const NoiseSch
}
#endif

static float resolve_guidance_scale(const GenerateParams & params, const DitConfig & config) {
if (params.guidance_scale > 0.0f) return params.guidance_scale;
return config.is_turbo ? TURBO_GUIDANCE : STANDARD_GUIDANCE;
}

template <typename EngineImpl>
static bool sample_dit_latent(EngineImpl & engine, const GenerateParams & params,
const GenerationState & state, EncoderConditioning & conditioning,
NoiseSchedule & noise, const StageReporter & report,
StageDump & dump, StageTimes & timing, std::vector<float> & latent) {
const DitConfig & config = engine.dit_cfg;
const float guidance = resolve_guidance_scale(params, config);
if (engine.opts.verbose) {
fprintf(stderr, "[acestep-engine] DiT: turbo=%d steps=%d shift=%.2f T=%d task=%s\n",
(int) config.is_turbo, noise.steps, noise.shift,
fprintf(stderr, "[acestep-engine] DiT: turbo=%d steps=%d shift=%.2f guidance=%.2f T=%d task=%s\n",
(int) config.is_turbo, noise.steps, noise.shift, guidance,
conditioning.frames, state.task.type.c_str());
}
engine.ensure_dit();
Expand All @@ -1038,7 +1075,11 @@ static bool sample_dit_latent(EngineImpl & engine, const GenerateParams & params
sample.schedule = noise.schedule.data();
sample.num_steps = noise.steps;
sample.real_enc_S = &conditioning.sequence;
sample.dcw_enabled = params.dcw_enabled;
sample.guidance_scale = guidance;
sample.null_cond_emb = conditioning.null_emb.empty() ? nullptr : conditioning.null_emb.data();
// DCW is a turbo-preset correction: the official UI disables it for
// base/sft models, and base-model quality is validated without it.
sample.dcw_enabled = params.dcw_enabled && config.is_turbo;
sample.dcw_scaler = params.dcw_scaler;
sample.dcw_high_scaler = params.dcw_high_scaler;
sample.on_step = [&](int step, int total) { return report("dit", step, total); };
Expand Down
Loading
Loading