diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 6946b607f5..4675fb9a97 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -18,6 +18,8 @@ add_library(mtmd mtmd-image.cpp mtmd.h mtmd-helper.cpp + mtmd-helper-gen.cpp + mtmd-helper-common.h mtmd-helper.h clip.cpp clip.h diff --git a/tools/mtmd/README-dev.md b/tools/mtmd/README-dev.md index 3a08915876..83256a4565 100644 --- a/tools/mtmd/README-dev.md +++ b/tools/mtmd/README-dev.md @@ -28,6 +28,40 @@ A typical pipeline of the core libmtmd is as follows: - Single image or batch is encoded, via `mtmd_encode()` or `mtmd_batch_encode()` - Get the output embeddings +## Audio generation support + +Audio generation is added to mtmd in PR [#26254](https://github.com/ggml-org/llama.cpp/pull/26254) + +Currently, we support the 3-stage pipeline below which should cover most TTS models: +1. (Optional) an audio encoder model that converts reference voice into codes or features +2. A backbone model that accepts text prompt and reference voice as input +3. A feature generator model that takes the hidden state from backbone and generate audio features (usually as audio codes or mel-spectrogram) +4. A model that converts audio features to the final PCM waveform + +For example, Qwen3-TTS: +1. Reference voice is encoded using ECAPA-TDNN speaker encoder (`speaker_encoder`) +2. Text prompt and reference voice are processed via a backbone (`talker.model`) +3. A model converts sampled semantic token and hidden state from stage 2 into a list of 15 acoustic codes (`talker.code_predictor`) +4. 16 generated codes are converted into waveform (`code2wav`) + +### API design constraint + +Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system is designed to be flexible and reusable by new models. + +`mtmd_gen_audio` is split into 2 main API: +- Core API `mtmd.h`: handles main inference. Important: the API surface must be stateless; caller must handle state management and audio frame accumulation. +- Helper API `mtmd-helper.h`: provides a model-agnostic stateful API. Usage example can be found in the `tools/tts` directory. + +### Checklist for porting new audio generation models to mtmd + +1. Establish a list of reusable and missing components from the current mtmd implementation. +2. Sidecar models (code2wav, bigvgan, etc) must live inside the same GGUF file (but can be in different `clip_context` if necessary) +3. Make sure most of the changes happen inside `mtmd-helper-gen.cpp`. A good PR looks like this: + - 10-20% changes is to add new backbone (text) model and conversion + - 60% changes inside `mtmd-helper-gen.cpp` + - 10% changes inside `libmtmd` and `clip.cpp` systems + - The rest downstream code (CLI, server) should have no changes at all + ## Helper We provide a set of helper functions via `mtmd_helper` to make using libmtmd easier. The helper provides: diff --git a/tools/mtmd/mtmd-helper-common.h b/tools/mtmd/mtmd-helper-common.h new file mode 100644 index 0000000000..3083e21eea --- /dev/null +++ b/tools/mtmd/mtmd-helper-common.h @@ -0,0 +1,180 @@ +#pragma once + +// shared internal utilities for the mtmd-helper-*.cpp translation units +// (mtmd-helper.cpp, mtmd-helper-gen.cpp) +// NOT part of the public mtmd-helper.h API + +#include "ggml.h" +#include "llama.h" +#include "mtmd.h" + +#include +#include +#include +#include + +// +// logging +// + +struct mtmd_helper_logger { + ggml_log_callback default_callback = [](ggml_log_level level, const char * text, void * user_data) { + (void) level; + (void) user_data; + fputs(text, stderr); + fflush(stderr); + }; + + ggml_log_callback log_callback = default_callback; + void * log_callback_user_data; + + void log_v(enum ggml_log_level level, const char * format, va_list args) { + if (format == NULL) { + return; + } + va_list args_copy; + va_copy(args_copy, args); + char buffer[128]; + int len = vsnprintf(buffer, 128, format, args); + if (len < 128) { + log_callback(level, buffer, log_callback_user_data); + } else { + char * buffer2 = (char *) calloc(len + 1, sizeof(char)); + vsnprintf(buffer2, len + 1, format, args_copy); + buffer2[len] = 0; + log_callback(level, buffer2, log_callback_user_data); + free(buffer2); + } + va_end(args_copy); + } + + void log(enum ggml_log_level level, const char * format, ...) { + va_list args; + va_start(args, format); + log_v(level, format, args); + va_end(args); + } +}; + +// inline (C++17): one shared instance across every TU that includes this header +inline mtmd_helper_logger g_logger; + +#define LOG_DBG(...) g_logger.log(GGML_LOG_LEVEL_DEBUG, __VA_ARGS__) +#define LOG_INF(...) g_logger.log(GGML_LOG_LEVEL_INFO, __VA_ARGS__) +#define LOG_WRN(...) g_logger.log(GGML_LOG_LEVEL_WARN, __VA_ARGS__) +#define LOG_ERR(...) g_logger.log(GGML_LOG_LEVEL_ERROR, __VA_ARGS__) + +// +// embd batch +// + +// helper struct to make working with embd batch easier +// note: this will be removed after llama_batch_ext refactoring +struct decode_embd_batch { + int n_pos_per_embd; + int n_mmproj_embd; + std::vector pos; + std::vector pos_view; // used by mrope + std::vector n_seq_id; + std::vector seq_id_0; + std::vector seq_ids; + std::vector logits; + llama_batch batch; + decode_embd_batch(float * embd, int32_t n_tokens, int n_pos_per_embd, int n_mmproj_embd) : n_pos_per_embd(n_pos_per_embd), n_mmproj_embd(n_mmproj_embd) { + GGML_ASSERT(n_tokens > 0 && n_pos_per_embd > 0 && n_mmproj_embd > 0); + pos .resize(n_tokens * n_pos_per_embd); + n_seq_id.resize(n_tokens); + seq_ids .resize(n_tokens + 1); + logits .resize(n_tokens); + seq_id_0.resize(1); + seq_ids [n_tokens] = nullptr; + batch = { + /*n_tokens =*/ n_tokens, + /*tokens =*/ nullptr, + /*embd =*/ embd, + /*pos =*/ pos.data(), + /*n_seq_id =*/ n_seq_id.data(), + /*seq_id =*/ seq_ids.data(), + /*logits =*/ logits.data(), + }; + } + + void set_position_normal(llama_pos pos_0, llama_seq_id seq_id) { + seq_id_0[0] = seq_id; + for (int i = 0; i < batch.n_tokens; i++) { + batch.pos [i] = pos_0 + i; + batch.n_seq_id[i] = 1; + batch.seq_id [i] = seq_id_0.data(); + batch.logits [i] = false; + } + } + + // M-RoPE for image + void set_position_mrope_2d(const std::vector & rel_pos, llama_seq_id seq_id) { + GGML_ASSERT(n_pos_per_embd == 4); + GGML_ASSERT(!rel_pos.empty() && (int32_t)rel_pos.size() == batch.n_tokens); + seq_id_0[0] = seq_id; + for (int32_t i = 0; i < batch.n_tokens; i++) { + pos[i ] = rel_pos[i].t; + pos[i + batch.n_tokens ] = rel_pos[i].y; + pos[i + batch.n_tokens * 2] = rel_pos[i].x; + pos[i + batch.n_tokens * 3] = rel_pos[i].z; + } + for (int i = 0; i < batch.n_tokens; i++) { + batch.n_seq_id[i] = 1; + batch.seq_id [i] = seq_id_0.data(); + batch.logits [i] = false; + } + } + + // M-RoPE for audio + void set_position_mrope_1d(llama_pos pos_0, llama_seq_id seq_id) { + GGML_ASSERT(n_pos_per_embd == 4); + seq_id_0[0] = seq_id; + for (int i = 0; i < batch.n_tokens; i++) { + pos[i ] = pos_0 + i; + pos[i + batch.n_tokens ] = pos_0 + i; + pos[i + batch.n_tokens * 2] = pos_0 + i; + pos[i + batch.n_tokens * 3] = pos_0 + i; + } + for (int i = 0; i < batch.n_tokens; i++) { + batch.n_seq_id[i] = 1; + batch.seq_id [i] = seq_id_0.data(); + batch.logits [i] = false; + } + } + + llama_batch get_view(int offset, int n_tokens) { + GGML_ASSERT(offset >= 0 && n_tokens > 0 && offset + n_tokens <= batch.n_tokens); + llama_pos * pos_ptr; + pos_view.clear(); + pos_view.reserve(n_tokens * n_pos_per_embd); + if (n_pos_per_embd > 1) { + // mrope + // for example, with layout of src: 1234...1234...1234...1234... + // offset 2 will give us dst: 34...34...34...34... + for (int i = 0; i < n_pos_per_embd; i++) { + // assume n_tokens is less than or equal to batch.n_tokens + // batch.n_tokens is number of **total** tokens + // n_tokens is number of viewed token + size_t src_idx = i * batch.n_tokens + offset; + pos_view.insert(pos_view.end(), + pos.data() + src_idx, + pos.data() + src_idx + n_tokens); + } + pos_ptr = pos_view.data(); + } else { + // normal + pos_ptr = pos.data() + offset; + } + return { + /*n_tokens =*/ n_tokens, + /*tokens =*/ nullptr, + /*embd =*/ batch.embd + offset * n_mmproj_embd, + /*pos =*/ pos_ptr, + /*n_seq_id =*/ batch.n_seq_id + offset, + /*seq_id =*/ batch.seq_id + offset, + /*logits =*/ batch.logits + offset, + }; + } +}; diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp new file mode 100644 index 0000000000..88ca7e8e8e --- /dev/null +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -0,0 +1,383 @@ +#include "mtmd.h" +#include "mtmd-helper.h" +#include "mtmd-helper-common.h" +#include "llama.h" +#include "../src/llama-ext.h" + +#include +#include +#include +#include + +#ifdef MTMD_INTERNAL_HEADER +#error "mtmd-helper is a public library outside of mtmd. it must not include internal headers" +#endif + +// +// Audio generation helpers +// +// model-specific pipeline logic is dispatched on mtmd_gen_audio_get_info(mctx).type; +// the public surface (init/free/reset/set_input/step/get_output) stays model-agnostic +// + +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 void write_wav16(std::vector & buf, const std::vector & pcm, int32_t rate) { + 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); + } +} + +struct mtmd_helper_gen_audio { + llama_context * lctx; + mtmd_context * mctx; + const llama_model * model; + const llama_vocab * vocab; + int n_embd = 0; + mtmd_gen_audio_info info; + + // qwen3tts: 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 tok_embd; // whole token embedding matrix, n_vocab * n_embd + + // matches hparams.wav_tfm_sliding_window hardcoded in clip.cpp; code2wav + // batches exactly this many frames per call + size_t window_frames = 72; + + // per-generation state, cleared by reset() + bool mrope = false; + int pos = 0; + int32_t top_k = 50; + float top_p = 1.0f; + std::vector codes_buf; + std::vector c2w_state; + std::vector audio_pcm; + std::vector> overlay; + size_t overlay_idx = 0; + std::vector h_state_buf; + mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; + std::vector out_buf; +}; + +static bool qwen3tts_ensure_cache(mtmd_helper_gen_audio * ctx) { + if (ctx->specials_ok) { + return true; + } + ctx->codec_0 = find_special_token(ctx->vocab, "<|codec_0|>"); + ctx->codec_bos = find_special_token(ctx->vocab, "<|codec_bos|>"); + ctx->codec_eos = find_special_token(ctx->vocab, "<|codec_eos_token|>"); + ctx->codec_pad = find_special_token(ctx->vocab, "<|codec_pad|>"); + ctx->c_think = find_special_token(ctx->vocab, "<|codec_think|>"); + ctx->c_think_b = find_special_token(ctx->vocab, "<|codec_think_bos|>"); + ctx->c_think_e = find_special_token(ctx->vocab, "<|codec_think_eos|>"); + ctx->tts_pad = find_special_token(ctx->vocab, ""); + ctx->tts_bos = find_special_token(ctx->vocab, ""); + ctx->tts_eos = find_special_token(ctx->vocab, ""); + for (llama_token t : { ctx->codec_0, ctx->codec_bos, ctx->codec_eos, ctx->codec_pad, + ctx->c_think, ctx->c_think_b, ctx->c_think_e, + ctx->tts_pad, ctx->tts_bos, ctx->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(ctx->model, nullptr); + if (n_tok_embd == 0) { + LOG_ERR("mtmd_helper_gen_audio: model has no token embeddings\n"); + return false; + } + ctx->tok_embd.resize(n_tok_embd); + llama_model_get_tok_embd(ctx->model, ctx->tok_embd.data()); + ctx->specials_ok = true; + return true; +} + +// encodes a reference wav (already loaded as a bitmap) through the mmproj's +// speaker encoder, returning the single x-vector embedding row it produces +static bool qwen3tts_encode_speaker(mtmd_helper_gen_audio * ctx, mtmd_bitmap * bitmap, std::vector & out) { + if (!mtmd_support_audio(ctx->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(ctx->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(ctx->mctx, chunk) != 0) { + LOG_ERR("mtmd_helper_gen_audio: speaker encode failed\n"); + break; + } + const float * embd = mtmd_get_output_embd(ctx->mctx); + const size_t n = (size_t) llama_model_n_embd_inp(ctx->model) * mtmd_input_chunk_get_n_tokens(chunk); + out.assign(embd, embd + n); + ok = true; + break; + } + } + mtmd_input_chunks_free(chunks); + return ok; +} + +// runs one CODE2WAV process() call on whatever is currently buffered, carrying +// the persisted state (KV cache + conv left-context) across batches +static bool qwen3tts_flush_c2w(mtmd_helper_gen_audio * ctx) { + if (ctx->codes_buf.empty()) { + return true; + } + mtmd_gen_inp inp{}; + inp.type = MTMD_GEN_PROCESS_TYPE_CODE2WAV; + inp.codes = ctx->codes_buf.data(); + inp.n_codes = ctx->codes_buf.size(); + inp.state_data = ctx->c2w_state.empty() ? nullptr : (const char *) ctx->c2w_state.data(); + inp.state_size = ctx->c2w_state.size(); + mtmd_gen_out out{}; + if (mtmd_gen_audio_process(ctx->mctx, &inp, &out) != 0) { + LOG_ERR("mtmd_helper_gen_audio: code2wav process failed\n"); + return false; + } + ctx->audio_pcm.insert(ctx->audio_pcm.end(), out.audio, out.audio + out.n_samples); + ctx->c2w_state.assign(out.state_data, out.state_data + out.state_size); + ctx->codes_buf.clear(); + return true; +} + +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->lctx = lctx; + ctx->mctx = mctx; + ctx->model = llama_get_model(lctx); + ctx->vocab = llama_model_get_vocab(ctx->model); + ctx->n_embd = llama_model_n_embd(ctx->model); + ctx->info = mtmd_gen_audio_get_info(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) { + ctx->pos = 0; + ctx->codes_buf.clear(); + ctx->c2w_state.clear(); + ctx->audio_pcm.clear(); + ctx->overlay.clear(); + ctx->overlay_idx = 0; + ctx->h_state_buf.clear(); + ctx->out_buf.clear(); +} + +int32_t mtmd_helper_gen_audio_set_input(mtmd_helper_gen_audio * ctx, const mtmd_helper_gen_audio_inp * inp) { + mtmd_helper_gen_audio_reset(ctx); + + if (ctx->info.type != MTMD_GEN_AUDIO_TYPE_QWEN3TTS) { + LOG_ERR("mtmd_helper_gen_audio: unsupported or missing gen-audio pipeline\n"); + return 1; + } + if (!qwen3tts_ensure_cache(ctx)) { + return 1; + } + + const std::string lang = inp->lang ? inp->lang : "english"; + const llama_token c_lang = find_special_token(ctx->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 speaker_embd; + if (inp->speaker_ref) { + if (!qwen3tts_encode_speaker(ctx, inp->speaker_ref, speaker_embd)) { + return 1; + } + } + + const int n_embd = ctx->n_embd; + auto row = [&](llama_token t) { + return std::vector(ctx->tok_embd.begin() + (size_t) t * n_embd, + ctx->tok_embd.begin() + (size_t) (t + 1) * n_embd); + }; + auto sum_row = [&](llama_token a, llama_token b) { + std::vector va = row(a), vb = row(b); + for (int i = 0; i < n_embd; i++) va[(size_t) i] += vb[(size_t) i]; + return va; + }; + auto sum_vec = [&](llama_token a, const std::vector & vb) { + std::vector va = row(a); + for (int i = 0; i < n_embd; 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 ids(full.size() + 16); + int n_ids = llama_tokenize(ctx->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> prompt; + for (int i = 0; i < 3; i++) prompt.push_back(row(ids[(size_t) i])); + prompt.push_back(sum_row(ctx->tts_pad, ctx->c_think)); + prompt.push_back(sum_row(ctx->tts_pad, ctx->c_think_b)); + prompt.push_back(sum_row(ctx->tts_pad, c_lang)); + prompt.push_back(sum_row(ctx->tts_pad, ctx->c_think_e)); + if (!speaker_embd.empty()) prompt.push_back(sum_vec(ctx->tts_pad, speaker_embd)); + prompt.push_back(sum_row(ctx->tts_bos, ctx->codec_pad)); + for (int i = 3; i < n_ids - 5; i++) prompt.push_back(sum_row(ids[(size_t) i], ctx->codec_pad)); + prompt.push_back(sum_row(ctx->tts_eos, ctx->codec_pad)); + prompt.push_back(sum_row(ctx->tts_pad, ctx->codec_bos)); + + const int n_prompt = (int) prompt.size(); + + // the talker rides the qwen3vl interleaved mrope: positions carry + // n_pos_per_embd sections laid out [section * n_tokens + i], all + // equal for a pure text/codec stream + ctx->mrope = llama_model_rope_type(ctx->model) == LLAMA_ROPE_TYPE_MROPE || + llama_model_rope_type(ctx->model) == LLAMA_ROPE_TYPE_IMROPE; + const int n_pos_per_embd = ctx->mrope ? 4 : 1; + + std::vector embd_buf((size_t) n_prompt * (size_t) n_embd); + for (int i = 0; i < n_prompt; i++) { + memcpy(embd_buf.data() + (size_t) i * n_embd, prompt[(size_t) i].data(), (size_t) n_embd * sizeof(float)); + } + + decode_embd_batch batch_embd(embd_buf.data(), n_prompt, n_pos_per_embd, n_embd); + if (ctx->mrope) batch_embd.set_position_mrope_1d(0, 0); + else batch_embd.set_position_normal(0, 0); + batch_embd.batch.logits[n_prompt - 1] = 1; + + if (llama_decode(ctx->lctx, batch_embd.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: prefill decode failed\n"); + return 1; + } + + ctx->pos = n_prompt; + ctx->top_k = inp->top_k > 0 ? inp->top_k : 50; + ctx->top_p = inp->top_p > 0 ? inp->top_p : 1.0f; + ctx->out_type = inp->out_type; + + // the text stream keeps flowing during generation: the input after + // frame k adds trailing text row k on top of the codes embedding, + // then tts_eos, then tts_pad once the utterance is spent + for (int i = 3; i < n_ids - 5; i++) ctx->overlay.push_back(row(ids[(size_t) i])); + ctx->overlay.push_back(row(ctx->tts_eos)); + ctx->overlay.push_back(row(ctx->tts_pad)); + + return 0; +} + +int32_t mtmd_helper_gen_audio_step(mtmd_helper_gen_audio * ctx, llama_token sampled, + const float * h_state_in, const float ** h_state_out) { + if (ctx->info.type != MTMD_GEN_AUDIO_TYPE_QWEN3TTS) { + return 1; + } + + mtmd_gen_inp inp{}; + inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; + inp.code0 = sampled - ctx->codec_0; + inp.embd = const_cast(h_state_in); + inp.top_k = ctx->top_k; + inp.top_p = ctx->top_p; + mtmd_gen_out out{}; + if (mtmd_gen_audio_process(ctx->mctx, &inp, &out) != 0) { + LOG_ERR("mtmd_helper_gen_audio: gen_code process failed\n"); + return 1; + } + + ctx->codes_buf.insert(ctx->codes_buf.end(), out.codes, out.codes + out.n_codes); + if (out.n_codes > 0 && ctx->codes_buf.size() / out.n_codes >= ctx->window_frames) { + if (!qwen3tts_flush_c2w(ctx)) { + return 1; + } + } + + std::vector fb(out.embd, out.embd + ctx->n_embd); + const auto & ov = ctx->overlay[std::min(ctx->overlay_idx, ctx->overlay.size() - 1)]; + for (int i = 0; i < ctx->n_embd; i++) fb[(size_t) i] += ov[(size_t) i]; + ctx->overlay_idx++; + + const int n_pos_per_embd = ctx->mrope ? 4 : 1; + decode_embd_batch batch_embd(fb.data(), 1, n_pos_per_embd, ctx->n_embd); + if (ctx->mrope) batch_embd.set_position_mrope_1d(ctx->pos, 0); + else batch_embd.set_position_normal(ctx->pos, 0); + batch_embd.batch.logits[0] = 1; + ctx->pos++; + + if (llama_decode(ctx->lctx, batch_embd.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: decode failed\n"); + return 1; + } + + const float * he = llama_get_embeddings_ith(ctx->lctx, -1); + ctx->h_state_buf.assign(he, he + ctx->n_embd); + *h_state_out = ctx->h_state_buf.data(); + + return 0; +} + +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) { + if (!qwen3tts_flush_c2w(ctx)) { + return 1; + } + + *out_sample_rate = ctx->info.sample_rate; + + if (ctx->out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) { + *out_data = (const char *) ctx->audio_pcm.data(); + *out_data_len = ctx->audio_pcm.size() * sizeof(float); + return 0; + } + + ctx->out_buf.clear(); + write_wav16(ctx->out_buf, ctx->audio_pcm, ctx->info.sample_rate); + *out_data = ctx->out_buf.data(); + *out_data_len = ctx->out_buf.size(); + return 0; +} diff --git a/tools/mtmd/mtmd-helper.cpp b/tools/mtmd/mtmd-helper.cpp index 90451d02eb..8d29360c86 100644 --- a/tools/mtmd/mtmd-helper.cpp +++ b/tools/mtmd/mtmd-helper.cpp @@ -9,6 +9,7 @@ #include "mtmd.h" #include "mtmd-helper.h" +#include "mtmd-helper-common.h" #include "llama.h" #include @@ -45,50 +46,6 @@ // internal logging functions // -struct mtmd_helper_logger { - ggml_log_callback default_callback = [](ggml_log_level level, const char * text, void * user_data) { - (void) level; - (void) user_data; - fputs(text, stderr); - fflush(stderr); - }; - - ggml_log_callback log_callback = default_callback; - void * log_callback_user_data; - - void log_v(enum ggml_log_level level, const char * format, va_list args) { - if (format == NULL) { - return; - } - va_list args_copy; - va_copy(args_copy, args); - char buffer[128]; - int len = vsnprintf(buffer, 128, format, args); - if (len < 128) { - log_callback(level, buffer, log_callback_user_data); - } else { - char * buffer2 = (char *) calloc(len + 1, sizeof(char)); - vsnprintf(buffer2, len + 1, format, args_copy); - buffer2[len] = 0; - log_callback(level, buffer2, log_callback_user_data); - free(buffer2); - } - va_end(args_copy); - } - - void log(enum ggml_log_level level, const char * format, ...) { - va_list args; - va_start(args, format); - log_v(level, format, args); - va_end(args); - } -} g_logger; - -#define LOG_DBG(...) g_logger.log(GGML_LOG_LEVEL_DEBUG, __VA_ARGS__) -#define LOG_INF(...) g_logger.log(GGML_LOG_LEVEL_INFO, __VA_ARGS__) -#define LOG_WRN(...) g_logger.log(GGML_LOG_LEVEL_WARN, __VA_ARGS__) -#define LOG_ERR(...) g_logger.log(GGML_LOG_LEVEL_ERROR, __VA_ARGS__) - void mtmd_helper_log_set(ggml_log_callback log_callback, void * user_data) { if (log_callback == nullptr) { log_callback = g_logger.default_callback; @@ -127,117 +84,6 @@ void mtmd_helper_image_get_decoder_pos(const mtmd_image_tokens * chunks, llama_p } } -// helper struct to make working with embd batch easier -// note: this will be removed after llama_batch_ext refactoring -struct decode_embd_batch { - int n_pos_per_embd; - int n_mmproj_embd; - std::vector pos; - std::vector pos_view; // used by mrope - std::vector n_seq_id; - std::vector seq_id_0; - std::vector seq_ids; - std::vector logits; - llama_batch batch; - decode_embd_batch(float * embd, int32_t n_tokens, int n_pos_per_embd, int n_mmproj_embd) : n_pos_per_embd(n_pos_per_embd), n_mmproj_embd(n_mmproj_embd) { - GGML_ASSERT(n_tokens > 0 && n_pos_per_embd > 0 && n_mmproj_embd > 0); - pos .resize(n_tokens * n_pos_per_embd); - n_seq_id.resize(n_tokens); - seq_ids .resize(n_tokens + 1); - logits .resize(n_tokens); - seq_id_0.resize(1); - seq_ids [n_tokens] = nullptr; - batch = { - /*n_tokens =*/ n_tokens, - /*tokens =*/ nullptr, - /*embd =*/ embd, - /*pos =*/ pos.data(), - /*n_seq_id =*/ n_seq_id.data(), - /*seq_id =*/ seq_ids.data(), - /*logits =*/ logits.data(), - }; - } - - void set_position_normal(llama_pos pos_0, llama_seq_id seq_id) { - seq_id_0[0] = seq_id; - for (int i = 0; i < batch.n_tokens; i++) { - batch.pos [i] = pos_0 + i; - batch.n_seq_id[i] = 1; - batch.seq_id [i] = seq_id_0.data(); - batch.logits [i] = false; - } - } - - // M-RoPE for image - void set_position_mrope_2d(const std::vector & rel_pos, llama_seq_id seq_id) { - GGML_ASSERT(n_pos_per_embd == 4); - GGML_ASSERT(!rel_pos.empty() && (int32_t)rel_pos.size() == batch.n_tokens); - seq_id_0[0] = seq_id; - for (int32_t i = 0; i < batch.n_tokens; i++) { - pos[i ] = rel_pos[i].t; - pos[i + batch.n_tokens ] = rel_pos[i].y; - pos[i + batch.n_tokens * 2] = rel_pos[i].x; - pos[i + batch.n_tokens * 3] = rel_pos[i].z; - } - for (int i = 0; i < batch.n_tokens; i++) { - batch.n_seq_id[i] = 1; - batch.seq_id [i] = seq_id_0.data(); - batch.logits [i] = false; - } - } - - // M-RoPE for audio - void set_position_mrope_1d(llama_pos pos_0, llama_seq_id seq_id) { - GGML_ASSERT(n_pos_per_embd == 4); - seq_id_0[0] = seq_id; - for (int i = 0; i < batch.n_tokens; i++) { - pos[i ] = pos_0 + i; - pos[i + batch.n_tokens ] = pos_0 + i; - pos[i + batch.n_tokens * 2] = pos_0 + i; - pos[i + batch.n_tokens * 3] = pos_0 + i; - } - for (int i = 0; i < batch.n_tokens; i++) { - batch.n_seq_id[i] = 1; - batch.seq_id [i] = seq_id_0.data(); - batch.logits [i] = false; - } - } - - llama_batch get_view(int offset, int n_tokens) { - GGML_ASSERT(offset >= 0 && n_tokens > 0 && offset + n_tokens <= batch.n_tokens); - llama_pos * pos_ptr; - pos_view.clear(); - pos_view.reserve(n_tokens * n_pos_per_embd); - if (n_pos_per_embd > 1) { - // mrope - // for example, with layout of src: 1234...1234...1234...1234... - // offset 2 will give us dst: 34...34...34...34... - for (int i = 0; i < n_pos_per_embd; i++) { - // assume n_tokens is less than or equal to batch.n_tokens - // batch.n_tokens is number of **total** tokens - // n_tokens is number of viewed token - size_t src_idx = i * batch.n_tokens + offset; - pos_view.insert(pos_view.end(), - pos.data() + src_idx, - pos.data() + src_idx + n_tokens); - } - pos_ptr = pos_view.data(); - } else { - // normal - pos_ptr = pos.data() + offset; - } - return { - /*n_tokens =*/ n_tokens, - /*tokens =*/ nullptr, - /*embd =*/ batch.embd + offset * n_mmproj_embd, - /*pos =*/ pos_ptr, - /*n_seq_id =*/ batch.n_seq_id + offset, - /*seq_id =*/ batch.seq_id + offset, - /*logits =*/ batch.logits + offset, - }; - } -}; - // Helper class to set non-causal attention via RAII class scope_non_causal { public: diff --git a/tools/mtmd/mtmd-helper.h b/tools/mtmd/mtmd-helper.h index 680a2317df..f7974fffba 100644 --- a/tools/mtmd/mtmd-helper.h +++ b/tools/mtmd/mtmd-helper.h @@ -157,6 +157,59 @@ MTMD_API int32_t mtmd_helper_video_read_next(mtmd_helper_video * ctx, mtmd_bitmap ** out_bitmap, char ** out_text); +// +// Audio generation helpers +// (early-stage experimental, subjected to breaking changes) +// + +// audio generation helper context +// contains accumulator for generated audio features and PCM audio +struct mtmd_helper_gen_audio; +typedef struct mtmd_helper_gen_audio mtmd_helper_gen_audio; + +enum mtmd_helper_gen_audio_outtype { + MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM, // raw PCM + MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV, // WAV PCM 16-bit LE, mono +}; +struct mtmd_helper_gen_audio_inp { + const char * prompt; + size_t prompt_len; + + mtmd_bitmap * speaker_ref; // optional, can be NULL + const char * lang; // optional, can be NULL + + int32_t top_k; + float top_p; + + mtmd_helper_gen_audio_outtype out_type; +}; + +MTMD_API mtmd_helper_gen_audio * mtmd_helper_gen_audio_init( + struct llama_context * lctx, + struct mtmd_context * mctx); + +MTMD_API void mtmd_helper_gen_audio_free(mtmd_helper_gen_audio * ctx); + +MTMD_API void mtmd_helper_gen_audio_reset(mtmd_helper_gen_audio * ctx); + +MTMD_API int32_t mtmd_helper_gen_audio_set_input( + mtmd_helper_gen_audio * ctx, + const struct mtmd_helper_gen_audio_inp * inp); + +// h_state_out is valid until next step() or reset() call +MTMD_API int32_t mtmd_helper_gen_audio_step( + mtmd_helper_gen_audio * ctx, + llama_token sampled, + const float * h_state_in, + const float ** h_state_out); + +// out_data valid until next get_output() or reset() call +MTMD_API 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); + #ifdef __cplusplus } // extern "C" #endif @@ -177,6 +230,28 @@ struct mtmd_helper_video_deleter { }; using video_ptr = std::unique_ptr; +// audio generation-related C++ wrappers +struct mtmd_helper_gen_audio_deleter { + void operator()(mtmd_helper_gen_audio * val) { mtmd_helper_gen_audio_free(val); } +}; +using gen_audio_ptr = std::unique_ptr; +struct gen_audio { + gen_audio_ptr ctx; + gen_audio(struct llama_context * lctx, struct mtmd_context * mctx) : ctx(mtmd_helper_gen_audio_init(lctx, mctx)) {} + void reset() { + mtmd_helper_gen_audio_reset(ctx.get()); + } + int32_t set_input(const struct mtmd_helper_gen_audio_inp * inp) { + return mtmd_helper_gen_audio_set_input(ctx.get(), inp); + } + int32_t step(llama_token sampled, const float * h_state, const float ** h_state_out) { + return mtmd_helper_gen_audio_step(ctx.get(), sampled, h_state, h_state_out); + } + int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len) { + return mtmd_helper_gen_audio_get_output(ctx.get(), out_sample_rate, out_data, out_data_len); + } +}; + } // namespace mtmd_helper #endif diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 2bcd58104e..508f75e1f2 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1570,16 +1570,22 @@ float * mtmd_get_output_embd(mtmd_context * ctx) { // audio generation // -mtmd_gen_audio_type mtmd_gen_audio_get_type(const mtmd_context * ctx) { +mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) { + mtmd_gen_audio_info info; if (!ctx->ctx_gen_a) { - return MTMD_GEN_AUDIO_TYPE_NONE; + info.type = MTMD_GEN_AUDIO_TYPE_NONE; + return info; } switch (clip_get_projector_type(ctx->ctx_gen_a)) { case PROJECTOR_TYPE_QWEN3TTS_GEN: - return MTMD_GEN_AUDIO_TYPE_MTP; + info.type = MTMD_GEN_AUDIO_TYPE_QWEN3TTS; + info.sample_rate = 24000; + break; default: - return MTMD_GEN_AUDIO_TYPE_NONE; + info.type = MTMD_GEN_AUDIO_TYPE_NONE; + break; } + return info; } static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) { diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 6056ac56f5..5a99b0dbec 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -330,11 +330,16 @@ MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname); ///////////////////////////////////////// // EXPERIMENTAL API for audio generation, subjected to breaking changes +// represent the pipeline type enum mtmd_gen_audio_type { MTMD_GEN_AUDIO_TYPE_NONE, // not supported - MTMD_GEN_AUDIO_TYPE_MTP, // qwen3tts style, with MTP-like generation head + MTMD_GEN_AUDIO_TYPE_QWEN3TTS, }; -MTMD_API mtmd_gen_audio_type mtmd_gen_audio_get_type(const mtmd_context * ctx); +struct mtmd_gen_audio_info { + mtmd_gen_audio_type type; + int32_t sample_rate; // in Hz, for example 24000 for qwen3tts +}; +MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx); enum mtmd_gen_process_type { MTMD_GEN_PROCESS_TYPE_GEN_CODE, // h_state to codes @@ -370,6 +375,7 @@ struct mtmd_gen_out { const char * state_data; size_t state_size; }; +// note: this API is stateless, caller must handle state management and audio frame accumulation MTMD_API int32_t mtmd_gen_audio_process(mtmd_context * ctx, const struct mtmd_gen_inp * inp, struct mtmd_gen_out * out); diff --git a/tools/tts/qwen3-tts.cpp b/tools/tts/qwen3-tts.cpp index 422b582033..ebca346756 100644 --- a/tools/tts/qwen3-tts.cpp +++ b/tools/tts/qwen3-tts.cpp @@ -1,203 +1,41 @@ -// Qwen3-TTS end to end throwaway test driver, adapted from Pascal's tts-qwen3.cpp -// reference to use the split mtmd_gen_audio_process() API (GEN_CODE / CODE2WAV) -// instead of the fused mtmd_gen_audio(). Quick manual test only, not polished, -// will be removed once the real accumulation helper lands. -// -// Codes are accumulated here (in this file) frame by frame; once 24 frames -// are buffered, a CODE2WAV process() call turns them into a batch of PCM. +// Qwen3-TTS end to end throwaway test driver, using the mtmd_helper_gen_audio_* +// helper (tools/mtmd/mtmd-helper.h) which owns all the Qwen3TTS-specific prompt +// construction, code/audio accumulation and batching. Quick manual test only, +// not polished, will be removed once a real CLI tool lands. #include "llama.h" #include "mtmd.h" #include "mtmd-helper.h" #include "common.h" #include "log.h" -#include "ggml.h" -#include "gguf.h" -#include -#include #include #include #include #include -// Read one dequantized row of a 2D tensor from a GGUF file. -struct gguf_row_reader { - struct gguf_context * gguf = nullptr; - struct ggml_context * meta = nullptr; - FILE * f = nullptr; - size_t data_off = 0; - - bool open(const char * path) { - struct ggml_init_params ip = { 0, nullptr, true }; - struct gguf_init_params gp = { true, &meta }; - gguf = gguf_init_from_file(path, gp); - if (!gguf) { - return false; - } - data_off = gguf_get_data_offset(gguf); - f = fopen(path, "rb"); - return f != nullptr; - } - - bool read_row(const char * tensor_name, int64_t row, std::vector & out) { - const int64_t idx = gguf_find_tensor(gguf, tensor_name); - if (idx < 0) { - return false; - } - struct ggml_tensor * t = ggml_get_tensor(meta, tensor_name); - if (!t || row < 0 || row >= t->ne[1]) { - return false; - } - const size_t row_bytes = ggml_row_size(t->type, t->ne[0]); - std::vector raw(row_bytes); - if (fseek(f, (long) (data_off + gguf_get_tensor_offset(gguf, idx) + (size_t) row * row_bytes), SEEK_SET) != 0) { - return false; - } - if (fread(raw.data(), 1, row_bytes, f) != row_bytes) { - return false; - } - out.resize((size_t) t->ne[0]); - if (t->type == GGML_TYPE_F32) { - memcpy(out.data(), raw.data(), row_bytes); - } else { - const auto * traits = ggml_get_type_traits(t->type); - traits->to_float(raw.data(), out.data(), t->ne[0]); - } - return true; - } - - ~gguf_row_reader() { - if (f) fclose(f); - if (gguf) gguf_free(gguf); - if (meta) ggml_free(meta); - } -}; - -static llama_token find_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 void save_wav16(const char * path, const std::vector & pcm, int rate) { - FILE * f = fopen(path, "wb"); - if (!f) { - LOG_ERR("failed to open %s\n", path); - return; - } - 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; - fwrite("RIFF", 1, 4, f); fwrite(&riff_sz, 4, 1, f); fwrite("WAVE", 1, 4, f); - fwrite("fmt ", 1, 4, f); fwrite(&fmt_sz, 4, 1, f); - fwrite(&fmt, 2, 1, f); fwrite(&ch, 2, 1, f); fwrite(&rate32, 4, 1, f); - fwrite(&byte_rate, 4, 1, f); fwrite(&align, 2, 1, f); fwrite(&bits, 2, 1, f); - fwrite("data", 1, 4, f); fwrite(&data_sz, 4, 1, f); - for (float v : pcm) { - int16_t s = (int16_t) (std::max(-1.0f, std::min(1.0f, v)) * 32767.0f); - fwrite(&s, 2, 1, f); - } - fclose(f); -} - -// loads a reference wav and runs it through the mmproj's speaker encoder (ECAPA-TDNN, -// ctx_a in mtmd terms), returning the single x-vector embedding row it produces -static bool encode_speaker_wav(mtmd_context * mctx, const llama_model * model, const char * path, std::vector & out_embd) { - if (!mtmd_support_audio(mctx)) { - LOG_ERR("mmproj has no audio encoder, can't use --speaker\n"); - return false; - } - mtmd_helper_bitmap_wrapper wrapper = mtmd_helper_bitmap_init_from_file(mctx, path, false); - if (!wrapper.bitmap) { - LOG_ERR("failed to load %s\n", path); - 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 * bitmap = wrapper.bitmap; - bool ok = mtmd_tokenize(mctx, chunks, &text, &bitmap, 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("speaker encode failed\n"); - break; - } - const float * embd = mtmd_get_output_embd(mctx); - const size_t n_embd_out = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk); - out_embd.assign(embd, embd + n_embd_out); - ok = true; - break; - } - } - mtmd_input_chunks_free(chunks); - mtmd_bitmap_free(wrapper.bitmap); - return ok; -} - -// runs one CODE2WAV process() call on a batch of frames' codes (frame-major; -// the model wants exactly one window's worth, clip.cpp front-pads a shorter -// batch), carrying the persisted state (KV cache + conv left-context) across -// batches; appends the resulting PCM to audio_out and updates state for the -// next batch -static bool code2wav_step(mtmd_context * mctx, const std::vector & codes, std::vector & state, - std::vector & audio_out) { - mtmd_gen_inp inp{}; - inp.type = MTMD_GEN_PROCESS_TYPE_CODE2WAV; - inp.codes = const_cast(codes.data()); - inp.n_codes = codes.size(); - inp.state_data = state.empty() ? nullptr : (const char *) state.data(); - inp.state_size = state.size(); - - mtmd_gen_out out{}; - const auto t0 = std::chrono::steady_clock::now(); - const int rc = mtmd_gen_audio_process(mctx, &inp, &out); - const auto t1 = std::chrono::steady_clock::now(); - const double ms = std::chrono::duration(t1 - t0).count(); - LOG_INF("code2wav: %zu codes -> %zu samples in %.1f ms\n", codes.size() / 16, out.n_samples, ms); - if (rc != 0) { - LOG_ERR("code2wav process failed\n"); - return false; - } - audio_out.insert(audio_out.end(), out.audio, out.audio + out.n_samples); - state.assign(out.state_data, out.state_data + out.state_size); - return true; -} - int main(int argc, char ** argv) { - const char * model_path = nullptr; - const char * mmproj_path = nullptr; - const char * out_path = "output.wav"; + const char * model_path = nullptr; + const char * mmproj_path = nullptr; + const char * out_path = "output.wav"; const char * speaker_path = nullptr; std::string text; - std::string lang = "english"; - int max_new = 512; - int n_gpu = 999; + std::string lang = "english"; + int max_new = 512; + int n_gpu = 999; for (int i = 1; i < argc; i++) { auto next = [&](const char * flag) -> const char * { if (i + 1 >= argc) { fprintf(stderr, "missing value for %s\n", flag); exit(1); } return argv[++i]; }; - if (!strcmp(argv[i], "-m")) model_path = next("-m"); - else if (!strcmp(argv[i], "--mmproj")) mmproj_path = next("--mmproj"); - else if (!strcmp(argv[i], "-p")) text = next("-p"); - else if (!strcmp(argv[i], "-o")) out_path = next("-o"); - else if (!strcmp(argv[i], "--lang")) lang = next("--lang"); - else if (!strcmp(argv[i], "--max-new")) max_new = atoi(next("--max-new")); - else if (!strcmp(argv[i], "-ngl")) n_gpu = atoi(next("-ngl")); + if (!strcmp(argv[i], "-m")) model_path = next("-m"); + else if (!strcmp(argv[i], "--mmproj")) mmproj_path = next("--mmproj"); + else if (!strcmp(argv[i], "-p")) text = next("-p"); + else if (!strcmp(argv[i], "-o")) out_path = next("-o"); + else if (!strcmp(argv[i], "--lang")) lang = next("--lang"); + else if (!strcmp(argv[i], "--max-new")) max_new = atoi(next("--max-new")); + else if (!strcmp(argv[i], "-ngl")) n_gpu = atoi(next("-ngl")); else if (!strcmp(argv[i], "--speaker")) speaker_path = next("--speaker"); else { fprintf(stderr, @@ -217,8 +55,7 @@ int main(int argc, char ** argv) { mparams.n_gpu_layers = n_gpu; llama_model * model = llama_model_load_from_file(model_path, mparams); if (!model) { LOG_ERR("failed to load %s\n", model_path); return 1; } - const llama_vocab * vocab = llama_model_get_vocab(model); - const int n_embd = llama_model_n_embd(model); + const llama_vocab * vocab = llama_model_get_vocab(model); llama_context_params cparams = llama_context_default_params(); cparams.n_ctx = 4096; @@ -230,190 +67,76 @@ int main(int argc, char ** argv) { mtmd_context_params mtmd_params = mtmd_context_params_default(); mtmd_context * mctx = mtmd_init_from_file(mmproj_path, model, mtmd_params); if (!mctx) { LOG_ERR("failed to load %s\n", mmproj_path); return 1; } - if (mtmd_gen_audio_get_type(mctx) == MTMD_GEN_AUDIO_TYPE_NONE) { + if (mtmd_gen_audio_get_info(mctx).type == MTMD_GEN_AUDIO_TYPE_NONE) { LOG_ERR("mmproj does not support audio generation\n"); return 1; } - std::vector speaker_embd; + mtmd_helper_bitmap_wrapper speaker_wrapper{ nullptr, nullptr }; if (speaker_path) { - if (!encode_speaker_wav(mctx, model, speaker_path, speaker_embd)) return 1; - LOG_INF("speaker: encoded %s into a %zu-dim x-vector\n", speaker_path, speaker_embd.size()); + speaker_wrapper = mtmd_helper_bitmap_init_from_file(mctx, speaker_path, false); + if (!speaker_wrapper.bitmap) { LOG_ERR("failed to load %s\n", speaker_path); return 1; } } - // vocab landmarks: the codec rows sit after the text vocab - const llama_token codec_0 = find_token(vocab, "<|codec_0|>"); - const llama_token codec_bos = find_token(vocab, "<|codec_bos|>"); - const llama_token codec_eos = find_token(vocab, "<|codec_eos_token|>"); - const llama_token codec_pad = find_token(vocab, "<|codec_pad|>"); - const llama_token c_think = find_token(vocab, "<|codec_think|>"); - const llama_token c_think_b = find_token(vocab, "<|codec_think_bos|>"); - const llama_token c_think_e = find_token(vocab, "<|codec_think_eos|>"); - const llama_token c_lang = find_token(vocab, ("<|codec_language_" + lang + "|>").c_str()); - const llama_token tts_pad = find_token(vocab, ""); - const llama_token tts_bos = find_token(vocab, ""); - const llama_token tts_eos = find_token(vocab, ""); - for (llama_token t : { codec_0, codec_bos, codec_eos, codec_pad, c_think, c_think_b, c_think_e, c_lang, - tts_pad, tts_bos, tts_eos }) { - if (t == LLAMA_TOKEN_NULL) { - LOG_ERR("missing special token in vocab (lang '%s'?)\n", lang.c_str()); - return 1; - } + mtmd_helper::gen_audio gen(lctx, mctx); + mtmd_helper_gen_audio_inp inp{}; + inp.prompt = text.c_str(); + inp.prompt_len = text.size(); + inp.speaker_ref = speaker_wrapper.bitmap; + inp.lang = lang.c_str(); + inp.top_k = 40; + inp.top_p = 0.95f; + inp.out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; + if (gen.set_input(&inp) != 0) { LOG_ERR("set_input failed\n"); return 1; } + mtmd_bitmap_free(speaker_wrapper.bitmap); + + // vocab landmarks: only the backbone-level sampling policy (which tokens + // are valid codec_0 candidates, and which one means EOS) stays here, since + // that's ordinary LLM sampling, not part of the audio-generation pipeline + llama_token codec_0_tok = LLAMA_TOKEN_NULL; + llama_token codec_eos_tok = LLAMA_TOKEN_NULL; + for (llama_token t = 0; t < llama_vocab_n_tokens(vocab); t++) { + const char * piece = llama_vocab_get_text(vocab, t); + if (!strcmp(piece, "<|codec_0|>")) codec_0_tok = t; + else if (!strcmp(piece, "<|codec_eos_token|>")) codec_eos_tok = t; + } + if (codec_0_tok == LLAMA_TOKEN_NULL || codec_eos_tok == LLAMA_TOKEN_NULL) { + LOG_ERR("missing codec special tokens in vocab\n"); + return 1; } - // embedding rows straight from the gguf: the prompt sums two rows - // per position, which tokens cannot express - gguf_row_reader rows; - if (!rows.open(model_path)) { LOG_ERR("failed to open %s for row reads\n", model_path); return 1; } - const char * EMBD = "token_embd.weight"; - auto row = [&](llama_token t) { - std::vector v; - if (!rows.read_row(EMBD, t, v)) { LOG_ERR("row read failed for token %d\n", t); exit(1); } - return v; - }; - auto sum_row = [&](llama_token a, llama_token b) { - std::vector va = row(a), vb = row(b); - for (size_t i = 0; i < va.size(); i++) va[i] += vb[i]; - return va; - }; - auto sum_vec = [&](llama_token a, const std::vector & vb) { - std::vector va = row(a); - for (size_t i = 0; i < va.size(); i++) va[i] += vb[i]; - return va; - }; - - // upstream wrap, then slices: [0:3] role, [3:-5] utterance body - const std::string full = "<|im_start|>assistant\n" + text + "<|im_end|>\n<|im_start|>assistant\n"; - std::vector 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("tokenization failed\n"); return 1; } - ids.resize((size_t) n_ids); - - std::vector> 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)); - - const int n_prompt = (int) prompt.size(); - LOG_INF("prompt: %d positions (%d text tokens)\n", n_prompt, n_ids); - - // the talker rides the qwen3vl interleaved mrope: positions carry - // n_pos_per_embd sections laid out [section * n_tokens + i], all - // equal for a pure text/codec stream - const bool mrope = llama_model_rope_type(model) == LLAMA_ROPE_TYPE_MROPE || - llama_model_rope_type(model) == LLAMA_ROPE_TYPE_IMROPE; - const int n_pos_sec = mrope ? 4 : 1; - std::vector pos_buf((size_t) n_pos_sec * (size_t) n_prompt); - - // prefill as one embd batch, logits on the last position - std::vector embd_buf((size_t) n_prompt * (size_t) n_embd); - for (int i = 0; i < n_prompt; i++) { - memcpy(embd_buf.data() + (size_t) i * n_embd, prompt[(size_t) i].data(), (size_t) n_embd * sizeof(float)); - } - llama_batch batch = llama_batch_init(n_prompt, n_embd, 1); - batch.n_tokens = n_prompt; - batch.pos = pos_buf.data(); - memcpy(batch.embd, embd_buf.data(), embd_buf.size() * sizeof(float)); - for (int i = 0; i < n_prompt; i++) { - for (int sec = 0; sec < n_pos_sec; sec++) { - pos_buf[(size_t) sec * n_prompt + (size_t) i] = i; - } - batch.n_seq_id[i] = 1; - batch.seq_id[i][0] = 0; - batch.logits[i] = (int8_t) (i == n_prompt - 1); - } - if (llama_decode(lctx, batch) != 0) { LOG_ERR("prefill decode failed\n"); return 1; } - - // the text stream keeps flowing during generation: the input after - // frame k adds trailing text row k on top of the codes embedding, - // then tts_eos, then tts_pad once the utterance is spent - std::vector> overlay; - for (int i = 3; i < n_ids - 5; i++) overlay.push_back(row(ids[(size_t) i])); - overlay.push_back(row(tts_eos)); - overlay.push_back(row(tts_pad)); - - // matches hparams.wav_tfm_sliding_window hardcoded in clip.cpp; code2wav - // batches exactly this many frames per call - const size_t C2W_WINDOW_FRAMES = 72; - - // AR loop: sample c0 among the semantic codec rows plus eos, hand the - // hidden state to the code predictor (GEN_CODE), buffer the 16 codes it - // returns. Once a full window is buffered, run CODE2WAV on it, carrying - // its state (KV cache + conv left-context) across batches. - std::vector audio; - std::vector c2w_state; - std::vector codes_buf; - std::vector h((size_t) n_embd), fb((size_t) n_embd); - int n_frames = 0; - int pos = n_prompt; - - for (; n_frames < max_new; n_frames++) { + auto sample_codec0 = [&]() -> llama_token { const float * logits = llama_get_logits_ith(lctx, -1); - llama_token best = codec_eos; - float bestv = logits[codec_eos]; - for (llama_token t = codec_0; t < codec_0 + 2048; t++) { + llama_token best = codec_eos_tok; + float bestv = logits[codec_eos_tok]; + for (llama_token t = codec_0_tok; t < codec_0_tok + 2048; t++) { if (logits[t] > bestv) { bestv = logits[t]; best = t; } } - if (best == codec_eos) { - break; - } + return best; + }; - const float * he = llama_get_embeddings_ith(lctx, -1); - memcpy(h.data(), he, (size_t) n_embd * sizeof(float)); - - mtmd_gen_inp inp{}; - inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; - inp.code0 = best - codec_0; - inp.embd = h.data(); - inp.top_k = 50; - inp.top_p = 1.0f; - mtmd_gen_out out{}; - if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) { LOG_ERR("gen_code process failed\n"); return 1; } - - codes_buf.insert(codes_buf.end(), out.codes, out.codes + out.n_codes); - memcpy(fb.data(), out.embd, (size_t) n_embd * sizeof(float)); - - if (codes_buf.size() / 16 >= C2W_WINDOW_FRAMES) { - if (!code2wav_step(mctx, codes_buf, c2w_state, audio)) return 1; - codes_buf.clear(); - } - - const auto & ov = overlay[std::min((size_t) n_frames, overlay.size() - 1)]; - for (int i = 0; i < n_embd; i++) { - fb[(size_t) i] += ov[(size_t) i]; - } - - batch.n_tokens = 1; - memcpy(batch.embd, fb.data(), (size_t) n_embd * sizeof(float)); - for (int sec = 0; sec < n_pos_sec; sec++) { - pos_buf[(size_t) sec] = pos; - } - pos++; - batch.n_seq_id[0] = 1; - batch.seq_id[0][0] = 0; - batch.logits[0] = 1; - if (llama_decode(lctx, batch) != 0) { LOG_ERR("decode failed at frame %d\n", n_frames); return 1; } + int n_frames = 0; + llama_token sampled = sample_codec0(); + const float * h_state = llama_get_embeddings_ith(lctx, -1); + for (; n_frames < max_new && sampled != codec_eos_tok; n_frames++) { + const float * h_next = nullptr; + if (gen.step(sampled, h_state, &h_next) != 0) { LOG_ERR("step failed at frame %d\n", n_frames); return 1; } + h_state = h_next; + sampled = sample_codec0(); } - // flush whatever's left, less than a full window (front-padded with code 0 by clip.cpp) - if (!codes_buf.empty()) { - if (!code2wav_step(mctx, codes_buf, c2w_state, audio)) return 1; - } + int32_t sample_rate = 0; + const char * data = nullptr; + size_t data_len = 0; + if (gen.get_output(&sample_rate, &data, &data_len) != 0) { LOG_ERR("get_output failed\n"); return 1; } - LOG_INF("generated %d frames, %zu samples (%.2f s)\n", n_frames, audio.size(), (double) audio.size() / 24000.0); - save_wav16(out_path, audio, 24000); + LOG_INF("generated %d frames, %zu bytes of WAV audio (%d Hz)\n", n_frames, data_len, sample_rate); + FILE * f = fopen(out_path, "wb"); + if (!f) { LOG_ERR("failed to open %s\n", out_path); return 1; } + fwrite(data, 1, data_len, f); + fclose(f); LOG_INF("wrote %s\n", out_path); - batch.pos = nullptr; - llama_batch_free(batch); mtmd_free(mctx); llama_free(lctx); llama_model_free(model);