From 2b38981c4654cc52a102ab020bf3c02efe2a47a1 Mon Sep 17 00:00:00 2001 From: Jeff Stein Date: Fri, 10 Jul 2026 16:46:44 +0000 Subject: [PATCH 01/11] feat: add POST /score endpoint for lyric alignment scoring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1 Add a C++ implementation of lyric alignment scoring that captures cross-attention matrices from the DiT forward pass and computes coverage, monotonicity, and confidence metrics via DTW pathfinding. This eliminates the need for a Python sidecar — the entire scoring pipeline runs in-process with the ace-server binary. Ported from the Python ACE-Step scoring modules (pinned to ace-step/ACE-Step-1.5@82252c24): - acestep/core/scoring/_dtw.py (DTW + median filter) - acestep/core/scoring/dit_score.py (MusicLyricScorer metrics) New files: - src/dtw-score.h: Pure C++ port of DTW, median filter, token type mask generation, attention preprocessing (head selection, averaging, min-max normalization), and the full scoring pipeline (coverage, monotonicity, confidence -> lyrics_score = cov^2 * mono^2 * conf). Zero external dependencies beyond standard library headers. Each function has comments referencing the exact Python source file and line numbers for traceability. - tests/test-dtw-score.cpp: 9 unit tests (58 assertions) covering DTW pathfinding on diagonal/horizontal matrices, median filter spike removal with reflect padding, token type mask generation (bracket detection), full scoring pipeline with perfect/poor/tagged alignment, multi-head attention averaging, and DTW path monotonicity. All pass. - .github/workflows/scoring-drift-check.yml: Weekly CI check that fetches the upstream Python scoring files at the pinned commit and compares SHA256 hashes. Fails with a warning if the Python reference has drifted, signaling that the C++ port needs re-evaluation. Modified files: - src/dit-graph.h: dit_attn_f32() and dit_ggml_build_cross_attn() accept an optional capture_scores parameter that forces the f32 attention path (instead of flash_attn_ext) and marks the softmax attention weights as a graph output via ggml_set_output(). dit_ggml_build_graph() accepts score_layers[] to capture cross-attention from specific layers, naming each captured tensor "cross_attn_scores_L{layer}" for later retrieval. - src/dit-sampler.h: New dit_ggml_score_forward() function that builds a DiT graph with score capture, runs a single forward pass at t=0 (no denoising loop), and reads back the attention matrices for each captured layer. Returns [enc_S, S, Nh] per layer (batch 0 only). - src/pipeline-synth-impl.h: Added per_lyric_ids field to SynthState to persist lyric token IDs from text encoding through to scoring. - src/pipeline-synth-ops.cpp/h: ops_encode_text() now stores lyric token IDs in SynthState. New ops_score_forward() op runs the scoring forward pass, stacks attention from all captured layers into [n_layers, n_heads, enc_S, S] layout, decodes token IDs to strings via the BPE tokenizer for type mask generation, and calls calculate_lyric_score() to compute metrics. - src/pipeline-synth.h/cpp: New ace_synth_score() public API and LyricScoreResult struct (defined in dtw-score.h). Runs the same phase 1 setup as ace_synth_job_run_dit (text encoding, context build, noise init) but calls ops_score_forward() instead of the full denoising loop. Currently supports text2music task only. - tools/ace-server.cpp: New POST /score endpoint with score_worker and handle_score. Accepts the same JSON request format as /synth (plain JSON or multipart), creates an async job, and returns a JSON array of per-request scores: [{"coverage":1.0,"monotonicity":0.95,"confidence":0.24, "lyrics_score":0.21}, ...] Uses the existing job queue system (POST /score -> job ID -> GET /job?id=N -> GET /job?id=N&result=1). - tests/CMakeLists.txt: Added test-dtw-score build target (links libm, no ggml dependency). Verified end-to-end on Strix Halo (Radeon 8060S, Vulkan backend): - 58/58 unit tests pass - ace-server builds and runs with /score endpoint - Single scoring forward pass: ~106ms for 10s audio (T=250, S=125) - Cross-attention captured from 5 layers x 32 heads x [100 x 125] - JSON response correctly returned via job queue Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/scoring-drift-check.yml | 54 +++ src/dit-graph.h | 70 ++- src/dit-sampler.h | 180 ++++++++ src/dtw-score.h | 492 ++++++++++++++++++++++ src/pipeline-synth-impl.h | 5 + src/pipeline-synth-ops.cpp | 140 ++++++ src/pipeline-synth-ops.h | 17 + src/pipeline-synth.cpp | 56 +++ src/pipeline-synth.h | 16 + tests/CMakeLists.txt | 7 + tests/test-dtw-score.cpp | 340 +++++++++++++++ tools/ace-server.cpp | 180 ++++++++ 12 files changed, 1545 insertions(+), 12 deletions(-) create mode 100644 .github/workflows/scoring-drift-check.yml create mode 100644 src/dtw-score.h create mode 100644 tests/test-dtw-score.cpp diff --git a/.github/workflows/scoring-drift-check.yml b/.github/workflows/scoring-drift-check.yml new file mode 100644 index 00000000..ed26da34 --- /dev/null +++ b/.github/workflows/scoring-drift-check.yml @@ -0,0 +1,54 @@ +# Check if the upstream Python scoring reference has drifted from the +# commit pinned in src/dtw-score.h. The C++ port is a snapshot of: +# ace-step/ACE-Step-1.5@82252c24 +# acestep/core/scoring/_dtw.py +# acestep/core/scoring/dit_score.py +# +# If the upstream files change, this workflow fails and the C++ port +# needs re-evaluation. The expected SHA256 hashes are embedded below +# so there's no separate file to maintain — update both hashes and +# the pin comment in dtw-score.h when re-syncing. +name: Scoring Drift Check + +on: + schedule: + - cron: '0 8 * * 1' # weekly: Monday 08:00 UTC + workflow_dispatch: + +jobs: + check: + runs-on: ubuntu-latest + steps: + - name: Fetch upstream Python scoring files + run: | + BASE="https://raw.githubusercontent.com/ace-step/ACE-Step-1.5/82252c24" + curl -fsSL "$BASE/acestep/core/scoring/_dtw.py" -o _dtw.py + curl -fsSL "$BASE/acestep/core/scoring/dit_score.py" -o dit_score.py + + - name: Verify hashes match pinned snapshot + run: | + # Expected SHA256 of the Python files at commit 82252c24. + # Update these (and the pin comment in src/dtw-score.h) when + # re-syncing the C++ port after an upstream algorithm change. + EXPECTED_DTW="2d2252e7108f296cd26722c7b622e1c1a1ed47c71459a8ff873164c4c1ada962" + EXPECTED_SCORE="73700ff976549339dad3d6c500f8cbe9a8cce08d1979bbb85e3e3ed53e0605a1" + + ACTUAL_DTW=$(sha256sum _dtw.py | cut -d' ' -f1) + ACTUAL_SCORE=$(sha256sum dit_score.py | cut -d' ' -f1) + + echo "_dtw.py: expected=$EXPECTED_DTW actual=$ACTUAL_DTW" + echo "dit_score.py: expected=$EXPECTED_SCORE actual=$ACTUAL_SCORE" + + if [ "$ACTUAL_DTW" != "$EXPECTED_DTW" ]; then + echo "::warning::_dtw.py has drifted from the pinned commit (82252c24)." + echo "::warning::The C++ port in src/dtw-score.h may need updating." + exit 1 + fi + + if [ "$ACTUAL_SCORE" != "$EXPECTED_SCORE" ]; then + echo "::warning::dit_score.py has drifted from the pinned commit (82252c24)." + echo "::warning::The C++ port in src/dtw-score.h may need updating." + exit 1 + fi + + echo "No drift detected. C++ port matches pinned Python reference." diff --git a/src/dit-graph.h b/src/dit-graph.h index ef537266..ef9e2630 100644 --- a/src/dit-graph.h +++ b/src/dit-graph.h @@ -120,16 +120,26 @@ static struct ggml_tensor * dit_ggml_build_temb(struct ggml_context * ctx, // Q: [D, S, Nh], K: [D, S_kv, Nkv], V: [D, S_kv, Nkv] // mask: [S_kv, S] F16 or NULL, scale: 1/sqrt(D) // Returns: [D, Nh, S] (same layout as flash_attn_ext output) +// +// When capture_scores is non-NULL, the softmax scores tensor (attention weights) +// is marked as a graph output and stored at *capture_scores. Shape: [S_kv, S, Nh, N]. +// This is used by the /score endpoint to extract cross-attention matrices for +// lyric alignment scoring. Ported from Python ACE-Step output_attentions=True. static struct ggml_tensor * dit_attn_f32(struct ggml_context * ctx, struct ggml_tensor * q, struct ggml_tensor * k, struct ggml_tensor * v, struct ggml_tensor * mask, - float scale) { + float scale, + struct ggml_tensor ** capture_scores = nullptr) { struct ggml_tensor * scores = ggml_mul_mat(ctx, k, q); scores = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f); - struct ggml_tensor * vt = ggml_cont(ctx, ggml_transpose(ctx, v)); - struct ggml_tensor * out = ggml_mul_mat(ctx, vt, scores); + if (capture_scores) { + *capture_scores = scores; + ggml_set_output(scores); + } + struct ggml_tensor * vt = ggml_cont(ctx, ggml_transpose(ctx, v)); + struct ggml_tensor * out = ggml_mul_mat(ctx, vt, scores); return ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3)); } @@ -271,6 +281,9 @@ static struct ggml_tensor * dit_ggml_build_mlp(struct ggml_context * ctx, // Build cross-attention sub-graph for a single layer. // norm_ca: [H, S, N] pre-normalized hidden state (Q source) // enc: [H, enc_S, N] condition-embedded encoder states (K/V source) +// When capture_scores is non-NULL, forces the f32 attention path and stores +// the softmax attention weights tensor at *capture_scores (marked as graph +// output). Shape: [enc_S, S, Nh, N]. Used by /score for lyric alignment. // Returns: output [H, S, N] (NOT added to residual yet) static struct ggml_tensor * dit_ggml_build_cross_attn(struct ggml_context * ctx, DiTGGML * m, @@ -281,7 +294,8 @@ static struct ggml_tensor * dit_ggml_build_cross_attn(struct ggml_context * ctx, struct ggml_tensor * mask, // [enc_S, S, 1, N] F16 or NULL int S, int enc_S, - int N) { + int N, + struct ggml_tensor ** capture_scores = nullptr) { DiTGGMLConfig & c = m->cfg; int D = c.head_dim; int Nh = c.n_heads; @@ -334,9 +348,13 @@ static struct ggml_tensor * dit_ggml_build_cross_attn(struct ggml_context * ctx, // mask blocks padding positions in encoder hidden states float scale = 1.0f / sqrtf((float) D); + // When capturing scores, force the f32 path (flash_attn_ext does not + // materialize the attention weights as a separate tensor). + bool use_fa = m->use_flash_attn && !capture_scores; + // K/V come in F32 from mul_mat (no KV cache here). Cast to F16 before FA, // mirroring llama.cpp build_attn_mha for graphs without a KV cache. - if (m->use_flash_attn) { + if (use_fa) { if (k->type == GGML_TYPE_F32) { k = ggml_cast(ctx, k, GGML_TYPE_F16); } @@ -345,9 +363,9 @@ static struct ggml_tensor * dit_ggml_build_cross_attn(struct ggml_context * ctx, } } - struct ggml_tensor * attn = m->use_flash_attn ? ggml_flash_attn_ext(ctx, q, k, v, mask, scale, 0.0f, 0.0f) : - dit_attn_f32(ctx, q, k, v, mask, scale); - if (m->use_flash_attn) { + struct ggml_tensor * attn = use_fa ? ggml_flash_attn_ext(ctx, q, k, v, mask, scale, 0.0f, 0.0f) : + dit_attn_f32(ctx, q, k, v, mask, scale, capture_scores); + if (use_fa) { ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32); } @@ -375,7 +393,8 @@ static struct ggml_tensor * dit_ggml_build_layer(struct ggml_context * ctx, struct ggml_tensor * ca_mask, // [enc_S, S, 1, N] or NULL int S, int enc_S, - int N) { + int N, + struct ggml_tensor ** capture_scores = nullptr) { DiTGGMLConfig & c = m->cfg; DiTGGMLLayer * ly = &m->layers[layer_idx]; int H = c.hidden_size; @@ -428,7 +447,7 @@ static struct ggml_tensor * dit_ggml_build_layer(struct ggml_context * ctx, if (enc) { struct ggml_tensor * norm_ca = dit_ggml_rms_norm_weighted(ctx, hidden, ly->cross_attn_norm, c.rms_norm_eps); struct ggml_tensor * ca_out = - dit_ggml_build_cross_attn(ctx, m, ly, norm_ca, enc, positions, ca_mask, S, enc_S, N); + dit_ggml_build_cross_attn(ctx, m, ly, norm_ca, enc, positions, ca_mask, S, enc_S, N, capture_scores); hidden = ggml_add(ctx, hidden, ca_out); } @@ -462,13 +481,20 @@ static struct ggml_tensor * dit_ggml_build_layer(struct ggml_context * ctx, // // Graph outputs: // "velocity" [out_channels, T, N] predicted flow velocity +// +// When score_layers is non-NULL, cross-attention scores are captured for the +// specified layers (forced f32 path). Each captured layer produces a named +// output tensor "cross_attn_scores_L{layer}" of shape [enc_S, S, Nh, N]. +// These are retrieved after graph compute for lyric alignment scoring. static struct ggml_cgraph * dit_ggml_build_graph(DiTGGML * m, struct ggml_context * ctx, int T, // temporal length (before patching) int enc_S, // encoder sequence length int N, // batch size struct ggml_tensor ** p_input, // [out] input tensor to fill - struct ggml_tensor ** p_output) { // [out] output tensor to read + struct ggml_tensor ** p_output, // [out] output tensor to read + const int * score_layers = nullptr, // layers to capture, or NULL + int n_score_layers = 0) { DiTGGMLConfig & c = m->cfg; int S = T / c.patch_size; // sequence length after patching @@ -573,7 +599,27 @@ static struct ggml_cgraph * dit_ggml_build_graph(DiTGGML * m, for (int i = 0; i < c.n_layers; i++) { // layer_type=0 (sliding window): sa_mask_sw, layer_type=1 (full): unmasked struct ggml_tensor * sa_mask = (m->layers[i].layer_type == 0) ? sa_mask_sw : nullptr; - hidden = dit_ggml_build_layer(ctx, m, i, hidden, tproj, enc, positions, sa_mask, ca_mask, S, enc_S, N); + + // Check if this layer should capture cross-attention scores + struct ggml_tensor * capture = nullptr; + if (score_layers) { + for (int sl = 0; sl < n_score_layers; sl++) { + if (score_layers[sl] == i) { + capture = (struct ggml_tensor *) 1; // sentinel: non-NULL triggers capture + break; + } + } + } + + hidden = dit_ggml_build_layer(ctx, m, i, hidden, tproj, enc, positions, sa_mask, ca_mask, S, enc_S, N, + capture ? &capture : nullptr); + + // Name captured cross-attention scores for later retrieval + if (capture && capture != (struct ggml_tensor *) 1) { + char score_name[64]; + snprintf(score_name, sizeof(score_name), "cross_attn_scores_L%d", i); + ggml_set_name(capture, score_name); + } // Debug dumps at key layers: 0, 6, 12, 18, last if (i == 0 || i == 6 || i == 12 || i == 18 || i == c.n_layers - 1) { char lname[64]; diff --git a/src/dit-sampler.h b/src/dit-sampler.h index 88b49995..ee33ec2c 100644 --- a/src/dit-sampler.h +++ b/src/dit-sampler.h @@ -696,3 +696,183 @@ static int dit_ggml_generate(DiTGGML * model, ggml_free(ctx); return 0; } + +// ============================================================================ +// Scoring forward pass: single DiT forward with cross-attention capture. +// +// Matches Python ACE-Step lyric_score.py:get_lyric_score() which runs a +// single DiT forward with output_attentions=True to extract cross-attention +// matrices for lyric alignment scoring. +// +// Builds the same DiT graph as generation but forces the f32 attention path +// for the specified score_layers, marking the softmax attention weights as +// graph outputs. Runs one forward pass at t=0 (pure noise input) and reads +// back the cross-attention scores. +// +// The caller is responsible for: +// - text encoding (enc_hidden, enc_S, real_enc_S) +// - context build (context_latents) +// - noise generation (noise) +// - running calculate_lyric_score() on the returned attention matrices +// +// attention_out: filled with [n_score_layers][Nh][enc_S][S] attention weights +// (one vector per captured layer, each enc_S*S*Nh floats, row-major). +// Only the conditional (first) batch sample is extracted. +// Returns 0 on success, -1 on error. +static int dit_ggml_score_forward(DiTGGML * model, + const float * noise, + const float * context_latents, + const float * enc_hidden_data, + int enc_S, + const int * real_enc_S, + int T, + int N, + const int * score_layers, + int n_score_layers, + std::vector> & attention_out) { + DiTGGMLConfig & c = model->cfg; + int Oc = c.out_channels; + int ctx_ch = c.in_channels - Oc; + int in_ch = c.in_channels; + int S = T / c.patch_size; + int n_per = T * Oc; + + // Scoring uses a single forward pass (no CFG, no sampling loop). + // The Python reference runs batch=2 (noise + partial) but we only need + // the conditional pass — the attention patterns from the conditional + // branch are what the scorer consumes. + int N_graph = N; + + fprintf(stderr, "[DiT-Score] Forward: N=%d, T=%d, S=%d, enc_S=%d, %d score layers\n", + N, T, S, enc_S, n_score_layers); + + // Graph context + size_t ctx_size = ggml_tensor_overhead() * 8192 + ggml_graph_overhead_custom(8192, false); + std::vector ctx_buf(ctx_size); + struct ggml_init_params gparams = { + /*.mem_size =*/ctx_size, + /*.mem_buffer =*/ctx_buf.data(), + /*.no_alloc =*/true, + }; + struct ggml_context * ctx = ggml_init(gparams); + + struct ggml_tensor * t_input = NULL; + struct ggml_tensor * t_output = NULL; + struct ggml_cgraph * gf = dit_ggml_build_graph(model, ctx, T, enc_S, N_graph, &t_input, &t_output, + score_layers, n_score_layers); + + fprintf(stderr, "[DiT-Score] Graph: %d nodes\n", ggml_graph_n_nodes(gf)); + + struct ggml_tensor * t_enc = ggml_graph_get_tensor(gf, "enc_hidden"); + int H_enc = (int) t_enc->ne[0]; + + // Allocate compute buffers + ggml_backend_sched_reset(model->sched); + if (model->backend != model->cpu_backend) { + const char * input_names[] = { + "enc_hidden", "input_latents", "t", "t_r", "positions", "sa_mask_sw", "ca_mask" + }; + for (const char * iname : input_names) { + struct ggml_tensor * t = ggml_graph_get_tensor(gf, iname); + if (t) { + ggml_backend_sched_set_tensor_backend(model->sched, t, model->backend); + } + } + } + if (!ggml_backend_sched_alloc_graph(model->sched, gf)) { + fprintf(stderr, "[DiT-Score] FATAL: failed to allocate graph\n"); + ggml_free(ctx); + return -1; + } + + // Set timesteps to 0 (scoring uses the final denoised state) + struct ggml_tensor * t_t = ggml_graph_get_tensor(gf, "t"); + struct ggml_tensor * t_tr = ggml_graph_get_tensor(gf, "t_r"); + float t_zero = 0.0f; + ggml_backend_tensor_set(t_t, &t_zero, 0, sizeof(float)); + ggml_backend_tensor_set(t_tr, &t_zero, 0, sizeof(float)); + + // Positions + struct ggml_tensor * t_pos = ggml_graph_get_tensor(gf, "positions"); + std::vector pos_data(S * N_graph); + for (int b = 0; b < N_graph; b++) { + for (int i = 0; i < S; i++) { + pos_data[b * S + i] = i; + } + } + ggml_backend_tensor_set(t_pos, pos_data.data(), 0, S * N_graph * sizeof(int32_t)); + + // Self-attention mask (full attention for scoring — no sliding window) + struct ggml_tensor * t_sa_mask_sw = ggml_graph_get_tensor(gf, "sa_mask_sw"); + int win = c.sliding_window; + std::vector sa_sw_data(S * S * N_graph); + for (int b = 0; b < N_graph; b++) { + for (int qi = 0; qi < S; qi++) { + for (int ki = 0; ki < S; ki++) { + int dist = (qi > ki) ? (qi - ki) : (ki - qi); + bool in_win = (win <= 0) || (S <= win) || (dist <= win); + sa_sw_data[b * S * S + qi * S + ki] = ggml_fp32_to_fp16(in_win ? 0.0f : -INFINITY); + } + } + } + ggml_backend_tensor_set(t_sa_mask_sw, sa_sw_data.data(), 0, S * S * N_graph * sizeof(uint16_t)); + + // Cross-attention mask + struct ggml_tensor * t_ca_mask = ggml_graph_get_tensor(gf, "ca_mask"); + std::vector ca_data(enc_S * S * N_graph); + for (int b = 0; b < N; b++) { + int re = real_enc_S ? real_enc_S[b] : enc_S; + for (int qi = 0; qi < S; qi++) { + for (int ki = 0; ki < enc_S; ki++) { + float v = (ki < re) ? 0.0f : -INFINITY; + ca_data[b * enc_S * S + qi * enc_S + ki] = ggml_fp32_to_fp16(v); + } + } + } + ggml_backend_tensor_set(t_ca_mask, ca_data.data(), 0, enc_S * S * N_graph * sizeof(uint16_t)); + + // Encoder hidden states + ggml_backend_tensor_set(t_enc, enc_hidden_data, 0, H_enc * enc_S * N * sizeof(float)); + + // Input: context_latents + noise + std::vector input_buf(in_ch * T * N_graph); + for (int b = 0; b < N; b++) { + for (int t = 0; t < T; t++) { + memcpy(&input_buf[b * T * in_ch + t * in_ch], &context_latents[b * T * ctx_ch + t * ctx_ch], + ctx_ch * sizeof(float)); + memcpy(&input_buf[b * T * in_ch + t * in_ch + ctx_ch], &noise[b * n_per + t * Oc], Oc * sizeof(float)); + } + } + ggml_backend_tensor_set(t_input, input_buf.data(), 0, in_ch * T * N_graph * sizeof(float)); + + // Run forward pass + ggml_backend_sched_graph_compute(model->sched, gf); + + // Read back cross-attention scores for each captured layer + attention_out.resize(n_score_layers); + for (int sl = 0; sl < n_score_layers; sl++) { + char score_name[64]; + snprintf(score_name, sizeof(score_name), "cross_attn_scores_L%d", score_layers[sl]); + struct ggml_tensor * t_scores = ggml_graph_get_tensor(gf, score_name); + if (!t_scores) { + fprintf(stderr, "[DiT-Score] WARNING: could not find tensor %s\n", score_name); + continue; + } + + // Scores shape: [enc_S, S, Nh, N] (ne[0]=enc_S, ne[1]=S, ne[2]=Nh, ne[3]=N) + // Extract only batch 0: [enc_S, S, Nh] = enc_S * S * Nh floats + int layer_enc_S = (int) t_scores->ne[0]; + int layer_S = (int) t_scores->ne[1]; + int layer_Nh = (int) t_scores->ne[2]; + size_t batch0_elems = (size_t) layer_enc_S * layer_S * layer_Nh; + + attention_out[sl].resize(batch0_elems); + ggml_backend_tensor_get(t_scores, attention_out[sl].data(), 0, batch0_elems * sizeof(float)); + + fprintf(stderr, "[DiT-Score] Layer %d: %d heads, [%d x %d] attention matrix\n", + score_layers[sl], layer_Nh, layer_enc_S, layer_S); + } + + ggml_free(ctx); + return 0; +} diff --git a/src/dtw-score.h b/src/dtw-score.h new file mode 100644 index 00000000..5b867c43 --- /dev/null +++ b/src/dtw-score.h @@ -0,0 +1,492 @@ +#pragma once +// dtw-score.h: DTW pathfinding + lyric alignment scoring (pure C++) +// +// Direct port of the Python ACE-Step scoring modules: +// acestep/core/scoring/_dtw.py — DTW + median filter (numba-jitted) +// acestep/core/scoring/dit_score.py — MusicLyricScorer (coverage, monotonicity, confidence) +// +// Pinned to ace-step/ACE-Step-1.5 commit 82252c24 (2026-07-09). +// If the Python scoring algorithm changes upstream, this port will need +// re-evaluation. The file:line references in comments below map to that commit. +// +// No external dependencies beyond , , , . +// All compute is CPU-side on small matrices (tokens x frames), matching the +// Python reference which forces CPU ("the scoring matrices are small and this +// avoids occupying GPU VRAM that DiT / VAE / LM need" — dit_score.py:294). + +#include +#include +#include +#include +#include +#include +#include +#include + +// ============================================================================ +// DTW (Dynamic Time Warping) +// Ported from acestep/core/scoring/_dtw.py:dtw_cpu + _backtrace +// ============================================================================ + +// DTW backtrace: walk the trace matrix from (N, M) to (0, 0). +// Returns path as two parallel vectors: text_indices[i], time_indices[i]. +// Python returns shape (2, path_len); here we return a struct for clarity. +// +// trace values: 0 = diagonal, 1 = up (i-1), 2 = left (j-1) +struct DTWPath { + std::vector text_idx; // row index into cost matrix (lyric tokens) + std::vector time_idx; // col index into cost matrix (audio frames) +}; + +static DTWPath dtw_backtrace(std::vector & trace, int N, int M) { + // Boundary handling (matches _dtw.py:61-62) + // trace is (N+1) x (M+1), row-major + auto T = [&](int i, int j) -> float & { return trace[(size_t) i * (M + 1) + j]; }; + for (int j = 0; j <= M; j++) { + T(0, j) = 2; + } + for (int i = 0; i <= N; i++) { + T(i, 0) = 1; + } + + // Pre-allocate (max path length = N + M), fill from the end + int max_path_len = N + M; + std::vector path_text(max_path_len, 0); + std::vector path_time(max_path_len, 0); + + int i = N, j = M; + int path_idx = max_path_len - 1; + + while (i > 0 || j > 0) { + path_text[path_idx] = i - 1; // text index + path_time[path_idx] = j - 1; // time index + path_idx--; + + float t = T(i, j); + if (t == 0) { + i--; + j--; + } else if (t == 1) { + i--; + } else if (t == 2) { + j--; + } else { + break; + } + } + + int start = path_idx + 1; + DTWPath result; + result.text_idx.assign(path_text.begin() + start, path_text.end()); + result.time_idx.assign(path_time.begin() + start, path_time.end()); + return result; +} + +// DTW forward: compute cost matrix and backtrace the optimal path. +// x: cost matrix of shape [N, M] (row-major: x[i * M + j]). +// Returns the alignment path. +// +// Ported from acestep/core/scoring/_dtw.py:dtw_cpu (numba-jitted). +// The caller passes -calc_matrix (negated) so DTW finds the maximum-energy +// path (matches dit_score.py:255: dtw_cpu(-calc_matrix.astype(np.float32))). +static DTWPath dtw_cpu(const float * x, int N, int M) { + // cost and trace are (N+1) x (M+1), row-major + std::vector cost((size_t) (N + 1) * (M + 1), std::numeric_limits::infinity()); + std::vector trace((size_t) (N + 1) * (M + 1), -1.0f); + + auto C = [&](int i, int j) -> float & { return cost[(size_t) i * (M + 1) + j]; }; + + C(0, 0) = 0.0f; + + for (int j = 1; j <= M; j++) { + for (int i = 1; i <= N; i++) { + float c0 = C(i - 1, j - 1); // diagonal + float c1 = C(i - 1, j); // up + float c2 = C(i, j - 1); // left + + float c; + float t; + if (c0 < c1 && c0 < c2) { + c = c0; + t = 0; + } else if (c1 < c0 && c1 < c2) { + c = c1; + t = 1; + } else { + c = c2; + t = 2; + } + + C(i, j) = x[(size_t) (i - 1) * M + (j - 1)] + c; + trace[(size_t) i * (M + 1) + j] = t; + } + } + + return dtw_backtrace(trace, N, M); +} + +// ============================================================================ +// Median Filter +// Ported from acestep/core/scoring/_dtw.py:median_filter +// ============================================================================ + +// 1D median filter with reflect padding. +// x: input data, filter_width: window size (must be odd for symmetric padding). +// Returns filtered vector of the same length. +// +// Python uses F.pad(mode="reflect") then unfold + sort + pick middle. +// We replicate with a sliding window sort. +static std::vector median_filter_1d(const std::vector & x, int filter_width) { + int n = (int) x.size(); + int pad = filter_width / 2; + if (n <= pad) { + return x; + } + + // Reflect padding (matches torch F.pad mode="reflect") + std::vector padded(n + 2 * pad); + for (int i = 0; i < pad; i++) { + padded[i] = x[pad - i]; // reflect from left + } + for (int i = 0; i < n; i++) { + padded[pad + i] = x[i]; + } + for (int i = 0; i < pad; i++) { + padded[pad + n + i] = x[n - 2 - i]; // reflect from right + } + + std::vector result(n); + std::vector window(filter_width); + int mid = filter_width / 2; + + for (int i = 0; i < n; i++) { + for (int j = 0; j < filter_width; j++) { + window[j] = padded[i + j]; + } + std::nth_element(window.begin(), window.begin() + mid, window.end()); + result[i] = window[mid]; + } + return result; +} + +// 2D median filter applied along the last dimension (columns / frames). +// input: [rows, cols] row-major. filter_width applied along cols. +// Returns [rows, cols] row-major. +// +// Matches median_filter in _dtw.py:90-110 which operates on the last dim. +static std::vector median_filter_2d_rows(const std::vector & input, int rows, int cols, int filter_width) { + std::vector result((size_t) rows * cols); + for (int r = 0; r < rows; r++) { + std::vector row(input.begin() + (size_t) r * cols, input.begin() + (size_t) r * cols + cols); + std::vector filtered = median_filter_1d(row, filter_width); + memcpy(result.data() + (size_t) r * cols, filtered.data(), (size_t) cols * sizeof(float)); + } + return result; +} + +// ============================================================================ +// Lyric Alignment Scorer +// Ported from acestep/core/scoring/dit_score.py:MusicLyricScorer +// ============================================================================ + +// Default layer/head configuration for cross-attention extraction. +// Matches the Python default: dit_score.py uses custom_config with +// {2: [6], 3: [10, 11], 4: [3], 5: [8, 9], 6: [8]} +// 5 layers, 7 heads total. These are the cross-attention heads that +// best track lyric-to-audio alignment. +struct ScoreLayerHeadConfig { + int layer; + int head; +}; + +static const ScoreLayerHeadConfig DEFAULT_SCORE_HEADS[] = { + {2, 6}, {3, 10}, {3, 11}, {4, 3}, {5, 8}, {5, 9}, {6, 8}, +}; +static const int DEFAULT_SCORE_HEADS_COUNT = 7; + +// Token type mask: 1 = lyric token, 0 = structural tag (inside [...]). +// Ported from dit_score.py:_generate_token_type_mask (lines 32-55). +// +// The Python version uses tokenizer.decode([tid]) to get the string for each +// token. Here we accept a pre-decoded vector of token strings from the caller +// (the BPE tokenizer is available in the C++ pipeline but not in this header). +static std::vector generate_token_type_mask(const std::vector & decoded_tokens) { + int n = (int) decoded_tokens.size(); + std::vector mask(n, 1); + bool in_bracket = false; + + for (int i = 0; i < n; i++) { + const std::string & s = decoded_tokens[i]; + if (s.find('[') != std::string::npos) { + in_bracket = true; + } + if (in_bracket) { + mask[i] = 0; + } + if (s.find(']') != std::string::npos) { + in_bracket = false; + mask[i] = 0; + } + } + return mask; +} + +// Preprocess attention matrix: select heads, average, median filter, min-max +// normalize, square for DTW pathfinding. +// Ported from dit_score.py:_preprocess_attention (lines 57-115). +// +// attention: [n_layers, n_heads, tokens, frames] row-major +// = layers * n_heads * tokens * frames floats +// config: array of (layer, head) pairs to select +// config_count: number of entries in config +// medfilt_width: median filter window (1 = no filter) +// +// Outputs (via pointers): +// calc_matrix: squared energy for DTW pathfinding [tokens, frames] +// energy_matrix: normalized energy for scoring [tokens, frames] +// Returns false if no valid heads were found. +static bool preprocess_attention(const float * attention, + int n_layers, + int n_heads, + int tokens, + int frames, + const ScoreLayerHeadConfig * config, + int config_count, + int medfilt_width, + std::vector & calc_matrix, + std::vector & energy_matrix) { + // 1. Select heads and stack (matches dit_score.py:84-93) + std::vector> selected; + for (int c = 0; c < config_count; c++) { + int layer = config[c].layer; + int head = config[c].head; + if (layer < n_layers && head < n_heads) { + const float * ptr = + attention + (size_t) layer * n_heads * tokens * frames + (size_t) head * tokens * frames; + selected.emplace_back(ptr, ptr + (size_t) tokens * frames); + } + } + + if (selected.empty()) { + return false; + } + + // 2. Average across selected heads (dit_score.py:96) + std::vector avg_weights((size_t) tokens * frames, 0.0f); + for (const auto & s : selected) { + for (size_t i = 0; i < avg_weights.size(); i++) { + avg_weights[i] += s[i]; + } + } + float inv_count = 1.0f / (float) selected.size(); + for (auto & v : avg_weights) { + v *= inv_count; + } + + // 3. Median filter (dit_score.py:101) + if (medfilt_width > 1) { + energy_matrix = median_filter_2d_rows(avg_weights, tokens, frames, medfilt_width); + } else { + energy_matrix = avg_weights; + } + + // 4. Min-Max normalization (dit_score.py:104-109) + float e_min = std::numeric_limits::max(); + float e_max = std::numeric_limits::lowest(); + for (float v : energy_matrix) { + if (v < e_min) e_min = v; + if (v > e_max) e_max = v; + } + + if (e_max - e_min > 1e-9f) { + float range = e_max - e_min; + for (auto & v : energy_matrix) { + v = (v - e_min) / range; + } + } else { + std::fill(energy_matrix.begin(), energy_matrix.end(), 0.0f); + } + + // 5. Contrast enhancement for DTW (dit_score.py:113) + calc_matrix.resize((size_t) tokens * frames); + for (size_t i = 0; i < energy_matrix.size(); i++) { + calc_matrix[i] = energy_matrix[i] * energy_matrix[i]; + } + + return true; +} + +// Alignment metrics: coverage, monotonicity, path confidence. +// Ported from dit_score.py:_compute_alignment_metrics (lines 117-214). +// +// energy_matrix: [rows, cols] normalized energy (row-major) +// path: DTW path (text_idx, time_idx pairs) +// type_mask: [rows] — 1 = lyric token, 0 = structural tag +// time_weight: minimum energy threshold for centroid computation (default 0.01) +// overlap_frames: allowed backward movement for monotonicity (default 9.0) +// instrumental_weight: weight for non-lyric path steps (default 1.0) +struct AlignmentMetrics { + double coverage; + double monotonicity; + double confidence; +}; + +static AlignmentMetrics compute_alignment_metrics(const std::vector & energy_matrix, + int rows, + int cols, + const DTWPath & path, + const std::vector & type_mask, + double time_weight = 0.01, + double overlap_frames = 9.0, + double instrumental_weight = 1.0) { + auto E = [&](int i, int j) -> double { return (double) energy_matrix[(size_t) i * cols + j]; }; + + // ================= A. Coverage Score (dit_score.py:150-161) ================= + int total_sung_rows = 0; + int valid_sung_rows = 0; + const double coverage_threshold = 0.1; + + for (int i = 0; i < rows; i++) { + if (type_mask[i] == 1) { + total_sung_rows++; + double row_max = 0.0; + for (int j = 0; j < cols; j++) { + double e = E(i, j); + if (e > row_max) row_max = e; + } + if (row_max > coverage_threshold) { + valid_sung_rows++; + } + } + } + + double coverage = (total_sung_rows > 0) ? (double) valid_sung_rows / total_sung_rows : 1.0; + + // ================= B. Monotonicity Score (dit_score.py:163-191) ================= + // Compute centroid (energy-weighted mean column) for each lyric row. + std::vector centroids(rows, -1.0); + for (int i = 0; i < rows; i++) { + double sum_w = 0.0; + double sum_t = 0.0; + for (int j = 0; j < cols; j++) { + double e = E(i, j); + if (e > time_weight) { + sum_w += e; + sum_t += e * (double) j; + } + } + if (sum_w > 1e-9) { + centroids[i] = sum_t / sum_w; + } + } + + // Collect centroids for lyric rows with valid centroid + std::vector sung_centroids; + for (int i = 0; i < rows; i++) { + if (type_mask[i] == 1 && centroids[i] >= 0.0) { + sung_centroids.push_back(centroids[i]); + } + } + + double monotonicity; + int cnt = (int) sung_centroids.size(); + if (cnt > 1) { + double non_decreasing = 0.0; + for (int k = 0; k < cnt - 1; k++) { + if (sung_centroids[k + 1] >= (sung_centroids[k] - overlap_frames)) { + non_decreasing += 1.0; + } + } + monotonicity = non_decreasing / (double) (cnt - 1); + } else { + monotonicity = 1.0; + } + + // ================= C. Path Confidence (dit_score.py:193-212) ================= + double path_confidence; + int path_len = (int) path.text_idx.size(); + if (path_len > 0) { + double total_energy = 0.0; + double total_steps = 0.0; + for (int k = 0; k < path_len; k++) { + int r = path.text_idx[k]; + int c = path.time_idx[k]; + double pe = E(r, c); + double sw = (type_mask[r] == 0) ? instrumental_weight : 1.0; + total_energy += pe * sw; + total_steps += sw; + } + path_confidence = (total_steps > 0.0) ? total_energy / total_steps : 0.0; + } else { + path_confidence = 0.0; + } + + return { coverage, monotonicity, path_confidence }; +} + +// Full scoring pipeline: preprocess attention -> DTW -> metrics -> final score. +// Ported from dit_score.py:lyrics_alignment_info + calculate_score. +// +// attention: [n_layers, n_heads, tokens, frames] row-major +// decoded_tokens: per-token decoded strings for type mask generation +// config / config_count: which (layer, head) pairs to use +// medfilt_width: median filter window (1 = disabled) +// +// Returns the final lyrics_score = cov^2 * mono^2 * conf, clipped to [0, 1]. +struct LyricScoreResult { + double coverage; + double monotonicity; + double confidence; + double lyrics_score; // cov^2 * mono^2 * conf, clipped [0,1] + DTWPath path; +}; + +static LyricScoreResult calculate_lyric_score(const float * attention, + int n_layers, + int n_heads, + int tokens, + int frames, + const std::vector & decoded_tokens, + const ScoreLayerHeadConfig * config = DEFAULT_SCORE_HEADS, + int config_count = DEFAULT_SCORE_HEADS_COUNT, + int medfilt_width = 1) { + LyricScoreResult result = { 0, 0, 0, 0, {} }; + + // 1. Preprocess attention (dit_score.py:237-239) + std::vector calc_matrix; + std::vector energy_matrix; + if (!preprocess_attention(attention, n_layers, n_heads, tokens, frames, config, config_count, medfilt_width, + calc_matrix, energy_matrix)) { + return result; + } + + // 2. Generate token type mask (dit_score.py:248) + std::vector type_mask = generate_token_type_mask(decoded_tokens); + + // Safety check for shape mismatch (dit_score.py:251-252) + if ((int) type_mask.size() != tokens) { + type_mask.assign(tokens, 1); + } + + // 3. DTW pathfinding on negated calc_matrix (dit_score.py:255) + // Negate so DTW finds the maximum-energy path (minimum cost = maximum energy) + std::vector neg_calc((size_t) tokens * frames); + for (size_t i = 0; i < calc_matrix.size(); i++) { + neg_calc[i] = -calc_matrix[i]; + } + result.path = dtw_cpu(neg_calc.data(), tokens, frames); + + // 4. Compute metrics (dit_score.py:313-320) + AlignmentMetrics m = compute_alignment_metrics(energy_matrix, tokens, frames, result.path, type_mask); + result.coverage = m.coverage; + result.monotonicity = m.monotonicity; + result.confidence = m.confidence; + + // 5. Final score: cov^2 * mono^2 * conf (dit_score.py:324-325) + double final_score = (m.coverage * m.coverage) * (m.monotonicity * m.monotonicity) * m.confidence; + if (final_score < 0.0) final_score = 0.0; + if (final_score > 1.0) final_score = 1.0; + result.lyrics_score = final_score; + + return result; +} diff --git a/src/pipeline-synth-impl.h b/src/pipeline-synth-impl.h index b7199e53..4722903c 100644 --- a/src/pipeline-synth-impl.h +++ b/src/pipeline-synth-impl.h @@ -112,6 +112,11 @@ struct SynthState { std::vector per_enc_S_nc; bool need_enc_switch; + // lyric token IDs (per-batch) — persisted for /score endpoint. + // Set by ops_encode_text, used by ops_score_forward to generate the + // token type mask (lyric vs structural tag) for DTW alignment scoring. + std::vector> per_lyric_ids; + // stacked encoder hidden states int max_enc_S; std::vector enc_hidden; diff --git a/src/pipeline-synth-ops.cpp b/src/pipeline-synth-ops.cpp index 58bcd109..fcab01ce 100644 --- a/src/pipeline-synth-ops.cpp +++ b/src/pipeline-synth-ops.cpp @@ -6,7 +6,9 @@ #include "pipeline-synth-ops.h" +#include "bpe.h" #include "dit-sampler.h" +#include "dtw-score.h" #include "philox.h" #include "pipeline-synth-impl.h" #include "task-types.h" @@ -440,6 +442,9 @@ int ops_encode_text(const AceSynth * ctx, const AceRequest * reqs, int batch_n, s.need_enc_switch = s.use_source_context && !s.is_repaint && !s.is_lego_region && s.rr.audio_cover_strength < 1.0f; + // Persist lyric token IDs for /score endpoint (per-batch). + s.per_lyric_ids.resize(batch_n); + BPETokenizer * bpe = store_bpe(ctx->store, ctx->params.text_encoder_path); if (!bpe) { fprintf(stderr, "[Encode-Text] FATAL: store_bpe failed\n"); @@ -474,6 +479,8 @@ int ops_encode_text(const AceSynth * ctx, const AceRequest * reqs, int batch_n, int S_text = (int) text_ids.size(); int S_lyric = (int) lyric_ids.size(); + s.per_lyric_ids[b] = lyric_ids; // persist for /score + main_fwd[b].S_text = S_text; main_fwd[b].S_lyric = S_lyric; main_fwd[b].text_hidden.resize((size_t) H_text * S_text); @@ -958,3 +965,136 @@ int ops_vae_decode(const AceSynth * ctx, } return 0; } + +// Scoring forward pass: extract cross-attention matrices and compute lyric +// alignment metrics. Matches Python ACE-Step dit_score.py:calculate_score. +// +// The Python reference runs the full DiT forward with output_attentions=True +// on the generated audio + lyrics, then extracts cross-attention from +// specific layers/heads to score lyric alignment quality. +// +// Here we run a single DiT forward on the current noise/context/encoder +// state, capture cross-attention from the default score layers, and compute +// coverage/monotonicity/confidence via dtw-score.h. +int ops_score_forward(const AceSynth * ctx, + int batch_n, + std::vector & out_scores, + SynthState & s) { + DiTGGML * dit = store_require_dit(ctx->store, ctx->dit_key); + if (!dit) { + fprintf(stderr, "[Score] FATAL: store_require_dit failed\n"); + return -1; + } + ModelHandle dit_guard(ctx->store, dit); + + if (!ctx->params.use_fa) { + dit->use_flash_attn = false; + } + + // Default score layers from Python: {2, 3, 4, 5, 6} + // Heads are selected per-layer in dtw-score.h:DEFAULT_SCORE_HEADS + static const int score_layers[] = { 2, 3, 4, 5, 6 }; + static const int n_score_layers = 5; + + fprintf(stderr, "[Score] Starting: T=%d, S=%d, enc_S=%d, batch=%d, %d score layers\n", + s.T, s.S, s.enc_S, batch_n, n_score_layers); + + // Run scoring forward pass with cross-attention capture + std::vector> layer_attentions; + s.timer.reset(); + int rc = dit_ggml_score_forward(dit, s.noise.data(), s.context.data(), s.enc_hidden.data(), + s.enc_S, s.per_enc_S.data(), s.T, batch_n, + score_layers, n_score_layers, layer_attentions); + if (rc != 0) { + fprintf(stderr, "[Score] FATAL: dit_ggml_score_forward failed\n"); + return -1; + } + fprintf(stderr, "[Score] DiT forward: %.1f ms\n", s.timer.ms()); + + // Get BPE tokenizer for decoding token IDs to strings (type mask) + BPETokenizer * bpe = store_bpe(ctx->store, ctx->params.text_encoder_path); + if (!bpe) { + fprintf(stderr, "[Score] WARNING: store_bpe failed, using all-lyric mask\n"); + } + + // DiT config + DiTGGMLConfig & c = dit->cfg; + int Nh = c.n_heads; + int S = s.S; + int enc_S = s.enc_S; + + // Compute scores per batch item + out_scores.resize(batch_n); + for (int b = 0; b < batch_n; b++) { + // Build decoded token strings for the type mask + std::vector decoded_tokens; + if (b < (int) s.per_lyric_ids.size() && !s.per_lyric_ids[b].empty()) { + int n_lyric = (int) s.per_lyric_ids[b].size(); + decoded_tokens.resize(n_lyric); + if (bpe) { + for (int i = 0; i < n_lyric; i++) { + int tid = s.per_lyric_ids[b][i]; + if (tid >= 0 && tid < bpe->n_vocab) { + decoded_tokens[i] = bpe->id_to_str[tid]; + } + } + } else { + // No tokenizer: treat all as lyric tokens + std::fill(decoded_tokens.begin(), decoded_tokens.end(), std::string("lyric")); + } + } else { + // No lyric IDs: single dummy token + decoded_tokens = { "lyric" }; + } + + // Stack attention from all captured layers into [n_layers, Nh, enc_S, S] + // Each layer_attentions[sl] is [enc_S, S, Nh] (batch 0 only) + // We need [n_layers, n_heads, tokens, frames] for calculate_lyric_score + // But the attention matrix is [enc_S (KV=tokens), S (Q=frames), Nh] + // So: tokens = enc_S, frames = S, n_heads = Nh, n_layers = n_score_layers + // + // Note: the Python reference uses per-lyric-token attention (enc_S = lyric_tokens). + // Here enc_S is the full encoder sequence (text + lyric). We use the full + // sequence and the type mask to distinguish lyric from structural tokens. + int n_heads_total = Nh; + int n_layers_total = n_score_layers; + + // Stack: [n_layers * n_heads * enc_S * S] + std::vector stacked((size_t) n_layers_total * n_heads_total * enc_S * S, 0.0f); + + for (int sl = 0; sl < n_score_layers; sl++) { + if (sl >= (int) layer_attentions.size() || layer_attentions[sl].empty()) { + continue; + } + // layer_attentions[sl] is [enc_S, S, Nh] row-major + // We need [n_layers, n_heads, enc_S, S] = layer * (Nh * enc_S * S) + head * (enc_S * S) + ... + const float * src = layer_attentions[sl].data(); + for (int h = 0; h < Nh; h++) { + for (int t = 0; t < enc_S; t++) { + for (int f = 0; f < S; f++) { + // src index: [enc_S, S, Nh] = t * (S * Nh) + f * Nh + h + size_t src_idx = (size_t) t * S * Nh + (size_t) f * Nh + h; + // dst index: [n_layers, n_heads, enc_S, S] = sl * (Nh * enc_S * S) + h * (enc_S * S) + t * S + f + size_t dst_idx = (size_t) sl * n_heads_total * enc_S * S + + (size_t) h * enc_S * S + + (size_t) t * S + f; + if (src_idx < layer_attentions[sl].size()) { + stacked[dst_idx] = src[src_idx]; + } + } + } + } + } + + // Calculate lyric score + LyricScoreResult result = calculate_lyric_score( + stacked.data(), n_layers_total, n_heads_total, enc_S, S, decoded_tokens); + + fprintf(stderr, "[Score] Batch %d: coverage=%.4f monotonicity=%.4f confidence=%.4f lyrics_score=%.4f\n", + b, result.coverage, result.monotonicity, result.confidence, result.lyrics_score); + + out_scores[b] = result; + } + + return 0; +} diff --git a/src/pipeline-synth-ops.h b/src/pipeline-synth-ops.h index 1267376b..e2a05393 100644 --- a/src/pipeline-synth-ops.h +++ b/src/pipeline-synth-ops.h @@ -8,6 +8,8 @@ #include "pipeline-synth.h" +#include + struct AceSynth; struct SynthState; @@ -68,3 +70,18 @@ int ops_vae_decode(const AceSynth * ctx, SynthState & s, bool (*cancel)(void *), void * cancel_data); + +// Scoring primitive: single DiT forward with cross-attention capture. +// Runs one forward pass on the current noise + context + encoder states, +// extracts cross-attention matrices from the configured score layers, +// and computes lyric alignment metrics (coverage, monotonicity, confidence). +// +// Requires ops_encode_text (for s.per_lyric_ids) and ops_build_context + +// ops_init_noise (for noise + context latents) to have been run first. +// +// out_scores: filled with one LyricScoreResult per batch item. +// Returns 0 on success, -1 on error. +int ops_score_forward(const AceSynth * ctx, + int batch_n, + std::vector & out_scores, + SynthState & s); diff --git a/src/pipeline-synth.cpp b/src/pipeline-synth.cpp index 9571bc1b..d7794149 100644 --- a/src/pipeline-synth.cpp +++ b/src/pipeline-synth.cpp @@ -695,3 +695,59 @@ void ace_synth_free(AceSynth * ctx) { } delete ctx; } + +// Scoring: run phase 1 setup (encode, context, noise) then a single DiT +// forward with cross-attention capture instead of the full denoising loop. +// Matches Python ACE-Step dit_score.py:get_lyric_score flow. +int ace_synth_score(AceSynth * ctx, + const AceRequest * reqs, + int batch_n, + std::vector & out_scores) { + if (!ctx || !reqs || batch_n < 1 || batch_n > 9) { + return -1; + } + + // Scoring only supports text2music (pure generation) — the Python reference + // also scores on the text2music forward pass. + const std::string & task = reqs[0].task_type; + if (task != TASK_TEXT2MUSIC) { + fprintf(stderr, "[Score] ERROR: scoring only supports task 'text2music', got '%s'\n", task.c_str()); + return -1; + } + + AceSynthJob * job = alloc_job(ctx, reqs, batch_n); + SynthState & s = job->state; + s.use_source_context = !reqs[0].audio_codes.empty(); + s.instruction_str = s.use_source_context ? DIT_INSTR_COVER : DIT_INSTR_TEXT2MUSIC; + + // Phase 1 setup (same as run_text2music minus the DiT generate) + if (!pinned_encode_src_and_timbre(ctx, NULL, 0, NULL, 0, NULL, 0, NULL, 0, s)) { + delete job; + return -1; + } + if (ops_resolve_params(ctx, reqs, batch_n, s) != 0) { + delete job; + return -1; + } + if (ops_resolve_T(ctx, s) != 0) { + delete job; + return -1; + } + ops_build_schedule(s); + if (ops_encode_text(ctx, reqs, batch_n, s) != 0) { + delete job; + return -1; + } + if (ops_build_context(ctx, reqs, batch_n, s) != 0) { + delete job; + return -1; + } + ops_build_context_silence(ctx, batch_n, s); + ops_init_noise(ctx, reqs, batch_n, s); + + // Scoring forward pass (single DiT forward with cross-attention capture) + int rc = ops_score_forward(ctx, batch_n, out_scores, s); + + delete job; + return rc; +} diff --git a/src/pipeline-synth.h b/src/pipeline-synth.h index 4f231610..c90e385e 100644 --- a/src/pipeline-synth.h +++ b/src/pipeline-synth.h @@ -8,9 +8,11 @@ // caches everything across calls. DiT weight swap between phases is an // invisible consequence of the store, not an orchestration concern anymore. +#include "dtw-score.h" #include "request.h" #include +#include struct AceSynth; struct AceSynthJob; @@ -99,3 +101,17 @@ void ace_synth_job_free(AceSynthJob * job); void ace_audio_free(AceAudio * audio); void ace_synth_free(AceSynth * ctx); + +// LyricScoreResult is defined in dtw-score.h (included above). + +// Phase 1 scoring: run a single DiT forward pass with cross-attention capture +// and compute lyric alignment metrics. Requires the same phase 1 setup as +// ace_synth_job_run_dit (text encoding, context build, noise init) but does +// NOT run the full denoising loop — just one forward pass to extract attention. +// +// Returns one LyricScoreResult per batch item in out_scores. +// Returns 0 on success, -1 on error. +int ace_synth_score(AceSynth * ctx, + const AceRequest * reqs, + int batch_n, + std::vector & out_scores); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e940df51..b2ca206e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -15,3 +15,10 @@ link_ggml_backends(test-model-store) add_executable(test-lm-prompt test-lm-prompt.cpp) target_link_libraries(test-lm-prompt PRIVATE acestep-core) link_ggml_backends(test-lm-prompt) + +# test-dtw-score: DTW pathfinding + lyric alignment scoring (pure C++, no ggml). +add_executable(test-dtw-score test-dtw-score.cpp) +target_include_directories(test-dtw-score PRIVATE ${CMAKE_SOURCE_DIR}/src) +if(NOT MSVC) + target_link_libraries(test-dtw-score PRIVATE m) +endif() diff --git a/tests/test-dtw-score.cpp b/tests/test-dtw-score.cpp new file mode 100644 index 00000000..33194413 --- /dev/null +++ b/tests/test-dtw-score.cpp @@ -0,0 +1,340 @@ +// test-dtw-score.cpp: unit tests for DTW pathfinding + lyric alignment scoring +// +// Verifies the C++ port against known-good values from the Python reference +// (acestep/core/scoring/_dtw.py and dit_score.py). No external dependencies +// beyond dtw-score.h itself — all test cases are self-contained. +// +// Usage: +// ./test-dtw-score +// +// Exit code 0 = all tests passed, 1 = at least one failed. + +#include "dtw-score.h" + +#include +#include +#include +#include + +static int g_pass = 0; +static int g_fail = 0; + +#define CHECK(cond, msg) \ + do { \ + if (cond) { \ + g_pass++; \ + } else { \ + g_fail++; \ + fprintf(stderr, "FAIL: %s (line %d)\n", msg, __LINE__); \ + } \ + } while (0) + +#define CHECK_NEAR(a, b, eps, msg) \ + do { \ + if (fabs((double) (a) - (double) (b)) < (eps)) { \ + g_pass++; \ + } else { \ + g_fail++; \ + fprintf(stderr, "FAIL: %s — expected %.8g, got %.8g (line %d)\n", \ + msg, (double) (b), (double) (a), __LINE__); \ + } \ + } while (0) + +// ============================================================================ +// Test 1: DTW on a simple diagonal cost matrix +// A diagonal matrix should produce a path that goes straight from (0,0) to (N-1,M-1) +// ============================================================================ +static void test_dtw_diagonal() { + // 3x3 cost matrix where the diagonal is cheapest + // Cost: [[0, 9, 9], + // [9, 0, 9], + // [9, 9, 0]] + // DTW on this should find the diagonal path. + float cost[] = { 0, 9, 9, 9, 0, 9, 9, 9, 0 }; + DTWPath path = dtw_cpu(cost, 3, 3); + + // Path should visit (0,0), (1,1), (2,2) + CHECK(path.text_idx.size() == 3, "diagonal DTW path length == 3"); + for (size_t k = 0; k < path.text_idx.size(); k++) { + CHECK(path.text_idx[k] == (int) k, "diagonal DTW text_idx[k] == k"); + CHECK(path.time_idx[k] == (int) k, "diagonal DTW time_idx[k] == k"); + } +} + +// ============================================================================ +// Test 2: DTW on a matrix that forces a horizontal step +// ============================================================================ +static void test_dtw_horizontal_step() { + // 2x3 cost matrix: + // [[0, 0, 0], + // [9, 9, 0]] + // The optimal path should go right along row 0, then diagonal to (1,2). + float cost[] = { 0, 0, 0, 9, 9, 0 }; + DTWPath path = dtw_cpu(cost, 2, 3); + + // Path must start at (0,0) and end at (1,2) + CHECK(!path.text_idx.empty(), "horizontal DTW path non-empty"); + CHECK(path.text_idx.front() == 0, "horizontal DTW starts at text=0"); + CHECK(path.time_idx.front() == 0, "horizontal DTW starts at time=0"); + CHECK(path.text_idx.back() == 1, "horizontal DTW ends at text=1"); + CHECK(path.time_idx.back() == 2, "horizontal DTW ends at time=2"); +} + +// ============================================================================ +// Test 3: Median filter — simple cases +// ============================================================================ +static void test_median_filter() { + // Window=3 on a smooth ramp: [1, 2, 3, 4, 5] -> [1, 2, 3, 4, 5] + // (reflect padding keeps edges intact for monotonic data) + std::vector ramp = { 1, 2, 3, 4, 5 }; + auto filtered = median_filter_1d(ramp, 3); + CHECK(filtered.size() == 5, "median filter preserves length"); + CHECK_NEAR(filtered[2], 3.0f, 1e-6, "median filter middle of ramp"); + + // Window=3 removes a single spike: [1, 2, 100, 4, 5] + // Reflect padding: [2,1,2,100,4,5,4] + // Windows: [2,1,2]=2, [1,2,100]=2, [2,100,4]=4, [100,4,5]=5, [4,5,4]=4 + std::vector spike = { 1, 2, 100, 4, 5 }; + auto fs = median_filter_1d(spike, 3); + CHECK_NEAR(fs[0], 2.0f, 1e-6, "median filter spike removal [0]"); + CHECK_NEAR(fs[1], 2.0f, 1e-6, "median filter spike removal [1]"); + CHECK_NEAR(fs[2], 4.0f, 1e-6, "median filter spike removal [2]"); + CHECK_NEAR(fs[3], 5.0f, 1e-6, "median filter spike removal [3]"); + CHECK_NEAR(fs[4], 4.0f, 1e-6, "median filter spike removal [4]"); + + // Window=1 = no-op + auto noop = median_filter_1d(ramp, 1); + CHECK(noop == ramp, "median filter width=1 is no-op"); +} + +// ============================================================================ +// Test 4: Token type mask generation +// ============================================================================ +static void test_token_type_mask() { + // Tokens: "hello", "[", "Verse", "]", "world" + // Mask: 1, 0, 0, 0, 1 + std::vector tokens = { "hello", "[", "Verse", "]", "world" }; + auto mask = generate_token_type_mask(tokens); + CHECK(mask.size() == 5, "token type mask size"); + CHECK(mask[0] == 1, "token type mask: lyric before bracket"); + CHECK(mask[1] == 0, "token type mask: opening bracket"); + CHECK(mask[2] == 0, "token type mask: inside bracket"); + CHECK(mask[3] == 0, "token type mask: closing bracket"); + CHECK(mask[4] == 1, "token type mask: lyric after bracket"); + + // Multi-token bracket: "[", "Intro", "Guitar", "]" + std::vector tokens2 = { "[", "Intro", "Guitar", "]", "sing" }; + auto mask2 = generate_token_type_mask(tokens2); + CHECK(mask2[0] == 0, "multi-token bracket: open"); + CHECK(mask2[1] == 0, "multi-token bracket: inside 1"); + CHECK(mask2[2] == 0, "multi-token bracket: inside 2"); + CHECK(mask2[3] == 0, "multi-token bracket: close"); + CHECK(mask2[4] == 1, "multi-token bracket: after"); +} + +// ============================================================================ +// Test 5: Full scoring pipeline with a synthetic attention matrix +// A perfect diagonal attention pattern should yield a high score. +// ============================================================================ +static void test_scoring_perfect_alignment() { + // Create a 4-token, 4-frame attention matrix where the diagonal has + // high energy and off-diagonal is near zero. With 1 layer, 1 head. + // This simulates perfect lyric-to-audio alignment. + int tokens = 4; + int frames = 4; + std::vector attn((size_t) tokens * frames, 0.01f); + // Diagonal high energy + for (int i = 0; i < tokens; i++) { + attn[(size_t) i * frames + i] = 1.0f; + } + + // All tokens are lyrics (no brackets) + std::vector decoded(tokens, "lyric"); + + // Use a custom config: just layer 0, head 0 + ScoreLayerHeadConfig config[] = { { 0, 0 } }; + + LyricScoreResult result = + calculate_lyric_score(attn.data(), 1, 1, tokens, frames, decoded, config, 1, 1); + + // With perfect diagonal alignment: + // - Coverage: all lyric rows have max energy 1.0 > 0.1, so coverage = 1.0 + // - Monotonicity: centroids are [0, 1, 2, 3], strictly increasing, so mono = 1.0 + // - Confidence: DTW tie-breaking (prefers left/up over diagonal when costs + // are equal) creates a 7-step zigzag path through the 4x4 matrix. + // 4 diagonal cells have energy 1.0, 3 off-diagonal have 0.0. + // confidence = 4/7 ≈ 0.571 + // - lyrics_score = 1^2 * 1^2 * (4/7) = 4/7 ≈ 0.571 + CHECK_NEAR(result.coverage, 1.0, 1e-6, "perfect alignment coverage"); + CHECK_NEAR(result.monotonicity, 1.0, 1e-6, "perfect alignment monotonicity"); + CHECK_NEAR(result.confidence, 4.0 / 7.0, 0.01, "perfect alignment confidence"); + CHECK_NEAR(result.lyrics_score, 4.0 / 7.0, 0.01, "perfect alignment final score"); +} + +// ============================================================================ +// Test 6: Full scoring pipeline with poor alignment +// Uniform attention should yield low confidence and low score. +// ============================================================================ +static void test_scoring_uniform_attention() { + int tokens = 4; + int frames = 8; + // Uniform attention — no structure + std::vector attn((size_t) tokens * frames, 0.5f); + + std::vector decoded(tokens, "lyric"); + ScoreLayerHeadConfig config[] = { { 0, 0 } }; + + LyricScoreResult result = + calculate_lyric_score(attn.data(), 1, 1, tokens, frames, decoded, config, 1, 1); + + // With uniform attention, min-max normalization zeros everything out + // (e_max - e_min < 1e-9), so energy_matrix = 0, confidence = 0, score = 0 + CHECK_NEAR(result.lyrics_score, 0.0, 1e-6, "uniform attention score is 0"); +} + +// ============================================================================ +// Test 7: Scoring with structural tags (non-lyric tokens) +// ============================================================================ +static void test_scoring_with_tags() { + // 5 tokens: lyric, [Intro], lyric, lyric, lyric + // 5 frames: perfect diagonal for lyric tokens, tag token gets low energy + int tokens = 5; + int frames = 5; + std::vector attn((size_t) tokens * frames, 0.01f); + + // Diagonal energy for all tokens (including tag) + for (int i = 0; i < tokens; i++) { + attn[(size_t) i * frames + i] = 1.0f; + } + + // Token 1 is a structural tag + std::vector decoded = { "sing", "[", "more", "singing", "here" }; + ScoreLayerHeadConfig config[] = { { 0, 0 } }; + + LyricScoreResult result = + calculate_lyric_score(attn.data(), 1, 1, tokens, frames, decoded, config, 1, 1); + + // Coverage: 4 lyric tokens, all with max energy 1.0 > 0.1 -> coverage = 1.0 + // Monotonicity: lyric centroids [0, 2, 3, 4] strictly increasing -> mono = 1.0 + // Confidence: DTW tie-breaking creates a 9-step zigzag through the 5x5 + // matrix. 5 diagonal cells have energy 1.0, 4 off-diagonal have 0.0. + // confidence = 5/9 ≈ 0.556 + // lyrics_score = 1 * 1 * (5/9) = 5/9 ≈ 0.556 + CHECK_NEAR(result.coverage, 1.0, 1e-6, "tagged alignment coverage"); + CHECK_NEAR(result.monotonicity, 1.0, 1e-6, "tagged alignment monotonicity"); + CHECK_NEAR(result.lyrics_score, 5.0 / 9.0, 0.01, "tagged alignment final score"); +} + +// ============================================================================ +// Test 8: Multi-head attention averaging +// ============================================================================ +static void test_multi_head_averaging() { + int tokens = 3; + int frames = 3; + int n_layers = 2; + int n_heads = 2; + + // 2 layers x 2 heads x 3 tokens x 3 frames + std::vector attn((size_t) n_layers * n_heads * tokens * frames, 0.0f); + + // Layer 0, Head 0: diagonal + // Layer 0, Head 1: anti-diagonal + // Layer 1, Head 0: diagonal + // Layer 1, Head 1: diagonal + auto A = [&](int l, int h, int t, int f) -> float & { + return attn[(size_t) ((l * n_heads + h) * tokens + t) * frames + f]; + }; + + // 3 of 4 heads have diagonal, 1 has anti-diagonal + for (int t = 0; t < tokens; t++) { + A(0, 0, t, t) = 1.0f; + A(0, 1, t, tokens - 1 - t) = 1.0f; // anti-diagonal + A(1, 0, t, t) = 1.0f; + A(1, 1, t, t) = 1.0f; + } + + // Select all 4 heads + ScoreLayerHeadConfig config[] = { { 0, 0 }, { 0, 1 }, { 1, 0 }, { 1, 1 } }; + + std::vector calc_matrix, energy_matrix; + bool ok = preprocess_attention(attn.data(), n_layers, n_heads, tokens, frames, config, 4, 1, + calc_matrix, energy_matrix); + CHECK(ok, "multi-head preprocess succeeded"); + + // Average: 3/4 diagonal + 1/4 anti-diagonal + // avg = [[0.75, 0, 0.25], [0, 1, 0], [0.25, 0, 0.75]] + // min=0, max=1, range=1 -> normalization is identity + // energy[0,0] = 0.75, energy[0,2] = 0.25 + CHECK_NEAR(energy_matrix[0], 0.75f, 1e-5, "multi-head averaged energy [0,0]"); + CHECK_NEAR(energy_matrix[2], 0.25f, 1e-5, "multi-head averaged energy [0,2]"); +} + +// ============================================================================ +// Test 9: DTW path monotonicity (paths must be non-decreasing in both dims) +// ============================================================================ +static void test_dtw_path_monotonic() { + // Random-ish cost matrix + int N = 5, M = 7; + std::vector cost((size_t) N * M); + for (int i = 0; i < N; i++) { + for (int j = 0; j < M; j++) { + // Cheaper near the diagonal (i/N * j/M) + float di = (float) i / N; + float dj = (float) j / M; + cost[(size_t) i * M + j] = fabsf(di - dj) + 0.01f * ((i * 7 + j * 3) % 5); + } + } + + DTWPath path = dtw_cpu(cost.data(), N, M); + + // Path must be non-decreasing in both text and time + for (size_t k = 1; k < path.text_idx.size(); k++) { + CHECK(path.text_idx[k] >= path.text_idx[k - 1], "DTW path text non-decreasing"); + CHECK(path.time_idx[k] >= path.time_idx[k - 1], "DTW path time non-decreasing"); + } + + // Path must start at (0,0) and end at (N-1, M-1) + CHECK(path.text_idx.front() == 0, "DTW path starts at text 0"); + CHECK(path.time_idx.front() == 0, "DTW path starts at time 0"); + CHECK(path.text_idx.back() == N - 1, "DTW path ends at text N-1"); + CHECK(path.time_idx.back() == M - 1, "DTW path ends at time M-1"); +} + +// ============================================================================ + +int main(int argc, char ** argv) { + (void) argc; + (void) argv; + + fprintf(stderr, "=== DTW + Lyric Score Tests ===\n\n"); + + test_dtw_diagonal(); + fprintf(stderr, "test_dtw_diagonal: done\n"); + + test_dtw_horizontal_step(); + fprintf(stderr, "test_dtw_horizontal_step: done\n"); + + test_median_filter(); + fprintf(stderr, "test_median_filter: done\n"); + + test_token_type_mask(); + fprintf(stderr, "test_token_type_mask: done\n"); + + test_scoring_perfect_alignment(); + fprintf(stderr, "test_scoring_perfect_alignment: done\n"); + + test_scoring_uniform_attention(); + fprintf(stderr, "test_scoring_uniform_attention: done\n"); + + test_scoring_with_tags(); + fprintf(stderr, "test_scoring_with_tags: done\n"); + + test_multi_head_averaging(); + fprintf(stderr, "test_multi_head_averaging: done\n"); + + test_dtw_path_monotonic(); + fprintf(stderr, "test_dtw_path_monotonic: done\n"); + + fprintf(stderr, "\n=== Results: %d passed, %d failed ===\n", g_pass, g_fail); + return g_fail > 0 ? 1 : 0; +} diff --git a/tools/ace-server.cpp b/tools/ace-server.cpp index 68dffeb0..2ed807ce 100644 --- a/tools/ace-server.cpp +++ b/tools/ace-server.cpp @@ -1367,6 +1367,185 @@ static void encode_worker(std::shared_ptr job, AceRequest ace_req, float * } // POST /vae +// score worker: load DiT + Text-Enc, run scoring forward pass, store JSON result. +static void score_worker(std::shared_ptr job, + std::vector ace_reqs) { + if (job->cancel.load()) { + job->status.store(JobStatus::CANCELLED); + return; + } + + // Resolve DiT + Text-Enc (VAE not needed for scoring, but the pipeline + // requires it in the registry to build the context). + std::string dit_name = resolve_name(g_registry.dit, ace_reqs[0].synth_model, g_loaded_dit); + const ModelEntry * dit = registry_find(g_registry.dit, dit_name.c_str()); + if (!dit) { + fprintf(stderr, "[Server] DiT not found: %s\n", dit_name.c_str()); + job->status.store(JobStatus::FAILED); + return; + } + if (g_registry.text_enc.empty()) { + fprintf(stderr, "[Server] Missing Text-Enc in registry\n"); + job->status.store(JobStatus::FAILED); + return; + } + std::string vae_name = resolve_name(g_registry.vae, ace_reqs[0].vae, g_loaded_vae); + const ModelEntry * vae = registry_find(g_registry.vae, vae_name.c_str()); + if (!vae) { + fprintf(stderr, "[Server] VAE not found: %s\n", vae_name.c_str()); + job->status.store(JobStatus::FAILED); + return; + } + + AceSynthParams p = g_synth_params; + p.text_encoder_path = g_registry.text_enc[0].path.c_str(); + p.dit_path = dit->path.c_str(); + p.vae_path = vae->path.c_str(); + p.adapter_path = nullptr; + p.adapter_scale = 1.0f; + + fprintf(stderr, "[Server] Loading score: DiT=%s\n", dit_name.c_str()); + AceSynth * ctx = ace_synth_load(g_store, &p); + if (!ctx) { + fprintf(stderr, "[Server] FATAL: synth load failed for scoring\n"); + job->status.store(JobStatus::FAILED); + return; + } + + std::vector scores; + int rc = ace_synth_score(ctx, ace_reqs.data(), (int) ace_reqs.size(), scores); + ace_synth_free(ctx); + + if (rc != 0) { + job->status.store(job->cancel.load() ? JobStatus::CANCELLED : JobStatus::FAILED); + return; + } + + // Build JSON result + // Format: {"scores":[{"coverage":0.95,"monotonicity":0.88,"confidence":0.72,"lyrics_score":0.61},...]} + yyjson_mut_doc * doc = yyjson_mut_doc_new(NULL); + yyjson_mut_val * root = yyjson_mut_arr(doc); + yyjson_mut_doc_set_root(doc, root); + for (size_t i = 0; i < scores.size(); i++) { + yyjson_mut_val * obj = yyjson_mut_obj(doc); + yyjson_mut_obj_add_real(doc, obj, "coverage", scores[i].coverage); + yyjson_mut_obj_add_real(doc, obj, "monotonicity", scores[i].monotonicity); + yyjson_mut_obj_add_real(doc, obj, "confidence", scores[i].confidence); + yyjson_mut_obj_add_real(doc, obj, "lyrics_score", scores[i].lyrics_score); + yyjson_mut_arr_append(root, obj); + } + + yyjson_write_err err; + size_t len = 0; + char * json_str = yyjson_mut_write_opts(doc, 0, NULL, &len, &err); + yyjson_mut_doc_free(doc); + if (!json_str) { + fprintf(stderr, "[Server] FATAL: JSON serialization failed: %s\n", err.msg); + job->status.store(JobStatus::FAILED); + return; + } + + job->result_body = std::string(json_str, len); + job->result_mime = "application/json"; + free(json_str); + + if (g_keep_loaded) { + g_loaded_dit = dit_name; + g_loaded_vae = vae_name; + } + + fprintf(stderr, "[Server] Job %s done (score)\n", job->id.c_str()); + job->status.store(JobStatus::DONE); +} + +// POST /score +// Accepts the same JSON request format as /synth (plain JSON or array), +// but only uses caption, lyrics, duration, and seed to run a single DiT +// forward pass with cross-attention capture for lyric alignment scoring. +// No audio is generated — the result is a JSON array of per-request scores. +// returns: JSON {"id":"N"} immediately. Result is JSON array of +// {"coverage","monotonicity","confidence","lyrics_score"} per request. +static void handle_score(const httplib::Request & req, httplib::Response & res) { + if (g_registry.dit.empty() || g_registry.text_enc.empty() || g_registry.vae.empty()) { + json_error(res, 501, "No score models in registry (need dit + text-encoder + vae)"); + return; + } + + // Parse request: plain JSON (single or array), same as /synth minus audio parts. + std::vector ace_reqs; + + if (req.is_multipart_form_data()) { + AceRequest ace_req; + std::string json_body; + if (req.form.has_file("request")) { + json_body = req.form.get_file("request").content; + } else if (req.form.has_field("request")) { + json_body = req.form.get_field("request"); + } else { + json_error(res, 400, "Multipart: missing 'request' part"); + return; + } + if (!request_parse_json(&ace_req, json_body.c_str())) { + json_error(res, 400, "Multipart: invalid JSON in 'request' part"); + return; + } + ace_reqs.push_back(std::move(ace_req)); + } else { + // Plain JSON: single object or array + if (req.body.empty()) { + json_error(res, 400, "Empty request body"); + return; + } + AceRequest ace_req; + if (!request_parse_json(&ace_req, req.body.c_str())) { + // Try as array + yyjson_doc * doc = yyjson_read(req.body.c_str(), req.body.size(), 0); + if (!doc || !yyjson_is_arr(yyjson_doc_get_root(doc))) { + yyjson_doc_free(doc); + json_error(res, 400, "Invalid JSON (expected object or array)"); + return; + } + size_t n = yyjson_arr_size(yyjson_doc_get_root(doc)); + for (size_t i = 0; i < n; i++) { + AceRequest r; + request_init(&r); + yyjson_val * obj = yyjson_arr_get(yyjson_doc_get_root(doc), i); + if (!request_parse_json(&r, yyjson_val_write(obj, 0, nullptr))) { + // Fallback: use the raw object directly + char * obj_str = yyjson_val_write(obj, 0, nullptr); + request_parse_json(&r, obj_str); + free(obj_str); + } + ace_reqs.push_back(std::move(r)); + } + yyjson_doc_free(doc); + } else { + ace_reqs.push_back(std::move(ace_req)); + } + } + + if (ace_reqs.empty()) { + json_error(res, 400, "Empty request"); + return; + } + if (ace_reqs[0].caption.empty()) { + json_error(res, 400, "Caption is required"); + return; + } + + // Create job, spawn worker, return ID + auto job = job_create(); + fprintf(stderr, "[Server] Job %s created (score, %d requests)\n", job->id.c_str(), (int) ace_reqs.size()); + + work_push([job, reqs = std::move(ace_reqs)]() mutable { + score_worker(job, std::move(reqs)); + }); + + // Return job ID immediately + std::string body = "{\"id\":\"" + job->id + "\"}"; + res.set_content(body, "application/json"); +} + // multipart/form-data: single VAE entrypoint, dispatches on which side is // supplied in the request body. Symmetric with /synth and /understand on // the 'audio or src_latents' input contract, except here they are mutually @@ -1764,6 +1943,7 @@ int main(int argc, char ** argv) { svr.Post("/synth", handle_synth); svr.Post("/understand", handle_understand); svr.Post("/vae", handle_vae); + svr.Post("/score", handle_score); svr.Get("/health", [](const httplib::Request &, httplib::Response & res) { res.set_content("{\"status\":\"ok\"}", "application/json"); }); From 7084db9375cb54955a097ec5006f4b8de0f54058 Mon Sep 17 00:00:00 2001 From: Jeff Stein Date: Fri, 10 Jul 2026 17:26:08 +0000 Subject: [PATCH 02/11] fix: address code review issues in /score endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove vacuous comment about LyricScoreResult include location - Fix misleading t=0 comment ("final denoised state" → "matches Python reference") - Replace UB pointer sentinel (ggml_tensor*)1 with a plain bool flag - Fix memory leak: yyjson_val_write in array-parse if-condition was never freed; collapse to a single alloc+free - Use ::error:: instead of ::warning:: in drift-check workflow (was emitting a warning annotation on a failing job) - Document that batch_n > 9 limit matches ace_synth_job_run Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01Bapse2pC9C1iJyEFFQ4svj --- .github/workflows/scoring-drift-check.yml | 8 ++++---- src/dit-graph.h | 9 +++++---- src/dit-sampler.h | 2 +- src/pipeline-synth.cpp | 2 +- src/pipeline-synth.h | 2 -- tools/ace-server.cpp | 11 ++++------- 6 files changed, 15 insertions(+), 19 deletions(-) diff --git a/.github/workflows/scoring-drift-check.yml b/.github/workflows/scoring-drift-check.yml index ed26da34..69c5dfcf 100644 --- a/.github/workflows/scoring-drift-check.yml +++ b/.github/workflows/scoring-drift-check.yml @@ -40,14 +40,14 @@ jobs: echo "dit_score.py: expected=$EXPECTED_SCORE actual=$ACTUAL_SCORE" if [ "$ACTUAL_DTW" != "$EXPECTED_DTW" ]; then - echo "::warning::_dtw.py has drifted from the pinned commit (82252c24)." - echo "::warning::The C++ port in src/dtw-score.h may need updating." + echo "::error::_dtw.py has drifted from the pinned commit (82252c24)." + echo "::error::The C++ port in src/dtw-score.h may need updating." exit 1 fi if [ "$ACTUAL_SCORE" != "$EXPECTED_SCORE" ]; then - echo "::warning::dit_score.py has drifted from the pinned commit (82252c24)." - echo "::warning::The C++ port in src/dtw-score.h may need updating." + echo "::error::dit_score.py has drifted from the pinned commit (82252c24)." + echo "::error::The C++ port in src/dtw-score.h may need updating." exit 1 fi diff --git a/src/dit-graph.h b/src/dit-graph.h index ef9e2630..a836f9c7 100644 --- a/src/dit-graph.h +++ b/src/dit-graph.h @@ -601,21 +601,22 @@ static struct ggml_cgraph * dit_ggml_build_graph(DiTGGML * m, struct ggml_tensor * sa_mask = (m->layers[i].layer_type == 0) ? sa_mask_sw : nullptr; // Check if this layer should capture cross-attention scores - struct ggml_tensor * capture = nullptr; + bool capture_this_layer = false; + struct ggml_tensor * capture = nullptr; if (score_layers) { for (int sl = 0; sl < n_score_layers; sl++) { if (score_layers[sl] == i) { - capture = (struct ggml_tensor *) 1; // sentinel: non-NULL triggers capture + capture_this_layer = true; break; } } } hidden = dit_ggml_build_layer(ctx, m, i, hidden, tproj, enc, positions, sa_mask, ca_mask, S, enc_S, N, - capture ? &capture : nullptr); + capture_this_layer ? &capture : nullptr); // Name captured cross-attention scores for later retrieval - if (capture && capture != (struct ggml_tensor *) 1) { + if (capture) { char score_name[64]; snprintf(score_name, sizeof(score_name), "cross_attn_scores_L%d", i); ggml_set_name(capture, score_name); diff --git a/src/dit-sampler.h b/src/dit-sampler.h index ee33ec2c..7d623da7 100644 --- a/src/dit-sampler.h +++ b/src/dit-sampler.h @@ -785,7 +785,7 @@ static int dit_ggml_score_forward(DiTGGML * model, return -1; } - // Set timesteps to 0 (scoring uses the final denoised state) + // Set timesteps to 0 (matches Python reference which scores at t=0) struct ggml_tensor * t_t = ggml_graph_get_tensor(gf, "t"); struct ggml_tensor * t_tr = ggml_graph_get_tensor(gf, "t_r"); float t_zero = 0.0f; diff --git a/src/pipeline-synth.cpp b/src/pipeline-synth.cpp index d7794149..d5b61098 100644 --- a/src/pipeline-synth.cpp +++ b/src/pipeline-synth.cpp @@ -703,7 +703,7 @@ int ace_synth_score(AceSynth * ctx, const AceRequest * reqs, int batch_n, std::vector & out_scores) { - if (!ctx || !reqs || batch_n < 1 || batch_n > 9) { + if (!ctx || !reqs || batch_n < 1 || batch_n > 9) { // same limit as ace_synth_job_run return -1; } diff --git a/src/pipeline-synth.h b/src/pipeline-synth.h index c90e385e..572269ac 100644 --- a/src/pipeline-synth.h +++ b/src/pipeline-synth.h @@ -102,8 +102,6 @@ void ace_audio_free(AceAudio * audio); void ace_synth_free(AceSynth * ctx); -// LyricScoreResult is defined in dtw-score.h (included above). - // Phase 1 scoring: run a single DiT forward pass with cross-attention capture // and compute lyric alignment metrics. Requires the same phase 1 setup as // ace_synth_job_run_dit (text encoding, context build, noise init) but does diff --git a/tools/ace-server.cpp b/tools/ace-server.cpp index 2ed807ce..f669316b 100644 --- a/tools/ace-server.cpp +++ b/tools/ace-server.cpp @@ -1509,13 +1509,10 @@ static void handle_score(const httplib::Request & req, httplib::Response & res) for (size_t i = 0; i < n; i++) { AceRequest r; request_init(&r); - yyjson_val * obj = yyjson_arr_get(yyjson_doc_get_root(doc), i); - if (!request_parse_json(&r, yyjson_val_write(obj, 0, nullptr))) { - // Fallback: use the raw object directly - char * obj_str = yyjson_val_write(obj, 0, nullptr); - request_parse_json(&r, obj_str); - free(obj_str); - } + yyjson_val * obj = yyjson_arr_get(yyjson_doc_get_root(doc), i); + char * obj_str = yyjson_val_write(obj, 0, nullptr); + request_parse_json(&r, obj_str); + free(obj_str); ace_reqs.push_back(std::move(r)); } yyjson_doc_free(doc); From 6f0caeee8585e1b8bd538f8616a257e13490d015 Mon Sep 17 00:00:00 2001 From: Jeff Stein Date: Fri, 10 Jul 2026 12:25:05 -0600 Subject: [PATCH 03/11] fix: make lyric scoring reference-compatible Require generated latents, score pure-noise and regressed states, and correct per-batch attention and layer selection. Harden model metadata, request parsing, drift checks, and CI test coverage. --- .github/workflows/ci-build.yml | 3 + .github/workflows/scoring-drift-check.yml | 34 ++- CMakeLists.txt | 2 + README.md | 9 +- docs/ARCHITECTURE.md | 26 ++- src/dit-sampler.h | 91 ++++---- src/dtw-score.h | 252 +++++++++++++++++----- src/model-store.cpp | 90 ++++++++ src/model-store.h | 19 +- src/pipeline-synth-impl.h | 9 +- src/pipeline-synth-ops.cpp | 220 ++++++++++--------- src/pipeline-synth-ops.h | 18 +- src/pipeline-synth.cpp | 61 ++++-- src/pipeline-synth.h | 26 ++- tests/CMakeLists.txt | 3 + tests/test-dtw-score.cpp | 188 +++++++++++----- tools/ace-server.cpp | 229 ++++++++++++-------- 17 files changed, 874 insertions(+), 406 deletions(-) diff --git a/.github/workflows/ci-build.yml b/.github/workflows/ci-build.yml index 69ee9af1..3a7f69e9 100644 --- a/.github/workflows/ci-build.yml +++ b/.github/workflows/ci-build.yml @@ -41,6 +41,9 @@ jobs: ./build/ace-synth --help 2>&1 | head -5 ./build/quantize --help 2>&1 | head -3 + - name: Unit tests + run: ctest --test-dir build --output-on-failure + lint: name: Lint & Static Analysis runs-on: ubuntu-latest diff --git a/.github/workflows/scoring-drift-check.yml b/.github/workflows/scoring-drift-check.yml index 69c5dfcf..6a4426ae 100644 --- a/.github/workflows/scoring-drift-check.yml +++ b/.github/workflows/scoring-drift-check.yml @@ -3,10 +3,12 @@ # ace-step/ACE-Step-1.5@82252c24 # acestep/core/scoring/_dtw.py # acestep/core/scoring/dit_score.py +# acestep/core/generation/handler/lyric_score.py +# acestep/core/generation/handler/lyric_alignment_common.py # # If the upstream files change, this workflow fails and the C++ port # needs re-evaluation. The expected SHA256 hashes are embedded below -# so there's no separate file to maintain — update both hashes and +# so there's no separate file to maintain — update all four hashes and # the pin comment in dtw-score.h when re-syncing. name: Scoring Drift Check @@ -15,15 +17,21 @@ on: - cron: '0 8 * * 1' # weekly: Monday 08:00 UTC workflow_dispatch: +permissions: {} + jobs: check: runs-on: ubuntu-latest + timeout-minutes: 5 steps: - name: Fetch upstream Python scoring files run: | - BASE="https://raw.githubusercontent.com/ace-step/ACE-Step-1.5/82252c24" - curl -fsSL "$BASE/acestep/core/scoring/_dtw.py" -o _dtw.py - curl -fsSL "$BASE/acestep/core/scoring/dit_score.py" -o dit_score.py + BASE="https://raw.githubusercontent.com/ace-step/ACE-Step-1.5/main" + CURL=(curl -fsSL --max-time 30 --retry 2) + "${CURL[@]}" "$BASE/acestep/core/scoring/_dtw.py" -o _dtw.py + "${CURL[@]}" "$BASE/acestep/core/scoring/dit_score.py" -o dit_score.py + "${CURL[@]}" "$BASE/acestep/core/generation/handler/lyric_score.py" -o lyric_score.py + "${CURL[@]}" "$BASE/acestep/core/generation/handler/lyric_alignment_common.py" -o lyric_alignment_common.py - name: Verify hashes match pinned snapshot run: | @@ -32,12 +40,18 @@ jobs: # re-syncing the C++ port after an upstream algorithm change. EXPECTED_DTW="2d2252e7108f296cd26722c7b622e1c1a1ed47c71459a8ff873164c4c1ada962" EXPECTED_SCORE="73700ff976549339dad3d6c500f8cbe9a8cce08d1979bbb85e3e3ed53e0605a1" + EXPECTED_HANDLER="92ccffbe079482b9e0e7a7fc6304a408fa7d996a7f9e88c7b3f819ee5d463510" + EXPECTED_COMMON="ec35e392a2354696e9def4f98db27c415f09763e839609dd368532971ec0eb08" ACTUAL_DTW=$(sha256sum _dtw.py | cut -d' ' -f1) ACTUAL_SCORE=$(sha256sum dit_score.py | cut -d' ' -f1) + ACTUAL_HANDLER=$(sha256sum lyric_score.py | cut -d' ' -f1) + ACTUAL_COMMON=$(sha256sum lyric_alignment_common.py | cut -d' ' -f1) echo "_dtw.py: expected=$EXPECTED_DTW actual=$ACTUAL_DTW" echo "dit_score.py: expected=$EXPECTED_SCORE actual=$ACTUAL_SCORE" + echo "lyric_score.py: expected=$EXPECTED_HANDLER actual=$ACTUAL_HANDLER" + echo "lyric_alignment_common.py: expected=$EXPECTED_COMMON actual=$ACTUAL_COMMON" if [ "$ACTUAL_DTW" != "$EXPECTED_DTW" ]; then echo "::error::_dtw.py has drifted from the pinned commit (82252c24)." @@ -51,4 +65,16 @@ jobs: exit 1 fi + if [ "$ACTUAL_HANDLER" != "$EXPECTED_HANDLER" ]; then + echo "::error::lyric_score.py has drifted from the pinned commit (82252c24)." + echo "::error::The C++ scoring forward path may need updating." + exit 1 + fi + + if [ "$ACTUAL_COMMON" != "$EXPECTED_COMMON" ]; then + echo "::error::lyric_alignment_common.py has drifted from the pinned commit (82252c24)." + echo "::error::Lyric slicing or model-specific head configuration may need updating." + exit 1 + fi + echo "No drift detected. C++ port matches pinned Python reference." diff --git a/CMakeLists.txt b/CMakeLists.txt index a8792bb3..1f490374 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,6 +1,8 @@ cmake_minimum_required(VERSION 3.14) project(acestep-ggml LANGUAGES C CXX) +include(CTest) + set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) diff --git a/README.md b/README.md index 62c9c622..021e925b 100644 --- a/README.md +++ b/README.md @@ -114,7 +114,7 @@ Debug:
API endpoints -The server exposes four POST endpoints and two GET endpoints: +The server exposes five POST endpoints and two GET endpoints: **POST /lm** - Generate lyrics and audio codes from a caption. Returns JSON. @@ -126,6 +126,13 @@ win over audio when both are sent on the same side). **POST /understand** - Reverse pipeline: audio in, metadata + lyrics + codes out. Multipart only (source audio or pre-encoded latents required, optional request JSON for params). +**POST /score** - Score lyric alignment for generated DiT latents. Multipart +only: send `request` as one AceRequest or an array and `pred_latents` as raw +f32 `[batch, T, 64]` data (`latent` is accepted as an alias). A single request +can reuse the latent part returned by `/synth`; batched latents are concatenated +in request order. Results contain separate pure-noise (`lm`) and regressed +generated-latent (`dit`) metrics. + **POST /vae** - Standalone VAE entrypoint: send `audio` to encode (latents out), send `src_latents` to decode (audio out). Multipart only, the two inputs are mutually exclusive. Lets the webui cache a latent on an diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 38cb5fa4..7934d10d 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -821,10 +821,10 @@ in one GPU pass. HTTP server exposing the same pipelines as `ace-lm`, `ace-synth`, and `ace-understand`. One binary, one port. -POST /lm, POST /synth, POST /understand and POST /vae are all **asynchronous**: they -return a job ID immediately, push the request to a FIFO queue, and the single -worker thread processes jobs in order. Clients poll GET /job?id=N for status -and fetch results with GET /job?id=N&result=1. +POST /lm, POST /synth, POST /understand, POST /score and POST /vae are all +**asynchronous**: they return a job ID immediately, push the request to a FIFO +queue, and the single worker thread processes jobs in order. Clients poll +GET /job?id=N for status and fetch results with GET /job?id=N&result=1. Cancel: POST /job?id=N&cancel=1 stops a specific job. `--models` scans a directory for GGUF files and classifies each by its @@ -846,6 +846,7 @@ quantisation raise both columns. | LM | Qwen3 LM + KV cache | ~2-3 GB | ~2-3 GB | | Synth | Qwen3 text-enc, cond-enc, DiT, VAE enc, VAE dec, FSQ tok/detok | ~2-3 GB (DiT or VAE tiles) | ~3-4 GB + tiles | | Understand | Qwen3 LM, VAE enc, FSQ tok | ~2-3 GB (LM or VAE tiles) | ~2-3 GB + tiles | +| Score | Qwen3 text-enc, cond-enc, DiT | ~2-3 GB (DiT) | ~2-3 GB | VAE tile activations scale with `--vae-chunk` and `--vae-overlap`. Bigger tiles process audio faster with fewer seams but cost more transient VRAM @@ -913,11 +914,18 @@ POST /understand Submit understand, returns job ID body: multipart/form-data (audio or src_latents required, optional request JSON) response: {"id":"3"} +POST /score Submit generated-latent lyric scoring, returns job ID + body: multipart/form-data + 'request': AceRequest or [AceRequest, ...] + 'pred_latents': raw f32 [batch, T, 64] ('latent' alias accepted) + batched requests must share duration, inference steps, model and adapter + response: {"id":"4"} + POST /vae Submit VAE encode or decode, returns job ID body: multipart/form-data (exactly one of 'audio' or 'src_latents') 'audio' -> encode path (latents out) 'src_latents' -> decode path (audio out) - response: {"id":"4"} + response: {"id":"5"} GET /job?id=N Poll job status response: {"status":"running|done|failed|cancelled"} @@ -926,6 +934,7 @@ GET /job?id=N&result=1 Fetch job result lm: application/json [AceRequest, ...] synth: multipart/mixed (one audio part + one latent part per track, paired) understand: multipart/mixed (one json part + one latent part for the source) + score: application/json [{"lm_score":...,"dit_score":...,"lm":{...},"dit":{...}}, ...] vae encode: application/octet-stream (raw .vae bytes, no audio echo: client already has it) vae decode: audio/mpeg or audio/wav (raw, no latent echo: client already has it) @@ -944,10 +953,13 @@ GET /logs SSE stream of server stderr GET / Embedded WebUI (gzipped HTML) ``` -Latent payload format (src_latents, ref_latents, synth/understand response latent parts, /vae encode response body): +Latent payload format (src_latents, ref_latents, pred_latents, synth/understand +response latent parts, /vae encode response body): raw f32 little-endian, flat [T, 64], no header. T = size / 256. Same byte layout neural-codec writes as `.vae` files. Hard cap T <= 15000 frames -(matches the silence_latent buffer baked into the DiT GGUF), 413 over. +(matches the silence_latent buffer baked into the DiT GGUF), 413 over. For a +batched score request, concatenate one equal-length `[T, 64]` tensor per JSON +request in the same order. `lm_model`, `synth_model`, `adapter`, `adapter_scale` fields in the JSON body select which model and adapter to load. `lm_mode` picks the LM instruction diff --git a/src/dit-sampler.h b/src/dit-sampler.h index 7d623da7..01fb4a16 100644 --- a/src/dit-sampler.h +++ b/src/dit-sampler.h @@ -706,29 +706,29 @@ static int dit_ggml_generate(DiTGGML * model, // // Builds the same DiT graph as generation but forces the f32 attention path // for the specified score_layers, marking the softmax attention weights as -// graph outputs. Runs one forward pass at t=0 (pure noise input) and reads -// back the cross-attention scores. +// graph outputs. Runs one forward pass at the supplied flow-matching timestep +// and reads back the cross-attention scores for every batch item. // // The caller is responsible for: // - text encoding (enc_hidden, enc_S, real_enc_S) // - context build (context_latents) -// - noise generation (noise) +// - constructing the latent state for the supplied timestep // - running calculate_lyric_score() on the returned attention matrices // -// attention_out: filled with [n_score_layers][Nh][enc_S][S] attention weights -// (one vector per captured layer, each enc_S*S*Nh floats, row-major). -// Only the conditional (first) batch sample is extracted. +// attention_out: one vector per captured layer. Each vector uses GGML layout +// [enc_S, S, Nh, N], where enc_S is the contiguous dimension. // Returns 0 on success, -1 on error. -static int dit_ggml_score_forward(DiTGGML * model, - const float * noise, - const float * context_latents, - const float * enc_hidden_data, - int enc_S, - const int * real_enc_S, - int T, - int N, - const int * score_layers, - int n_score_layers, +static int dit_ggml_score_forward(DiTGGML * model, + const float * latents, + const float * context_latents, + const float * enc_hidden_data, + int enc_S, + const int * real_enc_S, + int T, + int N, + float timestep, + const int * score_layers, + int n_score_layers, std::vector> & attention_out) { DiTGGMLConfig & c = model->cfg; int Oc = c.out_channels; @@ -737,18 +737,17 @@ static int dit_ggml_score_forward(DiTGGML * model, int S = T / c.patch_size; int n_per = T * Oc; - // Scoring uses a single forward pass (no CFG, no sampling loop). - // The Python reference runs batch=2 (noise + partial) but we only need - // the conditional pass — the attention patterns from the conditional - // branch are what the scorer consumes. + // Scoring uses a single conditional forward pass (no CFG or sampling + // loop). The caller invokes this for the pure-noise and regressed-latent + // states used by the Python reference. int N_graph = N; - fprintf(stderr, "[DiT-Score] Forward: N=%d, T=%d, S=%d, enc_S=%d, %d score layers\n", - N, T, S, enc_S, n_score_layers); + fprintf(stderr, "[DiT-Score] Forward: N=%d, T=%d, S=%d, enc_S=%d, %d score layers\n", N, T, S, enc_S, + n_score_layers); // Graph context - size_t ctx_size = ggml_tensor_overhead() * 8192 + ggml_graph_overhead_custom(8192, false); - std::vector ctx_buf(ctx_size); + size_t ctx_size = ggml_tensor_overhead() * 8192 + ggml_graph_overhead_custom(8192, false); + std::vector ctx_buf(ctx_size); struct ggml_init_params gparams = { /*.mem_size =*/ctx_size, /*.mem_buffer =*/ctx_buf.data(), @@ -758,8 +757,8 @@ static int dit_ggml_score_forward(DiTGGML * model, struct ggml_tensor * t_input = NULL; struct ggml_tensor * t_output = NULL; - struct ggml_cgraph * gf = dit_ggml_build_graph(model, ctx, T, enc_S, N_graph, &t_input, &t_output, - score_layers, n_score_layers); + struct ggml_cgraph * gf = + dit_ggml_build_graph(model, ctx, T, enc_S, N_graph, &t_input, &t_output, score_layers, n_score_layers); fprintf(stderr, "[DiT-Score] Graph: %d nodes\n", ggml_graph_n_nodes(gf)); @@ -785,12 +784,11 @@ static int dit_ggml_score_forward(DiTGGML * model, return -1; } - // Set timesteps to 0 (matches Python reference which scores at t=0) + // Python passes t_r=t, so the reference-time embedding receives zero. struct ggml_tensor * t_t = ggml_graph_get_tensor(gf, "t"); struct ggml_tensor * t_tr = ggml_graph_get_tensor(gf, "t_r"); - float t_zero = 0.0f; - ggml_backend_tensor_set(t_t, &t_zero, 0, sizeof(float)); - ggml_backend_tensor_set(t_tr, &t_zero, 0, sizeof(float)); + ggml_backend_tensor_set(t_t, ×tep, 0, sizeof(float)); + ggml_backend_tensor_set(t_tr, ×tep, 0, sizeof(float)); // Positions struct ggml_tensor * t_pos = ggml_graph_get_tensor(gf, "positions"); @@ -802,15 +800,16 @@ static int dit_ggml_score_forward(DiTGGML * model, } ggml_backend_tensor_set(t_pos, pos_data.data(), 0, S * N_graph * sizeof(int32_t)); - // Self-attention mask (full attention for scoring — no sliding window) + // Self-attention mask for the model's sliding-window layers. Full-attention + // layers do not consume this tensor, matching the generation graph. struct ggml_tensor * t_sa_mask_sw = ggml_graph_get_tensor(gf, "sa_mask_sw"); - int win = c.sliding_window; + int win = c.sliding_window; std::vector sa_sw_data(S * S * N_graph); for (int b = 0; b < N_graph; b++) { for (int qi = 0; qi < S; qi++) { for (int ki = 0; ki < S; ki++) { - int dist = (qi > ki) ? (qi - ki) : (ki - qi); - bool in_win = (win <= 0) || (S <= win) || (dist <= win); + int dist = (qi > ki) ? (qi - ki) : (ki - qi); + bool in_win = (win <= 0) || (S <= win) || (dist <= win); sa_sw_data[b * S * S + qi * S + ki] = ggml_fp32_to_fp16(in_win ? 0.0f : -INFINITY); } } @@ -824,7 +823,7 @@ static int dit_ggml_score_forward(DiTGGML * model, int re = real_enc_S ? real_enc_S[b] : enc_S; for (int qi = 0; qi < S; qi++) { for (int ki = 0; ki < enc_S; ki++) { - float v = (ki < re) ? 0.0f : -INFINITY; + float v = (ki < re) ? 0.0f : -INFINITY; ca_data[b * enc_S * S + qi * enc_S + ki] = ggml_fp32_to_fp16(v); } } @@ -834,13 +833,13 @@ static int dit_ggml_score_forward(DiTGGML * model, // Encoder hidden states ggml_backend_tensor_set(t_enc, enc_hidden_data, 0, H_enc * enc_S * N * sizeof(float)); - // Input: context_latents + noise + // Input: context_latents + latent state at this timestep. std::vector input_buf(in_ch * T * N_graph); for (int b = 0; b < N; b++) { for (int t = 0; t < T; t++) { memcpy(&input_buf[b * T * in_ch + t * in_ch], &context_latents[b * T * ctx_ch + t * ctx_ch], ctx_ch * sizeof(float)); - memcpy(&input_buf[b * T * in_ch + t * in_ch + ctx_ch], &noise[b * n_per + t * Oc], Oc * sizeof(float)); + memcpy(&input_buf[b * T * in_ch + t * in_ch + ctx_ch], &latents[b * n_per + t * Oc], Oc * sizeof(float)); } } ggml_backend_tensor_set(t_input, input_buf.data(), 0, in_ch * T * N_graph * sizeof(float)); @@ -860,17 +859,17 @@ static int dit_ggml_score_forward(DiTGGML * model, } // Scores shape: [enc_S, S, Nh, N] (ne[0]=enc_S, ne[1]=S, ne[2]=Nh, ne[3]=N) - // Extract only batch 0: [enc_S, S, Nh] = enc_S * S * Nh floats - int layer_enc_S = (int) t_scores->ne[0]; - int layer_S = (int) t_scores->ne[1]; - int layer_Nh = (int) t_scores->ne[2]; - size_t batch0_elems = (size_t) layer_enc_S * layer_S * layer_Nh; + int layer_enc_S = (int) t_scores->ne[0]; + int layer_S = (int) t_scores->ne[1]; + int layer_Nh = (int) t_scores->ne[2]; + int layer_N = (int) t_scores->ne[3]; + size_t layer_elems = (size_t) layer_enc_S * layer_S * layer_Nh * layer_N; - attention_out[sl].resize(batch0_elems); - ggml_backend_tensor_get(t_scores, attention_out[sl].data(), 0, batch0_elems * sizeof(float)); + attention_out[sl].resize(layer_elems); + ggml_backend_tensor_get(t_scores, attention_out[sl].data(), 0, layer_elems * sizeof(float)); - fprintf(stderr, "[DiT-Score] Layer %d: %d heads, [%d x %d] attention matrix\n", - score_layers[sl], layer_Nh, layer_enc_S, layer_S); + fprintf(stderr, "[DiT-Score] Layer %d: %d heads, %d batch, [%d x %d] attention matrix\n", score_layers[sl], + layer_Nh, layer_N, layer_enc_S, layer_S); } ggml_free(ctx); diff --git a/src/dtw-score.h b/src/dtw-score.h index 5b867c43..ae0e49e0 100644 --- a/src/dtw-score.h +++ b/src/dtw-score.h @@ -4,6 +4,8 @@ // Direct port of the Python ACE-Step scoring modules: // acestep/core/scoring/_dtw.py — DTW + median filter (numba-jitted) // acestep/core/scoring/dit_score.py — MusicLyricScorer (coverage, monotonicity, confidence) +// Integration parity follows generation/handler/lyric_score.py and +// lyric_alignment_common.py for latent states, lyric slicing, and model heads. // // Pinned to ace-step/ACE-Step-1.5 commit 82252c24 (2026-07-09). // If the Python scoring algorithm changes upstream, this port will need @@ -41,7 +43,9 @@ struct DTWPath { static DTWPath dtw_backtrace(std::vector & trace, int N, int M) { // Boundary handling (matches _dtw.py:61-62) // trace is (N+1) x (M+1), row-major - auto T = [&](int i, int j) -> float & { return trace[(size_t) i * (M + 1) + j]; }; + auto T = [&](int i, int j) -> float & { + return trace[(size_t) i * (M + 1) + j]; + }; for (int j = 0; j <= M; j++) { T(0, j) = 2; } @@ -50,7 +54,7 @@ static DTWPath dtw_backtrace(std::vector & trace, int N, int M) { } // Pre-allocate (max path length = N + M), fill from the end - int max_path_len = N + M; + int max_path_len = N + M; std::vector path_text(max_path_len, 0); std::vector path_time(max_path_len, 0); @@ -75,7 +79,7 @@ static DTWPath dtw_backtrace(std::vector & trace, int N, int M) { } } - int start = path_idx + 1; + int start = path_idx + 1; DTWPath result; result.text_idx.assign(path_text.begin() + start, path_text.end()); result.time_idx.assign(path_time.begin() + start, path_time.end()); @@ -94,7 +98,9 @@ static DTWPath dtw_cpu(const float * x, int N, int M) { std::vector cost((size_t) (N + 1) * (M + 1), std::numeric_limits::infinity()); std::vector trace((size_t) (N + 1) * (M + 1), -1.0f); - auto C = [&](int i, int j) -> float & { return cost[(size_t) i * (M + 1) + j]; }; + auto C = [&](int i, int j) -> float & { + return cost[(size_t) i * (M + 1) + j]; + }; C(0, 0) = 0.0f; @@ -117,7 +123,7 @@ static DTWPath dtw_cpu(const float * x, int N, int M) { t = 2; } - C(i, j) = x[(size_t) (i - 1) * M + (j - 1)] + c; + C(i, j) = x[(size_t) (i - 1) * M + (j - 1)] + c; trace[(size_t) i * (M + 1) + j] = t; } } @@ -137,7 +143,7 @@ static DTWPath dtw_cpu(const float * x, int N, int M) { // Python uses F.pad(mode="reflect") then unfold + sort + pick middle. // We replicate with a sliding window sort. static std::vector median_filter_1d(const std::vector & x, int filter_width) { - int n = (int) x.size(); + int n = (int) x.size(); int pad = filter_width / 2; if (n <= pad) { return x; @@ -146,7 +152,7 @@ static std::vector median_filter_1d(const std::vector & x, int fil // Reflect padding (matches torch F.pad mode="reflect") std::vector padded(n + 2 * pad); for (int i = 0; i < pad; i++) { - padded[i] = x[pad - i]; // reflect from left + padded[i] = x[pad - i]; // reflect from left } for (int i = 0; i < n; i++) { padded[pad + i] = x[i]; @@ -157,7 +163,7 @@ static std::vector median_filter_1d(const std::vector & x, int fil std::vector result(n); std::vector window(filter_width); - int mid = filter_width / 2; + int mid = filter_width / 2; for (int i = 0; i < n; i++) { for (int j = 0; j < filter_width; j++) { @@ -174,7 +180,10 @@ static std::vector median_filter_1d(const std::vector & x, int fil // Returns [rows, cols] row-major. // // Matches median_filter in _dtw.py:90-110 which operates on the last dim. -static std::vector median_filter_2d_rows(const std::vector & input, int rows, int cols, int filter_width) { +static std::vector median_filter_2d_rows(const std::vector & input, + int rows, + int cols, + int filter_width) { std::vector result((size_t) rows * cols); for (int r = 0; r < rows; r++) { std::vector row(input.begin() + (size_t) r * cols, input.begin() + (size_t) r * cols + cols); @@ -189,20 +198,131 @@ static std::vector median_filter_2d_rows(const std::vector & input // Ported from acestep/core/scoring/dit_score.py:MusicLyricScorer // ============================================================================ -// Default layer/head configuration for cross-attention extraction. +// Default 2B layer/head configuration for cross-attention extraction. // Matches the Python default: dit_score.py uses custom_config with // {2: [6], 3: [10, 11], 4: [3], 5: [8, 9], 6: [8]} // 5 layers, 7 heads total. These are the cross-attention heads that -// best track lyric-to-audio alignment. +// best track lyric-to-audio alignment. Layer values are absolute DiT layer +// numbers. Callers that compact captured layers must remap them before +// preprocess_attention() indexes the compact attention buffer. struct ScoreLayerHeadConfig { int layer; int head; }; -static const ScoreLayerHeadConfig DEFAULT_SCORE_HEADS[] = { - {2, 6}, {3, 10}, {3, 11}, {4, 3}, {5, 8}, {5, 9}, {6, 8}, +static const ScoreLayerHeadConfig DEFAULT_2B_SCORE_HEADS[] = { + { 2, 6 }, + { 3, 10 }, + { 3, 11 }, + { 4, 3 }, + { 5, 8 }, + { 5, 9 }, + { 6, 8 }, }; -static const int DEFAULT_SCORE_HEADS_COUNT = 7; +static const int DEFAULT_2B_SCORE_HEADS_COUNT = 7; + +// Remap absolute DiT layer numbers to compact slots in captured_layers. +// Returns false when any configured layer/head is unavailable; silently +// scoring a partial or model-incompatible configuration produces misleading +// alignment values. +static bool remap_score_heads(const ScoreLayerHeadConfig * config, + int config_count, + const int * captured_layers, + int captured_count, + int n_heads, + std::vector & remapped) { + remapped.clear(); + if (!config || config_count <= 0 || !captured_layers || captured_count <= 0 || n_heads <= 0) { + return false; + } + + for (int c = 0; c < config_count; c++) { + if (config[c].head < 0 || config[c].head >= n_heads) { + remapped.clear(); + return false; + } + + int compact_layer = -1; + for (int i = 0; i < captured_count; i++) { + if (captured_layers[i] == config[c].layer) { + compact_layer = i; + break; + } + } + if (compact_layer < 0) { + remapped.clear(); + return false; + } + remapped.push_back({ compact_layer, config[c].head }); + } + return !remapped.empty(); +} + +// Extract one batch item's pure-lyric rows from a GGML attention output. +// source layout is [tokens, frames, heads, batches] with tokens (ne[0]) +// contiguous. destination layout is [heads, lyric_tokens, frames], matching +// the per-layer slice expected by preprocess_attention(). +static bool extract_attention_slice(const float * source, + size_t source_size, + int tokens, + int frames, + int heads, + int batches, + int batch_index, + int token_start, + int token_count, + std::vector & destination) { + destination.clear(); + if (!source || tokens <= 0 || frames <= 0 || heads <= 0 || batches <= 0 || batch_index < 0 || + batch_index >= batches || token_start < 0 || token_count <= 0 || token_start + token_count > tokens) { + return false; + } + + size_t expected = (size_t) tokens * frames * heads * batches; + if (source_size < expected) { + return false; + } + + destination.resize((size_t) heads * token_count * frames); + for (int h = 0; h < heads; h++) { + for (int t = 0; t < token_count; t++) { + for (int f = 0; f < frames; f++) { + size_t src_idx = (size_t) (token_start + t) + + (size_t) tokens * (f + (size_t) frames * (h + (size_t) heads * batch_index)); + size_t dst_idx = ((size_t) h * token_count + t) * frames + f; + destination[dst_idx] = source[src_idx]; + } + } + } + return true; +} + +struct LyricTokenSegment { + int start; + std::vector token_ids; +}; + +// Match Python _extract_lyric_segment(): strip the generated prompt header and +// stop at the first end-of-text token. The returned start is also the row +// offset into encoder_hidden_states because the condition encoder packs lyric +// tokens first. +static LyricTokenSegment extract_lyric_token_segment(const std::vector & raw_ids, int header_tokens, int eos_id) { + LyricTokenSegment result = { header_tokens, {} }; + if (header_tokens < 0 || header_tokens > (int) raw_ids.size()) { + result.start = -1; + return result; + } + + int end = (int) raw_ids.size(); + for (int i = header_tokens; i < end; i++) { + if (raw_ids[i] == eos_id) { + end = i; + break; + } + } + result.token_ids.assign(raw_ids.begin() + header_tokens, raw_ids.begin() + end); + return result; +} // Token type mask: 1 = lyric token, 0 = structural tag (inside [...]). // Ported from dit_score.py:_generate_token_type_mask (lines 32-55). @@ -211,9 +331,9 @@ static const int DEFAULT_SCORE_HEADS_COUNT = 7; // token. Here we accept a pre-decoded vector of token strings from the caller // (the BPE tokenizer is available in the C++ pipeline but not in this header). static std::vector generate_token_type_mask(const std::vector & decoded_tokens) { - int n = (int) decoded_tokens.size(); + int n = (int) decoded_tokens.size(); std::vector mask(n, 1); - bool in_bracket = false; + bool in_bracket = false; for (int i = 0; i < n; i++) { const std::string & s = decoded_tokens[i]; @@ -225,7 +345,7 @@ static std::vector generate_token_type_mask(const std::vector } if (s.find(']') != std::string::npos) { in_bracket = false; - mask[i] = 0; + mask[i] = 0; } } return mask; @@ -245,16 +365,16 @@ static std::vector generate_token_type_mask(const std::vector // calc_matrix: squared energy for DTW pathfinding [tokens, frames] // energy_matrix: normalized energy for scoring [tokens, frames] // Returns false if no valid heads were found. -static bool preprocess_attention(const float * attention, - int n_layers, - int n_heads, - int tokens, - int frames, +static bool preprocess_attention(const float * attention, + int n_layers, + int n_heads, + int tokens, + int frames, const ScoreLayerHeadConfig * config, - int config_count, - int medfilt_width, - std::vector & calc_matrix, - std::vector & energy_matrix) { + int config_count, + int medfilt_width, + std::vector & calc_matrix, + std::vector & energy_matrix) { // 1. Select heads and stack (matches dit_score.py:84-93) std::vector> selected; for (int c = 0; c < config_count; c++) { @@ -294,8 +414,12 @@ static bool preprocess_attention(const float * attention, float e_min = std::numeric_limits::max(); float e_max = std::numeric_limits::lowest(); for (float v : energy_matrix) { - if (v < e_min) e_min = v; - if (v > e_max) e_max = v; + if (v < e_min) { + e_min = v; + } + if (v > e_max) { + e_max = v; + } } if (e_max - e_min > 1e-9f) { @@ -332,18 +456,20 @@ struct AlignmentMetrics { }; static AlignmentMetrics compute_alignment_metrics(const std::vector & energy_matrix, - int rows, - int cols, - const DTWPath & path, - const std::vector & type_mask, - double time_weight = 0.01, - double overlap_frames = 9.0, - double instrumental_weight = 1.0) { - auto E = [&](int i, int j) -> double { return (double) energy_matrix[(size_t) i * cols + j]; }; + int rows, + int cols, + const DTWPath & path, + const std::vector & type_mask, + double time_weight = 0.01, + double overlap_frames = 9.0, + double instrumental_weight = 1.0) { + auto E = [&](int i, int j) -> double { + return (double) energy_matrix[(size_t) i * cols + j]; + }; // ================= A. Coverage Score (dit_score.py:150-161) ================= - int total_sung_rows = 0; - int valid_sung_rows = 0; + int total_sung_rows = 0; + int valid_sung_rows = 0; const double coverage_threshold = 0.1; for (int i = 0; i < rows; i++) { @@ -352,7 +478,9 @@ static AlignmentMetrics compute_alignment_metrics(const std::vector & ene double row_max = 0.0; for (int j = 0; j < cols; j++) { double e = E(i, j); - if (e > row_max) row_max = e; + if (e > row_max) { + row_max = e; + } } if (row_max > coverage_threshold) { valid_sung_rows++; @@ -389,7 +517,7 @@ static AlignmentMetrics compute_alignment_metrics(const std::vector & ene } double monotonicity; - int cnt = (int) sung_centroids.size(); + int cnt = (int) sung_centroids.size(); if (cnt > 1) { double non_decreasing = 0.0; for (int k = 0; k < cnt - 1; k++) { @@ -404,13 +532,13 @@ static AlignmentMetrics compute_alignment_metrics(const std::vector & ene // ================= C. Path Confidence (dit_score.py:193-212) ================= double path_confidence; - int path_len = (int) path.text_idx.size(); + int path_len = (int) path.text_idx.size(); if (path_len > 0) { - double total_energy = 0.0; - double total_steps = 0.0; + double total_energy = 0.0; + double total_steps = 0.0; for (int k = 0; k < path_len; k++) { - int r = path.text_idx[k]; - int c = path.time_idx[k]; + int r = path.text_idx[k]; + int c = path.time_idx[k]; double pe = E(r, c); double sw = (type_mask[r] == 0) ? instrumental_weight : 1.0; total_energy += pe * sw; @@ -432,23 +560,29 @@ static AlignmentMetrics compute_alignment_metrics(const std::vector & ene // config / config_count: which (layer, head) pairs to use // medfilt_width: median filter window (1 = disabled) // -// Returns the final lyrics_score = cov^2 * mono^2 * conf, clipped to [0, 1]. +// Returns the final lyrics_score = cov^2 * mono^2 * conf, clipped to [0, 1] +// and rounded to four decimal places like the Python API. struct LyricScoreResult { - double coverage; - double monotonicity; - double confidence; - double lyrics_score; // cov^2 * mono^2 * conf, clipped [0,1] + double coverage; + double monotonicity; + double confidence; + double lyrics_score; // cov^2 * mono^2 * conf, clipped [0,1] DTWPath path; }; -static LyricScoreResult calculate_lyric_score(const float * attention, +struct LyricScoreComparison { + LyricScoreResult lm; + LyricScoreResult dit; +}; + +static LyricScoreResult calculate_lyric_score(const float * attention, int n_layers, int n_heads, int tokens, int frames, const std::vector & decoded_tokens, - const ScoreLayerHeadConfig * config = DEFAULT_SCORE_HEADS, - int config_count = DEFAULT_SCORE_HEADS_COUNT, + const ScoreLayerHeadConfig * config, + int config_count, int medfilt_width = 1) { LyricScoreResult result = { 0, 0, 0, 0, {} }; @@ -477,16 +611,20 @@ static LyricScoreResult calculate_lyric_score(const float * result.path = dtw_cpu(neg_calc.data(), tokens, frames); // 4. Compute metrics (dit_score.py:313-320) - AlignmentMetrics m = compute_alignment_metrics(energy_matrix, tokens, frames, result.path, type_mask); + AlignmentMetrics m = compute_alignment_metrics(energy_matrix, tokens, frames, result.path, type_mask); result.coverage = m.coverage; result.monotonicity = m.monotonicity; result.confidence = m.confidence; // 5. Final score: cov^2 * mono^2 * conf (dit_score.py:324-325) double final_score = (m.coverage * m.coverage) * (m.monotonicity * m.monotonicity) * m.confidence; - if (final_score < 0.0) final_score = 0.0; - if (final_score > 1.0) final_score = 1.0; - result.lyrics_score = final_score; + if (final_score < 0.0) { + final_score = 0.0; + } + if (final_score > 1.0) { + final_score = 1.0; + } + result.lyrics_score = std::round(final_score * 10000.0) / 10000.0; return result; } diff --git a/src/model-store.cpp b/src/model-store.cpp index 75fda310..24acaedd 100644 --- a/src/model-store.cpp +++ b/src/model-store.cpp @@ -14,15 +14,104 @@ #include "gguf-weights.h" #include "timer.h" +#include "yyjson.h" #include #include +#include #include #include #include #include #include +static void dit_meta_set_default_alignment_config(DiTMeta * meta) { + if (meta->cfg.n_layers == 24 && meta->cfg.n_heads == 16) { + static const DiTScoreHeadConfig defaults[] = { + { 2, 6 }, + { 3, 10 }, + { 3, 11 }, + { 4, 3 }, + { 5, 8 }, + { 5, 9 }, + { 6, 8 }, + }; + meta->lyric_alignment_heads.assign(defaults, defaults + sizeof(defaults) / sizeof(defaults[0])); + return; + } + if (meta->cfg.n_layers == 32 && meta->cfg.n_heads == 32) { + static const DiTScoreHeadConfig defaults[] = { + { 3, 18 }, + { 3, 27 }, + { 4, 22 }, + { 5, 5 }, + { 5, 6 }, + { 5, 7 }, + { 6, 2 }, + { 6, 12 }, + { 6, 13 }, + { 7, 20 }, + { 7, 21 }, + }; + meta->lyric_alignment_heads.assign(defaults, defaults + sizeof(defaults) / sizeof(defaults[0])); + } +} + +static void dit_meta_load_alignment_config(DiTMeta * meta, const GGUFModel & gf) { + const char * config_json = gf_get_str(gf, "acestep.config_json"); + bool invalid = false; + + if (config_json && config_json[0]) { + yyjson_doc * doc = yyjson_read(config_json, strlen(config_json), 0); + yyjson_val * root = doc ? yyjson_doc_get_root(doc) : nullptr; + yyjson_val * config = root ? yyjson_obj_get(root, "lyric_alignment_layers_config") : nullptr; + if (config && yyjson_is_obj(config)) { + size_t idx, max; + yyjson_val * key; + yyjson_val * heads; + yyjson_obj_foreach(config, idx, max, key, heads) { + const char * layer_text = yyjson_get_str(key); + char * end = nullptr; + long layer = layer_text ? strtol(layer_text, &end, 10) : -1; + if (!layer_text || !end || *end != '\0' || layer < 0 || layer >= meta->cfg.n_layers || + !yyjson_is_arr(heads)) { + invalid = true; + break; + } + + size_t head_idx, head_max; + yyjson_val * head_val; + yyjson_arr_foreach(heads, head_idx, head_max, head_val) { + int head = yyjson_is_int(head_val) ? (int) yyjson_get_int(head_val) : -1; + if (head < 0 || head >= meta->cfg.n_heads) { + invalid = true; + break; + } + meta->lyric_alignment_heads.push_back({ (int) layer, head }); + } + if (invalid) { + break; + } + } + } + yyjson_doc_free(doc); + } + + if (invalid) { + fprintf(stderr, "[Store] WARNING: invalid lyric_alignment_layers_config; using architecture default\n"); + meta->lyric_alignment_heads.clear(); + } + if (meta->lyric_alignment_heads.empty()) { + dit_meta_set_default_alignment_config(meta); + } + if (meta->lyric_alignment_heads.empty()) { + fprintf(stderr, "[Store] WARNING: no lyric alignment head config for %dL/%dH DiT\n", meta->cfg.n_layers, + meta->cfg.n_heads); + } else { + fprintf(stderr, "[Store] Lyric alignment: %zu configured heads\n", meta->lyric_alignment_heads.size()); + } +} + namespace { // Key hashing. Only the fields relevant for this ModelKind participate, so @@ -544,6 +633,7 @@ const DiTMeta * store_dit_meta(ModelStore * s, const char * dit_path) { return nullptr; } meta->is_turbo = gf_get_bool(gf, "acestep.is_turbo"); + dit_meta_load_alignment_config(meta, gf); // silence_latent: [15000, 64] f32, also accessible via store_silence for // callers that only need the pointer. Cached here too so DiTMeta is diff --git a/src/model-store.h b/src/model-store.h index e008ba91..27ab9ff3 100644 --- a/src/model-store.h +++ b/src/model-store.h @@ -57,6 +57,7 @@ #include #include +#include struct ModelStore; @@ -87,13 +88,19 @@ enum EvictPolicy { EVICT_NEVER, // --keep-loaded: never evict, accumulate }; -// DiT metadata cached on the CPU: needed by text encoding and T resolution -// before the DiT itself is loaded on the GPU. +struct DiTScoreHeadConfig { + int layer; + int head; +}; + +// DiT metadata cached on the CPU: needed by text encoding, T resolution, and +// model-specific lyric scoring before the DiT itself is loaded on the GPU. struct DiTMeta { - DiTGGMLConfig cfg; - std::vector silence_full; // [15000, 64] f32, from silence_latent tensor - std::vector null_cond_cpu; // [hidden_size] f32, empty when the model has none - bool is_turbo; + DiTGGMLConfig cfg; + std::vector silence_full; // [15000, 64] f32, from silence_latent tensor + std::vector null_cond_cpu; // [hidden_size] f32, empty when the model has none + std::vector lyric_alignment_heads; + bool is_turbo; }; ModelStore * store_create(EvictPolicy policy); diff --git a/src/pipeline-synth-impl.h b/src/pipeline-synth-impl.h index 4722903c..5310d693 100644 --- a/src/pipeline-synth-impl.h +++ b/src/pipeline-synth-impl.h @@ -112,10 +112,11 @@ struct SynthState { std::vector per_enc_S_nc; bool need_enc_switch; - // lyric token IDs (per-batch) — persisted for /score endpoint. - // Set by ops_encode_text, used by ops_score_forward to generate the - // token type mask (lyric vs structural tag) for DTW alignment scoring. - std::vector> per_lyric_ids; + // Pure lyric token IDs and their row offset in the packed condition + // sequence (per batch). The prompt header and trailing end-of-text tokens + // are excluded so /score measures lyrics rather than metadata/caption. + std::vector> per_lyric_ids; + std::vector per_lyric_start; // stacked encoder hidden states int max_enc_S; diff --git a/src/pipeline-synth-ops.cpp b/src/pipeline-synth-ops.cpp index fcab01ce..bd957f4d 100644 --- a/src/pipeline-synth-ops.cpp +++ b/src/pipeline-synth-ops.cpp @@ -11,9 +11,11 @@ #include "dtw-score.h" #include "philox.h" #include "pipeline-synth-impl.h" +#include "sampling.h" #include "task-types.h" #include "vae-enc.h" +#include #include #include #include @@ -403,6 +405,11 @@ struct TextEncForward { int S_lyric; }; +static std::string build_lyric_header(const AceRequest & rb) { + const char * language = rb.vocal_language.empty() ? "unknown" : rb.vocal_language.c_str(); + return std::string("# Languages\n") + language + "\n\n# Lyric\n"; +} + // Build the text/lyric prompt pair that feeds the text encoder for one batch // element. instruction is the DiT instruction header (main or non-cover). static void build_prompt_strings(const AceRequest & rb, @@ -416,14 +423,12 @@ static void build_prompt_strings(const AceRequest & rb, } const char * keyscale_b = rb.keyscale.empty() ? "N/A" : rb.keyscale.c_str(); const char * timesig_b = rb.timesignature.empty() ? "N/A" : rb.timesignature.c_str(); - const char * language_b = rb.vocal_language.empty() ? "unknown" : rb.vocal_language.c_str(); - - char metas_b[512]; + char metas_b[512]; snprintf(metas_b, sizeof(metas_b), "- bpm: %s\n- timesignature: %s\n- keyscale: %s\n- duration: %d seconds\n", bpm_b, timesig_b, keyscale_b, (int) duration); text_out = std::string("# Instruction\n") + instruction + "\n\n" + "# Caption\n" + rb.caption + "\n\n" + "# Metas\n" + metas_b + "<|endoftext|>\n"; - lyric_out = std::string("# Languages\n") + language_b + "\n\n# Lyric\n" + rb.lyrics + "<|endoftext|>"; + lyric_out = build_lyric_header(rb) + rb.lyrics + "<|endoftext|>"; } int ops_encode_text(const AceSynth * ctx, const AceRequest * reqs, int batch_n, SynthState & s) { @@ -442,8 +447,9 @@ int ops_encode_text(const AceSynth * ctx, const AceRequest * reqs, int batch_n, s.need_enc_switch = s.use_source_context && !s.is_repaint && !s.is_lego_region && s.rr.audio_cover_strength < 1.0f; - // Persist lyric token IDs for /score endpoint (per-batch). + // Persist the pure lyric token range for /score (per-batch). s.per_lyric_ids.resize(batch_n); + s.per_lyric_start.assign(batch_n, -1); BPETokenizer * bpe = store_bpe(ctx->store, ctx->params.text_encoder_path); if (!bpe) { @@ -479,7 +485,10 @@ int ops_encode_text(const AceSynth * ctx, const AceRequest * reqs, int batch_n, int S_text = (int) text_ids.size(); int S_lyric = (int) lyric_ids.size(); - s.per_lyric_ids[b] = lyric_ids; // persist for /score + std::vector header_ids = bpe_encode(bpe, build_lyric_header(reqs[b]), false); + LyricTokenSegment segment = extract_lyric_token_segment(lyric_ids, (int) header_ids.size(), bpe->eos_id); + s.per_lyric_ids[b] = std::move(segment.token_ids); + s.per_lyric_start[b] = segment.start; main_fwd[b].S_text = S_text; main_fwd[b].S_lyric = S_lyric; @@ -966,20 +975,19 @@ int ops_vae_decode(const AceSynth * ctx, return 0; } -// Scoring forward pass: extract cross-attention matrices and compute lyric -// alignment metrics. Matches Python ACE-Step dit_score.py:calculate_score. -// -// The Python reference runs the full DiT forward with output_attentions=True -// on the generated audio + lyrics, then extracts cross-attention from -// specific layers/heads to score lyric alignment quality. -// -// Here we run a single DiT forward on the current noise/context/encoder -// state, capture cross-attention from the default score layers, and compute -// coverage/monotonicity/confidence via dtw-score.h. -int ops_score_forward(const AceSynth * ctx, - int batch_n, - std::vector & out_scores, - SynthState & s) { +// Scoring forward pass: build the two latent states from Python +// LyricScoreMixin.get_lyric_score(), capture their cross-attention, slice to +// pure lyric rows, and compute alignment metrics via dtw-score.h. +int ops_score_forward(const AceSynth * ctx, + int batch_n, + const float * pred_latents, + std::vector & out_scores, + SynthState & s) { + if (!pred_latents || batch_n <= 0 || s.num_steps <= 0) { + fprintf(stderr, "[Score] FATAL: predicted latents, batch, and inference steps are required\n"); + return -1; + } + DiTGGML * dit = store_require_dit(ctx->store, ctx->dit_key); if (!dit) { fprintf(stderr, "[Score] FATAL: store_require_dit failed\n"); @@ -991,107 +999,115 @@ int ops_score_forward(const AceSynth * ctx, dit->use_flash_attn = false; } - // Default score layers from Python: {2, 3, 4, 5, 6} - // Heads are selected per-layer in dtw-score.h:DEFAULT_SCORE_HEADS - static const int score_layers[] = { 2, 3, 4, 5, 6 }; - static const int n_score_layers = 5; + if (ctx->meta->lyric_alignment_heads.empty()) { + fprintf(stderr, "[Score] FATAL: model has no lyric alignment layer/head configuration\n"); + return -1; + } + + std::vector absolute_heads; + std::vector score_layers; + for (const DiTScoreHeadConfig & head : ctx->meta->lyric_alignment_heads) { + absolute_heads.push_back({ head.layer, head.head }); + if (std::find(score_layers.begin(), score_layers.end(), head.layer) == score_layers.end()) { + score_layers.push_back(head.layer); + } + } - fprintf(stderr, "[Score] Starting: T=%d, S=%d, enc_S=%d, batch=%d, %d score layers\n", - s.T, s.S, s.enc_S, batch_n, n_score_layers); + std::vector compact_heads; + if (!remap_score_heads(absolute_heads.data(), (int) absolute_heads.size(), score_layers.data(), + (int) score_layers.size(), dit->cfg.n_heads, compact_heads)) { + fprintf(stderr, "[Score] FATAL: invalid lyric alignment layer/head configuration\n"); + return -1; + } + + fprintf(stderr, "[Score] Starting: T=%d, S=%d, enc_S=%d, batch=%d, %d score layers\n", s.T, s.S, s.enc_S, batch_n, + (int) score_layers.size()); + + // Python scores both pure noise (LM score, t=1) and a latent regressed + // toward the generated result (DiT score, t=1/inference_steps). + float t_last = 1.0f / (float) s.num_steps; + size_t latent_count = (size_t) batch_n * s.T * s.Oc; + std::vector regressed_latents(latent_count); + for (size_t i = 0; i < latent_count; i++) { + regressed_latents[i] = t_last * s.noise[i] + (1.0f - t_last) * pred_latents[i]; + } - // Run scoring forward pass with cross-attention capture - std::vector> layer_attentions; + std::vector> lm_attentions; + std::vector> dit_attentions; s.timer.reset(); - int rc = dit_ggml_score_forward(dit, s.noise.data(), s.context.data(), s.enc_hidden.data(), - s.enc_S, s.per_enc_S.data(), s.T, batch_n, - score_layers, n_score_layers, layer_attentions); + int rc = + dit_ggml_score_forward(dit, s.noise.data(), s.context.data(), s.enc_hidden.data(), s.enc_S, s.per_enc_S.data(), + s.T, batch_n, 1.0f, score_layers.data(), (int) score_layers.size(), lm_attentions); if (rc != 0) { - fprintf(stderr, "[Score] FATAL: dit_ggml_score_forward failed\n"); + fprintf(stderr, "[Score] FATAL: pure-noise score forward failed\n"); return -1; } - fprintf(stderr, "[Score] DiT forward: %.1f ms\n", s.timer.ms()); + rc = dit_ggml_score_forward(dit, regressed_latents.data(), s.context.data(), s.enc_hidden.data(), s.enc_S, + s.per_enc_S.data(), s.T, batch_n, t_last, score_layers.data(), + (int) score_layers.size(), dit_attentions); + if (rc != 0) { + fprintf(stderr, "[Score] FATAL: regressed-latent score forward failed\n"); + return -1; + } + fprintf(stderr, "[Score] Attention forwards: %.1f ms\n", s.timer.ms()); - // Get BPE tokenizer for decoding token IDs to strings (type mask) BPETokenizer * bpe = store_bpe(ctx->store, ctx->params.text_encoder_path); if (!bpe) { - fprintf(stderr, "[Score] WARNING: store_bpe failed, using all-lyric mask\n"); + fprintf(stderr, "[Score] FATAL: store_bpe failed\n"); + return -1; } - // DiT config - DiTGGMLConfig & c = dit->cfg; - int Nh = c.n_heads; - int S = s.S; - int enc_S = s.enc_S; + int Nh = dit->cfg.n_heads; + int S = s.S; + int enc_S = s.enc_S; + + auto stack_attention = [&](const std::vector> & attentions, int batch_index, int token_start, + int token_count, std::vector & stacked) -> bool { + size_t layer_size = (size_t) Nh * token_count * S; + stacked.assign(score_layers.size() * layer_size, 0.0f); + for (size_t sl = 0; sl < score_layers.size(); sl++) { + if (sl >= attentions.size() || attentions[sl].empty()) { + return false; + } + std::vector slice; + if (!extract_attention_slice(attentions[sl].data(), attentions[sl].size(), enc_S, S, Nh, batch_n, + batch_index, token_start, token_count, slice)) { + return false; + } + std::copy(slice.begin(), slice.end(), stacked.begin() + (size_t) sl * layer_size); + } + return true; + }; - // Compute scores per batch item out_scores.resize(batch_n); for (int b = 0; b < batch_n; b++) { - // Build decoded token strings for the type mask - std::vector decoded_tokens; - if (b < (int) s.per_lyric_ids.size() && !s.per_lyric_ids[b].empty()) { - int n_lyric = (int) s.per_lyric_ids[b].size(); - decoded_tokens.resize(n_lyric); - if (bpe) { - for (int i = 0; i < n_lyric; i++) { - int tid = s.per_lyric_ids[b][i]; - if (tid >= 0 && tid < bpe->n_vocab) { - decoded_tokens[i] = bpe->id_to_str[tid]; - } - } - } else { - // No tokenizer: treat all as lyric tokens - std::fill(decoded_tokens.begin(), decoded_tokens.end(), std::string("lyric")); - } - } else { - // No lyric IDs: single dummy token - decoded_tokens = { "lyric" }; + if (b >= (int) s.per_lyric_ids.size() || b >= (int) s.per_lyric_start.size() || s.per_lyric_ids[b].empty() || + s.per_lyric_start[b] < 0) { + fprintf(stderr, "[Score] ERROR: batch %d has no pure lyric tokens\n", b); + return -1; } - // Stack attention from all captured layers into [n_layers, Nh, enc_S, S] - // Each layer_attentions[sl] is [enc_S, S, Nh] (batch 0 only) - // We need [n_layers, n_heads, tokens, frames] for calculate_lyric_score - // But the attention matrix is [enc_S (KV=tokens), S (Q=frames), Nh] - // So: tokens = enc_S, frames = S, n_heads = Nh, n_layers = n_score_layers - // - // Note: the Python reference uses per-lyric-token attention (enc_S = lyric_tokens). - // Here enc_S is the full encoder sequence (text + lyric). We use the full - // sequence and the type mask to distinguish lyric from structural tokens. - int n_heads_total = Nh; - int n_layers_total = n_score_layers; - - // Stack: [n_layers * n_heads * enc_S * S] - std::vector stacked((size_t) n_layers_total * n_heads_total * enc_S * S, 0.0f); - - for (int sl = 0; sl < n_score_layers; sl++) { - if (sl >= (int) layer_attentions.size() || layer_attentions[sl].empty()) { - continue; - } - // layer_attentions[sl] is [enc_S, S, Nh] row-major - // We need [n_layers, n_heads, enc_S, S] = layer * (Nh * enc_S * S) + head * (enc_S * S) + ... - const float * src = layer_attentions[sl].data(); - for (int h = 0; h < Nh; h++) { - for (int t = 0; t < enc_S; t++) { - for (int f = 0; f < S; f++) { - // src index: [enc_S, S, Nh] = t * (S * Nh) + f * Nh + h - size_t src_idx = (size_t) t * S * Nh + (size_t) f * Nh + h; - // dst index: [n_layers, n_heads, enc_S, S] = sl * (Nh * enc_S * S) + h * (enc_S * S) + t * S + f - size_t dst_idx = (size_t) sl * n_heads_total * enc_S * S + - (size_t) h * enc_S * S + - (size_t) t * S + f; - if (src_idx < layer_attentions[sl].size()) { - stacked[dst_idx] = src[src_idx]; - } - } - } - } + int lyric_tokens = (int) s.per_lyric_ids[b].size(); + std::vector decoded_tokens(lyric_tokens); + for (int i = 0; i < lyric_tokens; i++) { + decoded_tokens[i] = bpe_decode(*bpe, std::vector{ s.per_lyric_ids[b][i] }); + } + + std::vector lm_stacked; + std::vector dit_stacked; + if (!stack_attention(lm_attentions, b, s.per_lyric_start[b], lyric_tokens, lm_stacked) || + !stack_attention(dit_attentions, b, s.per_lyric_start[b], lyric_tokens, dit_stacked)) { + fprintf(stderr, "[Score] FATAL: failed to slice attention for batch %d\n", b); + return -1; } - // Calculate lyric score - LyricScoreResult result = calculate_lyric_score( - stacked.data(), n_layers_total, n_heads_total, enc_S, S, decoded_tokens); + LyricScoreComparison result; + result.lm = calculate_lyric_score(lm_stacked.data(), (int) score_layers.size(), Nh, lyric_tokens, S, + decoded_tokens, compact_heads.data(), (int) compact_heads.size()); + result.dit = calculate_lyric_score(dit_stacked.data(), (int) score_layers.size(), Nh, lyric_tokens, S, + decoded_tokens, compact_heads.data(), (int) compact_heads.size()); - fprintf(stderr, "[Score] Batch %d: coverage=%.4f monotonicity=%.4f confidence=%.4f lyrics_score=%.4f\n", - b, result.coverage, result.monotonicity, result.confidence, result.lyrics_score); + fprintf(stderr, "[Score] Batch %d: lm=%.4f dit=%.4f\n", b, result.lm.lyrics_score, result.dit.lyrics_score); out_scores[b] = result; } diff --git a/src/pipeline-synth-ops.h b/src/pipeline-synth-ops.h index e2a05393..9aff511a 100644 --- a/src/pipeline-synth-ops.h +++ b/src/pipeline-synth-ops.h @@ -71,17 +71,17 @@ int ops_vae_decode(const AceSynth * ctx, bool (*cancel)(void *), void * cancel_data); -// Scoring primitive: single DiT forward with cross-attention capture. -// Runs one forward pass on the current noise + context + encoder states, -// extracts cross-attention matrices from the configured score layers, -// and computes lyric alignment metrics (coverage, monotonicity, confidence). +// Scoring primitive: DiT cross-attention capture for the Python reference's +// pure-noise and regressed-latent states. pred_latents is the generated DiT +// output [batch_n, T, 64]. // // Requires ops_encode_text (for s.per_lyric_ids) and ops_build_context + // ops_init_noise (for noise + context latents) to have been run first. // -// out_scores: filled with one LyricScoreResult per batch item. +// out_scores: filled with one LM/DiT score comparison per batch item. // Returns 0 on success, -1 on error. -int ops_score_forward(const AceSynth * ctx, - int batch_n, - std::vector & out_scores, - SynthState & s); +int ops_score_forward(const AceSynth * ctx, + int batch_n, + const float * pred_latents, + std::vector & out_scores, + SynthState & s); diff --git a/src/pipeline-synth.cpp b/src/pipeline-synth.cpp index d5b61098..d696c92c 100644 --- a/src/pipeline-synth.cpp +++ b/src/pipeline-synth.cpp @@ -35,7 +35,7 @@ void ace_synth_default_params(AceSynthParams * p) { p->dump_dir = NULL; } -AceSynth * ace_synth_load(ModelStore * store, const AceSynthParams * params) { +static AceSynth * ace_synth_load_impl(ModelStore * store, const AceSynthParams * params, bool require_vae) { if (!store || !params) { fprintf(stderr, "[Synth-Load] ERROR: store and params are required\n"); return NULL; @@ -48,7 +48,7 @@ AceSynth * ace_synth_load(ModelStore * store, const AceSynthParams * params) { fprintf(stderr, "[Synth-Load] ERROR: text_encoder_path is NULL\n"); return NULL; } - if (!params->vae_path) { + if (require_vae && !params->vae_path) { fprintf(stderr, "[Synth-Load] ERROR: vae_path is NULL\n"); return NULL; } @@ -89,10 +89,10 @@ AceSynth * ace_synth_load(ModelStore * store, const AceSynthParams * params) { ctx->dit_key.adapter_scale = params->adapter_scale; ctx->vae_enc_key.kind = MODEL_VAE_ENC; - ctx->vae_enc_key.path = params->vae_path; + ctx->vae_enc_key.path = params->vae_path ? params->vae_path : ""; ctx->vae_dec_key.kind = MODEL_VAE_DEC; - ctx->vae_dec_key.path = params->vae_path; + ctx->vae_dec_key.path = params->vae_path ? params->vae_path : ""; fprintf(stderr, "[Synth-Load] Ready: turbo=%s, fa=%s, batch_cfg=%s\n", ctx->meta->is_turbo ? "yes" : "no", params->use_fa ? "yes" : "no", params->use_batch_cfg ? "yes" : "no"); @@ -106,6 +106,14 @@ AceSynth * ace_synth_load(ModelStore * store, const AceSynthParams * params) { return ctx; } +AceSynth * ace_synth_load(ModelStore * store, const AceSynthParams * params) { + return ace_synth_load_impl(store, params, true); +} + +AceSynth * ace_synth_load_score(ModelStore * store, const AceSynthParams * params) { + return ace_synth_load_impl(store, params, false); +} + // Allocate job and init the SynthState fields every task poses the same way. static AceSynthJob * alloc_job(AceSynth * ctx, const AceRequest * reqs, int batch_n) { AceSynthJob * job = new AceSynthJob(); @@ -696,27 +704,35 @@ void ace_synth_free(AceSynth * ctx) { delete ctx; } -// Scoring: run phase 1 setup (encode, context, noise) then a single DiT -// forward with cross-attention capture instead of the full denoising loop. -// Matches Python ACE-Step dit_score.py:get_lyric_score flow. -int ace_synth_score(AceSynth * ctx, - const AceRequest * reqs, - int batch_n, - std::vector & out_scores) { - if (!ctx || !reqs || batch_n < 1 || batch_n > 9) { // same limit as ace_synth_job_run +// Scoring: run phase 1 setup, then compare pure-noise and regressed generated +// latent attention as in Python LyricScoreMixin.get_lyric_score(). +int ace_synth_score(AceSynth * ctx, + const AceRequest * reqs, + int batch_n, + const float * pred_latents, + int pred_T_latent, + std::vector & out_scores) { + if (!ctx || !reqs || !pred_latents || batch_n < 1 || batch_n > 9 || pred_T_latent <= 0) { return -1; } // Scoring only supports text2music (pure generation) — the Python reference // also scores on the text2music forward pass. - const std::string & task = reqs[0].task_type; - if (task != TASK_TEXT2MUSIC) { - fprintf(stderr, "[Score] ERROR: scoring only supports task 'text2music', got '%s'\n", task.c_str()); - return -1; + for (int i = 0; i < batch_n; ++i) { + const std::string & task = reqs[i].task_type; + if (task != TASK_TEXT2MUSIC) { + fprintf(stderr, "[Score] ERROR: scoring only supports task 'text2music', got '%s' at batch %d\n", + task.c_str(), i); + return -1; + } + if (reqs[i].seed < 0) { + fprintf(stderr, "[Score] ERROR: seed must be resolved before scoring at batch %d\n", i); + return -1; + } } - AceSynthJob * job = alloc_job(ctx, reqs, batch_n); - SynthState & s = job->state; + AceSynthJob * job = alloc_job(ctx, reqs, batch_n); + SynthState & s = job->state; s.use_source_context = !reqs[0].audio_codes.empty(); s.instruction_str = s.use_source_context ? DIT_INSTR_COVER : DIT_INSTR_TEXT2MUSIC; @@ -733,7 +749,11 @@ int ace_synth_score(AceSynth * ctx, delete job; return -1; } - ops_build_schedule(s); + if (s.T != pred_T_latent) { + fprintf(stderr, "[Score] ERROR: predicted latent has %d frames, request resolves to %d\n", pred_T_latent, s.T); + delete job; + return -1; + } if (ops_encode_text(ctx, reqs, batch_n, s) != 0) { delete job; return -1; @@ -745,8 +765,7 @@ int ace_synth_score(AceSynth * ctx, ops_build_context_silence(ctx, batch_n, s); ops_init_noise(ctx, reqs, batch_n, s); - // Scoring forward pass (single DiT forward with cross-attention capture) - int rc = ops_score_forward(ctx, batch_n, out_scores, s); + int rc = ops_score_forward(ctx, batch_n, pred_latents, out_scores, s); delete job; return rc; diff --git a/src/pipeline-synth.h b/src/pipeline-synth.h index 572269ac..8ea91950 100644 --- a/src/pipeline-synth.h +++ b/src/pipeline-synth.h @@ -21,7 +21,7 @@ struct ModelStore; struct AceSynthParams { const char * text_encoder_path; // Qwen3 text encoder GGUF (required) const char * dit_path; // DiT GGUF (required) - const char * vae_path; // VAE GGUF (required) + const char * vae_path; // VAE GGUF (required except for score-only contexts) const char * adapter_path; // adapter safetensors or directory (NULL to disable) float adapter_scale; // user scale multiplier bool use_fa; // flash attention @@ -47,6 +47,9 @@ void ace_synth_default_params(AceSynthParams * p); // modules are acquired per op, never owned by the context. NULL on failure. AceSynth * ace_synth_load(ModelStore * store, const AceSynthParams * params); +// Score-only variant: text encoder and DiT are required, but VAE is not. +AceSynth * ace_synth_load_score(ModelStore * store, const AceSynthParams * params); + // Phase 1: encode sources, build context, run all DiT denoising steps. // Modules are acquired as needed: VAE encoder for source and timbre, FSQ // tokenizer and detokenizer for cover-mode roundtrip, text encoder + cond @@ -102,14 +105,17 @@ void ace_audio_free(AceAudio * audio); void ace_synth_free(AceSynth * ctx); -// Phase 1 scoring: run a single DiT forward pass with cross-attention capture -// and compute lyric alignment metrics. Requires the same phase 1 setup as -// ace_synth_job_run_dit (text encoding, context build, noise init) but does -// NOT run the full denoising loop — just one forward pass to extract attention. +// Lyric scoring against generated DiT latents. Runs the Python reference's +// pure-noise and regressed-latent attention passes, then computes LM and DiT +// alignment metrics on the pure lyric token range. // -// Returns one LyricScoreResult per batch item in out_scores. +// pred_latents: [batch_n, pred_T_latent, 64] generated latent tensor. +// Each request seed must be resolved (non-negative) before calling. +// Returns one LyricScoreComparison per batch item in out_scores. // Returns 0 on success, -1 on error. -int ace_synth_score(AceSynth * ctx, - const AceRequest * reqs, - int batch_n, - std::vector & out_scores); +int ace_synth_score(AceSynth * ctx, + const AceRequest * reqs, + int batch_n, + const float * pred_latents, + int pred_T_latent, + std::vector & out_scores); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b2ca206e..095943dc 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -22,3 +22,6 @@ target_include_directories(test-dtw-score PRIVATE ${CMAKE_SOURCE_DIR}/src) if(NOT MSVC) target_link_libraries(test-dtw-score PRIVATE m) endif() +if(BUILD_TESTING) + add_test(NAME dtw-score COMMAND test-dtw-score) +endif() diff --git a/tests/test-dtw-score.cpp b/tests/test-dtw-score.cpp index 33194413..626b4343 100644 --- a/tests/test-dtw-score.cpp +++ b/tests/test-dtw-score.cpp @@ -19,25 +19,25 @@ static int g_pass = 0; static int g_fail = 0; -#define CHECK(cond, msg) \ - do { \ - if (cond) { \ - g_pass++; \ - } else { \ - g_fail++; \ - fprintf(stderr, "FAIL: %s (line %d)\n", msg, __LINE__); \ - } \ +#define CHECK(cond, msg) \ + do { \ + if (cond) { \ + g_pass++; \ + } else { \ + g_fail++; \ + fprintf(stderr, "FAIL: %s (line %d)\n", msg, __LINE__); \ + } \ } while (0) -#define CHECK_NEAR(a, b, eps, msg) \ - do { \ - if (fabs((double) (a) - (double) (b)) < (eps)) { \ - g_pass++; \ - } else { \ - g_fail++; \ - fprintf(stderr, "FAIL: %s — expected %.8g, got %.8g (line %d)\n", \ - msg, (double) (b), (double) (a), __LINE__); \ - } \ +#define CHECK_NEAR(a, b, eps, msg) \ + do { \ + if (fabs((double) (a) - (double) (b)) < (eps)) { \ + g_pass++; \ + } else { \ + g_fail++; \ + fprintf(stderr, "FAIL: %s — expected %.8g, got %.8g (line %d)\n", msg, (double) (b), (double) (a), \ + __LINE__); \ + } \ } while (0) // ============================================================================ @@ -50,8 +50,8 @@ static void test_dtw_diagonal() { // [9, 0, 9], // [9, 9, 0]] // DTW on this should find the diagonal path. - float cost[] = { 0, 9, 9, 9, 0, 9, 9, 9, 0 }; - DTWPath path = dtw_cpu(cost, 3, 3); + float cost[] = { 0, 9, 9, 9, 0, 9, 9, 9, 0 }; + DTWPath path = dtw_cpu(cost, 3, 3); // Path should visit (0,0), (1,1), (2,2) CHECK(path.text_idx.size() == 3, "diagonal DTW path length == 3"); @@ -69,8 +69,8 @@ static void test_dtw_horizontal_step() { // [[0, 0, 0], // [9, 9, 0]] // The optimal path should go right along row 0, then diagonal to (1,2). - float cost[] = { 0, 0, 0, 9, 9, 0 }; - DTWPath path = dtw_cpu(cost, 2, 3); + float cost[] = { 0, 0, 0, 9, 9, 0 }; + DTWPath path = dtw_cpu(cost, 2, 3); // Path must start at (0,0) and end at (1,2) CHECK(!path.text_idx.empty(), "horizontal DTW path non-empty"); @@ -86,7 +86,7 @@ static void test_dtw_horizontal_step() { static void test_median_filter() { // Window=3 on a smooth ramp: [1, 2, 3, 4, 5] -> [1, 2, 3, 4, 5] // (reflect padding keeps edges intact for monotonic data) - std::vector ramp = { 1, 2, 3, 4, 5 }; + std::vector ramp = { 1, 2, 3, 4, 5 }; auto filtered = median_filter_1d(ramp, 3); CHECK(filtered.size() == 5, "median filter preserves length"); CHECK_NEAR(filtered[2], 3.0f, 1e-6, "median filter middle of ramp"); @@ -95,7 +95,7 @@ static void test_median_filter() { // Reflect padding: [2,1,2,100,4,5,4] // Windows: [2,1,2]=2, [1,2,100]=2, [2,100,4]=4, [100,4,5]=5, [4,5,4]=4 std::vector spike = { 1, 2, 100, 4, 5 }; - auto fs = median_filter_1d(spike, 3); + auto fs = median_filter_1d(spike, 3); CHECK_NEAR(fs[0], 2.0f, 1e-6, "median filter spike removal [0]"); CHECK_NEAR(fs[1], 2.0f, 1e-6, "median filter spike removal [1]"); CHECK_NEAR(fs[2], 4.0f, 1e-6, "median filter spike removal [2]"); @@ -114,7 +114,7 @@ static void test_token_type_mask() { // Tokens: "hello", "[", "Verse", "]", "world" // Mask: 1, 0, 0, 0, 1 std::vector tokens = { "hello", "[", "Verse", "]", "world" }; - auto mask = generate_token_type_mask(tokens); + auto mask = generate_token_type_mask(tokens); CHECK(mask.size() == 5, "token type mask size"); CHECK(mask[0] == 1, "token type mask: lyric before bracket"); CHECK(mask[1] == 0, "token type mask: opening bracket"); @@ -124,7 +124,7 @@ static void test_token_type_mask() { // Multi-token bracket: "[", "Intro", "Guitar", "]" std::vector tokens2 = { "[", "Intro", "Guitar", "]", "sing" }; - auto mask2 = generate_token_type_mask(tokens2); + auto mask2 = generate_token_type_mask(tokens2); CHECK(mask2[0] == 0, "multi-token bracket: open"); CHECK(mask2[1] == 0, "multi-token bracket: inside 1"); CHECK(mask2[2] == 0, "multi-token bracket: inside 2"); @@ -140,8 +140,8 @@ static void test_scoring_perfect_alignment() { // Create a 4-token, 4-frame attention matrix where the diagonal has // high energy and off-diagonal is near zero. With 1 layer, 1 head. // This simulates perfect lyric-to-audio alignment. - int tokens = 4; - int frames = 4; + int tokens = 4; + int frames = 4; std::vector attn((size_t) tokens * frames, 0.01f); // Diagonal high energy for (int i = 0; i < tokens; i++) { @@ -152,10 +152,11 @@ static void test_scoring_perfect_alignment() { std::vector decoded(tokens, "lyric"); // Use a custom config: just layer 0, head 0 - ScoreLayerHeadConfig config[] = { { 0, 0 } }; + ScoreLayerHeadConfig config[] = { + { 0, 0 } + }; - LyricScoreResult result = - calculate_lyric_score(attn.data(), 1, 1, tokens, frames, decoded, config, 1, 1); + LyricScoreResult result = calculate_lyric_score(attn.data(), 1, 1, tokens, frames, decoded, config, 1, 1); // With perfect diagonal alignment: // - Coverage: all lyric rows have max energy 1.0 > 0.1, so coverage = 1.0 @@ -176,16 +177,17 @@ static void test_scoring_perfect_alignment() { // Uniform attention should yield low confidence and low score. // ============================================================================ static void test_scoring_uniform_attention() { - int tokens = 4; - int frames = 8; + int tokens = 4; + int frames = 8; // Uniform attention — no structure std::vector attn((size_t) tokens * frames, 0.5f); std::vector decoded(tokens, "lyric"); - ScoreLayerHeadConfig config[] = { { 0, 0 } }; + ScoreLayerHeadConfig config[] = { + { 0, 0 } + }; - LyricScoreResult result = - calculate_lyric_score(attn.data(), 1, 1, tokens, frames, decoded, config, 1, 1); + LyricScoreResult result = calculate_lyric_score(attn.data(), 1, 1, tokens, frames, decoded, config, 1, 1); // With uniform attention, min-max normalization zeros everything out // (e_max - e_min < 1e-9), so energy_matrix = 0, confidence = 0, score = 0 @@ -198,8 +200,8 @@ static void test_scoring_uniform_attention() { static void test_scoring_with_tags() { // 5 tokens: lyric, [Intro], lyric, lyric, lyric // 5 frames: perfect diagonal for lyric tokens, tag token gets low energy - int tokens = 5; - int frames = 5; + int tokens = 5; + int frames = 5; std::vector attn((size_t) tokens * frames, 0.01f); // Diagonal energy for all tokens (including tag) @@ -208,11 +210,15 @@ static void test_scoring_with_tags() { } // Token 1 is a structural tag - std::vector decoded = { "sing", "[", "more", "singing", "here" }; - ScoreLayerHeadConfig config[] = { { 0, 0 } }; + std::vector decoded = { "sing", "[Intro]", "more", "singing", "here" }; + ScoreLayerHeadConfig config[] = { + { 0, 0 } + }; - LyricScoreResult result = - calculate_lyric_score(attn.data(), 1, 1, tokens, frames, decoded, config, 1, 1); + std::vector type_mask = generate_token_type_mask(decoded); + CHECK(type_mask == std::vector({ 1, 0, 1, 1, 1 }), "tagged alignment mask excludes only the tag"); + + LyricScoreResult result = calculate_lyric_score(attn.data(), 1, 1, tokens, frames, decoded, config, 1, 1); // Coverage: 4 lyric tokens, all with max energy 1.0 > 0.1 -> coverage = 1.0 // Monotonicity: lyric centroids [0, 2, 3, 4] strictly increasing -> mono = 1.0 @@ -229,8 +235,8 @@ static void test_scoring_with_tags() { // Test 8: Multi-head attention averaging // ============================================================================ static void test_multi_head_averaging() { - int tokens = 3; - int frames = 3; + int tokens = 3; + int frames = 3; int n_layers = 2; int n_heads = 2; @@ -247,18 +253,23 @@ static void test_multi_head_averaging() { // 3 of 4 heads have diagonal, 1 has anti-diagonal for (int t = 0; t < tokens; t++) { - A(0, 0, t, t) = 1.0f; + A(0, 0, t, t) = 1.0f; A(0, 1, t, tokens - 1 - t) = 1.0f; // anti-diagonal - A(1, 0, t, t) = 1.0f; - A(1, 1, t, t) = 1.0f; + A(1, 0, t, t) = 1.0f; + A(1, 1, t, t) = 1.0f; } // Select all 4 heads - ScoreLayerHeadConfig config[] = { { 0, 0 }, { 0, 1 }, { 1, 0 }, { 1, 1 } }; + ScoreLayerHeadConfig config[] = { + { 0, 0 }, + { 0, 1 }, + { 1, 0 }, + { 1, 1 } + }; std::vector calc_matrix, energy_matrix; - bool ok = preprocess_attention(attn.data(), n_layers, n_heads, tokens, frames, config, 4, 1, - calc_matrix, energy_matrix); + bool ok = + preprocess_attention(attn.data(), n_layers, n_heads, tokens, frames, config, 4, 1, calc_matrix, energy_matrix); CHECK(ok, "multi-head preprocess succeeded"); // Average: 3/4 diagonal + 1/4 anti-diagonal @@ -270,17 +281,79 @@ static void test_multi_head_averaging() { } // ============================================================================ -// Test 9: DTW path monotonicity (paths must be non-decreasing in both dims) +// Test 9: Remap absolute model layers to compact captured-layer slots +// ============================================================================ +static void test_score_head_remap() { + int captured_layers[] = { 2, 3, 4, 5, 6 }; + std::vector remapped; + bool ok = remap_score_heads(DEFAULT_2B_SCORE_HEADS, DEFAULT_2B_SCORE_HEADS_COUNT, captured_layers, 5, 16, remapped); + CHECK(ok, "score head remap succeeds"); + CHECK(remapped.size() == 7, "score head remap preserves every configured head"); + CHECK(remapped[0].layer == 0 && remapped[0].head == 6, "absolute layer 2 maps to compact layer 0"); + CHECK(remapped[1].layer == 1 && remapped[1].head == 10, "absolute layer 3 maps to compact layer 1"); + CHECK(remapped[6].layer == 4 && remapped[6].head == 8, "absolute layer 6 maps to compact layer 4"); + + int missing_layers[] = { 2, 3, 4 }; + CHECK(!remap_score_heads(DEFAULT_2B_SCORE_HEADS, DEFAULT_2B_SCORE_HEADS_COUNT, missing_layers, 3, 16, remapped), + "score head remap rejects a partial capture"); +} + +// ============================================================================ +// Test 10: GGML attention layout and per-batch pure-lyric slicing +// ============================================================================ +static void test_attention_slice_layout() { + const int tokens = 4; + const int frames = 2; + const int heads = 2; + const int batches = 2; + std::vector source((size_t) tokens * frames * heads * batches); + + for (int b = 0; b < batches; b++) { + for (int h = 0; h < heads; h++) { + for (int f = 0; f < frames; f++) { + for (int t = 0; t < tokens; t++) { + size_t index = (size_t) t + (size_t) tokens * (f + frames * (h + heads * b)); + source[index] = (float) (1000 * b + 100 * h + 10 * f + t); + } + } + } + } + + std::vector slice; + bool ok = extract_attention_slice(source.data(), source.size(), tokens, frames, heads, batches, 1, 1, 2, slice); + CHECK(ok, "attention slice succeeds"); + CHECK(slice.size() == 8, "attention slice has heads*lyric_tokens*frames elements"); + CHECK_NEAR(slice[0], 1001.0f, 1e-6, "attention slice batch 1, head 0, token 1, frame 0"); + CHECK_NEAR(slice[1], 1011.0f, 1e-6, "attention slice batch 1, head 0, token 1, frame 1"); + CHECK_NEAR(slice[6], 1102.0f, 1e-6, "attention slice batch 1, head 1, token 2, frame 0"); + CHECK_NEAR(slice[7], 1112.0f, 1e-6, "attention slice batch 1, head 1, token 2, frame 1"); +} + +// ============================================================================ +// Test 11: Strip lyric prompt headers and trailing end-of-text tokens +// ============================================================================ +static void test_lyric_token_segment() { + std::vector raw = { 10, 11, 20, 21, 151643, 151643 }; + LyricTokenSegment segment = extract_lyric_token_segment(raw, 2, 151643); + CHECK(segment.start == 2, "lyric segment preserves encoder row offset"); + CHECK(segment.token_ids == std::vector({ 20, 21 }), "lyric segment excludes header and end tokens"); + + LyricTokenSegment invalid = extract_lyric_token_segment(raw, 10, 151643); + CHECK(invalid.start == -1 && invalid.token_ids.empty(), "lyric segment rejects an invalid header length"); +} + +// ============================================================================ +// Test 12: DTW path monotonicity (paths must be non-decreasing in both dims) // ============================================================================ static void test_dtw_path_monotonic() { // Random-ish cost matrix - int N = 5, M = 7; + int N = 5, M = 7; std::vector cost((size_t) N * M); for (int i = 0; i < N; i++) { for (int j = 0; j < M; j++) { // Cheaper near the diagonal (i/N * j/M) - float di = (float) i / N; - float dj = (float) j / M; + float di = (float) i / N; + float dj = (float) j / M; cost[(size_t) i * M + j] = fabsf(di - dj) + 0.01f * ((i * 7 + j * 3) % 5); } } @@ -332,6 +405,15 @@ int main(int argc, char ** argv) { test_multi_head_averaging(); fprintf(stderr, "test_multi_head_averaging: done\n"); + test_score_head_remap(); + fprintf(stderr, "test_score_head_remap: done\n"); + + test_attention_slice_layout(); + fprintf(stderr, "test_attention_slice_layout: done\n"); + + test_lyric_token_segment(); + fprintf(stderr, "test_lyric_token_segment: done\n"); + test_dtw_path_monotonic(); fprintf(stderr, "test_dtw_path_monotonic: done\n"); diff --git a/tools/ace-server.cpp b/tools/ace-server.cpp index f669316b..47fe1ff7 100644 --- a/tools/ace-server.cpp +++ b/tools/ace-server.cpp @@ -27,6 +27,7 @@ // /lm LM // /synth DiT + Text-Enc + VAE // /understand LM + DiT + VAE +// /score DiT + Text-Enc #include "audio-io.h" #include "model-registry.h" @@ -265,6 +266,37 @@ static int latent_payload_validate(size_t size, int * http_code_out) { return T; } +// Validate a contiguous batch of [T, 64] f32 latent tensors. +static int latent_batch_payload_validate(size_t size, int batch_n, int * http_code_out) { + if (batch_n <= 0) { + if (http_code_out) { + *http_code_out = 400; + } + return -1; + } + size_t batch_frame_bytes = (size_t) batch_n * LATENT_FRAME_BYTES; + if (size == 0 || (size % batch_frame_bytes) != 0) { + if (http_code_out) { + *http_code_out = 400; + } + return -1; + } + int T = (int) (size / batch_frame_bytes); + if (T <= 0) { + if (http_code_out) { + *http_code_out = 400; + } + return -1; + } + if (T > MAX_T_LATENT) { + if (http_code_out) { + *http_code_out = 413; + } + return -1; + } + return T; +} + // Build a multipart/mixed body that bundles the primary payload with its // latents. The audio variant pairs one audio part with one latent part per // track, the JSON variant carries a single payload and one optional latent. @@ -1366,17 +1398,17 @@ static void encode_worker(std::shared_ptr job, AceRequest ace_req, float * fprintf(stderr, "[Server] Job %s done (encode)\n", job->id.c_str()); } -// POST /vae -// score worker: load DiT + Text-Enc, run scoring forward pass, store JSON result. +// Score worker: load DiT + Text-Enc, run reference-compatible attention +// scoring against generated latents, and store the JSON result. static void score_worker(std::shared_ptr job, - std::vector ace_reqs) { + std::vector ace_reqs, + std::vector pred_latents, + int pred_T_latent) { if (job->cancel.load()) { job->status.store(JobStatus::CANCELLED); return; } - // Resolve DiT + Text-Enc (VAE not needed for scoring, but the pipeline - // requires it in the registry to build the context). std::string dit_name = resolve_name(g_registry.dit, ace_reqs[0].synth_model, g_loaded_dit); const ModelEntry * dit = registry_find(g_registry.dit, dit_name.c_str()); if (!dit) { @@ -1389,31 +1421,34 @@ static void score_worker(std::shared_ptr job, job->status.store(JobStatus::FAILED); return; } - std::string vae_name = resolve_name(g_registry.vae, ace_reqs[0].vae, g_loaded_vae); - const ModelEntry * vae = registry_find(g_registry.vae, vae_name.c_str()); - if (!vae) { - fprintf(stderr, "[Server] VAE not found: %s\n", vae_name.c_str()); - job->status.store(JobStatus::FAILED); - return; - } - AceSynthParams p = g_synth_params; p.text_encoder_path = g_registry.text_enc[0].path.c_str(); p.dit_path = dit->path.c_str(); - p.vae_path = vae->path.c_str(); + p.vae_path = nullptr; p.adapter_path = nullptr; p.adapter_scale = 1.0f; + if (!ace_reqs[0].adapter.empty()) { + const AdapterEntry * adapter = registry_find_adapter(g_registry, ace_reqs[0].adapter.c_str()); + if (!adapter) { + fprintf(stderr, "[Server] Adapter not found: %s\n", ace_reqs[0].adapter.c_str()); + job->status.store(JobStatus::FAILED); + return; + } + p.adapter_path = adapter->path.c_str(); + p.adapter_scale = ace_reqs[0].adapter_scale; + } - fprintf(stderr, "[Server] Loading score: DiT=%s\n", dit_name.c_str()); - AceSynth * ctx = ace_synth_load(g_store, &p); + fprintf(stderr, "[Server] Loading score: DiT=%s%s%s\n", dit_name.c_str(), + ace_reqs[0].adapter.empty() ? "" : " Adapter=", ace_reqs[0].adapter.c_str()); + AceSynth * ctx = ace_synth_load_score(g_store, &p); if (!ctx) { fprintf(stderr, "[Server] FATAL: synth load failed for scoring\n"); job->status.store(JobStatus::FAILED); return; } - std::vector scores; - int rc = ace_synth_score(ctx, ace_reqs.data(), (int) ace_reqs.size(), scores); + std::vector scores; + int rc = ace_synth_score(ctx, ace_reqs.data(), (int) ace_reqs.size(), pred_latents.data(), pred_T_latent, scores); ace_synth_free(ctx); if (rc != 0) { @@ -1421,22 +1456,29 @@ static void score_worker(std::shared_ptr job, return; } - // Build JSON result - // Format: {"scores":[{"coverage":0.95,"monotonicity":0.88,"confidence":0.72,"lyrics_score":0.61},...]} - yyjson_mut_doc * doc = yyjson_mut_doc_new(NULL); + // Bare array with one LM/DiT comparison per request. + yyjson_mut_doc * doc = yyjson_mut_doc_new(NULL); yyjson_mut_val * root = yyjson_mut_arr(doc); yyjson_mut_doc_set_root(doc, root); + auto build_metrics = [&](const LyricScoreResult & score) { + yyjson_mut_val * metrics = yyjson_mut_obj(doc); + yyjson_mut_obj_add_real(doc, metrics, "coverage", score.coverage); + yyjson_mut_obj_add_real(doc, metrics, "monotonicity", score.monotonicity); + yyjson_mut_obj_add_real(doc, metrics, "confidence", score.confidence); + yyjson_mut_obj_add_real(doc, metrics, "lyrics_score", score.lyrics_score); + return metrics; + }; for (size_t i = 0; i < scores.size(); i++) { yyjson_mut_val * obj = yyjson_mut_obj(doc); - yyjson_mut_obj_add_real(doc, obj, "coverage", scores[i].coverage); - yyjson_mut_obj_add_real(doc, obj, "monotonicity", scores[i].monotonicity); - yyjson_mut_obj_add_real(doc, obj, "confidence", scores[i].confidence); - yyjson_mut_obj_add_real(doc, obj, "lyrics_score", scores[i].lyrics_score); + yyjson_mut_obj_add_real(doc, obj, "lm_score", scores[i].lm.lyrics_score); + yyjson_mut_obj_add_real(doc, obj, "dit_score", scores[i].dit.lyrics_score); + yyjson_mut_obj_add_val(doc, obj, "lm", build_metrics(scores[i].lm)); + yyjson_mut_obj_add_val(doc, obj, "dit", build_metrics(scores[i].dit)); yyjson_mut_arr_append(root, obj); } yyjson_write_err err; - size_t len = 0; + size_t len = 0; char * json_str = yyjson_mut_write_opts(doc, 0, NULL, &len, &err); yyjson_mut_doc_free(doc); if (!json_str) { @@ -1445,13 +1487,18 @@ static void score_worker(std::shared_ptr job, return; } - job->result_body = std::string(json_str, len); - job->result_mime = "application/json"; + job->result_body = std::string(json_str, len); + job->result_mime = "application/json"; free(json_str); if (g_keep_loaded) { - g_loaded_dit = dit_name; - g_loaded_vae = vae_name; + g_loaded_dit = dit_name; + g_loaded_adapter = ace_reqs[0].adapter; + g_loaded_adapter_scale = ace_reqs[0].adapter_scale; + } else { + g_loaded_dit.clear(); + g_loaded_adapter.clear(); + g_loaded_adapter_scale = 1.0f; } fprintf(stderr, "[Server] Job %s done (score)\n", job->id.c_str()); @@ -1459,86 +1506,96 @@ static void score_worker(std::shared_ptr job, } // POST /score -// Accepts the same JSON request format as /synth (plain JSON or array), -// but only uses caption, lyrics, duration, and seed to run a single DiT -// forward pass with cross-attention capture for lyric alignment scoring. -// No audio is generated — the result is a JSON array of per-request scores. -// returns: JSON {"id":"N"} immediately. Result is JSON array of -// {"coverage","monotonicity","confidence","lyrics_score"} per request. +// multipart/form-data: +// part "request": JSON object or array with caption/lyrics/generation params +// part "pred_latents": generated raw f32 latents [batch, T, 64] +// ("latent" is accepted as an alias) +// Returns a job ID immediately. The result is a bare JSON array containing +// LM and DiT score metrics for each request. static void handle_score(const httplib::Request & req, httplib::Response & res) { - if (g_registry.dit.empty() || g_registry.text_enc.empty() || g_registry.vae.empty()) { - json_error(res, 501, "No score models in registry (need dit + text-encoder + vae)"); + if (g_registry.dit.empty() || g_registry.text_enc.empty()) { + json_error(res, 501, "No score models in registry (need dit + text-encoder)"); + return; + } + if (!req.is_multipart_form_data()) { + json_error(res, 400, "Score endpoint requires multipart/form-data with request and pred_latents parts"); + return; + } + + std::string json_body; + if (req.form.has_file("request")) { + json_body = req.form.get_file("request").content; + } else if (req.form.has_field("request")) { + json_body = req.form.get_field("request"); + } else { + json_error(res, 400, "Multipart: missing 'request' part"); return; } - // Parse request: plain JSON (single or array), same as /synth minus audio parts. std::vector ace_reqs; + if (!request_parse_json_array(json_body.c_str(), &ace_reqs)) { + json_error(res, 400, "Multipart: invalid JSON object or array in 'request' part"); + return; + } + if (ace_reqs.size() > 9) { + json_error(res, 400, "Score batch size exceeds maximum of 9"); + return; + } - if (req.is_multipart_form_data()) { - AceRequest ace_req; - std::string json_body; - if (req.form.has_file("request")) { - json_body = req.form.get_file("request").content; - } else if (req.form.has_field("request")) { - json_body = req.form.get_field("request"); - } else { - json_error(res, 400, "Multipart: missing 'request' part"); + for (size_t i = 0; i < ace_reqs.size(); i++) { + AceRequest & current = ace_reqs[i]; + if (current.caption.empty()) { + json_error(res, 400, "Caption is required for every score request"); return; } - if (!request_parse_json(&ace_req, json_body.c_str())) { - json_error(res, 400, "Multipart: invalid JSON in 'request' part"); + if (current.lyrics.empty() || current.lyrics == "[Instrumental]") { + json_error(res, 400, "Non-instrumental lyrics are required for scoring"); return; } - ace_reqs.push_back(std::move(ace_req)); - } else { - // Plain JSON: single object or array - if (req.body.empty()) { - json_error(res, 400, "Empty request body"); + if (current.task_type != TASK_TEXT2MUSIC) { + json_error(res, 400, "Scoring only supports task_type 'text2music'"); return; } - AceRequest ace_req; - if (!request_parse_json(&ace_req, req.body.c_str())) { - // Try as array - yyjson_doc * doc = yyjson_read(req.body.c_str(), req.body.size(), 0); - if (!doc || !yyjson_is_arr(yyjson_doc_get_root(doc))) { - yyjson_doc_free(doc); - json_error(res, 400, "Invalid JSON (expected object or array)"); - return; - } - size_t n = yyjson_arr_size(yyjson_doc_get_root(doc)); - for (size_t i = 0; i < n; i++) { - AceRequest r; - request_init(&r); - yyjson_val * obj = yyjson_arr_get(yyjson_doc_get_root(doc), i); - char * obj_str = yyjson_val_write(obj, 0, nullptr); - request_parse_json(&r, obj_str); - free(obj_str); - ace_reqs.push_back(std::move(r)); - } - yyjson_doc_free(doc); - } else { - ace_reqs.push_back(std::move(ace_req)); + if (i > 0 && + (current.duration != ace_reqs[0].duration || current.inference_steps != ace_reqs[0].inference_steps || + current.custom_timesteps != ace_reqs[0].custom_timesteps || + current.synth_model != ace_reqs[0].synth_model || current.adapter != ace_reqs[0].adapter || + current.adapter_scale != ace_reqs[0].adapter_scale)) { + json_error(res, 400, "Batched score requests must share duration, steps, model, and adapter"); + return; } + request_resolve_seed(¤t); } - if (ace_reqs.empty()) { - json_error(res, 400, "Empty request"); + std::string latent_content; + if (req.form.has_file("pred_latents")) { + latent_content = req.form.get_file("pred_latents").content; + } else if (req.form.has_file("latent")) { + latent_content = req.form.get_file("latent").content; + } else { + json_error(res, 400, "Multipart: missing 'pred_latents' part"); return; } - if (ace_reqs[0].caption.empty()) { - json_error(res, 400, "Caption is required"); + + int http_code = 0; + int pred_T_latent = latent_batch_payload_validate(latent_content.size(), (int) ace_reqs.size(), &http_code); + if (pred_T_latent < 0) { + json_error( + res, http_code, + http_code == 413 ? "pred_latents exceeds max frames" : "pred_latents size must be batch*T*64*4 bytes"); return; } + std::vector pred_latents(latent_content.size() / sizeof(float)); + memcpy(pred_latents.data(), latent_content.data(), latent_content.size()); - // Create job, spawn worker, return ID auto job = job_create(); - fprintf(stderr, "[Server] Job %s created (score, %d requests)\n", job->id.c_str(), (int) ace_reqs.size()); + fprintf(stderr, "[Server] Job %s created (score, %d requests, %d latent frames)\n", job->id.c_str(), + (int) ace_reqs.size(), pred_T_latent); - work_push([job, reqs = std::move(ace_reqs)]() mutable { - score_worker(job, std::move(reqs)); + work_push([job, reqs = std::move(ace_reqs), latents = std::move(pred_latents), pred_T_latent]() mutable { + score_worker(job, std::move(reqs), std::move(latents), pred_T_latent); }); - // Return job ID immediately std::string body = "{\"id\":\"" + job->id + "\"}"; res.set_content(body, "application/json"); } From 080d708c85e59a6d8ac337917a54dc5f789f6dbc Mon Sep 17 00:00:00 2001 From: Jeff Stein Date: Fri, 10 Jul 2026 13:04:09 -0600 Subject: [PATCH 04/11] docs: document scoring helpers --- src/model-store.cpp | 4 ++++ src/pipeline-synth-ops.cpp | 2 ++ src/pipeline-synth.cpp | 2 ++ 3 files changed, 8 insertions(+) diff --git a/src/model-store.cpp b/src/model-store.cpp index 24acaedd..7b8998d6 100644 --- a/src/model-store.cpp +++ b/src/model-store.cpp @@ -25,6 +25,8 @@ #include #include +// Populate the official lyric-alignment layer/head defaults for supported +// 2B and XL DiT architectures when model metadata does not provide them. static void dit_meta_set_default_alignment_config(DiTMeta * meta) { if (meta->cfg.n_layers == 24 && meta->cfg.n_heads == 16) { static const DiTScoreHeadConfig defaults[] = { @@ -57,6 +59,8 @@ static void dit_meta_set_default_alignment_config(DiTMeta * meta) { } } +// Load and validate lyric-alignment heads from the DiT GGUF config JSON, +// falling back to the architecture defaults when metadata is absent or invalid. static void dit_meta_load_alignment_config(DiTMeta * meta, const GGUFModel & gf) { const char * config_json = gf_get_str(gf, "acestep.config_json"); bool invalid = false; diff --git a/src/pipeline-synth-ops.cpp b/src/pipeline-synth-ops.cpp index bd957f4d..27c9cff3 100644 --- a/src/pipeline-synth-ops.cpp +++ b/src/pipeline-synth-ops.cpp @@ -405,6 +405,8 @@ struct TextEncForward { int S_lyric; }; +// Build the exact language-and-lyric prefix used for both text encoding and +// pure-lyric token-range extraction so their encoder row offsets stay aligned. static std::string build_lyric_header(const AceRequest & rb) { const char * language = rb.vocal_language.empty() ? "unknown" : rb.vocal_language.c_str(); return std::string("# Languages\n") + language + "\n\n# Lyric\n"; diff --git a/src/pipeline-synth.cpp b/src/pipeline-synth.cpp index d696c92c..87a5b56e 100644 --- a/src/pipeline-synth.cpp +++ b/src/pipeline-synth.cpp @@ -35,6 +35,8 @@ void ace_synth_default_params(AceSynthParams * p) { p->dump_dir = NULL; } +// Validate model paths and construct the shared lightweight synth context. +// Score-only contexts may omit the VAE while generation contexts require it. static AceSynth * ace_synth_load_impl(ModelStore * store, const AceSynthParams * params, bool require_vae) { if (!store || !params) { fprintf(stderr, "[Synth-Load] ERROR: store and params are required\n"); From 515759ce3dcb5ac6c110b0562fe97f43a6790646 Mon Sep 17 00:00:00 2001 From: Jeff Stein Date: Fri, 10 Jul 2026 13:22:22 -0600 Subject: [PATCH 05/11] fix: reject malformed alignment metadata --- src/model-store.cpp | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/src/model-store.cpp b/src/model-store.cpp index 7b8998d6..446dc7b9 100644 --- a/src/model-store.cpp +++ b/src/model-store.cpp @@ -60,16 +60,26 @@ static void dit_meta_set_default_alignment_config(DiTMeta * meta) { } // Load and validate lyric-alignment heads from the DiT GGUF config JSON, -// falling back to the architecture defaults when metadata is absent or invalid. +// falling back to architecture defaults only when the metadata key is absent. static void dit_meta_load_alignment_config(DiTMeta * meta, const GGUFModel & gf) { const char * config_json = gf_get_str(gf, "acestep.config_json"); bool invalid = false; + bool has_config = false; if (config_json && config_json[0]) { yyjson_doc * doc = yyjson_read(config_json, strlen(config_json), 0); yyjson_val * root = doc ? yyjson_doc_get_root(doc) : nullptr; - yyjson_val * config = root ? yyjson_obj_get(root, "lyric_alignment_layers_config") : nullptr; - if (config && yyjson_is_obj(config)) { + yyjson_val * config = nullptr; + if (!doc || !yyjson_is_obj(root)) { + invalid = true; + } else { + config = yyjson_obj_get(root, "lyric_alignment_layers_config"); + has_config = config != nullptr; + if (has_config && !yyjson_is_obj(config)) { + invalid = true; + } + } + if (!invalid && has_config) { size_t idx, max; yyjson_val * key; yyjson_val * heads; @@ -98,14 +108,15 @@ static void dit_meta_load_alignment_config(DiTMeta * meta, const GGUFModel & gf) } } } - yyjson_doc_free(doc); + if (doc) { + yyjson_doc_free(doc); + } } if (invalid) { - fprintf(stderr, "[Store] WARNING: invalid lyric_alignment_layers_config; using architecture default\n"); + fprintf(stderr, "[Store] WARNING: invalid lyric_alignment_layers_config; ignoring scoring heads\n"); meta->lyric_alignment_heads.clear(); - } - if (meta->lyric_alignment_heads.empty()) { + } else if (!has_config && meta->lyric_alignment_heads.empty()) { dit_meta_set_default_alignment_config(meta); } if (meta->lyric_alignment_heads.empty()) { From 8e2adb2c1d77ed2cc7e062886470e2e20180a492 Mon Sep 17 00:00:00 2001 From: Jeff Stein Date: Fri, 10 Jul 2026 13:25:36 -0600 Subject: [PATCH 06/11] fix: propagate score cancellation --- src/pipeline-synth-ops.cpp | 16 +++++++++++++++- src/pipeline-synth-ops.h | 7 +++++-- src/pipeline-synth.cpp | 10 ++++++++-- src/pipeline-synth.h | 7 +++++-- tools/ace-server.cpp | 20 ++++++++++++++++++-- 5 files changed, 51 insertions(+), 9 deletions(-) diff --git a/src/pipeline-synth-ops.cpp b/src/pipeline-synth-ops.cpp index 27c9cff3..c14f2a9f 100644 --- a/src/pipeline-synth-ops.cpp +++ b/src/pipeline-synth-ops.cpp @@ -984,11 +984,17 @@ int ops_score_forward(const AceSynth * ctx, int batch_n, const float * pred_latents, std::vector & out_scores, - SynthState & s) { + SynthState & s, + bool (*cancel)(void *), + void * cancel_data) { if (!pred_latents || batch_n <= 0 || s.num_steps <= 0) { fprintf(stderr, "[Score] FATAL: predicted latents, batch, and inference steps are required\n"); return -1; } + if (cancel && cancel(cancel_data)) { + fprintf(stderr, "[Score] Cancelled before attention forwards\n"); + return -1; + } DiTGGML * dit = store_require_dit(ctx->store, ctx->dit_key); if (!dit) { @@ -1044,6 +1050,10 @@ int ops_score_forward(const AceSynth * ctx, fprintf(stderr, "[Score] FATAL: pure-noise score forward failed\n"); return -1; } + if (cancel && cancel(cancel_data)) { + fprintf(stderr, "[Score] Cancelled between attention forwards\n"); + return -1; + } rc = dit_ggml_score_forward(dit, regressed_latents.data(), s.context.data(), s.enc_hidden.data(), s.enc_S, s.per_enc_S.data(), s.T, batch_n, t_last, score_layers.data(), (int) score_layers.size(), dit_attentions); @@ -1051,6 +1061,10 @@ int ops_score_forward(const AceSynth * ctx, fprintf(stderr, "[Score] FATAL: regressed-latent score forward failed\n"); return -1; } + if (cancel && cancel(cancel_data)) { + fprintf(stderr, "[Score] Cancelled after attention forwards\n"); + return -1; + } fprintf(stderr, "[Score] Attention forwards: %.1f ms\n", s.timer.ms()); BPETokenizer * bpe = store_bpe(ctx->store, ctx->params.text_encoder_path); diff --git a/src/pipeline-synth-ops.h b/src/pipeline-synth-ops.h index 9aff511a..3e9c8680 100644 --- a/src/pipeline-synth-ops.h +++ b/src/pipeline-synth-ops.h @@ -79,9 +79,12 @@ int ops_vae_decode(const AceSynth * ctx, // ops_init_noise (for noise + context latents) to have been run first. // // out_scores: filled with one LM/DiT score comparison per batch item. -// Returns 0 on success, -1 on error. +// cancel/cancel_data: abort callback, polled before and between DiT forwards. +// Returns 0 on success, -1 on error or cancellation. int ops_score_forward(const AceSynth * ctx, int batch_n, const float * pred_latents, std::vector & out_scores, - SynthState & s); + SynthState & s, + bool (*cancel)(void *), + void * cancel_data); diff --git a/src/pipeline-synth.cpp b/src/pipeline-synth.cpp index 87a5b56e..3df014e7 100644 --- a/src/pipeline-synth.cpp +++ b/src/pipeline-synth.cpp @@ -713,10 +713,16 @@ int ace_synth_score(AceSynth * ctx, int batch_n, const float * pred_latents, int pred_T_latent, - std::vector & out_scores) { + std::vector & out_scores, + bool (*cancel)(void *), + void * cancel_data) { if (!ctx || !reqs || !pred_latents || batch_n < 1 || batch_n > 9 || pred_T_latent <= 0) { return -1; } + if (cancel && cancel(cancel_data)) { + fprintf(stderr, "[Score] Cancelled before setup\n"); + return -1; + } // Scoring only supports text2music (pure generation) — the Python reference // also scores on the text2music forward pass. @@ -767,7 +773,7 @@ int ace_synth_score(AceSynth * ctx, ops_build_context_silence(ctx, batch_n, s); ops_init_noise(ctx, reqs, batch_n, s); - int rc = ops_score_forward(ctx, batch_n, pred_latents, out_scores, s); + int rc = ops_score_forward(ctx, batch_n, pred_latents, out_scores, s, cancel, cancel_data); delete job; return rc; diff --git a/src/pipeline-synth.h b/src/pipeline-synth.h index 8ea91950..919e1d85 100644 --- a/src/pipeline-synth.h +++ b/src/pipeline-synth.h @@ -112,10 +112,13 @@ void ace_synth_free(AceSynth * ctx); // pred_latents: [batch_n, pred_T_latent, 64] generated latent tensor. // Each request seed must be resolved (non-negative) before calling. // Returns one LyricScoreComparison per batch item in out_scores. -// Returns 0 on success, -1 on error. +// cancel/cancel_data: abort callback, polled before and between DiT forwards. +// Returns 0 on success, -1 on error or cancellation. int ace_synth_score(AceSynth * ctx, const AceRequest * reqs, int batch_n, const float * pred_latents, int pred_T_latent, - std::vector & out_scores); + std::vector & out_scores, + bool (*cancel)(void *) = nullptr, + void * cancel_data = nullptr); diff --git a/tools/ace-server.cpp b/tools/ace-server.cpp index 47fe1ff7..30f4e895 100644 --- a/tools/ace-server.cpp +++ b/tools/ace-server.cpp @@ -1448,11 +1448,16 @@ static void score_worker(std::shared_ptr job, } std::vector scores; - int rc = ace_synth_score(ctx, ace_reqs.data(), (int) ace_reqs.size(), pred_latents.data(), pred_T_latent, scores); + int rc = ace_synth_score(ctx, ace_reqs.data(), (int) ace_reqs.size(), pred_latents.data(), pred_T_latent, scores, + server_cancel_job, (void *) &job->cancel); ace_synth_free(ctx); + if (job->cancel.load()) { + job->status.store(JobStatus::CANCELLED); + return; + } if (rc != 0) { - job->status.store(job->cancel.load() ? JobStatus::CANCELLED : JobStatus::FAILED); + job->status.store(JobStatus::FAILED); return; } @@ -1486,6 +1491,11 @@ static void score_worker(std::shared_ptr job, job->status.store(JobStatus::FAILED); return; } + if (job->cancel.load()) { + free(json_str); + job->status.store(JobStatus::CANCELLED); + return; + } job->result_body = std::string(json_str, len); job->result_mime = "application/json"; @@ -1501,6 +1511,12 @@ static void score_worker(std::shared_ptr job, g_loaded_adapter_scale = 1.0f; } + if (job->cancel.load()) { + job->result_body.clear(); + job->result_mime.clear(); + job->status.store(JobStatus::CANCELLED); + return; + } fprintf(stderr, "[Server] Job %s done (score)\n", job->id.c_str()); job->status.store(JobStatus::DONE); } From 3e59699613dca5e3c1b5282c64f7e85d0749d964 Mon Sep 17 00:00:00 2001 From: Jeff Stein Date: Fri, 10 Jul 2026 13:30:45 -0600 Subject: [PATCH 07/11] docs: document changed scoring entry points --- src/model-store.cpp | 9 +++++---- src/pipeline-synth-ops.cpp | 5 +++-- src/pipeline-synth.cpp | 7 +++++-- tests/test-dtw-score.cpp | 3 +-- tools/ace-server.cpp | 1 + 5 files changed, 15 insertions(+), 10 deletions(-) diff --git a/src/model-store.cpp b/src/model-store.cpp index 446dc7b9..c3988d08 100644 --- a/src/model-store.cpp +++ b/src/model-store.cpp @@ -25,8 +25,8 @@ #include #include -// Populate the official lyric-alignment layer/head defaults for supported -// 2B and XL DiT architectures when model metadata does not provide them. +/// Populate the official lyric-alignment layer/head defaults for supported +/// 2B and XL DiT architectures when model metadata does not provide them. static void dit_meta_set_default_alignment_config(DiTMeta * meta) { if (meta->cfg.n_layers == 24 && meta->cfg.n_heads == 16) { static const DiTScoreHeadConfig defaults[] = { @@ -59,8 +59,8 @@ static void dit_meta_set_default_alignment_config(DiTMeta * meta) { } } -// Load and validate lyric-alignment heads from the DiT GGUF config JSON, -// falling back to architecture defaults only when the metadata key is absent. +/// Load and validate lyric-alignment heads from the DiT GGUF config JSON, +/// falling back to architecture defaults only when the metadata key is absent. static void dit_meta_load_alignment_config(DiTMeta * meta, const GGUFModel & gf) { const char * config_json = gf_get_str(gf, "acestep.config_json"); bool invalid = false; @@ -628,6 +628,7 @@ MetadataFSM * store_fsm(ModelStore * s, const char * lm_path, int vocab_size) { return fsm; } +/// Load and cache the CPU-resident DiT configuration and scoring metadata. const DiTMeta * store_dit_meta(ModelStore * s, const char * dit_path) { std::lock_guard lock(s->mtx); std::string key = dit_path ? dit_path : ""; diff --git a/src/pipeline-synth-ops.cpp b/src/pipeline-synth-ops.cpp index c14f2a9f..9b058d66 100644 --- a/src/pipeline-synth-ops.cpp +++ b/src/pipeline-synth-ops.cpp @@ -405,8 +405,8 @@ struct TextEncForward { int S_lyric; }; -// Build the exact language-and-lyric prefix used for both text encoding and -// pure-lyric token-range extraction so their encoder row offsets stay aligned. +/// Build the exact language-and-lyric prefix used for both text encoding and +/// pure-lyric token-range extraction so their encoder row offsets stay aligned. static std::string build_lyric_header(const AceRequest & rb) { const char * language = rb.vocal_language.empty() ? "unknown" : rb.vocal_language.c_str(); return std::string("# Languages\n") + language + "\n\n# Lyric\n"; @@ -433,6 +433,7 @@ static void build_prompt_strings(const AceRequest & rb, lyric_out = build_lyric_header(rb) + rb.lyrics + "<|endoftext|>"; } +/// Encode text and lyrics while preserving each request's pure-lyric token range for scoring. int ops_encode_text(const AceSynth * ctx, const AceRequest * reqs, int batch_n, SynthState & s) { // Per-batch text encoding in two GPU phases to keep EVICT_STRICT at one // module resident at a time: diff --git a/src/pipeline-synth.cpp b/src/pipeline-synth.cpp index 3df014e7..93687492 100644 --- a/src/pipeline-synth.cpp +++ b/src/pipeline-synth.cpp @@ -21,6 +21,7 @@ #include #include +/// Initialize synthesis parameters with the library's generation defaults. void ace_synth_default_params(AceSynthParams * p) { p->text_encoder_path = NULL; p->dit_path = NULL; @@ -35,8 +36,8 @@ void ace_synth_default_params(AceSynthParams * p) { p->dump_dir = NULL; } -// Validate model paths and construct the shared lightweight synth context. -// Score-only contexts may omit the VAE while generation contexts require it. +/// Validate model paths and construct the shared lightweight synth context. +/// Score-only contexts may omit the VAE while generation contexts require it. static AceSynth * ace_synth_load_impl(ModelStore * store, const AceSynthParams * params, bool require_vae) { if (!store || !params) { fprintf(stderr, "[Synth-Load] ERROR: store and params are required\n"); @@ -108,10 +109,12 @@ static AceSynth * ace_synth_load_impl(ModelStore * store, const AceSynthParams * return ctx; } +/// Create a generation-capable synthesis context that requires VAE weights. AceSynth * ace_synth_load(ModelStore * store, const AceSynthParams * params) { return ace_synth_load_impl(store, params, true); } +/// Create a score-only synthesis context without requiring VAE weights. AceSynth * ace_synth_load_score(ModelStore * store, const AceSynthParams * params) { return ace_synth_load_impl(store, params, false); } diff --git a/tests/test-dtw-score.cpp b/tests/test-dtw-score.cpp index 626b4343..0fa65448 100644 --- a/tests/test-dtw-score.cpp +++ b/tests/test-dtw-score.cpp @@ -373,8 +373,7 @@ static void test_dtw_path_monotonic() { CHECK(path.time_idx.back() == M - 1, "DTW path ends at time M-1"); } -// ============================================================================ - +/// Run the self-contained DTW and lyric-alignment regression suite. int main(int argc, char ** argv) { (void) argc; (void) argv; diff --git a/tools/ace-server.cpp b/tools/ace-server.cpp index 30f4e895..f83b15cb 100644 --- a/tools/ace-server.cpp +++ b/tools/ace-server.cpp @@ -1847,6 +1847,7 @@ static void usage(const char * prog) { prog, synth_d.vae_chunk, synth_d.vae_overlap, g_max_batch, lm_d.max_seq); } +/// Parse server options, register HTTP routes, and run the worker-backed API server. int main(int argc, char ** argv) { ace_lm_default_params(&g_lm_params); ace_synth_default_params(&g_synth_params); From 5c07cd5209506cd13feecf8d6ea23598193869ae Mon Sep 17 00:00:00 2001 From: Jeff Stein Date: Fri, 10 Jul 2026 14:14:52 -0600 Subject: [PATCH 08/11] ci: cap scoring drift retries --- .github/workflows/scoring-drift-check.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/scoring-drift-check.yml b/.github/workflows/scoring-drift-check.yml index 6a4426ae..e467f375 100644 --- a/.github/workflows/scoring-drift-check.yml +++ b/.github/workflows/scoring-drift-check.yml @@ -27,7 +27,7 @@ jobs: - name: Fetch upstream Python scoring files run: | BASE="https://raw.githubusercontent.com/ace-step/ACE-Step-1.5/main" - CURL=(curl -fsSL --max-time 30 --retry 2) + CURL=(curl -fsSL --max-time 30 --retry 2 --retry-max-time 60) "${CURL[@]}" "$BASE/acestep/core/scoring/_dtw.py" -o _dtw.py "${CURL[@]}" "$BASE/acestep/core/scoring/dit_score.py" -o dit_score.py "${CURL[@]}" "$BASE/acestep/core/generation/handler/lyric_score.py" -o lyric_score.py From c09f22fc71d6d03a088e9d2a100e07be235bdacc Mon Sep 17 00:00:00 2001 From: Jeff Stein Date: Fri, 10 Jul 2026 14:18:14 -0600 Subject: [PATCH 09/11] fix: correct DTW tie-breaking --- src/dtw-score.h | 7 +++++-- tests/test-dtw-score.cpp | 35 +++++++++++++++++++++++------------ 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/src/dtw-score.h b/src/dtw-score.h index ae0e49e0..662526c5 100644 --- a/src/dtw-score.h +++ b/src/dtw-score.h @@ -112,10 +112,13 @@ static DTWPath dtw_cpu(const float * x, int N, int M) { float c; float t; - if (c0 < c1 && c0 < c2) { + // The pinned Python snapshot uses strict comparisons here, which + // can select a larger c2 when c0 and c1 tie. Keep the intended + // diagonal -> up -> left priority while always choosing a minimum. + if (c0 <= c1 && c0 <= c2) { c = c0; t = 0; - } else if (c1 < c0 && c1 < c2) { + } else if (c1 <= c2) { c = c1; t = 1; } else { diff --git a/tests/test-dtw-score.cpp b/tests/test-dtw-score.cpp index 0fa65448..be9e9611 100644 --- a/tests/test-dtw-score.cpp +++ b/tests/test-dtw-score.cpp @@ -78,6 +78,16 @@ static void test_dtw_horizontal_step() { CHECK(path.time_idx.front() == 0, "horizontal DTW starts at time=0"); CHECK(path.text_idx.back() == 1, "horizontal DTW ends at text=1"); CHECK(path.time_idx.back() == 2, "horizontal DTW ends at time=2"); + + // Diagonal-first tie-breaking selects the minimum-cost three-step path: + // (0,0) -> (0,1) -> (1,2). The old strict comparisons instead selected + // the expensive four-step path through (1,0) and (1,1). + CHECK(path.text_idx.size() == 3, "horizontal DTW path length == 3"); + if (path.text_idx.size() == 3) { + CHECK(path.text_idx[0] == 0 && path.time_idx[0] == 0, "horizontal DTW step 0"); + CHECK(path.text_idx[1] == 0 && path.time_idx[1] == 1, "horizontal DTW step 1"); + CHECK(path.text_idx[2] == 1 && path.time_idx[2] == 2, "horizontal DTW step 2"); + } } // ============================================================================ @@ -161,15 +171,17 @@ static void test_scoring_perfect_alignment() { // With perfect diagonal alignment: // - Coverage: all lyric rows have max energy 1.0 > 0.1, so coverage = 1.0 // - Monotonicity: centroids are [0, 1, 2, 3], strictly increasing, so mono = 1.0 - // - Confidence: DTW tie-breaking (prefers left/up over diagonal when costs - // are equal) creates a 7-step zigzag path through the 4x4 matrix. - // 4 diagonal cells have energy 1.0, 3 off-diagonal have 0.0. - // confidence = 4/7 ≈ 0.571 - // - lyrics_score = 1^2 * 1^2 * (4/7) = 4/7 ≈ 0.571 + // - Diagonal-first tie-breaking keeps the four-step path on the diagonal. + // - Confidence and lyrics_score are therefore both 1.0. CHECK_NEAR(result.coverage, 1.0, 1e-6, "perfect alignment coverage"); CHECK_NEAR(result.monotonicity, 1.0, 1e-6, "perfect alignment monotonicity"); - CHECK_NEAR(result.confidence, 4.0 / 7.0, 0.01, "perfect alignment confidence"); - CHECK_NEAR(result.lyrics_score, 4.0 / 7.0, 0.01, "perfect alignment final score"); + CHECK_NEAR(result.confidence, 1.0, 1e-6, "perfect alignment confidence"); + CHECK_NEAR(result.lyrics_score, 1.0, 1e-6, "perfect alignment final score"); + CHECK(result.path.text_idx.size() == 4, "perfect alignment path length == 4"); + for (size_t i = 0; i < result.path.text_idx.size(); i++) { + CHECK(result.path.text_idx[i] == (int) i && result.path.time_idx[i] == (int) i, + "perfect alignment path stays diagonal"); + } } // ============================================================================ @@ -222,13 +234,12 @@ static void test_scoring_with_tags() { // Coverage: 4 lyric tokens, all with max energy 1.0 > 0.1 -> coverage = 1.0 // Monotonicity: lyric centroids [0, 2, 3, 4] strictly increasing -> mono = 1.0 - // Confidence: DTW tie-breaking creates a 9-step zigzag through the 5x5 - // matrix. 5 diagonal cells have energy 1.0, 4 off-diagonal have 0.0. - // confidence = 5/9 ≈ 0.556 - // lyrics_score = 1 * 1 * (5/9) = 5/9 ≈ 0.556 + // Diagonal-first tie-breaking keeps all five path steps on energy 1.0, + // so confidence and lyrics_score are both 1.0. CHECK_NEAR(result.coverage, 1.0, 1e-6, "tagged alignment coverage"); CHECK_NEAR(result.monotonicity, 1.0, 1e-6, "tagged alignment monotonicity"); - CHECK_NEAR(result.lyrics_score, 5.0 / 9.0, 0.01, "tagged alignment final score"); + CHECK_NEAR(result.confidence, 1.0, 1e-6, "tagged alignment confidence"); + CHECK_NEAR(result.lyrics_score, 1.0, 1e-6, "tagged alignment final score"); } // ============================================================================ From 519e859e6f917223c494b8709440506138240dfd Mon Sep 17 00:00:00 2001 From: Jeff Stein Date: Fri, 10 Jul 2026 14:19:46 -0600 Subject: [PATCH 10/11] fix: validate score head indices --- src/dtw-score.h | 2 +- tests/test-dtw-score.cpp | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/dtw-score.h b/src/dtw-score.h index 662526c5..955719ef 100644 --- a/src/dtw-score.h +++ b/src/dtw-score.h @@ -383,7 +383,7 @@ static bool preprocess_attention(const float * attention, for (int c = 0; c < config_count; c++) { int layer = config[c].layer; int head = config[c].head; - if (layer < n_layers && head < n_heads) { + if (layer >= 0 && layer < n_layers && head >= 0 && head < n_heads) { const float * ptr = attention + (size_t) layer * n_heads * tokens * frames + (size_t) head * tokens * frames; selected.emplace_back(ptr, ptr + (size_t) tokens * frames); diff --git a/tests/test-dtw-score.cpp b/tests/test-dtw-score.cpp index be9e9611..3a84af60 100644 --- a/tests/test-dtw-score.cpp +++ b/tests/test-dtw-score.cpp @@ -289,6 +289,19 @@ static void test_multi_head_averaging() { // energy[0,0] = 0.75, energy[0,2] = 0.25 CHECK_NEAR(energy_matrix[0], 0.75f, 1e-5, "multi-head averaged energy [0,0]"); CHECK_NEAR(energy_matrix[2], 0.25f, 1e-5, "multi-head averaged energy [0,2]"); + + ScoreLayerHeadConfig negative_layer[] = { + { -1, 0 } + }; + ScoreLayerHeadConfig negative_head[] = { + { 0, -1 } + }; + CHECK(!preprocess_attention(attn.data(), n_layers, n_heads, tokens, frames, negative_layer, 1, 1, calc_matrix, + energy_matrix), + "preprocess rejects a negative layer"); + CHECK(!preprocess_attention(attn.data(), n_layers, n_heads, tokens, frames, negative_head, 1, 1, calc_matrix, + energy_matrix), + "preprocess rejects a negative head"); } // ============================================================================ From 7f28afe08c01254f1611fc7e4dfcbc362536e1c9 Mon Sep 17 00:00:00 2001 From: Jeff Stein Date: Mon, 13 Jul 2026 09:19:25 -0600 Subject: [PATCH 11/11] docs: clarify scoring reference behavior --- .github/workflows/scoring-drift-check.yml | 3 ++- src/dit-graph.h | 15 +++++++-------- src/dtw-score.h | 12 +++++++----- tools/ace-server.cpp | 4 ++-- 4 files changed, 18 insertions(+), 16 deletions(-) diff --git a/.github/workflows/scoring-drift-check.yml b/.github/workflows/scoring-drift-check.yml index e467f375..4ed236e6 100644 --- a/.github/workflows/scoring-drift-check.yml +++ b/.github/workflows/scoring-drift-check.yml @@ -77,4 +77,5 @@ jobs: exit 1 fi - echo "No drift detected. C++ port matches pinned Python reference." + echo "No upstream source drift detected against the pinned snapshot." + echo "The C++ implementation may contain documented correctness fixes." diff --git a/src/dit-graph.h b/src/dit-graph.h index a836f9c7..9c813560 100644 --- a/src/dit-graph.h +++ b/src/dit-graph.h @@ -488,14 +488,13 @@ static struct ggml_tensor * dit_ggml_build_layer(struct ggml_context * ctx, // These are retrieved after graph compute for lyric alignment scoring. static struct ggml_cgraph * dit_ggml_build_graph(DiTGGML * m, struct ggml_context * ctx, - int T, // temporal length (before patching) - int enc_S, // encoder sequence length - int N, // batch size - struct ggml_tensor ** p_input, // [out] input tensor to fill - struct ggml_tensor ** p_output, // [out] output tensor to read - const int * score_layers = nullptr, // layers to capture, or NULL - int n_score_layers = 0) { - + int T, // temporal length (before patching) + int enc_S, // encoder sequence length + int N, // batch size + struct ggml_tensor ** p_input, // [out] input tensor to fill + struct ggml_tensor ** p_output, // [out] output tensor to read + const int * score_layers = nullptr, // layers to capture, or NULL + int n_score_layers = 0) { DiTGGMLConfig & c = m->cfg; int S = T / c.patch_size; // sequence length after patching int H = c.hidden_size; diff --git a/src/dtw-score.h b/src/dtw-score.h index 955719ef..256c9013 100644 --- a/src/dtw-score.h +++ b/src/dtw-score.h @@ -1,7 +1,7 @@ #pragma once // dtw-score.h: DTW pathfinding + lyric alignment scoring (pure C++) // -// Direct port of the Python ACE-Step scoring modules: +// C++ implementation based on the Python ACE-Step scoring modules: // acestep/core/scoring/_dtw.py — DTW + median filter (numba-jitted) // acestep/core/scoring/dit_score.py — MusicLyricScorer (coverage, monotonicity, confidence) // Integration parity follows generation/handler/lyric_score.py and @@ -9,7 +9,8 @@ // // Pinned to ace-step/ACE-Step-1.5 commit 82252c24 (2026-07-09). // If the Python scoring algorithm changes upstream, this port will need -// re-evaluation. The file:line references in comments below map to that commit. +// re-evaluation. The file:line references in comments below map to that commit; +// documented correctness fixes may intentionally differ from the pinned source. // // No external dependencies beyond , , , . // All compute is CPU-side on small matrices (tokens x frames), matching the @@ -112,9 +113,10 @@ static DTWPath dtw_cpu(const float * x, int N, int M) { float c; float t; - // The pinned Python snapshot uses strict comparisons here, which - // can select a larger c2 when c0 and c1 tie. Keep the intended - // diagonal -> up -> left priority while always choosing a minimum. + // Intentional correctness divergence from the pinned Python + // snapshot: its strict comparisons can select a larger c2 when + // c0 and c1 tie. Keep diagonal -> up -> left priority while + // always choosing a minimum. if (c0 <= c1 && c0 <= c2) { c = c0; t = 0; diff --git a/tools/ace-server.cpp b/tools/ace-server.cpp index f83b15cb..38257c0b 100644 --- a/tools/ace-server.cpp +++ b/tools/ace-server.cpp @@ -1398,8 +1398,8 @@ static void encode_worker(std::shared_ptr job, AceRequest ace_req, float * fprintf(stderr, "[Server] Job %s done (encode)\n", job->id.c_str()); } -// Score worker: load DiT + Text-Enc, run reference-compatible attention -// scoring against generated latents, and store the JSON result. +// Score worker: load DiT + Text-Enc, run lyric-alignment attention scoring +// against generated latents, and store the JSON result. static void score_worker(std::shared_ptr job, std::vector ace_reqs, std::vector pred_latents,