From 13bfdc667bfa65cb679768501f1b13a1242394fe Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Fri, 31 Jul 2026 16:45:20 +0200 Subject: [PATCH] wire up --- tools/mtmd/clip.cpp | 103 ++++++++++++++++++----------- tools/mtmd/clip.h | 20 ++++-- tools/mtmd/models/models.h | 16 +++-- tools/mtmd/models/qwen3tts-gen.cpp | 39 ++++++++--- tools/mtmd/mtmd.cpp | 82 ++++++++++++++++------- tools/mtmd/mtmd.h | 6 +- 6 files changed, 180 insertions(+), 86 deletions(-) diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index d3522d25d7..15200c2a36 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1056,9 +1056,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { + const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_CODE_GEN; const int top_k = params ? params->top_k : 50; const float top_p = params ? params->top_p : 1.0f; - builder = std::make_unique(ctx, img, top_k, top_p); + builder = std::make_unique(ctx, img, gen_process, top_k, top_p); } break; case PROJECTOR_TYPE_YOUTUVL: { @@ -4163,8 +4164,9 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } set_input_f32("inp_raw", inp_raw); - } else { - // audio input + } else if (!(ctx->proj_type() == PROJECTOR_TYPE_QWEN3TTS_GEN && params->gen_process == CLIP_GEN_PROCESS_CODE2WAV)) { + // audio input (code2wav has no hidden-state/raw input at all, its + // only input is the "inp_codes" tensor handled in the switch below) GGML_ASSERT(imgs.entries.size() == 1); const auto & mel_inp = imgs.entries[0]; @@ -4722,17 +4724,23 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { - std::vector code0 = { params->code0 }; - set_input_i32("inp_code0", code0); + if (params->gen_process == CLIP_GEN_PROCESS_CODE2WAV) { + GGML_ASSERT(params->codes != nullptr); + std::vector codes = *params->codes; + set_input_i32("inp_codes", codes); + } else { + std::vector code0 = { params->code0 }; + set_input_i32("inp_code0", code0); - // one uniform(0,1) draw per codebook, consumed by do_sampling()'s - // inverse-CDF token selection (inp_rand_0 .. inp_rand_{n_acoustic-1}) - static std::mt19937 rng{ std::random_device{}() }; - std::uniform_real_distribution dist(0.0f, 1.0f); - const int64_t n_acoustic = model.gen_code_head_w->ne[2]; - for (int64_t g = 0; g < n_acoustic; g++) { - std::vector r = { dist(rng) }; - set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r); + // one uniform(0,1) draw per codebook, consumed by do_sampling()'s + // inverse-CDF token selection (inp_rand_0 .. inp_rand_{n_acoustic-1}) + static std::mt19937 rng{ std::random_device{}() }; + std::uniform_real_distribution dist(0.0f, 1.0f); + const int64_t n_acoustic = model.gen_code_head_w->ne[2]; + for (int64_t g = 0; g < n_acoustic; g++) { + std::vector r = { dist(rng) }; + set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r); + } } } break; case PROJECTOR_TYPE_HUNYUANVL: @@ -5150,35 +5158,49 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { return false; } - // the last node is the embedding tensor - ggml_tensor * embeddings = ggml_graph_node(gf, -1); + // the last node is the embedding tensor (not produced by the code2wav + // sub-graph, which has no out_embd at all) + ggml_tensor * embeddings = params->out_embd ? ggml_graph_node(gf, -1) : nullptr; - // sanity check (assuming that all images in batch have the same number of tokens, so we only check the first one) - const int n_tokens_out = embeddings->ne[1]; - const int expected_n_tokens_out = clip_n_output_tokens(ctx, &imgs.entries[0]); - if (n_tokens_out != expected_n_tokens_out) { - LOG_ERR("%s: expected output %d tokens, got %d\n", __func__, expected_n_tokens_out, n_tokens_out); - GGML_ABORT("Invalid number of output tokens"); - } - - LOG_DBG("%s: output embedding shape [%d, %d, %d]\n", __func__, - (int)embeddings->ne[0], (int)embeddings->ne[1], (int)embeddings->ne[2]); - - // copy output to user buffer if provided - // if output is empty, skip the copy - auto & out_batch_embd = *params->out_embd; - if (!out_batch_embd.empty()) { - if (out_batch_embd.size() != (size_t)ggml_nelements(embeddings)) { - LOG_ERR("%s: output buffer has %zu elements but expected %zu\n", __func__, out_batch_embd.size(), (size_t)ggml_nelements(embeddings)); - GGML_ABORT("Output buffer size mismatch"); + if (embeddings != nullptr) { + // sanity check (assuming that all images in batch have the same number of tokens, so we only check the first one) + const int n_tokens_out = embeddings->ne[1]; + const int expected_n_tokens_out = clip_n_output_tokens(ctx, &imgs.entries[0]); + if (n_tokens_out != expected_n_tokens_out) { + LOG_ERR("%s: expected output %d tokens, got %d\n", __func__, expected_n_tokens_out, n_tokens_out); + GGML_ABORT("Invalid number of output tokens"); + } + + LOG_DBG("%s: output embedding shape [%d, %d, %d]\n", __func__, + (int)embeddings->ne[0], (int)embeddings->ne[1], (int)embeddings->ne[2]); + + // copy output to user buffer if provided + // if output is empty, skip the copy + auto & out_batch_embd = *params->out_embd; + if (!out_batch_embd.empty()) { + if (out_batch_embd.size() != (size_t)ggml_nelements(embeddings)) { + LOG_ERR("%s: output buffer has %zu elements but expected %zu\n", __func__, out_batch_embd.size(), (size_t)ggml_nelements(embeddings)); + GGML_ABORT("Output buffer size mismatch"); + } + ggml_backend_tensor_get(embeddings, out_batch_embd.data(), 0, ggml_nbytes(embeddings)); + } else { + LOG_WRN("%s: output buffer is empty, skipping copy\n", __func__); } - ggml_backend_tensor_get(embeddings, out_batch_embd.data(), 0, ggml_nbytes(embeddings)); - } else { - LOG_WRN("%s: output buffer is empty, skipping copy\n", __func__); } - // for audio gen: also copy out the decoded PCM samples - // auto-sized to whatever the graph produced (fixed per model, but not known up-front) + // + // for audio gen models + // + + if (params->out_codes != nullptr) { + ggml_tensor * codes = ggml_graph_get_tensor(gf, "out_codes"); + if (codes == nullptr) { + GGML_ABORT("out_codes requested but graph has no \"out_codes\" tensor"); + } + auto & out_codes = *params->out_codes; + out_codes.resize(ggml_nelements(codes)); + ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes)); + } if (params->out_audio != nullptr) { ggml_tensor * audio = ggml_graph_get_tensor(gf, "out_audio"); if (audio == nullptr) { @@ -5189,8 +5211,11 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { ggml_backend_tensor_get(audio, out_audio.data(), 0, ggml_nbytes(audio)); } + // // Debug: dump final embeddings if MTMD_DEBUG_EMBEDDINGS is set - if (ctx->debug_output_embeddings) { + // + + if (ctx->debug_output_embeddings && embeddings != nullptr) { const int64_t n_embd = embeddings->ne[0]; const int64_t n_tokens = embeddings->ne[1]; std::vector emb_data(ggml_nelements(embeddings)); diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index 6de8bb9ed2..8c1b96d666 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -86,18 +86,30 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx); bool clip_image_encode (struct clip_ctx * ctx, int n_threads, const clip_image_f32 * img, std::vector & out_vec); bool clip_image_batch_encode(struct clip_ctx * ctx, int n_threads, const struct clip_image_f32_batch * imgs, std::vector & out_batch_embd); +enum clip_gen_process_type { + CLIP_GEN_PROCESS_CODE_GEN, // h_state to codes + CLIP_GEN_PROCESS_CODE2WAV, // codes to raw PCM audio +}; struct clip_encode_params { int n_threads = 1; const clip_image_f32_batch * imgs = nullptr; std::vector * out_embd = nullptr; - // note: for audio gen, imgs has expectly one entry of size (n_text_embd, 1), it's the hidden state from backbone - // code0 is the sampled semantic code from backbone - // out_embd holds the embd to be fed back to backbone - // out_audio holds the generated audio samples (PCM float32) + // for audio gen, imgs has exactly one entry (unused content for CODE2WAV, + // for CODE_GEN it holds the hidden state from backbone, size (n_text_embd, 1)) + clip_gen_process_type gen_process = CLIP_GEN_PROCESS_CODE_GEN; + + // CODE_GEN: code0 is the sampled semantic code from backbone, out_codes + // receives this frame's 16 sampled codes, out_embd receives the embd to + // be fed back to the backbone for the next frame int32_t code0 = 0; int32_t top_k = 50; float top_p = 1.0f; + std::vector * out_codes = nullptr; + + // CODE2WAV: codes holds this frame's 16 RVQ codes, out_audio receives the + // decoded PCM samples (F32) + const std::vector * codes = nullptr; std::vector * out_audio = nullptr; }; bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params); diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 1331ac8ac3..81b2a94eb9 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -227,11 +227,14 @@ struct clip_graph_qwen3tts_spkenc : clip_graph { }; struct clip_graph_qwen3tts_gen : clip_graph { - clip_graph_qwen3tts_gen(clip_ctx * ctx, const clip_image_f32 & img, int top_k, float top_p) - : clip_graph(ctx, img), top_k(top_k), top_p(top_p) {} + clip_graph_qwen3tts_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int top_k, float top_p) + : clip_graph(ctx, img), gen_process(gen_process), top_k(top_k), top_p(top_p) {} ggml_cgraph * build() override; - // sampling params, fixed at graph-build time + // which sub-graph build() constructs, fixed at graph-build time + clip_gen_process_type gen_process; + + // sampling params, fixed at graph-build time (CODE_GEN only) int top_k; float top_p; @@ -242,6 +245,7 @@ struct clip_graph_qwen3tts_gen : clip_graph { struct code_gen : clip_graph { code_gen(const clip_graph & parent, int top_k, float top_p) : clip_graph(parent), top_k(top_k), top_p(top_p) {} + ggml_cgraph * build() override { GGML_ABORT("call prefill()/step() instead"); } int top_k; float top_p; @@ -293,14 +297,14 @@ struct clip_graph_qwen3tts_gen : clip_graph { ggml_tensor * causal_conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride) const; ggml_tensor * snake(ggml_tensor * x, ggml_tensor * alpha, ggml_tensor * beta) const; - ggml_tensor * quant_decode(ggml_tensor * out_code_cache) const; + ggml_tensor * quant_decode(ggml_tensor * inp_codes) const; ggml_tensor * tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * pos0, ggml_tensor * mask) const; ggml_tensor * convnext_block(ggml_tensor * x, const clip_code2wav::upsample_block & blk) const; ggml_tensor * dac_res_unit(ggml_tensor * x, const clip_code2wav::dac_res & res, int dilation) const; - // out_code_cache: [1, n_codes] I32 (as produced by prefill()/step()). + // inp_codes: [1, n_codes] I32, one frame's RVQ codes. // returns audio samples, [n_samples] F32, clamped to [-1, 1]. - ggml_tensor * decode(ggml_tensor * out_code_cache) const; + ggml_tensor * decode(ggml_tensor * inp_codes) const; }; }; diff --git a/tools/mtmd/models/qwen3tts-gen.cpp b/tools/mtmd/models/qwen3tts-gen.cpp index 895ee77016..bead164f96 100644 --- a/tools/mtmd/models/qwen3tts-gen.cpp +++ b/tools/mtmd/models/qwen3tts-gen.cpp @@ -357,17 +357,17 @@ ggml_tensor * clip_graph_qwen3tts_gen::code2wav::snake(ggml_tensor * x, ggml_ten // RVQ codebook decode: 16 codes -> 512-dim hidden (C-first, [512, 1]). // codebook 0 (semantic) and 1..15 (acoustic) are summed within their own // group, projected out_proj'd separately, then the two projections added. -ggml_tensor * clip_graph_qwen3tts_gen::code2wav::quant_decode(ggml_tensor * out_code_cache) const { +ggml_tensor * clip_graph_qwen3tts_gen::code2wav::quant_decode(ggml_tensor * inp_codes) const { const auto & c2w = model.c2w; - ggml_tensor * code0 = ggml_view_1d(ctx0, out_code_cache, 1, 0); + ggml_tensor * code0 = ggml_view_1d(ctx0, inp_codes, 1, 0); ggml_tensor * sem = ggml_get_rows(ctx0, c2w.quant_first_cb_w, code0); // [256, 1] ggml_tensor * sem_out = ggml_mul_mat(ctx0, c2w.quant_first_out_w, sem); // [512, 1] ggml_tensor * acc = nullptr; const int64_t n_acoustic = c2w.quant_rest_cb_w->ne[2]; for (int g = 1; g <= n_acoustic; g++) { - ggml_tensor * codeg = ggml_view_1d(ctx0, out_code_cache, 1, (size_t) g * out_code_cache->nb[1]); + ggml_tensor * codeg = ggml_view_1d(ctx0, inp_codes, 1, (size_t) g * inp_codes->nb[1]); ggml_tensor * cb_g = ggml_view_2d(ctx0, c2w.quant_rest_cb_w, c2w.quant_rest_cb_w->ne[0], c2w.quant_rest_cb_w->ne[1], c2w.quant_rest_cb_w->nb[1], (size_t) (g - 1) * c2w.quant_rest_cb_w->nb[2]); ggml_tensor * embd = ggml_get_rows(ctx0, cb_g, codeg); // [256, 1] @@ -466,11 +466,11 @@ ggml_tensor * clip_graph_qwen3tts_gen::code2wav::dac_res_unit(ggml_tensor * x, c } // RVQ codes -> raw PCM. Single frame only: no cross-call state. -ggml_tensor * clip_graph_qwen3tts_gen::code2wav::decode(ggml_tensor * out_code_cache) const { +ggml_tensor * clip_graph_qwen3tts_gen::code2wav::decode(ggml_tensor * inp_codes) const { const auto & c2w = model.c2w; // 1. quantizer decode: 16 codes -> [512, 1] (C-first) - ggml_tensor * hidden = quant_decode(out_code_cache); + ggml_tensor * hidden = quant_decode(inp_codes); // 2. pre_conv: [512, 1] -> T-first [1, 512] -> causal conv k=3 -> [1, 1024] ggml_tensor * x = ggml_cont(ctx0, ggml_transpose(ctx0, hidden)); // [1, 512] @@ -534,9 +534,28 @@ ggml_tensor * clip_graph_qwen3tts_gen::code2wav::decode(ggml_tensor * out_code_c return x; } +// master build(): switches on gen_process to construct either the code_gen +// sub-graph (backbone hidden state -> 16 RVQ codes + next-step embd) or the +// code2wav sub-graph (16 RVQ codes -> raw PCM), both hosted in this one clip_ctx. ggml_cgraph * clip_graph_qwen3tts_gen::build() { GGML_ASSERT(n_batch == 1); // this module only ever processes one frame at a time + if (gen_process == CLIP_GEN_PROCESS_CODE2WAV) { + const int64_t n_acoustic = model.gen_code_head_w->ne[2]; // 15 + const int n_codes = (int) n_acoustic + 1; // 16 + + ggml_tensor * inp_codes = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_codes); + ggml_set_name(inp_codes, "inp_codes"); + ggml_set_input(inp_codes); + + ggml_tensor * out_audio = code2wav(*this).decode(inp_codes); + ggml_set_name(out_audio, "out_audio"); + ggml_set_output(out_audio); + ggml_build_forward_expand(gf, out_audio); + return gf; + } + + // CLIP_GEN_PROCESS_CODE_GEN ggml_tensor * h_state = build_inp_raw(1); h_state = ggml_reshape_1d(ctx0, h_state, h_state->ne[0]); cb(h_state, "inp_h_state", -1); @@ -582,11 +601,11 @@ ggml_cgraph * clip_graph_qwen3tts_gen::build() { out_code_cache = cg.step(k_cache, v_cache, out_code_cache, inp_rand, g); } - // output 1: raw PCM audio for this frame, decoded from the 16 sampled codes - ggml_tensor * out_audio = code2wav(*this).decode(out_code_cache); - ggml_set_name(out_audio, "out_audio"); - ggml_set_output(out_audio); - ggml_build_forward_expand(gf, out_audio); + // output 1: this frame's 16 sampled codes, for the caller's code2wav window + ggml_tensor * out_codes = ggml_cont(ctx0, out_code_cache); + ggml_set_name(out_codes, "out_codes"); + ggml_set_output(out_codes); + ggml_build_forward_expand(gf, out_codes); // output 2 (last node, read by clip_encode()): the sum of all 16 // codebook embeddings, fed back to the talker backbone for the next frame diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index fb4d7357dc..889b5e2b3d 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -264,7 +264,9 @@ struct mtmd_context { // generation context struct clip_ctx * ctx_gen_a; // audio - std::vector gen_out_audio; // decoded PCM samples for the current frame + std::vector gen_out_codes; // this frame's 16 sampled codes (CODE_GEN) + std::vector gen_out_embd; // next-step hidden state fed back to backbone (CODE_GEN) + std::vector gen_out_audio; // decoded PCM samples for the current frame (CODE2WAV) bool print_timings; int n_threads; @@ -1586,42 +1588,72 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in return 1; } - const size_t n_embd = (size_t) clip_n_mmproj_embd(ctx_clip); - if (inp->n_embd != n_embd) { - LOG_ERR("%s: n_embd mismatch: model expects %zu, got %zu\n", __func__, n_embd, inp->n_embd); - return 1; + if (inp->type == MTMD_GEN_PROCESS_TYPE_GEN_CODE) { + const size_t n_embd = (size_t) clip_n_mmproj_embd(ctx_clip); + + clip_image_f32 hidden_state; + hidden_state.set_size({(int) n_embd, 1}, false, true); + hidden_state.cpy_buf(std::vector(inp->embd, inp->embd + n_embd)); + + clip_image_f32_batch batch; + batch.is_audio = true; + batch.entries.push_back(std::move(hidden_state)); + + std::vector out_embd(n_embd); + std::vector out_codes; + + clip_encode_params params; + params.imgs = &batch; + params.n_threads = ctx->n_threads; + params.gen_process = CLIP_GEN_PROCESS_CODE_GEN; + params.out_embd = &out_embd; + params.out_codes = &out_codes; + params.code0 = inp->code0; + params.top_k = inp->top_k; + params.top_p = inp->top_p; + + if (!clip_encode(ctx_clip, ¶ms)) { + LOG_ERR("%s: clip_encode failed (gen_code)\n", __func__); + return 1; + } + + ctx->gen_out_embd = std::move(out_embd); + ctx->gen_out_codes = std::move(out_codes); + + out->embd = ctx->gen_out_embd.data(); + out->codes = ctx->gen_out_codes.data(); + out->n_codes = ctx->gen_out_codes.size(); + return 0; } - clip_image_f32 hidden_state; - hidden_state.set_size({(int) n_embd, 1}, false, true); - hidden_state.cpy_buf(std::vector(inp->embd, inp->embd + n_embd)); + // MTMD_GEN_PROCESS_TYPE_CODE2WAV + if (!inp->codes || inp->n_codes == 0) { + LOG_ERR("%s: codes required for code2wav\n", __func__); + return 1; + } + std::vector in_codes(inp->codes, inp->codes + inp->n_codes); + + // code2wav has no hidden-state input, the batch entry is an unused placeholder + clip_image_f32 dummy; + dummy.set_size({1, 1}, false, true); + dummy.cpy_buf(std::vector(1, 0.0f)); clip_image_f32_batch batch; batch.is_audio = true; - batch.entries.push_back(std::move(hidden_state)); + batch.entries.push_back(std::move(dummy)); - std::vector out_embd(n_embd); - ctx->gen_out_audio.clear(); clip_encode_params params; - params.imgs = &batch; - params.n_threads = ctx->n_threads; - params.out_embd = &out_embd; - params.out_audio = &ctx->gen_out_audio; - params.code0 = inp->code0; - params.top_k = inp->top_k; - params.top_p = inp->top_p; + params.imgs = &batch; + params.n_threads = ctx->n_threads; + params.gen_process = CLIP_GEN_PROCESS_CODE2WAV; + params.codes = &in_codes; + params.out_audio = &ctx->gen_out_audio; if (!clip_encode(ctx_clip, ¶ms)) { - LOG_ERR("%s: clip_encode failed\n", __func__); + LOG_ERR("%s: clip_encode failed (code2wav)\n", __func__); return 1; } - if (!out->embd || out->n_embd != out_embd.size()) { - LOG_ERR("%s: output buffer size mismatch: expected %zu, got %zu\n", __func__, out_embd.size(), out->n_embd); - return 1; - } - std::copy(out_embd.begin(), out_embd.end(), out->embd); - out->audio = ctx->gen_out_audio.data(); out->n_samples = ctx->gen_out_audio.size(); diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 29fd274860..35b8170a08 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -350,13 +350,15 @@ struct mtmd_gen_inp { float top_p; // for MTMD_GEN_PROCESS_TYPE_CODE2WAV - int32_t * codes; // the sampled codebook entries, must have n_codes elements - size_t n_codes; + int32_t * codes; + size_t n_codes; }; struct mtmd_gen_out { // note: output memory is allocated by the context, valid until next process() call // for MTMD_GEN_PROCESS_TYPE_GEN_CODE + const int32_t * codes; + size_t n_codes; const float * embd; // the generated hidden state, to be fed back to backbone // it must have n_text_embd elements