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:
co-authored by
Pascal
parent
8d274dd7c6
commit
6e62ba5384
+77
-16
@@ -477,6 +477,7 @@ struct mtmd_context {
|
||||
// generation context
|
||||
struct clip_ctx * ctx_gen_a; // audio
|
||||
std::vector<int32_t> gen_out_codes; // this frame's 16 sampled codes (GEN_CODE)
|
||||
std::vector<float> gen_out_feats; // this frame's continuous features, if any (GEN_CODE)
|
||||
std::vector<float> gen_out_embd; // next-step hidden state fed back to backbone (GEN_CODE)
|
||||
std::vector<float> gen_out_audio; // decoded PCM samples for the current frame (GEN_WAV)
|
||||
std::vector<uint8_t> gen_out_state; // state to feed into the next GEN_WAV call
|
||||
@@ -979,6 +980,10 @@ struct mtmd_context {
|
||||
{
|
||||
audio_preproc = std::make_unique<mtmd_audio_preprocessor_qwen3tts_spk>(ctx_a);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
|
||||
{
|
||||
audio_preproc = std::make_unique<mtmd_audio_preprocessor_pockettts>(ctx_a);
|
||||
} break;
|
||||
default:
|
||||
throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj));
|
||||
}
|
||||
@@ -1798,16 +1803,22 @@ float * mtmd_get_output_embd(mtmd_context * ctx) {
|
||||
//
|
||||
|
||||
mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
|
||||
mtmd_gen_audio_info info;
|
||||
mtmd_gen_audio_info info{};
|
||||
info.model_variant = "";
|
||||
if (!ctx->ctx_gen_a) {
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_NONE;
|
||||
return info;
|
||||
}
|
||||
info.model_variant = clip_get_hparams(ctx->ctx_gen_a)->gen_model_variant.c_str();
|
||||
switch (clip_get_projector_type(ctx->ctx_gen_a)) {
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_QWEN3TTS;
|
||||
info.sample_rate = 24000;
|
||||
break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN:
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS;
|
||||
info.sample_rate = 24000;
|
||||
break;
|
||||
default:
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_NONE;
|
||||
break;
|
||||
@@ -1815,6 +1826,33 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
|
||||
return info;
|
||||
}
|
||||
|
||||
mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) {
|
||||
mtmd_gen_inp inp{};
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
|
||||
inp.seed = UINT32_MAX;
|
||||
if (!ctx->ctx_gen_a) {
|
||||
return inp;
|
||||
}
|
||||
|
||||
switch (clip_get_projector_type(ctx->ctx_gen_a)) {
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
// https://huggingface.co/Qwen/Qwen3-TTS-12Hz-1.7B-Base/blob/main/generation_config.json
|
||||
inp.top_k = 50;
|
||||
inp.top_p = 1.0f;
|
||||
inp.temp = 0.9f; // TODO: handle this on graph
|
||||
break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN:
|
||||
// https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/default_parameters.py
|
||||
inp.top_k = 50;
|
||||
inp.top_p = 1.0f;
|
||||
inp.temp = 0.7f;
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
return inp;
|
||||
}
|
||||
|
||||
static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) {
|
||||
clip_ctx * ctx_clip = ctx->ctx_gen_a;
|
||||
if (!ctx_clip) {
|
||||
@@ -1822,6 +1860,8 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
|
||||
return 1;
|
||||
}
|
||||
|
||||
*out = {};
|
||||
|
||||
if (inp->type == MTMD_GEN_PROCESS_TYPE_GEN_CODE) {
|
||||
const size_t n_embd = (size_t) clip_n_mmproj_embd(ctx_clip);
|
||||
|
||||
@@ -1835,16 +1875,22 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
|
||||
|
||||
std::vector<float> out_embd(n_embd);
|
||||
std::vector<int32_t> out_codes;
|
||||
std::vector<float> out_feats;
|
||||
bool is_eos = false;
|
||||
|
||||
clip_encode_params params;
|
||||
params.imgs = &batch;
|
||||
params.n_threads = ctx->n_threads;
|
||||
params.gen_process = CLIP_GEN_PROCESS_GEN_CODE;
|
||||
params.out_embd = &out_embd;
|
||||
params.out_codes = &out_codes;
|
||||
params.code0 = inp->code0;
|
||||
params.top_k = inp->top_k;
|
||||
params.top_p = inp->top_p;
|
||||
params.imgs = &batch;
|
||||
params.n_threads = ctx->n_threads;
|
||||
params.gen_process = CLIP_GEN_PROCESS_GEN_CODE;
|
||||
params.out_embd = &out_embd;
|
||||
params.out_codes = &out_codes;
|
||||
params.out_feats = &out_feats;
|
||||
params.code0 = inp->code0;
|
||||
params.top_k = inp->top_k;
|
||||
params.top_p = inp->top_p;
|
||||
params.seed = inp->seed;
|
||||
params.temp = inp->temp;
|
||||
params.out_is_eos = &is_eos;
|
||||
|
||||
if (!clip_encode(ctx_clip, ¶ms)) {
|
||||
LOG_ERR("%s: clip_encode failed (gen_code)\n", __func__);
|
||||
@@ -1853,19 +1899,31 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
|
||||
|
||||
ctx->gen_out_embd = std::move(out_embd);
|
||||
ctx->gen_out_codes = std::move(out_codes);
|
||||
ctx->gen_out_feats = std::move(out_feats);
|
||||
|
||||
out->embd = ctx->gen_out_embd.data();
|
||||
out->codes = ctx->gen_out_codes.data();
|
||||
out->n_codes = ctx->gen_out_codes.size();
|
||||
out->embd = ctx->gen_out_embd.data();
|
||||
out->codes = ctx->gen_out_codes.data();
|
||||
out->n_codes = ctx->gen_out_codes.size();
|
||||
out->feats = ctx->gen_out_feats.data();
|
||||
out->n_feats = ctx->gen_out_feats.size();
|
||||
out->is_eos = is_eos;
|
||||
return 0;
|
||||
}
|
||||
|
||||
// MTMD_GEN_PROCESS_TYPE_GEN_WAV
|
||||
if (!inp->codes || inp->n_codes == 0) {
|
||||
LOG_ERR("%s: codes required for gen_wav\n", __func__);
|
||||
const bool has_codes = inp->codes && inp->n_codes > 0;
|
||||
const bool has_feats = inp->feats && inp->n_feats > 0;
|
||||
if (has_codes == has_feats) {
|
||||
LOG_ERR("%s: gen_wav requires exactly one of codes or feats\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
std::vector<int32_t> in_codes(inp->codes, inp->codes + inp->n_codes);
|
||||
std::vector<int32_t> in_codes;
|
||||
std::vector<float> in_feats;
|
||||
if (has_codes) {
|
||||
in_codes.assign(inp->codes, inp->codes + inp->n_codes);
|
||||
} else {
|
||||
in_feats.assign(inp->feats, inp->feats + inp->n_feats);
|
||||
}
|
||||
std::vector<uint8_t> in_state;
|
||||
if (inp->state_data) {
|
||||
in_state.assign(inp->state_data, inp->state_data + inp->state_size);
|
||||
@@ -1885,7 +1943,10 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
|
||||
params.imgs = &batch;
|
||||
params.n_threads = ctx->n_threads;
|
||||
params.gen_process = CLIP_GEN_PROCESS_GEN_WAV;
|
||||
params.codes = &in_codes;
|
||||
// gen_wav draws no randomness, but keep the seed so it does not reseed mid-generation
|
||||
params.seed = inp->seed;
|
||||
params.codes = has_codes ? &in_codes : nullptr;
|
||||
params.feats = has_feats ? &in_feats : nullptr;
|
||||
params.out_audio = &ctx->gen_out_audio;
|
||||
params.state_in = inp->state_data ? &in_state : nullptr;
|
||||
params.state_out = &ctx->gen_out_state;
|
||||
|
||||
Reference in New Issue
Block a user