From f55a236349d03272cdc092705a75486b6adf5886 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Wed, 5 Aug 2026 02:00:52 +0200 Subject: [PATCH] adapt the api --- tools/mtmd/clip-model.h | 3 ++ tools/mtmd/clip.cpp | 34 ++++++++++++++++++--- tools/mtmd/clip.h | 5 ++++ tools/mtmd/mtmd-helper-gen.cpp | 34 +++++++++++++++++---- tools/mtmd/mtmd-helper.h | 16 ++++++---- tools/mtmd/mtmd.cpp | 54 ++++++++++++++++++++++++---------- tools/mtmd/mtmd.h | 10 ++++++- tools/tts/tts.cpp | 15 ++++++---- 8 files changed, 133 insertions(+), 38 deletions(-) diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 101f49cd18..f3533f7eb4 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -136,6 +136,9 @@ struct clip_hparams { int32_t rvq_num_quantizers = 0; std::vector rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17) + // threshold for the "out_eos_score" graph output + float gen_eos_threshold = 0.0f; + // qwen3tts code2wav int32_t wav_tfm_n_layer = 0; int32_t wav_tfm_n_embd = 0; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index d6670030ff..fcf022b8d2 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -174,6 +174,10 @@ struct clip_ctx { bool support_batch = false; + // for audio gen, reseeded only when the caller asks for another seed + std::mt19937 rng{std::random_device{}()}; + uint32_t rng_seed = UINT32_MAX; + clip_ctx(clip_context_params & ctx_params) { flash_attn_type = ctx_params.flash_attn_type; no_alloc = ctx_params.no_alloc; @@ -4080,6 +4084,11 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { clip_model_loader::warmup(*ctx, *params->imgs); } + if (params->seed != ctx->rng_seed) { + ctx->rng_seed = params->seed; + ctx->rng.seed(params->seed == UINT32_MAX ? std::random_device{}() : params->seed); + } + // build the inference graph ggml_backend_sched_reset(ctx->sched.get()); ggml_cgraph * gf = clip_get_graph_builder(ctx, imgs, params)->build(); @@ -4786,11 +4795,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { set_input_i32("inp_code0", code0); // one uniform(0,1) draw per codebook, used by do_sampling() - 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) }; + std::vector r = { dist(ctx->rng) }; set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r); } } @@ -5252,6 +5260,24 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { out_codes.resize(ggml_nelements(codes)); ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes)); } + // optional outputs, a missing tensor is not an error + if (params->out_feats != nullptr) { + ggml_tensor * feats = ggml_graph_get_tensor(gf, "out_feats"); + if (feats != nullptr) { + auto & out_feats = *params->out_feats; + out_feats.resize(ggml_nelements(feats)); + ggml_backend_tensor_get(feats, out_feats.data(), 0, ggml_nbytes(feats)); + } + } + if (params->out_is_eos != nullptr) { + ggml_tensor * eos = ggml_graph_get_tensor(gf, "out_eos_score"); + if (eos != nullptr) { + GGML_ASSERT(ggml_nelements(eos) == 1); + float score = 0.0f; + ggml_backend_tensor_get(eos, &score, 0, sizeof(float)); + *params->out_is_eos = score > hparams.gen_eos_threshold; + } + } if (params->out_audio != nullptr) { ggml_tensor * audio = ggml_graph_get_tensor(gf, "out_audio"); if (audio == nullptr) { @@ -5262,9 +5288,9 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { ggml_backend_tensor_get(audio, out_audio.data(), 0, ggml_nbytes(audio)); // drop the tail audio that comes from the code-0 rear padding - const int64_t n_codes = model.gen_code_head_w->ne[2] + 1; + const int64_t n_codes = params->codes ? model.gen_code_head_w->ne[2] + 1 : 0; const int64_t n_frames_w = hparams.wav_tfm_swa; - const int64_t n_frames = (int64_t) params->codes->size() / n_codes; + const int64_t n_frames = params->codes ? (int64_t) params->codes->size() / n_codes : n_frames_w; if (n_frames < n_frames_w) { const size_t hop = out_audio.size() / n_frames_w; out_audio.resize((size_t) n_frames * hop); diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index 7f706d976e..97de0606db 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -104,9 +104,14 @@ struct clip_encode_params { int32_t top_k = 50; float top_p = 1.0f; std::vector * out_codes = nullptr; // this frame's 16 sampled codes + std::vector * out_feats = nullptr; // continuous counterpart of out_codes + uint32_t seed = UINT32_MAX; // UINT32_MAX for random + int32_t n_steps = -1; // integration steps, for flow-matching decoders + bool * out_is_eos = nullptr; // GEN_WAV const std::vector * codes = nullptr; // this frame's 16 RVQ codes + const std::vector * feats = nullptr; // continuous counterpart of codes std::vector * out_audio = nullptr; // decoded PCM samples, F32 const std::vector * state_in = nullptr; // state from previous call, null or wrong size means cold start std::vector * state_out = nullptr; // state for the next call diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index b52dc8e5a3..bc83f5ba23 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -87,7 +87,8 @@ public: virtual int32_t step_prompt(int32_t n_batch) = 0; // sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token, // those read what they need from h_state_in instead - virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) = 0; + // set out_stop on end-of-speech, h_state_out must be null if no frame is generated + virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) = 0; virtual int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) = 0; protected: @@ -203,6 +204,7 @@ public: pos = 0; top_k = inp->top_k > 0 ? inp->top_k : 50; top_p = inp->top_p > 0 ? inp->top_p : 1.0f; + seed = inp->seed; out_type = inp->out_type; // the text stream keeps flowing during generation: after frame k, the input adds @@ -244,13 +246,26 @@ public: return n_prompt - prompt_pos; } - int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) override { + int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override { + if (sampled == LLAMA_TOKEN_NULL) { + LOG_ERR("mtmd_helper_gen_audio: qwen3tts requires a token sampled from the backbone\n"); + return 1; + } + + // backbone signals end-of-speech with a token, no frame for this step + if (sampled == codec_eos || llama_vocab_is_eog(vocab, sampled)) { + *out_stop = true; + *h_state_out = nullptr; + return 0; + } + mtmd_gen_inp inp{}; inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; inp.code0 = sampled - codec_0; inp.embd = const_cast(h_state_in); inp.top_k = top_k; inp.top_p = top_p; + inp.seed = seed; mtmd_gen_out out{}; if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) { LOG_ERR("mtmd_helper_gen_audio: gen_code process failed\n"); @@ -432,8 +447,9 @@ private: std::unique_ptr prompt_batch; int n_prompt = 0; int prompt_pos = 0; - int32_t top_k = 50; - float top_p = 1.0f; + int32_t top_k = 50; + float top_p = 1.0f; + uint32_t seed = UINT32_MAX; std::vector codes_buf; std::vector c2w_state; std::vector audio_pcm; @@ -489,11 +505,17 @@ int32_t mtmd_helper_gen_audio_step_prompt(mtmd_helper_gen_audio * ctx, int32_t n } int32_t mtmd_helper_gen_audio_step_gen(mtmd_helper_gen_audio * ctx, llama_token sampled, - const float * h_state_in, const float ** h_state_out) { + const float * h_state_in, const float ** h_state_out, + bool * out_stop) { if (!ctx->pipeline) { return 1; } - return ctx->pipeline->step_gen(sampled, h_state_in, h_state_out); + bool stop = false; + const int32_t ret = ctx->pipeline->step_gen(sampled, h_state_in, h_state_out, &stop); + if (out_stop) { + *out_stop = stop; + } + return ret; } int32_t mtmd_helper_gen_audio_get_output(mtmd_helper_gen_audio * ctx, int32_t * out_sample_rate, diff --git a/tools/mtmd/mtmd-helper.h b/tools/mtmd/mtmd-helper.h index 7e5cf9b509..832f7171ac 100644 --- a/tools/mtmd/mtmd-helper.h +++ b/tools/mtmd/mtmd-helper.h @@ -183,8 +183,9 @@ struct mtmd_helper_gen_audio_inp { mtmd_bitmap * speaker_ref; // optional, can be NULL const char * lang; // optional, can be NULL - int32_t top_k; - float top_p; + int32_t top_k; + float top_p; + uint32_t seed; // UINT32_MAX for random (default: random) enum mtmd_helper_gen_audio_outtype out_type; }; @@ -208,12 +209,15 @@ MTMD_API int32_t mtmd_helper_gen_audio_step_prompt( int32_t n_batch); // generates one frame; must only be called after step_prompt() has returned 0 -// h_state_out is valid until next step_gen() or reset() call +// sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token +// out_stop (optional) is set on end-of-speech, the caller must then stop the loop +// h_state_out is valid until next step_gen() or reset() call, null if no frame is generated MTMD_API int32_t mtmd_helper_gen_audio_step_gen( mtmd_helper_gen_audio * ctx, llama_token sampled, const float * h_state_in, - const float ** h_state_out); + const float ** h_state_out, + bool * out_stop); // out_data valid until next get_output() or reset() call // out_n_samples (optional, can be NULL) receives the number of generated PCM samples @@ -261,8 +265,8 @@ struct gen_audio { int32_t step_prompt(int32_t n_batch) { return mtmd_helper_gen_audio_step_prompt(ctx.get(), n_batch); } - int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out) { - return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out); + int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out, bool * out_stop = nullptr) { + return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out, out_stop); } int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples = nullptr) { return mtmd_helper_gen_audio_get_output(ctx.get(), out_sample_rate, out_data, out_data_len, out_n_samples); diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index ff90d6818c..6a5b3f5a60 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -265,6 +265,7 @@ struct mtmd_context { // generation context struct clip_ctx * ctx_gen_a; // audio std::vector gen_out_codes; // this frame's 16 sampled codes (GEN_CODE) + std::vector gen_out_feats; // this frame's continuous features, if any (GEN_CODE) std::vector gen_out_embd; // next-step hidden state fed back to backbone (GEN_CODE) std::vector gen_out_audio; // decoded PCM samples for the current frame (GEN_WAV) std::vector gen_out_state; // state to feed into the next GEN_WAV call @@ -1580,7 +1581,7 @@ float * mtmd_get_output_embd(mtmd_context * ctx) { // mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) { - mtmd_gen_audio_info info; + mtmd_gen_audio_info info{}; if (!ctx->ctx_gen_a) { info.type = MTMD_GEN_AUDIO_TYPE_NONE; return info; @@ -1604,6 +1605,8 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in return 1; } + *out = {}; + if (inp->type == MTMD_GEN_PROCESS_TYPE_GEN_CODE) { const size_t n_embd = (size_t) clip_n_mmproj_embd(ctx_clip); @@ -1617,16 +1620,22 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in std::vector out_embd(n_embd); std::vector out_codes; + std::vector out_feats; + bool is_eos = false; clip_encode_params params; - params.imgs = &batch; - params.n_threads = ctx->n_threads; - params.gen_process = CLIP_GEN_PROCESS_GEN_CODE; - 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; + params.imgs = &batch; + params.n_threads = ctx->n_threads; + params.gen_process = CLIP_GEN_PROCESS_GEN_CODE; + params.out_embd = &out_embd; + params.out_codes = &out_codes; + params.out_feats = &out_feats; + params.code0 = inp->code0; + params.top_k = inp->top_k; + params.top_p = inp->top_p; + params.seed = inp->seed; + params.n_steps = inp->n_steps; + params.out_is_eos = &is_eos; if (!clip_encode(ctx_clip, ¶ms)) { LOG_ERR("%s: clip_encode failed (gen_code)\n", __func__); @@ -1635,19 +1644,31 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in ctx->gen_out_embd = std::move(out_embd); ctx->gen_out_codes = std::move(out_codes); + ctx->gen_out_feats = std::move(out_feats); - out->embd = ctx->gen_out_embd.data(); - out->codes = ctx->gen_out_codes.data(); - out->n_codes = ctx->gen_out_codes.size(); + out->embd = ctx->gen_out_embd.data(); + out->codes = ctx->gen_out_codes.data(); + out->n_codes = ctx->gen_out_codes.size(); + out->feats = ctx->gen_out_feats.data(); + out->n_feats = ctx->gen_out_feats.size(); + out->is_eos = is_eos; return 0; } // MTMD_GEN_PROCESS_TYPE_GEN_WAV - if (!inp->codes || inp->n_codes == 0) { - LOG_ERR("%s: codes required for gen_wav\n", __func__); + const bool has_codes = inp->codes && inp->n_codes > 0; + const bool has_feats = inp->feats && inp->n_feats > 0; + if (has_codes == has_feats) { + LOG_ERR("%s: gen_wav requires exactly one of codes or feats\n", __func__); return 1; } - std::vector in_codes(inp->codes, inp->codes + inp->n_codes); + std::vector in_codes; + std::vector in_feats; + if (has_codes) { + in_codes.assign(inp->codes, inp->codes + inp->n_codes); + } else { + in_feats.assign(inp->feats, inp->feats + inp->n_feats); + } std::vector in_state; if (inp->state_data) { in_state.assign(inp->state_data, inp->state_data + inp->state_size); @@ -1667,7 +1688,8 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in params.imgs = &batch; params.n_threads = ctx->n_threads; params.gen_process = CLIP_GEN_PROCESS_GEN_WAV; - params.codes = &in_codes; + params.codes = has_codes ? &in_codes : nullptr; + params.feats = has_feats ? &in_feats : nullptr; params.out_audio = &ctx->gen_out_audio; params.state_in = inp->state_data ? &in_state : nullptr; params.state_out = &ctx->gen_out_state; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 84651f8dcd..32d40dca32 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -354,10 +354,15 @@ struct mtmd_gen_inp { float * embd; // the hidden state from backbone, must have n_text_embd elements int32_t top_k; float top_p; + uint32_t seed; // UINT32_MAX for random + int32_t n_steps; // integration steps, for flow-matching decoders (-1 for default) // for MTMD_GEN_PROCESS_TYPE_GEN_WAV + // pass either codes (discrete) or feats (continuous), depending on the pipeline int32_t * codes; size_t n_codes; + const float * feats; + size_t n_feats; const char * state_data; size_t state_size; }; @@ -366,9 +371,12 @@ struct mtmd_gen_out { // for MTMD_GEN_PROCESS_TYPE_GEN_CODE const int32_t * codes; - size_t n_codes; + size_t n_codes; + const float * feats; // continuous counterpart of codes + size_t n_feats; const float * embd; // the generated hidden state, to be fed back to backbone // it must have n_text_embd elements + bool is_eos; // only set by pipelines having the EOS head inside mmproj // for MTMD_GEN_PROCESS_TYPE_GEN_WAV const float * audio; diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index b68edcaf57..9e52505d3b 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -119,6 +119,7 @@ int main(int argc, char ** argv) { inp.lang = params.tts_lang.c_str(); inp.top_k = params.sampling.top_k; inp.top_p = params.sampling.top_p; + inp.seed = params.sampling.seed; inp.out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; // @@ -143,8 +144,7 @@ int main(int argc, char ** argv) { } } - const llama_vocab * vocab = llama_model_get_vocab(model); - + // note: some pipelines ignore this token and use the hidden state instead auto sample_semantic_code = [&]() -> llama_token { llama_token t = common_sampler_sample(smpl, lctx, -1); common_sampler_accept(smpl, t, true); @@ -159,19 +159,24 @@ int main(int argc, char ** argv) { tts_timings timings; const int64_t t_gen_start_us = ggml_time_us(); - for (; n_frames < max_new && !llama_vocab_is_eog(vocab, sampled); n_frames++) { + bool stop = false; + while (!stop && n_frames < max_new) { const float * h_next = nullptr; // stage 2+3: semantic --> acoustic details --> audio waveform // step_gen() runs both stages and returns new h_state for next step - if (gen.step_gen(sampled, h_state, &h_next) != 0) { + if (gen.step_gen(sampled, h_state, &h_next, &stop) != 0) { LOG_ERR("step_gen failed at frame %d\n", n_frames); return 1; } + if (!h_next) { + break; // stopped without generating a frame + } + n_frames++; h_state = h_next; sampled = sample_semantic_code(); - timings.report(n_frames + 1); + timings.report(n_frames); } const double t_gen_s = (ggml_time_us() - t_gen_start_us) / 1e6;