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>
This commit is contained in:
Xuan-Son Nguyen
co-authored by
Pascal
Pascal
parent
8d274dd7c6
commit
6e62ba5384
+573
-10
@@ -5,6 +5,8 @@
|
||||
#include "../src/llama-ext.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cctype>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
@@ -87,7 +89,8 @@ public:
|
||||
virtual int32_t step_prompt(int32_t n_batch) = 0;
|
||||
// sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token,
|
||||
// those read what they need from h_state_in instead
|
||||
virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) = 0;
|
||||
// set out_stop on end-of-speech, h_state_out must be null if no frame is generated
|
||||
virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) = 0;
|
||||
virtual int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) = 0;
|
||||
|
||||
protected:
|
||||
@@ -200,8 +203,10 @@ public:
|
||||
prompt_pos = 0;
|
||||
|
||||
pos = 0;
|
||||
top_k = inp->top_k > 0 ? inp->top_k : 50;
|
||||
top_p = inp->top_p > 0 ? inp->top_p : 1.0f;
|
||||
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
|
||||
@@ -241,13 +246,26 @@ public:
|
||||
return n_prompt - prompt_pos;
|
||||
}
|
||||
|
||||
int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) override {
|
||||
mtmd_gen_inp inp{};
|
||||
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");
|
||||
@@ -384,10 +402,11 @@ private:
|
||||
if (codes_buf.empty()) {
|
||||
return true;
|
||||
}
|
||||
mtmd_gen_inp inp{};
|
||||
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{};
|
||||
@@ -427,8 +446,9 @@ private:
|
||||
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;
|
||||
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;
|
||||
@@ -438,10 +458,547 @@ private:
|
||||
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;
|
||||
}
|
||||
@@ -483,11 +1040,17 @@ int32_t mtmd_helper_gen_audio_step_prompt(mtmd_helper_gen_audio * ctx, int32_t n
|
||||
}
|
||||
|
||||
int32_t mtmd_helper_gen_audio_step_gen(mtmd_helper_gen_audio * ctx, llama_token sampled,
|
||||
const float * h_state_in, const float ** h_state_out) {
|
||||
const float * h_state_in, const float ** h_state_out,
|
||||
bool * out_stop) {
|
||||
if (!ctx->pipeline) {
|
||||
return 1;
|
||||
}
|
||||
return ctx->pipeline->step_gen(sampled, h_state_in, h_state_out);
|
||||
bool stop = false;
|
||||
const int32_t ret = ctx->pipeline->step_gen(sampled, h_state_in, h_state_out, &stop);
|
||||
if (out_stop) {
|
||||
*out_stop = stop;
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
int32_t mtmd_helper_gen_audio_get_output(mtmd_helper_gen_audio * ctx, int32_t * out_sample_rate,
|
||||
|
||||
Reference in New Issue
Block a user