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

[freddy311082]

Summary

Adds the ACE-Step Multi-Track (lego) task to the audiogen engine: given a source song and a target track name, the DiT generates a new isolated instrument layer that follows the source (tempo, key, groove), returned as a separate stem for mixing.

This required unlocking the base/sft DiT sampling path, which was deferred until now: classifier-free guidance via APG (Adaptive Projected Guidance) with the null condition embedding, momentum -0.75, per-channel norm clipping at 2.5, and orthogonal projection — matching the official implementation.

Changes

• dit_ggml: per-step unconditional forward + dit_apg_guide combine; guidance_scale and null_cond_emb sampling params.
• engine: task_type: "lego" with the official per-track instruction (uppercase track name), base/sft enforcement (turbo rejected with a clear error), guidance auto-resolution (turbo 1.0 / base 7.0), Haar DCW applied on turbo only (official base preset disables it), and lyric-stream language defaulting to "unknown" for lego (reference parity for stem tasks).
• API: new GenerateParams::guidance_scale and GenerateParams::track; 12 official track names validated.
• music-cli: --track, --guidance, and request JSON fields.
• Tests (all in-repo): lego task validation and generation-plan units, APG golden-value unit (hand-derived from the reference math), and an integration scenario verifying the DiT context dump (full 1.0 mask plane + source acoustic plane) with a graceful skip on turbo fixtures.

Validation

• 881/881 unit checks pass; full non-GPU ctest suite green.
• Numerical parity against the validated acestep.cpp prototype on identical inputs: conditioning tensors bit-identical, final latent correlation 0.999 after 50 CFG steps, audio correlation 0.996.
• Full-length (3:54) base-model lego generation validated by listening against the prototype's approved output (waveform correlation 0.91).

🤖 Generated with Claude Code

[github-actions]

Review Status

Current Status: ❌ PENDING
Approvals so far: none

Pending reviews: Needs 1 Management or Team Lead, and 1 more from Management, Team Lead, or Member.

[GustavoA1604]

[critical] engines/audiogen/src/acestep/engine.cpp — line 963 (if (is_lego_task(task.type) && config.is_turbo))

This permits LEGO generation with SFT weights, but ACE-Step defines LEGO as a Base-only task; SFT supports CFG but not the Base-exclusive extract/lego/complete tasks. Require an actual Base model identity rather than merely !is_turbo, and add an SFT rejection test.

[medium] engines/audiogen/test/test_acestep_integration.cpp — line 466 (CHECK(result.pcm.size() >= params.source_audio.size());)

LEGO promises a stem aligned to the source, but source latent frames are rounded to the DiT patch size and the decoded PCM is never trimmed. Arbitrary WAV lengths can therefore produce a longer stem that cannot be mixed sample-for-sample. Preserve the original source sample count, trim the result to it, and assert equality here.

[medium] engines/audiogen/src/acestep/engine.cpp — line 1048 (if (params.guidance_scale > 0.0f) return params.guidance_scale;)

An explicit guidance_scale > 1 enables APG on Turbo even though Turbo is guidance-distilled and CFG is unsupported. Force guidance to 1.0 for Turbo or reject such an override; otherwise --guidance 7 can run an invalid second forward and degrade output.

[medium] engines/audiogen/test/test_acestep_integration.cpp — line 468 (catch (const std::invalid_argument & error))

The integration test treats a Turbo fixture as a successful skip, so the green PR checks never execute LEGO generation. Add a Base-model CI lane

[medium][coding-standards] engines/audiogen/src/acestep/engine.cpp — line 1047 (static float resolve_guidance_scale(...))

The new guidance defaults, explicit override, and model-specific DCW gate have no unit coverage; the APG golden test does not exercise this policy. Add tests for Turbo/Base automatic guidance, explicit overrides, and DCW being enabled only for Turbo.

[low] engines/audiogen/README.md — line 92 (DiT CFG and APG ... are deferred)

This is now false, and the README also omits LEGO usage, --track, --guidance, and the new request fields. Update both this engine README and the repository root README as required for new public functionality.

[low][coding-standards] engines/audiogen/src/acestep/dit_ggml.cpp — line 909 (for (size_t i = 0; i < diff.size(); i++))

dit_apg_guide contains three separate loops/actions, contrary to the team rule that each loop be isolated in its own function. Extract the difference calculation, per-batch projection, and guided-update application into named helpers.

[low][coding-standards] engines/audiogen/src/acestep/dit_ggml.cpp — line 833 (// APG (Adaptive Projected Guidance)...)

This newly added explanatory block narrates implementation details, which the team standard prohibits. The helper and constant names already express the behavior; remove the narration or encode any missing intent in names.

[freddy311082]

Addressed the full review in 1aabecc — thanks for the sharp catches, especially the sft gap:

• [critical] sft passed the lego gate — the GGUF's general.name (stamped by convert.py with the checkpoint name) now feeds DitConfig::is_sft; lego requires an actual base DiT and rejects turbo and sft with distinct messages. Unit-covered in test_lego_model_policy.
• [medium] stem length — the decoded stem is trimmed to the source sample count (patch-alignment overshoot removed), and the integration scenario asserts result.pcm.size() == source_audio.size().
• [medium] guidance override on turbo — resolve_guidance_scale clamps any override to 1.0 on turbo (guidance-distilled); base/sft keep explicit overrides and the 7.0 auto default.
• [medium] base CI lane — the integration test now runs a mandatory-execution lego lane when AUDIOGEN_TEST_BASE_MODELS_DIR is set (a rejection there fails instead of skipping). Wiring a base fixture into the CI workflow is left as an infra follow-up.
• [medium] policy unit coverage — test_guidance_and_dcw_policy covers turbo/base auto defaults, explicit overrides, the turbo clamp, and DCW being turbo-only.
• [low] READMEs — engine README drops the stale 'CFG/APG deferred' note and documents lego / --track / --guidance with the per-variant defaults table; the root README model table now lists the base variant with stem support.
• [low] loop-per-function + narration — dit_apg_guide is now composition of apg_velocity_difference / apg_accumulate_momentum / apg_shape_batch_updates / apg_apply_guided_update, and the explanatory block is gone.

894/894 unit checks pass; E2E re-verified locally: turbo rejection message, base run with guidance 7, and stem length equal to the source.

[GustavoA1604]

Remove output files