From 6c197a77ea77ab77baaf5c3238beec924f373e64 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 10 Aug 2026 23:54:13 +0200 Subject: [PATCH] address security problems --- tools/mtmd/clip-model.h | 1 + tools/mtmd/clip.cpp | 12 +++++++++++ tools/mtmd/models/pockettts-gen.cpp | 6 ++++++ tools/mtmd/mtmd-audio.cpp | 8 ++++++++ tools/mtmd/mtmd-helper-gen.cpp | 32 +++++++++++++++-------------- 5 files changed, 44 insertions(+), 15 deletions(-) diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 206fc4be67..ad25c008e7 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -148,6 +148,7 @@ struct clip_hparams { std::string gen_model_variant; // pocket-tts + static constexpr int32_t pockettts_max_spk_seconds = 30; int32_t seanet_n_stage = 0; std::vector seanet_ratios; // encoder order (reversed compared to the config) int32_t mimi_downsample = 0; // encoder frame rate / model frame rate diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index cb03baabb1..a0c1c169ba 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1072,6 +1072,13 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE; const int n_step = params && params->n_steps > 0 ? params->n_steps : ctx->model.hparams.flow_n_step; const int64_t n_latent = ctx->model.gen_input_lin_w->ne[0]; + GGML_ASSERT(n_step > 0); + GGML_ASSERT(n_latent > 0); + // "inp_feats" takes the caller's buffer as-is, the graph must consume all of it + if (params && params->feats) { + GGML_ASSERT(params->feats->size() % (size_t) n_latent == 0); + GGML_ASSERT(params->feats->size() >= (size_t) n_latent); + } const int n_frames = params && params->feats ? (int) (params->feats->size() / n_latent) : 1; builder = std::make_unique(ctx, img, gen_process, n_step, n_frames); } break; @@ -4341,12 +4348,17 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { // the mask is causal with a sliding window, see _build_attention_mask() in the reference auto set_pockettts_tfm_inputs = [&]() { const int64_t n_pos = ggml_nelements(get_inp_tensor("inp_pos")); + GGML_ASSERT(n_pos > 0); std::vector positions((size_t) n_pos); for (int64_t i = 0; i < n_pos; i++) { positions[(size_t) i] = (int32_t) i; } set_input_i32("inp_pos", positions); + // the preprocessor truncates the waveform to keep this mask bounded + const int64_t max_pos = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate / 120; + GGML_ASSERT(n_pos <= max_pos && "pocket-tts speaker reference too long for a dense mask"); + const int64_t context = hparams.mimi_tfm_context; std::vector mask((size_t) n_pos * n_pos, -INFINITY); for (int64_t q = 0; q < n_pos; q++) { diff --git a/tools/mtmd/models/pockettts-gen.cpp b/tools/mtmd/models/pockettts-gen.cpp index 3266742150..3fd613e5f7 100644 --- a/tools/mtmd/models/pockettts-gen.cpp +++ b/tools/mtmd/models/pockettts-gen.cpp @@ -106,6 +106,12 @@ std::vector list_pockettts_state_slots(const clip_hparams & hpar } const auto & seanet = model.seanet; + // the slots below are sized from these + GGML_ASSERT(!model.gen_tfm_layers.empty()); + GGML_ASSERT((int) seanet.stages.size() >= hparams.seanet_n_stage); + GGML_ASSERT((int) hparams.seanet_ratios.size() >= hparams.seanet_n_stage); + GGML_ASSERT(hparams.mimi_tfm_context > 1 && hparams.mimi_downsample > 0); + slots.push_back({"tfm_pos", 1, 1}); const int64_t n_embd_a = model.gen_tfm_layers[0].q_w->ne[1]; diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index 67c533f40a..98a8c11ee9 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1440,6 +1440,14 @@ bool mtmd_audio_preprocessor_pockettts::preprocess(const float * return false; } + // the mimi transformer mask is dense, so cost is quadratic in the reference length + const int64_t max_samples = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate; + if ((int64_t) n_samples > max_samples) { + LOG_WRN("%s: speaker reference is %.1f s, truncating to the first %d s\n", __func__, + (double) n_samples / hparams.audio_sample_rate, clip_hparams::pockettts_max_spk_seconds); + n_samples = (size_t) max_samples; + } + const int64_t n_frames = (int64_t) (n_samples + frame_size - 1) / frame_size; const int64_t n_padded = n_frames * frame_size; diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 40b0ef7472..72106aefb2 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -549,16 +549,12 @@ public: } const int n_e = n_embd; - auto push_row = [&](llama_token t) { - prompt_embd_buf.insert(prompt_embd_buf.end(), - tok_embd.begin() + (size_t) t * n_e, - tok_embd.begin() + (size_t) (t + 1) * n_e); - }; // sequence order is voice, then text, then the audio BOS that starts generation if (!voice.empty()) { + GGML_ASSERT(voice.size() % (size_t) n_e == 0); if (bos_before_voice != LLAMA_TOKEN_NULL) { - push_row(bos_before_voice); + push_embd_row(prompt_embd_buf, bos_before_voice); } prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end()); } @@ -566,9 +562,9 @@ public: n_voice_pos = (int) (prompt_embd_buf.size() / (size_t) n_e); for (llama_token t : chunks[0]) { - push_row(t); + push_embd_row(prompt_embd_buf, t); } - push_row(audio_bos); + push_embd_row(prompt_embd_buf, audio_bos); arm_chunk_budget(0); n_prompt = (int) (prompt_embd_buf.size() / (size_t) n_e); @@ -712,10 +708,20 @@ private: LOG_ERR("mtmd_helper_gen_audio: token embedding copy failed\n"); return false; } + GGML_ASSERT(n_embd > 0 && n_tok_embd % (uint32_t) n_embd == 0); specials_ok = true; return true; } + // the table can be shorter than the vocab, so bound the row lookup + void push_embd_row(std::vector & dst, llama_token t) const { + const size_t n_rows = tok_embd.size() / (size_t) n_embd; + GGML_ASSERT(t >= 0 && (size_t) t < n_rows); + dst.insert(dst.end(), + tok_embd.begin() + (size_t) t * n_embd, + tok_embd.begin() + (size_t) (t + 1) * n_embd); + } + // token ids of the pieces the reference splits on, see split_into_best_sentences(). // the leading token is dropped, it is the tokenizer's dummy prefix std::vector punct_ids(const char * s) const { @@ -824,17 +830,13 @@ private: const int n_e = n_embd; prompt_embd_buf.clear(); for (llama_token t : chunks[chunk_idx]) { - prompt_embd_buf.insert(prompt_embd_buf.end(), - tok_embd.begin() + (size_t) t * n_e, - tok_embd.begin() + (size_t) (t + 1) * n_e); + push_embd_row(prompt_embd_buf, t); } - prompt_embd_buf.insert(prompt_embd_buf.end(), - tok_embd.begin() + (size_t) audio_bos * n_e, - tok_embd.begin() + (size_t) (audio_bos + 1) * n_e); + push_embd_row(prompt_embd_buf, audio_bos); arm_chunk_budget(chunk_idx); - // bounded by max_chunk_tokens + 1, so one decode is enough const int n_rows = (int) (prompt_embd_buf.size() / (size_t) n_e); + GGML_ASSERT(n_rows > 0); decode_embd_batch batch(prompt_embd_buf.data(), n_rows, 1, n_e); batch.set_position_normal(pos, seq_id); batch.batch.logits[n_rows - 1] = 1;