LLM-Inference/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-08-12 22:31:11 +04:00
Code Issues Packages Projects Releases Wiki Activity
Files
master
llama.cpp/tools/mtmd/mtmd-helper-gen.cpp
T
Xuan-Son NguyenGitHubPascal
 6e62ba5384 mtmd: support pocket-tts (#26871) 
* adapt the api

* text model ok

* working impl, need verify and clean up

* mtmd: build the pocket-tts transposed convolutions as GEMM + col2im

ggml_conv_transpose_1d has no grouped mode, so the depthwise upsample
was built as one convolution and one concat per channel, which floods
the graph with small nodes and makes kernel launches dominate the
decoder.

Fold both cases into the column form the seanet decoder already needs:
the general case reshapes the kernel to [IC, K * OC] and matmuls it
with the input, the depthwise case batches a matmul over the channels
so a step scales its own kernel. A single col2im_1d then scatter-adds
the columns back to the signal, with the same shape as before, so the
overlap-add tail, the streaming state and the bias are untouched.

Generation time per frame drops by 80% on CUDA and by 50% on CPU. The
output matches the previous implementation sample for sample, with a
correlation of 0.999994 and identical frame counts.

* flow_temp +  frames_after_eos

* chunking

* mtmd: carry the remaining pocket-tts per-pack settings

The language packs also tune the end-of-speech padding and the padding
of short prompts, next to the temperature already carried in the
mmproj: french_24l asks for 8 tail frames instead of the guessed 3,
english_2026-01 asks for short prompts to be padded with spaces.

Write both in the mmproj as clip.gen.audio.frames_after_eos and
clip.gen.audio.pad_short_text, keyed on the pack in the conversion
script like the temperature. The loader keeps them optional, so a
mmproj without them behaves as before. Map semicolons to commas for
every pack instead, the reference only asks for it on three of them and
it costs nothing elsewhere.

Existing mmproj files must be converted again to carry the two keys.

On a long french text the port now lands within 2% of the reference:
22.96s against 23.44s, with the same peak level and the same amount of
silence.

* clip.gen.audio.model_variant

* clean up code comments

* nit: drop the dead flow_temp hparam, the pack table holds the default

* update docs

* address security problems

* less invasive base.py

* lint

* add mtmd_gen_inp_default

* add docs

* rm gen_flow_temp

---------

Co-authored-by: Pascal <admin@serveurperso.com>
2026-08-11 14:18:30 +02:00

1063 lines
39 KiB
C++
Raw Permalink Blame History

#include "mtmd.h"
#include "mtmd-helper.h"
#include "mtmd-helper-common.h"
#include "llama.h"
#include "../src/llama-ext.h"
#include <algorithm>
#include <cctype>
#include <cmath>
#include <cstring>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#ifdef MTMD_INTERNAL_HEADER
#error "mtmd-helper is a public library outside of mtmd. it must not include internal headers"
#endif
//
// Audio generation helpers
//
// --tts-lang codes -> language names used by the codec_language special tokens
static const std::unordered_map<std::string, std::string> tts_lang_codes = {
{ "zh", "chinese" },
{ "en", "english" },
{ "de", "german" },
{ "it", "italian" },
{ "pt", "portuguese" },
{ "es", "spanish" },
{ "ja", "japanese" },
{ "ko", "korean" },
{ "fr", "french" },
{ "ru", "russian" },
};
static std::string tts_resolve_lang(const std::string & lang) {
auto it = tts_lang_codes.find(lang);
return it != tts_lang_codes.end() ? it->second : lang;
}
static llama_token find_special_token(const llama_vocab * vocab, const std::string & piece) {
const int32_t n = llama_vocab_n_tokens(vocab);
for (llama_token t = 0; t < n; t++) {
if (piece == llama_vocab_get_text(vocab, t)) {
return t;
}
}
return LLAMA_TOKEN_NULL;
}
static bool write_wav16(std::vector<char> & buf, const std::vector<float> & pcm, int32_t rate) {
// RIFF chunk sizes are 32-bit; refuse to emit a file with a truncated header
if (pcm.size() > ((size_t) UINT32_MAX - 36) / 2) {
return false;
}
const uint32_t data_sz = (uint32_t) (pcm.size() * 2);
const uint32_t riff_sz = 36 + data_sz;
const uint32_t fmt_sz = 16, byte_rate = (uint32_t) rate * 2;
const uint16_t fmt = 1, ch = 1, align = 2, bits = 16;
const uint32_t rate32 = (uint32_t) rate;
auto put = [&](const void * p, size_t n) {
const char * c = (const char *) p;
buf.insert(buf.end(), c, c + n);
};
put("RIFF", 4); put(&riff_sz, 4); put("WAVE", 4);
put("fmt ", 4); put(&fmt_sz, 4);
put(&fmt, 2); put(&ch, 2); put(&rate32, 4);
put(&byte_rate, 4); put(&align, 2); put(&bits, 2);
put("data", 4); put(&data_sz, 4);
for (float v : pcm) {
int16_t s = (int16_t) (std::max(-1.0f, std::min(1.0f, v)) * 32767.0f);
put(&s, 2);
}
return true;
}
class mtmd_gen_audio_pipeline {
public:
mtmd_gen_audio_pipeline(llama_context * lctx, mtmd_context * mctx)
: lctx(lctx), mctx(mctx), model(llama_get_model(lctx)), vocab(llama_model_get_vocab(model)),
n_embd(llama_model_n_embd(model)), info(mtmd_gen_audio_get_info(mctx)) {}
virtual ~mtmd_gen_audio_pipeline() = default;
virtual void reset() = 0;
virtual int32_t set_input(const mtmd_helper_gen_audio_inp * inp) = 0;
// decodes at most n_batch prompt tokens; returns remaining count (0 = done), <0 on error
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
// 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:
llama_context * lctx;
mtmd_context * mctx;
const llama_model * model;
const llama_vocab * vocab;
int n_embd;
mtmd_gen_audio_info info;
};
// Qwen3-TTS: backbone samples codec_0, code_predictor gives the other 15 codebooks,
// then code2wav decodes them to PCM
class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
public:
using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline;
void reset() override {
seq_id = 0;
pos = 0;
codes_buf.clear();
c2w_state.clear();
audio_pcm.clear();
overlay.clear();
h_state_buf.clear();
out_buf.clear();
prompt_embd_buf.clear();
prompt_batch.reset();
n_prompt = 0;
prompt_pos = 0;
}
int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
reset();
seq_id = inp->seq_id;
if (!ensure_cache()) {
return 1;
}
const std::string lang = tts_resolve_lang((inp->lang && inp->lang[0]) ? inp->lang : "english");
const llama_token c_lang = find_special_token(vocab, ("<|codec_language_" + lang + "|>").c_str());
if (c_lang == LLAMA_TOKEN_NULL) {
LOG_ERR("mtmd_helper_gen_audio: unknown language '%s'\n", lang.c_str());
return 1;
}
std::vector<float> speaker_embd;
if (inp->speaker_ref) {
if (!encode_speaker(inp->speaker_ref, speaker_embd)) {
return 1;
}
}
const int n_e = n_embd;
auto row = [&](llama_token t) {
return std::vector<float>(tok_embd.begin() + (size_t) t * n_e,
tok_embd.begin() + (size_t) (t + 1) * n_e);
};
auto sum_row = [&](llama_token a, llama_token b) {
std::vector<float> va = row(a), vb = row(b);
for (int i = 0; i < n_e; i++) va[(size_t) i] += vb[(size_t) i];
return va;
};
auto sum_vec = [&](llama_token a, const std::vector<float> & vb) {
std::vector<float> va = row(a);
for (int i = 0; i < n_e; i++) va[(size_t) i] += vb[(size_t) i];
return va;
};
// upstream chat wrap, then slices: [0:3] role, [3:-5] utterance body
const std::string full = "<|im_start|>assistant\n" + std::string(inp->prompt, inp->prompt_len) +
"<|im_end|>\n<|im_start|>assistant\n";
std::vector<llama_token> ids(full.size() + 16);
int n_ids = llama_tokenize(vocab, full.c_str(), (int32_t) full.size(), ids.data(), (int32_t) ids.size(),
false, true);
if (n_ids < 8) {
LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n");
return 1;
}
ids.resize((size_t) n_ids);
std::vector<std::vector<float>> prompt;
for (int i = 0; i < 3; i++) prompt.push_back(row(ids[(size_t) i]));
prompt.push_back(sum_row(tts_pad, c_think));
prompt.push_back(sum_row(tts_pad, c_think_b));
prompt.push_back(sum_row(tts_pad, c_lang));
prompt.push_back(sum_row(tts_pad, c_think_e));
if (!speaker_embd.empty()) prompt.push_back(sum_vec(tts_pad, speaker_embd));
prompt.push_back(sum_row(tts_bos, codec_pad));
for (int i = 3; i < n_ids - 5; i++) prompt.push_back(sum_row(ids[(size_t) i], codec_pad));
prompt.push_back(sum_row(tts_eos, codec_pad));
prompt.push_back(sum_row(tts_pad, codec_bos));
n_prompt = (int) prompt.size();
// the talker uses the qwen3vl interleaved mrope, all sections are equal for a text/codec stream
mrope = llama_model_rope_type(model) == LLAMA_ROPE_TYPE_MROPE ||
llama_model_rope_type(model) == LLAMA_ROPE_TYPE_IMROPE;
const int n_pos_per_embd = mrope ? 4 : 1;
prompt_embd_buf.resize((size_t) n_prompt * (size_t) n_e);
for (int i = 0; i < n_prompt; i++) {
memcpy(prompt_embd_buf.data() + (size_t) i * n_e, prompt[(size_t) i].data(), (size_t) n_e * sizeof(float));
}
prompt_batch.reset(new decode_embd_batch(prompt_embd_buf.data(), n_prompt, n_pos_per_embd, n_e));
if (mrope) prompt_batch->set_position_mrope_1d(0, seq_id);
else prompt_batch->set_position_normal (0, seq_id);
prompt_pos = 0;
pos = 0;
const mtmd_gen_inp def = mtmd_gen_inp_default(mctx);
top_k = inp->top_k > 0 ? inp->top_k : def.top_k;
top_p = inp->top_p > 0 ? inp->top_p : def.top_p;
seed = inp->seed;
out_type = inp->out_type;
// the prompt above holds the whole text stream up to tts_eos, so every generated
// frame adds tts_pad on top of the codes embedding
overlay = row(tts_pad);
return 0;
}
int32_t step_prompt(int32_t n_batch) override {
GGML_ASSERT(n_batch > 0);
if (prompt_pos >= n_prompt) {
return 0;
}
const int32_t n_tokens_batch = std::min(n_batch, n_prompt - prompt_pos);
llama_batch batch_view = prompt_batch->get_view(prompt_pos, n_tokens_batch);
const bool is_last_batch = (prompt_pos + n_tokens_batch) == n_prompt;
if (is_last_batch) {
batch_view.logits[n_tokens_batch - 1] = 1;
}
if (llama_decode(lctx, batch_view) != 0) {
LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n");
return -1;
}
pos += n_tokens_batch;
prompt_pos += n_tokens_batch;
if (prompt_pos >= n_prompt) {
// prompt fully processed, its embedding buffer is no longer needed
prompt_batch.reset();
prompt_embd_buf.clear();
return 0;
}
return n_prompt - prompt_pos;
}
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 = mtmd_gen_inp_default(mctx);
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
inp.code0 = sampled - codec_0;
inp.embd = const_cast<float *>(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");
return 1;
}
codes_buf.insert(codes_buf.end(), out.codes, out.codes + out.n_codes);
if (out.n_codes > 0 && codes_buf.size() / out.n_codes >= window_frames) {
if (!flush_gen_wav()) {
return 1;
}
}
std::vector<float> fb(out.embd, out.embd + n_embd);
for (int i = 0; i < n_embd; i++) fb[(size_t) i] += overlay[(size_t) i];
const int n_pos_per_embd = mrope ? 4 : 1;
decode_embd_batch batch_embd(fb.data(), 1, n_pos_per_embd, n_embd);
if (mrope) batch_embd.set_position_mrope_1d(pos, seq_id);
else batch_embd.set_position_normal (pos, seq_id);
batch_embd.batch.logits[0] = 1;
pos++;
if (llama_decode(lctx, batch_embd.batch) != 0) {
LOG_ERR("mtmd_helper_gen_audio: decode failed\n");
return 1;
}
const float * he = llama_get_embeddings_ith(lctx, -1);
h_state_buf.assign(he, he + n_embd);
*h_state_out = h_state_buf.data();
return 0;
}
int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override {
if (!flush_gen_wav()) {
return 1;
}
*out_sample_rate = info.sample_rate;
if (out_n_samples) {
*out_n_samples = (int64_t) audio_pcm.size();
}
if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
*out_data = (const char *) audio_pcm.data();
*out_data_len = audio_pcm.size() * sizeof(float);
return 0;
}
out_buf.clear();
if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
return 1;
}
*out_data = out_buf.data();
*out_data_len = out_buf.size();
return 0;
}
private:
bool ensure_cache() {
if (specials_ok) {
return true;
}
codec_0 = find_special_token(vocab, "<|codec_0|>");
codec_bos = find_special_token(vocab, "<|codec_bos|>");
codec_eos = find_special_token(vocab, "<|codec_eos_token|>");
codec_pad = find_special_token(vocab, "<|codec_pad|>");
c_think = find_special_token(vocab, "<|codec_think|>");
c_think_b = find_special_token(vocab, "<|codec_think_bos|>");
c_think_e = find_special_token(vocab, "<|codec_think_eos|>");
tts_pad = find_special_token(vocab, "<tts_pad>");
tts_bos = find_special_token(vocab, "<tts_text_bos>");
tts_eos = find_special_token(vocab, "<tts_text_eod>");
for (llama_token t : { codec_0, codec_bos, codec_eos, codec_pad,
c_think, c_think_b, c_think_e,
tts_pad, tts_bos, tts_eos }) {
if (t == LLAMA_TOKEN_NULL) {
LOG_ERR("mtmd_helper_gen_audio: missing a required special token in vocab\n");
return false;
}
}
const uint32_t n_tok_embd = llama_model_get_tok_embd(model, nullptr);
if (n_tok_embd == 0) {
LOG_ERR("mtmd_helper_gen_audio: model has no token embeddings\n");
return false;
}
tok_embd.resize(n_tok_embd);
if (llama_model_get_tok_embd(model, tok_embd.data()) != n_tok_embd) {
LOG_ERR("mtmd_helper_gen_audio: token embedding copy failed\n");
return false;
}
specials_ok = true;
return true;
}
// runs the reference wav through the speaker encoder, returns one x-vector embedding row
bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
if (!mtmd_support_audio(mctx)) {
LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n");
return false;
}
const std::string marker = mtmd_default_marker();
mtmd_input_text text{ marker.c_str(), marker.size(), false, true };
mtmd_input_chunks * chunks = mtmd_input_chunks_init();
const mtmd_bitmap * bptr = bitmap;
bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0;
if (ok) {
ok = false;
for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {
const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) {
continue;
}
if (mtmd_encode_chunk(mctx, chunk) != 0) {
LOG_ERR("mtmd_helper_gen_audio: speaker encode failed\n");
break;
}
const float * embd = mtmd_get_output_embd(mctx);
const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk);
out.assign(embd, embd + n);
ok = true;
break;
}
}
mtmd_input_chunks_free(chunks);
return ok;
}
// one GEN_WAV process() call over the buffered codes, state is carried across batches
bool flush_gen_wav() {
if (codes_buf.empty()) {
return true;
}
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
inp.codes = codes_buf.data();
inp.n_codes = codes_buf.size();
inp.seed = seed; // same seed as gen_code, else clip reseeds mid-generation
inp.state_data = c2w_state.empty() ? nullptr : (const char *) c2w_state.data();
inp.state_size = c2w_state.size();
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: gen_wav process failed\n");
return false;
}
audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples);
c2w_state.assign(out.state_data, out.state_data + out.state_size);
codes_buf.clear();
return true;
}
// vocab specials fixed across the whole session, looked up once
bool specials_ok = false;
llama_token codec_0 = LLAMA_TOKEN_NULL;
llama_token codec_bos = LLAMA_TOKEN_NULL;
llama_token codec_eos = LLAMA_TOKEN_NULL;
llama_token codec_pad = LLAMA_TOKEN_NULL;
llama_token c_think = LLAMA_TOKEN_NULL;
llama_token c_think_b = LLAMA_TOKEN_NULL;
llama_token c_think_e = LLAMA_TOKEN_NULL;
llama_token tts_pad = LLAMA_TOKEN_NULL;
llama_token tts_bos = LLAMA_TOKEN_NULL;
llama_token tts_eos = LLAMA_TOKEN_NULL;
std::vector<float> tok_embd; // whole token embedding matrix, n_vocab * n_embd
// must match hparams.wav_tfm_swa hardcoded in clip.cpp
size_t window_frames = 72;
// per-generation state, cleared by reset()
llama_seq_id seq_id = 0;
bool mrope = false;
int pos = 0;
// prompt decode state, consumed batch-by-batch by step_prompt()
std::vector<float> prompt_embd_buf;
std::unique_ptr<decode_embd_batch> prompt_batch;
int n_prompt = 0;
int prompt_pos = 0;
int32_t top_k = 50;
float top_p = 1.0f;
uint32_t seed = UINT32_MAX;
std::vector<int32_t> codes_buf;
std::vector<uint8_t> c2w_state;
std::vector<float> audio_pcm;
std::vector<float> overlay;
std::vector<float> h_state_buf;
mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
std::vector<char> out_buf;
};
// settings that only live in the reference's per-pack yaml, not in the checkpoint
// the english packs share the same shapes and tokenizer, but disagree on these
// all three are 0 / false when the pack does not tune them, the model default is then used
struct pockettts_pack_settings {
float temp = 0.0f;
int frames_after_eos = 0;
bool pad_short_text = false;
};
static pockettts_pack_settings pockettts_pack(const char * variant) {
static const std::unordered_map<std::string, pockettts_pack_settings> packs = {
{ "english", { 0.3f, 0, false } },
{ "english_2026-01", { 0.7f, 0, true } },
{ "english_2026-04", { 0.3f, 0, false } },
{ "french_24l", { 0.7f, 8, false } },
};
auto it = packs.find(variant ? variant : "");
if (it == packs.end()) {
LOG_WRN("mtmd_helper_gen_audio: no tuned settings for pocket-tts variant \"%s\"\n",
variant ? variant : "");
return {};
}
return it->second;
}
// pocket-tts: the backbone emits no token, the flow net turns each hidden state into a latent
// the end-of-speech head also lives in the mmproj
class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
public:
using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline;
void reset() override {
seq_id = 0;
pos = 0;
feats_buf.clear();
dec_state.clear();
audio_pcm.clear();
h_state_buf.clear();
out_buf.clear();
prompt_embd_buf.clear();
prompt_batch.reset();
n_prompt = 0;
prompt_pos = 0;
step_idx = 0;
eos_step = -1;
chunks.clear();
chunk_idx = 0;
n_voice_pos = 0;
chunk_budget = 0;
}
int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
reset();
seq_id = inp->seq_id;
if (!ensure_cache()) {
return 1;
}
std::vector<float> voice;
if (inp->speaker_ref) {
if (!encode_speaker(inp->speaker_ref, voice)) {
return 1;
}
}
pack = pockettts_pack(info.model_variant);
const std::string text = prepare_text(std::string(inp->prompt, inp->prompt_len),
pack.pad_short_text);
if (text.empty()) {
LOG_ERR("mtmd_helper_gen_audio: empty prompt\n");
return 1;
}
std::vector<llama_token> ids(text.size() + 16);
int n_ids = llama_tokenize(vocab, text.c_str(), (int32_t) text.size(), ids.data(),
(int32_t) ids.size(), false, false);
if (n_ids <= 0) {
LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n");
return 1;
}
ids.resize((size_t) n_ids);
// long inputs degrade badly, so each chunk restarts from the voice conditioning
// see split_into_best_sentences() in the reference
chunks = split_chunks(ids);
chunk_idx = 0;
if (chunks.size() > 1) {
LOG_INF("mtmd_helper_gen_audio: %d tokens split into %zu chunks\n", n_ids, chunks.size());
}
const int n_e = n_embd;
// 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_embd_row(prompt_embd_buf, bos_before_voice);
}
prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end());
}
// every later chunk rewinds to here and re-prompts, so the voice stays primed
n_voice_pos = (int) (prompt_embd_buf.size() / (size_t) n_e);
for (llama_token t : chunks[0]) {
push_embd_row(prompt_embd_buf, t);
}
push_embd_row(prompt_embd_buf, audio_bos);
arm_chunk_budget(0);
n_prompt = (int) (prompt_embd_buf.size() / (size_t) n_e);
prompt_batch.reset(new decode_embd_batch(prompt_embd_buf.data(), n_prompt, 1, n_e));
prompt_batch->set_position_normal(0, seq_id);
prompt_pos = 0;
seed = inp->seed;
out_type = inp->out_type;
return 0;
}
int32_t step_prompt(int32_t n_batch) override {
GGML_ASSERT(n_batch > 0);
if (prompt_pos >= n_prompt) {
return 0;
}
const int32_t n_tokens_batch = std::min(n_batch, n_prompt - prompt_pos);
llama_batch batch_view = prompt_batch->get_view(prompt_pos, n_tokens_batch);
if ((prompt_pos + n_tokens_batch) == n_prompt) {
batch_view.logits[n_tokens_batch - 1] = 1;
}
if (llama_decode(lctx, batch_view) != 0) {
LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n");
return -1;
}
pos += n_tokens_batch;
prompt_pos += n_tokens_batch;
if (prompt_pos >= n_prompt) {
prompt_batch.reset();
prompt_embd_buf.clear();
return 0;
}
return n_prompt - prompt_pos;
}
int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override {
(void) sampled; // the backbone output is continuous, there is no token to consume
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
inp.embd = const_cast<float *>(h_state_in);
// clip only reseeds when the seed changes, so pass the same one on every step
inp.seed = seed;
if (pack.temp > 0.0f) {
inp.temp = pack.temp;
}
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: flow decode failed\n");
return 1;
}
if (out.is_eos && eos_step < 0) {
eos_step = step_idx;
}
// the frame of the stopping step is discarded, matching _autoregressive_generation().
// the budget is the reference's fallback for a chunk whose eos head never fires
const bool chunk_done = (eos_step >= 0 && step_idx >= eos_step + frames_after_eos) ||
step_idx >= chunk_budget;
if (chunk_done) {
if (eos_step < 0) {
LOG_WRN("mtmd_helper_gen_audio: chunk %zu hit its budget without end-of-speech\n", chunk_idx);
}
return finish_chunk(h_state_out, out_stop);
}
feats_buf.insert(feats_buf.end(), out.feats, out.feats + out.n_feats);
step_idx++;
if (out.n_feats > 0 && feats_buf.size() / out.n_feats >= window_frames) {
if (!flush_gen_wav()) {
return 1;
}
}
decode_embd_batch batch_embd(const_cast<float *>(out.embd), 1, 1, n_embd);
batch_embd.set_position_normal(pos, seq_id);
batch_embd.batch.logits[0] = 1;
pos++;
if (llama_decode(lctx, batch_embd.batch) != 0) {
LOG_ERR("mtmd_helper_gen_audio: decode failed\n");
return 1;
}
const float * he = llama_get_embeddings_ith(lctx, -1);
h_state_buf.assign(he, he + n_embd);
*h_state_out = h_state_buf.data();
return 0;
}
int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override {
if (!flush_gen_wav()) {
return 1;
}
*out_sample_rate = info.sample_rate;
if (out_n_samples) {
*out_n_samples = (int64_t) audio_pcm.size();
}
if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
*out_data = (const char *) audio_pcm.data();
*out_data_len = audio_pcm.size() * sizeof(float);
return 0;
}
out_buf.clear();
if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
return 1;
}
*out_data = out_buf.data();
*out_data_len = out_buf.size();
return 0;
}
private:
bool ensure_cache() {
if (specials_ok) {
return true;
}
// bos_before_voice is optional, some packs do not insert it
bos_before_voice = find_special_token(vocab, "<|bos_before_voice|>");
audio_bos = find_special_token(vocab, "<|audio_bos|>");
if (audio_bos == LLAMA_TOKEN_NULL) {
LOG_ERR("mtmd_helper_gen_audio: missing <|audio_bos|> in vocab\n");
return false;
}
const uint32_t n_tok_embd = llama_model_get_tok_embd(model, nullptr);
if (n_tok_embd == 0) {
LOG_ERR("mtmd_helper_gen_audio: model has no token embeddings\n");
return false;
}
tok_embd.resize(n_tok_embd);
if (llama_model_get_tok_embd(model, tok_embd.data()) != n_tok_embd) {
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<float> & 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<llama_token> punct_ids(const char * s) const {
std::vector<llama_token> ids(16);
const int n = llama_tokenize(vocab, s, (int32_t) strlen(s), ids.data(), (int32_t) ids.size(), false, false);
if (n <= 1) {
return {};
}
return std::vector<llama_token>(ids.begin() + 1, ids.begin() + n);
}
// cut after runs of boundary tokens, so punctuation stays with the sentence it ends
static std::vector<std::vector<llama_token>> split_on(const std::vector<llama_token> & ids,
const std::vector<llama_token> & boundary) {
std::vector<std::vector<llama_token>> out;
size_t start = 0;
bool prev_was_boundary = false;
for (size_t i = 0; i < ids.size(); i++) {
const bool is_boundary = std::find(boundary.begin(), boundary.end(), ids[i]) != boundary.end();
if (!is_boundary && prev_was_boundary) {
out.emplace_back(ids.begin() + start, ids.begin() + i);
start = i;
}
prev_was_boundary = is_boundary;
}
out.emplace_back(ids.begin() + start, ids.end());
return out;
}
std::vector<std::vector<llama_token>> split_chunks(const std::vector<llama_token> & ids) const {
if ((int) ids.size() <= max_chunk_tokens) {
return { ids };
}
const std::vector<llama_token> eos_punct = punct_ids(".!...?");
const std::vector<llama_token> mid_punct = punct_ids(",;:");
// oversized sentences are split again on weaker punctuation, else words get skipped
std::vector<std::vector<llama_token>> segments;
for (auto & seg : split_on(ids, eos_punct)) {
if ((int) seg.size() <= max_chunk_tokens) {
segments.push_back(std::move(seg));
continue;
}
auto sub = split_on(seg, mid_punct);
if (sub.size() > 1) {
for (auto & s : sub) {
segments.push_back(std::move(s));
}
} else {
segments.push_back(std::move(seg));
}
}
std::vector<std::vector<llama_token>> out;
for (auto & seg : segments) {
if (seg.empty()) {
continue;
}
if (!out.empty() && (int) (out.back().size() + seg.size()) <= max_chunk_tokens) {
out.back().insert(out.back().end(), seg.begin(), seg.end());
} else {
out.push_back(std::move(seg));
}
}
if (out.empty()) {
out.push_back(ids);
}
for (const auto & c : out) {
if ((int) c.size() > max_chunk_tokens) {
LOG_WRN("mtmd_helper_gen_audio: chunk of %zu tokens exceeds the %d token budget, "
"generation may skip words\n", c.size(), max_chunk_tokens);
}
}
return out;
}
// _estimate_max_gen_len() plus the per-chunk tail guess, both in frames
void arm_chunk_budget(size_t idx) {
const int n_tok = (int) chunks[idx].size();
chunk_budget = (int) std::ceil((n_tok / 3.0 + 2.0) * frame_rate);
// the pack may pin the tail, else the reference guesses it from the word count
frames_after_eos = pack.frames_after_eos > 0 ? pack.frames_after_eos : (n_tok <= 6 ? 5 : 3);
step_idx = 0;
eos_step = -1;
}
// ends the current chunk and, if there is another, re-prompts it on top of the voice
int32_t finish_chunk(const float ** h_state_out, bool * out_stop) {
if (!flush_gen_wav()) {
return 1;
}
// the decoder restarts too, the next chunk's audio is not continuous with this one
dec_state.clear();
if (chunk_idx + 1 >= chunks.size()) {
*out_stop = true;
*h_state_out = nullptr;
return 0;
}
chunk_idx++;
// drop this chunk's text and audio, keep the voice conditioning
llama_memory_seq_rm(llama_get_memory(lctx), seq_id, n_voice_pos, -1);
pos = n_voice_pos;
const int n_e = n_embd;
prompt_embd_buf.clear();
for (llama_token t : chunks[chunk_idx]) {
push_embd_row(prompt_embd_buf, t);
}
push_embd_row(prompt_embd_buf, audio_bos);
arm_chunk_budget(chunk_idx);
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;
if (llama_decode(lctx, batch.batch) != 0) {
LOG_ERR("mtmd_helper_gen_audio: chunk prompt decode failed\n");
return 1;
}
pos += n_rows;
prompt_embd_buf.clear();
const float * he = llama_get_embeddings_ith(lctx, -1);
h_state_buf.assign(he, he + n_embd);
*h_state_out = h_state_buf.data();
*out_stop = false;
return 0;
}
// same normalization as prepare_text_prompt() in the reference, it affects quality
static std::string prepare_text(const std::string & in, bool pad_short) {
std::string s;
s.reserve(in.size() + 1);
for (char c : in) {
if (c == '\n' || c == '\r') {
s += ' ';
} else if (c == ';') {
s += ',';
} else {
s += c;
}
}
const size_t b = s.find_first_not_of(' ');
const size_t e = s.find_last_not_of(' ');
if (b == std::string::npos) {
return "";
}
s = s.substr(b, e - b + 1);
if (s[0] >= 'a' && s[0] <= 'z') {
s[0] = (char) (s[0] - 'a' + 'A');
}
const unsigned char last = (unsigned char) s.back();
if (std::isalnum(last)) {
s += '.';
}
if (pad_short && count_words(s) < 5) {
s = std::string(8, ' ') + s;
}
return s;
}
static int count_words(const std::string & s) {
int n = 0;
bool in_word = false;
for (char c : s) {
if (c == ' ') {
in_word = false;
} else if (!in_word) {
in_word = true;
n++;
}
}
return n;
}
// runs the reference wav through the mimi encoder, returns one row per 12.5Hz frame
bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
if (!mtmd_support_audio(mctx)) {
LOG_ERR("mtmd_helper_gen_audio: mmproj has no voice encoder\n");
return false;
}
const std::string marker = mtmd_default_marker();
mtmd_input_text text{ marker.c_str(), marker.size(), false, true };
mtmd_input_chunks * chunks = mtmd_input_chunks_init();
const mtmd_bitmap * bptr = bitmap;
bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0;
if (ok) {
ok = false;
for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {
const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) {
continue;
}
if (mtmd_encode_chunk(mctx, chunk) != 0) {
LOG_ERR("mtmd_helper_gen_audio: voice encode failed\n");
break;
}
const float * embd = mtmd_get_output_embd(mctx);
const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk);
out.assign(embd, embd + n);
ok = true;
break;
}
}
mtmd_input_chunks_free(chunks);
return ok;
}
// decodes the buffered latents, the mimi decoder state carries over between calls
bool flush_gen_wav() {
if (feats_buf.empty()) {
return true;
}
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
inp.feats = feats_buf.data();
inp.n_feats = feats_buf.size();
inp.seed = seed;
inp.state_data = dec_state.empty() ? nullptr : (const char *) dec_state.data();
inp.state_size = dec_state.size();
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: mimi decode failed\n");
return false;
}
audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples);
dec_state.assign(out.state_data, out.state_data + out.state_size);
feats_buf.clear();
return true;
}
pockettts_pack_settings pack;
bool specials_ok = false;
llama_token bos_before_voice = LLAMA_TOKEN_NULL;
llama_token audio_bos = LLAMA_TOKEN_NULL;
std::vector<float> tok_embd;
llama_seq_id seq_id = 0;
int pos = 0;
std::vector<float> prompt_embd_buf;
std::unique_ptr<decode_embd_batch> prompt_batch;
int n_prompt = 0;
int prompt_pos = 0;
uint32_t seed = UINT32_MAX;
// end-of-speech is latched, then a few more frames are generated as tail padding
int step_idx = 0;
int eos_step = -1;
int frames_after_eos = 3;
static constexpr int max_chunk_tokens = 50; // MAX_TOKEN_PER_CHUNK in the reference
static constexpr double frame_rate = 12.5;
std::vector<std::vector<llama_token>> chunks;
size_t chunk_idx = 0;
int n_voice_pos = 0; // KV positions held by the voice conditioning
int chunk_budget = 0;
// latents are decoded a window at a time, the decoder state bridges the windows
size_t window_frames = 8;
std::vector<float> feats_buf;
std::vector<uint8_t> dec_state;
std::vector<float> audio_pcm;
std::vector<float> h_state_buf;
mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
std::vector<char> out_buf;
};
static std::unique_ptr<mtmd_gen_audio_pipeline> make_pipeline(llama_context * lctx, mtmd_context * mctx) {
switch (mtmd_gen_audio_get_info(mctx).type) {
case MTMD_GEN_AUDIO_TYPE_QWEN3TTS:
return std::unique_ptr<mtmd_gen_audio_pipeline>(new qwen3tts_gen_audio_pipeline(lctx, mctx));
case MTMD_GEN_AUDIO_TYPE_POCKETTTS:
return std::unique_ptr<mtmd_gen_audio_pipeline>(new pockettts_gen_audio_pipeline(lctx, mctx));
default:
return nullptr;
}
}
struct mtmd_helper_gen_audio {
std::unique_ptr<mtmd_gen_audio_pipeline> pipeline;
};
mtmd_helper_gen_audio * mtmd_helper_gen_audio_init(struct llama_context * lctx, struct mtmd_context * mctx) {
auto * ctx = new mtmd_helper_gen_audio();
ctx->pipeline = make_pipeline(lctx, mctx);
return ctx;
}
void mtmd_helper_gen_audio_free(mtmd_helper_gen_audio * ctx) {
delete ctx;
}
void mtmd_helper_gen_audio_reset(mtmd_helper_gen_audio * ctx) {
if (ctx->pipeline) {
ctx->pipeline->reset();
}
}
int32_t mtmd_helper_gen_audio_set_input(mtmd_helper_gen_audio * ctx, const mtmd_helper_gen_audio_inp * inp) {
if (!ctx->pipeline) {
LOG_ERR("mtmd_helper_gen_audio: unsupported or missing gen-audio pipeline\n");
return 1;
}
return ctx->pipeline->set_input(inp);
}
int32_t mtmd_helper_gen_audio_step_prompt(mtmd_helper_gen_audio * ctx, int32_t n_batch) {
if (!ctx->pipeline) {
return -1;
}
return ctx->pipeline->step_prompt(n_batch);
}
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,
bool * out_stop) {
if (!ctx->pipeline) {
return 1;
}
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,
const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) {
if (!ctx->pipeline) {
return 1;
}
return ctx->pipeline->get_output(out_sample_rate, out_data, out_data_len, out_n_samples);
}
Reference in New Issue View Git Blame Copy Permalink