From f55a236349d03272cdc092705a75486b6adf5886 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Wed, 5 Aug 2026 02:00:52 +0200 Subject: [PATCH 01/17] adapt the api --- tools/mtmd/clip-model.h | 3 ++ tools/mtmd/clip.cpp | 34 ++++++++++++++++++--- tools/mtmd/clip.h | 5 +++ tools/mtmd/mtmd-helper-gen.cpp | 34 +++++++++++++++++---- tools/mtmd/mtmd-helper.h | 16 ++++++---- tools/mtmd/mtmd.cpp | 56 +++++++++++++++++++++++----------- tools/mtmd/mtmd.h | 10 +++++- tools/tts/tts.cpp | 15 ++++++--- 8 files changed, 134 insertions(+), 39 deletions(-) diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 101f49cd184..f3533f7eb49 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -136,6 +136,9 @@ struct clip_hparams { int32_t rvq_num_quantizers = 0; std::vector rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17) + // threshold for the "out_eos_score" graph output + float gen_eos_threshold = 0.0f; + // qwen3tts code2wav int32_t wav_tfm_n_layer = 0; int32_t wav_tfm_n_embd = 0; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index d6670030ff7..fcf022b8d28 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -174,6 +174,10 @@ struct clip_ctx { bool support_batch = false; + // for audio gen, reseeded only when the caller asks for another seed + std::mt19937 rng{std::random_device{}()}; + uint32_t rng_seed = UINT32_MAX; + clip_ctx(clip_context_params & ctx_params) { flash_attn_type = ctx_params.flash_attn_type; no_alloc = ctx_params.no_alloc; @@ -4080,6 +4084,11 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { clip_model_loader::warmup(*ctx, *params->imgs); } + if (params->seed != ctx->rng_seed) { + ctx->rng_seed = params->seed; + ctx->rng.seed(params->seed == UINT32_MAX ? std::random_device{}() : params->seed); + } + // build the inference graph ggml_backend_sched_reset(ctx->sched.get()); ggml_cgraph * gf = clip_get_graph_builder(ctx, imgs, params)->build(); @@ -4786,11 +4795,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { set_input_i32("inp_code0", code0); // one uniform(0,1) draw per codebook, used by do_sampling() - static std::mt19937 rng{ std::random_device{}() }; std::uniform_real_distribution dist(0.0f, 1.0f); const int64_t n_acoustic = model.gen_code_head_w->ne[2]; for (int64_t g = 0; g < n_acoustic; g++) { - std::vector r = { dist(rng) }; + std::vector r = { dist(ctx->rng) }; set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r); } } @@ -5252,6 +5260,24 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { out_codes.resize(ggml_nelements(codes)); ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes)); } + // optional outputs, a missing tensor is not an error + if (params->out_feats != nullptr) { + ggml_tensor * feats = ggml_graph_get_tensor(gf, "out_feats"); + if (feats != nullptr) { + auto & out_feats = *params->out_feats; + out_feats.resize(ggml_nelements(feats)); + ggml_backend_tensor_get(feats, out_feats.data(), 0, ggml_nbytes(feats)); + } + } + if (params->out_is_eos != nullptr) { + ggml_tensor * eos = ggml_graph_get_tensor(gf, "out_eos_score"); + if (eos != nullptr) { + GGML_ASSERT(ggml_nelements(eos) == 1); + float score = 0.0f; + ggml_backend_tensor_get(eos, &score, 0, sizeof(float)); + *params->out_is_eos = score > hparams.gen_eos_threshold; + } + } if (params->out_audio != nullptr) { ggml_tensor * audio = ggml_graph_get_tensor(gf, "out_audio"); if (audio == nullptr) { @@ -5262,9 +5288,9 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { ggml_backend_tensor_get(audio, out_audio.data(), 0, ggml_nbytes(audio)); // drop the tail audio that comes from the code-0 rear padding - const int64_t n_codes = model.gen_code_head_w->ne[2] + 1; + const int64_t n_codes = params->codes ? model.gen_code_head_w->ne[2] + 1 : 0; const int64_t n_frames_w = hparams.wav_tfm_swa; - const int64_t n_frames = (int64_t) params->codes->size() / n_codes; + const int64_t n_frames = params->codes ? (int64_t) params->codes->size() / n_codes : n_frames_w; if (n_frames < n_frames_w) { const size_t hop = out_audio.size() / n_frames_w; out_audio.resize((size_t) n_frames * hop); diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index 7f706d976eb..97de0606dbf 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -104,9 +104,14 @@ struct clip_encode_params { int32_t top_k = 50; float top_p = 1.0f; std::vector * out_codes = nullptr; // this frame's 16 sampled codes + std::vector * out_feats = nullptr; // continuous counterpart of out_codes + uint32_t seed = UINT32_MAX; // UINT32_MAX for random + int32_t n_steps = -1; // integration steps, for flow-matching decoders + bool * out_is_eos = nullptr; // GEN_WAV const std::vector * codes = nullptr; // this frame's 16 RVQ codes + const std::vector * feats = nullptr; // continuous counterpart of codes std::vector * out_audio = nullptr; // decoded PCM samples, F32 const std::vector * state_in = nullptr; // state from previous call, null or wrong size means cold start std::vector * state_out = nullptr; // state for the next call diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index b52dc8e5a34..bc83f5ba230 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -87,7 +87,8 @@ class mtmd_gen_audio_pipeline { 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: @@ -203,6 +204,7 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { pos = 0; top_k = inp->top_k > 0 ? inp->top_k : 50; top_p = inp->top_p > 0 ? inp->top_p : 1.0f; + seed = inp->seed; out_type = inp->out_type; // the text stream keeps flowing during generation: after frame k, the input adds @@ -244,13 +246,26 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { return n_prompt - prompt_pos; } - int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) override { + 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{}; inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; inp.code0 = sampled - codec_0; inp.embd = const_cast(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"); @@ -432,8 +447,9 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::unique_ptr 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 codes_buf; std::vector c2w_state; std::vector audio_pcm; @@ -489,11 +505,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, diff --git a/tools/mtmd/mtmd-helper.h b/tools/mtmd/mtmd-helper.h index 7e5cf9b5098..832f7171ac7 100644 --- a/tools/mtmd/mtmd-helper.h +++ b/tools/mtmd/mtmd-helper.h @@ -183,8 +183,9 @@ struct mtmd_helper_gen_audio_inp { mtmd_bitmap * speaker_ref; // optional, can be NULL const char * lang; // optional, can be NULL - int32_t top_k; - float top_p; + int32_t top_k; + float top_p; + uint32_t seed; // UINT32_MAX for random (default: random) enum mtmd_helper_gen_audio_outtype out_type; }; @@ -208,12 +209,15 @@ MTMD_API int32_t mtmd_helper_gen_audio_step_prompt( int32_t n_batch); // generates one frame; must only be called after step_prompt() has returned 0 -// h_state_out is valid until next step_gen() or reset() call +// sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token +// out_stop (optional) is set on end-of-speech, the caller must then stop the loop +// h_state_out is valid until next step_gen() or reset() call, null if no frame is generated MTMD_API 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_out, + bool * out_stop); // out_data valid until next get_output() or reset() call // out_n_samples (optional, can be NULL) receives the number of generated PCM samples @@ -261,8 +265,8 @@ struct gen_audio { int32_t step_prompt(int32_t n_batch) { return mtmd_helper_gen_audio_step_prompt(ctx.get(), n_batch); } - int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out) { - return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out); + int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out, bool * out_stop = nullptr) { + return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out, out_stop); } int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples = nullptr) { return mtmd_helper_gen_audio_get_output(ctx.get(), out_sample_rate, out_data, out_data_len, out_n_samples); diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index ff90d6818cb..6a5b3f5a606 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -265,6 +265,7 @@ struct mtmd_context { // generation context struct clip_ctx * ctx_gen_a; // audio std::vector gen_out_codes; // this frame's 16 sampled codes (GEN_CODE) + std::vector gen_out_feats; // this frame's continuous features, if any (GEN_CODE) std::vector gen_out_embd; // next-step hidden state fed back to backbone (GEN_CODE) std::vector gen_out_audio; // decoded PCM samples for the current frame (GEN_WAV) std::vector gen_out_state; // state to feed into the next GEN_WAV call @@ -1580,7 +1581,7 @@ 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{}; if (!ctx->ctx_gen_a) { info.type = MTMD_GEN_AUDIO_TYPE_NONE; return info; @@ -1604,6 +1605,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); @@ -1617,16 +1620,22 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in std::vector out_embd(n_embd); std::vector out_codes; + std::vector 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.n_steps = inp->n_steps; + params.out_is_eos = &is_eos; if (!clip_encode(ctx_clip, ¶ms)) { LOG_ERR("%s: clip_encode failed (gen_code)\n", __func__); @@ -1635,19 +1644,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); - - out->embd = ctx->gen_out_embd.data(); - out->codes = ctx->gen_out_codes.data(); - out->n_codes = ctx->gen_out_codes.size(); + 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->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 in_codes(inp->codes, inp->codes + inp->n_codes); + std::vector in_codes; + std::vector 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 in_state; if (inp->state_data) { in_state.assign(inp->state_data, inp->state_data + inp->state_size); @@ -1667,7 +1688,8 @@ 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; + 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; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 84651f8dcd0..32d40dca327 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -354,10 +354,15 @@ struct mtmd_gen_inp { float * embd; // the hidden state from backbone, must have n_text_embd elements int32_t top_k; float top_p; + uint32_t seed; // UINT32_MAX for random + int32_t n_steps; // integration steps, for flow-matching decoders (-1 for default) // for MTMD_GEN_PROCESS_TYPE_GEN_WAV + // pass either codes (discrete) or feats (continuous), depending on the pipeline int32_t * codes; size_t n_codes; + const float * feats; + size_t n_feats; const char * state_data; size_t state_size; }; @@ -366,9 +371,12 @@ struct mtmd_gen_out { // for MTMD_GEN_PROCESS_TYPE_GEN_CODE const int32_t * codes; - size_t n_codes; + size_t n_codes; + const float * feats; // continuous counterpart of codes + size_t n_feats; const float * embd; // the generated hidden state, to be fed back to backbone // it must have n_text_embd elements + bool is_eos; // only set by pipelines having the EOS head inside mmproj // for MTMD_GEN_PROCESS_TYPE_GEN_WAV const float * audio; diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index b68edcaf575..9e52505d3b1 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -119,6 +119,7 @@ int main(int argc, char ** argv) { inp.lang = params.tts_lang.c_str(); inp.top_k = params.sampling.top_k; inp.top_p = params.sampling.top_p; + inp.seed = params.sampling.seed; inp.out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; // @@ -143,8 +144,7 @@ int main(int argc, char ** argv) { } } - const llama_vocab * vocab = llama_model_get_vocab(model); - + // note: some pipelines ignore this token and use the hidden state instead auto sample_semantic_code = [&]() -> llama_token { llama_token t = common_sampler_sample(smpl, lctx, -1); common_sampler_accept(smpl, t, true); @@ -159,19 +159,24 @@ int main(int argc, char ** argv) { tts_timings timings; const int64_t t_gen_start_us = ggml_time_us(); - for (; n_frames < max_new && !llama_vocab_is_eog(vocab, sampled); n_frames++) { + bool stop = false; + while (!stop && n_frames < max_new) { const float * h_next = nullptr; // stage 2+3: semantic --> acoustic details --> audio waveform // step_gen() runs both stages and returns new h_state for next step - if (gen.step_gen(sampled, h_state, &h_next) != 0) { + if (gen.step_gen(sampled, h_state, &h_next, &stop) != 0) { LOG_ERR("step_gen failed at frame %d\n", n_frames); return 1; } + if (!h_next) { + break; // stopped without generating a frame + } + n_frames++; h_state = h_next; sampled = sample_semantic_code(); - timings.report(n_frames + 1); + timings.report(n_frames); } const double t_gen_s = (ggml_time_us() - t_gen_start_us) / 1e6; From a0c869f04380388b6c16b49f376fdf58a644316f Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Wed, 5 Aug 2026 02:41:43 +0200 Subject: [PATCH 02/17] text model ok --- conversion/__init__.py | 2 + conversion/base.py | 47 ++++++ conversion/pockettts.py | 312 ++++++++++++++++++++++++++++++++++++++ gguf-py/gguf/constants.py | 107 +++++++++++++ src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/llama-model.cpp | 3 + src/models/models.h | 13 ++ src/models/pockettts.cpp | 146 ++++++++++++++++++ 9 files changed, 632 insertions(+) create mode 100644 conversion/pockettts.py create mode 100644 src/models/pockettts.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index 06c2c50ad24..5e31a3d9184 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -210,6 +210,7 @@ "Qwen3MoeForCausalLM": "qwen", "Qwen3NextForCausalLM": "qwen", "Qwen3OmniMoeForConditionalGeneration": "qwen3vl", + "PocketTTSModel": "pockettts", "Qwen3TTSForConditionalGeneration": "qwen3tts", "Qwen3VLForConditionalGeneration": "qwen3vl", "Qwen3VLMoeForConditionalGeneration": "qwen3vl", @@ -305,6 +306,7 @@ "Qwen2_5_VLForConditionalGeneration": "qwenvl", "Qwen3ASRForConditionalGeneration": "qwen3vl", "Qwen3OmniMoeForConditionalGeneration": "qwen3vl", + "PocketTTSModel": "pockettts", "Qwen3TTSForConditionalGeneration": "qwen3tts", "Qwen3VLForConditionalGeneration": "qwen3vl", "Qwen3VLMoeForConditionalGeneration": "qwen3vl", diff --git a/conversion/base.py b/conversion/base.py index a7cd3fd904a..69f0fd9e657 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -1053,6 +1053,10 @@ def load_hparams(dir_model: Path, is_mistral_format: bool): config = AutoConfig.from_pretrained(dir_model, trust_remote_code=False).to_dict() except Exception as e: logger.warning(f"Failed to load model config from {dir_model}: {e}") + if not (dir_model / "config.json").is_file(): + config = load_hparams_non_hf(dir_model) + if config is not None: + return config logger.warning("Trying to load config.json instead") with open(dir_model / "config.json", "r", encoding="utf-8") as f: config = json.load(f) @@ -2618,6 +2622,49 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): LazyTorchTensor._dtype_str_map["F8_E8M0"] = torch.uint8 +def load_hparams_non_hf(dir_model: Path) -> dict[str, Any] | None: + # some models ship no config.json at all, their hparams are derived from the checkpoint + part_names = ModelBase.get_model_part_names(dir_model, "model", ".safetensors") + if len(part_names) != 1: + return None + with gguf.utility.SafetensorsLocal(dir_model / part_names[0]) as part: + shapes = {name: tuple(part[name].shape) for name in part.keys()} + + if "flow_lm.bos_emb" in shapes: + return _load_hparams_pockettts(shapes) + + return None + + +def _load_hparams_pockettts(shapes: dict[str, tuple[int, ...]]) -> dict[str, Any]: + logger.info("gguf: detected pocket-tts checkpoint, deriving hparams from tensor shapes") + n_vocab, n_embd = shapes["flow_lm.conditioner.embed.weight"] + n_layer = sum(1 for name in shapes if re.fullmatch(r"flow_lm\.transformer\.layers\.\d+\.norm1\.weight", name)) + n_layer_a = sum(1 for name in shapes if re.fullmatch(r"mimi\.encoder_transformer\.transformer\.layers\.\d+\.norm1\.weight", name)) + n_embd_a = shapes["mimi.encoder_transformer.transformer.layers.0.norm1.weight"][0] + return { + "architectures": ["PocketTTSModel"], + "model_type": "pockettts", + "num_hidden_layers": n_layer, + "hidden_size": n_embd, + "intermediate_size": shapes["flow_lm.transformer.layers.0.linear1.weight"][0], + # the transformer is fully causal with no context limit, this only bounds the KV cache + "max_position_embeddings": 4096, + # not stored anywhere in the checkpoint, but every released variant uses head_dim 64 + "num_attention_heads": n_embd // 64, + # 2 learned vectors are appended to the embedding table as extra tokens, see pockettts.py + "vocab_size": n_vocab + 2, + "rope_theta": 10000.0, + "layer_norm_eps": 1e-5, + "audio_config": { + "num_hidden_layers": n_layer_a, + "hidden_size": n_embd_a, + "intermediate_size": shapes["mimi.encoder_transformer.transformer.layers.0.linear1.weight"][0], + "num_attention_heads": n_embd_a // 64, + }, + } + + def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> str: # TODO @ngxson : this won't work correctly if the model has both audio & vision encoders # maybe we should fallback to text model's arch in that case, since not many models have both diff --git a/conversion/pockettts.py b/conversion/pockettts.py new file mode 100644 index 00000000000..4aa41f23010 --- /dev/null +++ b/conversion/pockettts.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +from typing import Iterable, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf + +# Pocket TTS is a CALM: an autoregressive backbone conditions a flow-matching decoder that +# generates one continuous 32-d latent per frame. There is no codebook anywhere in this model. +# +# The checkpoint ships no config.json, hparams are derived in base.load_hparams_non_hf(). +# +# Tricks being used to support this model via existing llama.cpp code paths: +# - bos_before_voice and bos_emb are learned input vectors, not tokens. they are appended to +# the embedding table as extra tokens so the helper can look them up like any other row. +# bos_emb lives in latent space, so input_linear is folded into it here +# - the backbone has no lm_head, the embedding table is reused as output so that a sampler +# can run over the (unused) logits +# +# pipeline stage mapping: +# mimi encoder + speaker_proj --> mapped to normal mtmd audio encoder +# flow_lm.transformer --> mapped to normal libllama text model (autoregressive) +# flow_lm.flow_net + out_eos --> MTMD_GEN_PROCESS_TYPE_GEN_CODE +# mimi decoder --> MTMD_GEN_PROCESS_TYPE_GEN_WAV + +# indices into mimi.encoder.model / mimi.decoder.model for stage i, see SEANetEncoder/SEANetDecoder +_ENC_RES_IDX = lambda i: 1 + 3 * i # noqa: E731 +_ENC_SCALE_IDX = lambda i: 3 + 3 * i # noqa: E731 +_DEC_SCALE_IDX = lambda i: 2 + 3 * i # noqa: E731 +_DEC_RES_IDX = lambda i: 3 + 3 * i # noqa: E731 + +_N_SEANET_STAGES = 3 +_SAMPLE_RATE = 24000 + + +@ModelBase.register("PocketTTSModel") +class PocketTTSModel(TextModel): + model_arch = gguf.MODEL_ARCH.POCKETTTS + + _LAYER_TENSOR_MAP = { + "norm1": gguf.MODEL_TENSOR.ATTN_NORM, + "norm2": gguf.MODEL_TENSOR.FFN_NORM, + "self_attn.out_proj": gguf.MODEL_TENSOR.ATTN_OUT, + "linear1": gguf.MODEL_TENSOR.FFN_UP, + "linear2": gguf.MODEL_TENSOR.FFN_DOWN, + } + + def set_vocab(self): + tokens, scores, toktypes = self._create_vocab_sentencepiece() + + # the last 3 rows of the embedding table are not sentencepiece pieces: the conditioner's + # padding row, then the two learned vectors appended by generate_extra_tensors() + extra = ["<|pad|>", "<|bos_before_voice|>", "<|audio_bos|>"] + for i, name in enumerate(extra): + tokens[len(tokens) - len(extra) + i] = name.encode("utf-8") + toktypes[len(tokens) - len(extra) + i] = SentencePieceTokenTypes.CONTROL + scores[len(tokens) - len(extra) + i] = -1000.0 + + self.gguf_writer.add_tokenizer_model("llama") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + self.gguf_writer.add_add_bos_token(False) + self.gguf_writer.add_add_eos_token(False) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if not name.startswith("flow_lm."): + return # mimi and the flow net go to the mmproj + + if name == "flow_lm.conditioner.embed.weight": + yield (self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD), self._embd_table(data_torch)) + return + + if name.startswith("flow_lm.out_norm."): + suffix = "." + name.rsplit(".", 1)[1] + yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT_NORM, suffix=suffix), data_torch) + return + + if name.startswith("flow_lm.transformer.layers."): + assert bid is not None + key_with_suffix = name.split(f"layers.{bid}.", 1)[1] + key, suffix = key_with_suffix.rsplit(".", 1) + + if key == "self_attn.in_proj": + q, k, v = data_torch.chunk(3, dim=0) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), q) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), k) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), v) + return + + tensor = self._LAYER_TENSOR_MAP.get(key) + if tensor is not None: + yield (self.format_tensor_name(tensor, bid, suffix="." + suffix), data_torch) + return + + return + + def _embd_table(self, embed: Tensor) -> Tensor: + bos_before_voice = self.model_tensors["flow_lm.bos_before_voice"]().reshape(1, -1) + # bos_emb is a latent, it only enters the backbone through input_linear + bos_emb = self.model_tensors["flow_lm.bos_emb"]() + input_linear = self.model_tensors["flow_lm.input_linear.weight"]() + audio_bos = torch.nn.functional.linear(bos_emb.float(), input_linear.float()).reshape(1, -1) + + return torch.cat([embed, bos_before_voice.to(embed.dtype), audio_bos.to(embed.dtype)], dim=0) + + +@ModelBase.register("PocketTTSModel") +class PocketTTSMmprojModel(MmprojModel): + has_audio_encoder = True + has_vision_encoder = False + + _MIMI_TFM_MAP = { + "norm1": (gguf.MODEL_TENSOR.A_ENC_INPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM), + "norm2": (gguf.MODEL_TENSOR.A_ENC_OUTPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM), + "self_attn.out_proj": (gguf.MODEL_TENSOR.A_ENC_OUTPUT, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT), + "linear1": (gguf.MODEL_TENSOR.A_ENC_FFN_UP, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP), + "linear2": (gguf.MODEL_TENSOR.A_ENC_FFN_DOWN, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN), + "layer_scale_1.scale": (gguf.MODEL_TENSOR.A_ENC_ATTN_SCALE, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE), + "layer_scale_2.scale": (gguf.MODEL_TENSOR.A_ENC_FFN_SCALE, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE), + } + _MIMI_TFM_QKV = ( + (gguf.MODEL_TENSOR.A_ENC_ATTN_Q, gguf.MODEL_TENSOR.A_ENC_ATTN_K, gguf.MODEL_TENSOR.A_ENC_ATTN_V), + (gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V), + ) + + def set_gguf_parameters(self): + self.gguf_writer.add_file_type(self.ftype) + assert self.hparams_audio is not None + + # voice-prompt encoder: mimi encoder + speaker_proj + self.gguf_writer.add_clip_has_audio_encoder(True) + # note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models + self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC) + self.gguf_writer.add_audio_projection_dim(self.n_embd_text) + self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"]) + self.gguf_writer.add_audio_embedding_length(self.hparams_audio["hidden_size"]) + self.gguf_writer.add_audio_feed_forward_length(self.hparams_audio["intermediate_size"]) + self.gguf_writer.add_audio_head_count(self.hparams_audio["num_attention_heads"]) + self.gguf_writer.add_audio_attention_layernorm_eps(1e-5) + + # generation: flow-matching decoder + mimi decoder + # note: the SEANet and flow net hparams are hardcoded on the clip.cpp side for now + self.gguf_writer.add_clip_has_gen_audio_encoder(True) + self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_GEN) + self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text) + self.gguf_writer.add_gen_audio_embedding_length(self.hparams_audio["hidden_size"]) + self.gguf_writer.add_gen_audio_feed_forward_length(self.hparams_audio["intermediate_size"]) + self.gguf_writer.add_gen_audio_block_count(self.hparams_audio["num_hidden_layers"]) + self.gguf_writer.add_gen_audio_head_count(self.hparams_audio["num_attention_heads"]) + self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + del name, bid, n_dims + # conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path + if ".seanet." in new_name or new_name in ("a.downsample.conv.weight", "a.gen.wav.upsample.weight"): + return gguf.GGMLQuantizationType.F16 + return False + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # the block index of the mimi transformers is parsed here, not by the base class + T = gguf.MODEL_TENSOR + + if name in ("flow_lm.bos_emb", "flow_lm.bos_before_voice", "flow_lm.conditioner.embed.weight"): + return # folded into the backbone embedding table + if name.startswith("flow_lm.transformer.") or name.startswith("flow_lm.out_norm."): + return # backbone + + if name == "flow_lm.speaker_proj_weight": + yield (self.format_tensor_name(T.A_ENC_SPEAKER_PROJ), data_torch) + return + if name == "flow_lm.input_linear.weight": + yield (self.format_tensor_name(T.A_GEN_INPUT_LINEAR), data_torch) + return + if name == "flow_lm.emb_mean": + yield (self.format_tensor_name(T.A_GEN_EMB_MEAN, suffix=""), data_torch) + return + if name == "flow_lm.emb_std": + yield (self.format_tensor_name(T.A_GEN_EMB_STD, suffix=""), data_torch) + return + if name.startswith("flow_lm.out_eos."): + suffix = "." + name.rsplit(".", 1)[1] + yield (self.format_tensor_name(T.A_GEN_OUT_EOS, suffix=suffix), data_torch) + return + + if name.startswith("flow_lm.flow_net."): + yield from self._flow_net_tensor(name, data_torch) + return + + if name == "mimi.downsample.conv.conv.weight": + yield (self.format_tensor_name(T.A_ENC_DOWNSAMPLE_CONV), data_torch) + return + if name == "mimi.upsample.convtr.convtr.weight": + yield (self.format_tensor_name(T.A_GEN_WAV_UPSAMPLE), data_torch) + return + if name == "mimi.quantizer.output_proj.weight": + yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_OUT), data_torch.squeeze(-1)) + return + + if "_transformer.transformer.layers." in name: + yield from self._mimi_tfm_tensor(name, data_torch) + return + + if name.startswith("mimi.encoder.model.") or name.startswith("mimi.decoder.model."): + yield from self._seanet_tensor(name, data_torch) + return + + return + + def _flow_net_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: + T = gguf.MODEL_TENSOR + key = name.split("flow_lm.flow_net.", 1)[1] + suffix = "." + key.rsplit(".", 1)[1] + + simple = { + "input_proj": T.A_GEN_FLOW_INPUT_PROJ, + "cond_embed": T.A_GEN_FLOW_COND_EMBD, + "final_layer.linear": T.A_GEN_FLOW_FINAL_PROJ, + "final_layer.adaLN_modulation.1": T.A_GEN_FLOW_FINAL_ADA, + } + tensor = simple.get(key.rsplit(".", 1)[0]) + if tensor is not None: + yield (self.format_tensor_name(tensor, suffix=suffix), data_torch) + return + + if key.startswith("time_embed."): + bid = int(key.split(".")[1]) + rest = key.split(f"time_embed.{bid}.", 1)[1] + time_map = { + "freqs": (T.A_GEN_FLOW_TIME_FREQS, ""), + "mlp.0": (T.A_GEN_FLOW_TIME_UP, suffix), + "mlp.2": (T.A_GEN_FLOW_TIME_DOWN, suffix), + "mlp.3.alpha": (T.A_GEN_FLOW_TIME_NORM, ""), + } + entry = time_map.get(rest) or time_map.get(rest.rsplit(".", 1)[0]) + if entry is not None: + yield (self.format_tensor_name(entry[0], bid, suffix=entry[1]), data_torch) + return + + if key.startswith("res_blocks."): + bid = int(key.split(".")[1]) + rest = key.split(f"res_blocks.{bid}.", 1)[1].rsplit(".", 1)[0] + blk_map = { + "in_ln": T.A_GEN_FLOW_BLK_NORM, + "mlp.0": T.A_GEN_FLOW_BLK_UP, + "mlp.2": T.A_GEN_FLOW_BLK_DOWN, + "adaLN_modulation.1": T.A_GEN_FLOW_BLK_ADA, + } + tensor = blk_map.get(rest) + if tensor is not None: + yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch) + return + + def _mimi_tfm_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: + is_decoder = name.startswith("mimi.decoder_transformer.") + bid = int(name.split("_transformer.transformer.layers.", 1)[1].split(".")[0]) + key_with_suffix = name.split(f".layers.{bid}.", 1)[1] + + if key_with_suffix == "self_attn.in_proj.weight": + q, k, v = data_torch.chunk(3, dim=0) + names = self._MIMI_TFM_QKV[1 if is_decoder else 0] + for tensor, part in zip(names, (q, k, v)): + yield (self.format_tensor_name(tensor, bid), part) + return + + key, suffix = key_with_suffix.rsplit(".", 1) + entry = self._MIMI_TFM_MAP.get(key) or self._MIMI_TFM_MAP.get(key_with_suffix) + if entry is None: + return + tensor = entry[1 if is_decoder else 0] + # layer_scale is stored without a .weight/.bias suffix + suffix = "" if key_with_suffix.endswith(".scale") else "." + suffix + yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch) + + def _seanet_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: + T = gguf.MODEL_TENSOR + is_decoder = name.startswith("mimi.decoder.") + idx = int(name.split(".model.", 1)[1].split(".")[0]) + suffix = "." + name.rsplit(".", 1)[1] + + conv_in, conv_out, res1, res2, scale = ( + (T.A_GEN_WAV_SEANET_CONV_IN, T.A_GEN_WAV_SEANET_CONV_OUT, T.A_GEN_WAV_SEANET_RES_CONV1, + T.A_GEN_WAV_SEANET_RES_CONV2, T.A_GEN_WAV_SEANET_SCALE_CONV) + if is_decoder else + (T.A_ENC_SEANET_CONV_IN, T.A_ENC_SEANET_CONV_OUT, T.A_ENC_SEANET_RES_CONV1, + T.A_ENC_SEANET_RES_CONV2, T.A_ENC_SEANET_SCALE_CONV) + ) + + if idx == 0: + yield (self.format_tensor_name(conv_in, suffix=suffix), data_torch) + return + if idx == 3 * _N_SEANET_STAGES + 2: + yield (self.format_tensor_name(conv_out, suffix=suffix), data_torch) + return + + for stage in range(_N_SEANET_STAGES): + res_idx = _DEC_RES_IDX(stage) if is_decoder else _ENC_RES_IDX(stage) + scale_idx = _DEC_SCALE_IDX(stage) if is_decoder else _ENC_SCALE_IDX(stage) + if idx == scale_idx: + yield (self.format_tensor_name(scale, stage, suffix=suffix), data_torch) + return + if idx == res_idx: + # block.1 is the dilated conv, block.3 the pointwise one (0 and 2 are ELU) + inner = int(name.split(".block.", 1)[1].split(".")[0]) + tensor = res1 if inner == 1 else res2 + yield (self.format_tensor_name(tensor, stage, suffix=suffix), data_torch) + return diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 8516222cccb..445144e13f5 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -572,6 +572,7 @@ class MODEL_ARCH(IntEnum): MELLUM = auto() NANBEIGE = auto() QWEN3TTS = auto() + POCKETTTS = auto() class VISION_PROJECTOR_TYPE(IntEnum): @@ -1031,6 +1032,37 @@ class MODEL_TENSOR(IntEnum): A_GEN_WAV_DAC_RES_CONV2 = auto() # DAC residual unit, pointwise causal conv A_GEN_WAV_DAC_POST_SNAKE = auto() # DAC final SnakeBeta A_GEN_WAV_DAC_POST_CONV = auto() # DAC conv_post -> 1-channel PCM + # pocket-tts: SEANet encoder (speaker path) and decoder (a.gen.wav path) + A_ENC_SEANET_CONV_IN = auto() + A_ENC_SEANET_CONV_OUT = auto() + A_ENC_SEANET_RES_CONV1 = auto() # residual unit, dilated conv + A_ENC_SEANET_RES_CONV2 = auto() # residual unit, pointwise conv + A_ENC_SEANET_SCALE_CONV = auto() # strided downsample conv + A_ENC_ATTN_SCALE = auto() # layer scale (gamma) on the attn output + A_ENC_SPEAKER_PROJ = auto() # voice latent -> backbone embd + A_GEN_FLOW_INPUT_PROJ = auto() + A_GEN_FLOW_COND_EMBD = auto() + A_GEN_FLOW_TIME_FREQS = auto() # timestep embedder, stored cos/sin frequencies + A_GEN_FLOW_TIME_UP = auto() + A_GEN_FLOW_TIME_DOWN = auto() + A_GEN_FLOW_TIME_NORM = auto() # RMSNorm alpha + A_GEN_FLOW_BLK_NORM = auto() # AdaLN res block, in_ln + A_GEN_FLOW_BLK_UP = auto() + A_GEN_FLOW_BLK_DOWN = auto() + A_GEN_FLOW_BLK_ADA = auto() # AdaLN modulation, -> shift/scale/gate + A_GEN_FLOW_FINAL_ADA = auto() # final layer AdaLN modulation, -> shift/scale + A_GEN_FLOW_FINAL_PROJ = auto() + A_GEN_OUT_EOS = auto() # end-of-speech head on the backbone hidden state + A_GEN_INPUT_LINEAR = auto() # generated latent -> backbone embd + A_GEN_EMB_MEAN = auto() # latent denormalization stats + A_GEN_EMB_STD = auto() + A_GEN_WAV_QUANT_OUT = auto() # DummyQuantizer output_proj, latent -> decoder dim + A_GEN_WAV_UPSAMPLE = auto() # frame rate -> encoder frame rate, depthwise convtr + A_GEN_WAV_SEANET_CONV_IN = auto() + A_GEN_WAV_SEANET_CONV_OUT = auto() # -> 1-channel PCM + A_GEN_WAV_SEANET_RES_CONV1 = auto() + A_GEN_WAV_SEANET_RES_CONV2 = auto() + A_GEN_WAV_SEANET_SCALE_CONV = auto() # strided upsample convtr A_MMPROJ = auto() A_MMPROJ_FC = auto() A_MM_NORM_PRE = auto() @@ -1244,6 +1276,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.MELLUM: "mellum", MODEL_ARCH.NANBEIGE: "nanbeige", MODEL_ARCH.QWEN3TTS: "qwen3tts", + MODEL_ARCH.POCKETTTS: "pockettts", } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { @@ -1698,6 +1731,36 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2: "a.gen.wav.dac.blk.{bid}.res.{xid}.conv2", MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE: "a.gen.wav.dac.post_snake", MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV: "a.gen.wav.dac.post_conv", + MODEL_TENSOR.A_ENC_SEANET_CONV_IN: "a.seanet.conv_in", + MODEL_TENSOR.A_ENC_SEANET_CONV_OUT: "a.seanet.conv_out", + MODEL_TENSOR.A_ENC_SEANET_RES_CONV1: "a.seanet.blk.{bid}.res_conv1", + MODEL_TENSOR.A_ENC_SEANET_RES_CONV2: "a.seanet.blk.{bid}.res_conv2", + MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV: "a.seanet.blk.{bid}.scale_conv", + MODEL_TENSOR.A_ENC_ATTN_SCALE: "a.blk.{bid}.attn_scale", + MODEL_TENSOR.A_ENC_SPEAKER_PROJ: "a.speaker_proj", + MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ: "a.gen.flow.input_proj", + MODEL_TENSOR.A_GEN_FLOW_COND_EMBD: "a.gen.flow.cond_embd", + MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS: "a.gen.flow.time.{bid}.freqs", + MODEL_TENSOR.A_GEN_FLOW_TIME_UP: "a.gen.flow.time.{bid}.up", + MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN: "a.gen.flow.time.{bid}.down", + MODEL_TENSOR.A_GEN_FLOW_TIME_NORM: "a.gen.flow.time.{bid}.norm", + MODEL_TENSOR.A_GEN_FLOW_BLK_NORM: "a.gen.flow.blk.{bid}.norm", + MODEL_TENSOR.A_GEN_FLOW_BLK_UP: "a.gen.flow.blk.{bid}.up", + MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN: "a.gen.flow.blk.{bid}.down", + MODEL_TENSOR.A_GEN_FLOW_BLK_ADA: "a.gen.flow.blk.{bid}.ada", + MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA: "a.gen.flow.final.ada", + MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ: "a.gen.flow.final.proj", + MODEL_TENSOR.A_GEN_OUT_EOS: "a.gen.out_eos", + MODEL_TENSOR.A_GEN_INPUT_LINEAR: "a.gen.input_linear", + MODEL_TENSOR.A_GEN_EMB_MEAN: "a.gen.emb_mean", + MODEL_TENSOR.A_GEN_EMB_STD: "a.gen.emb_std", + MODEL_TENSOR.A_GEN_WAV_QUANT_OUT: "a.gen.wav.quant_out", + MODEL_TENSOR.A_GEN_WAV_UPSAMPLE: "a.gen.wav.upsample", + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN: "a.gen.wav.seanet.conv_in", + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT: "a.gen.wav.seanet.conv_out", + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1: "a.gen.wav.seanet.blk.{bid}.res_conv1", + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2: "a.gen.wav.seanet.blk.{bid}.res_conv2", + MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV: "a.gen.wav.seanet.blk.{bid}.scale_conv", MODEL_TENSOR.A_MMPROJ: "mm.a.mlp.{bid}", MODEL_TENSOR.A_MMPROJ_FC: "mm.a.fc", MODEL_TENSOR.A_MM_NORM_PRE: "mm.a.norm_pre", @@ -2009,6 +2072,36 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2, MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE, MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV, + MODEL_TENSOR.A_ENC_SEANET_CONV_IN, + MODEL_TENSOR.A_ENC_SEANET_CONV_OUT, + MODEL_TENSOR.A_ENC_SEANET_RES_CONV1, + MODEL_TENSOR.A_ENC_SEANET_RES_CONV2, + MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV, + MODEL_TENSOR.A_ENC_ATTN_SCALE, + MODEL_TENSOR.A_ENC_SPEAKER_PROJ, + MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ, + MODEL_TENSOR.A_GEN_FLOW_COND_EMBD, + MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS, + MODEL_TENSOR.A_GEN_FLOW_TIME_UP, + MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN, + MODEL_TENSOR.A_GEN_FLOW_TIME_NORM, + MODEL_TENSOR.A_GEN_FLOW_BLK_NORM, + MODEL_TENSOR.A_GEN_FLOW_BLK_UP, + MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN, + MODEL_TENSOR.A_GEN_FLOW_BLK_ADA, + MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA, + MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ, + MODEL_TENSOR.A_GEN_OUT_EOS, + MODEL_TENSOR.A_GEN_INPUT_LINEAR, + MODEL_TENSOR.A_GEN_EMB_MEAN, + MODEL_TENSOR.A_GEN_EMB_STD, + MODEL_TENSOR.A_GEN_WAV_QUANT_OUT, + MODEL_TENSOR.A_GEN_WAV_UPSAMPLE, + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN, + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT, + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1, + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2, + MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV, MODEL_TENSOR.A_ENC_CONV_NORM_MEAN, MODEL_TENSOR.A_ENC_CONV_NORM_VAR, MODEL_TENSOR.A_ENC_MEL_FILTERS, @@ -4852,6 +4945,18 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, ], + MODEL_ARCH.POCKETTTS: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], } # tensors that will not be serialized @@ -5128,6 +5233,8 @@ class VisionProjectorType: NEMOTRON_V2_VL = "nemotron_v2_vl" QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor + POCKETTTS_SPKENC = "pockettts_spkenc" # audio: mimi encoder as voice-prompt encoder + POCKETTTS_GEN = "pockettts_gen" # audio generation: flow-matching decoder + mimi decoder HUNYUANVL = "hunyuanvl" PARAKEET = "parakeet" # audio MINIMAXM3 = "minimax_m3" diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 836cfade226..360bd319ff2 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -145,6 +145,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_MELLUM, "mellum" }, { LLM_ARCH_NANBEIGE, "nanbeige" }, { LLM_ARCH_QWEN3TTS, "qwen3tts" }, + { LLM_ARCH_POCKETTTS, "pockettts" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; diff --git a/src/llama-arch.h b/src/llama-arch.h index 49c2a6ac399..81947579070 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -150,6 +150,7 @@ enum llm_arch { LLM_ARCH_DFLASH, LLM_ARCH_NANBEIGE, LLM_ARCH_QWEN3TTS, + LLM_ARCH_POCKETTTS, LLM_ARCH_UNKNOWN, }; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index dda311c47bb..6dc47ff742a 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -114,6 +114,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_qwen3vlmoe(params); case LLM_ARCH_QWEN3TTS: return new llama_model_qwen3tts(params); + case LLM_ARCH_POCKETTTS: + return new llama_model_pockettts(params); case LLM_ARCH_PHI2: return new llama_model_phi2(params); case LLM_ARCH_PHI3: @@ -2610,6 +2612,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_MAINCODER: case LLM_ARCH_GLM_DSA: case LLM_ARCH_NANBEIGE: + case LLM_ARCH_POCKETTTS: return LLAMA_ROPE_TYPE_NORM; // the pairs of head values are offset by n_rot/2 diff --git a/src/models/models.h b/src/models/models.h index ad3dadaf393..ae6ecca4eda 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -697,6 +697,19 @@ struct llama_model_gpt2 : public llama_model_base { }; +struct llama_model_pockettts : public llama_model_base { + llama_model_pockettts(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_codeshell : public llama_model_base { llama_model_codeshell(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/src/models/pockettts.cpp b/src/models/pockettts.cpp new file mode 100644 index 00000000000..1b3bb6c648a --- /dev/null +++ b/src/models/pockettts.cpp @@ -0,0 +1,146 @@ +#include "models.h" + +// backbone of the pocket-tts CALM pipeline: the "text" side of a flow language model. +// it has no lm_head, the audio latents are produced by the flow net inside the mmproj + +void llama_model_pockettts::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); + + switch (hparams.n_layer()) { + case 6: type = LLM_TYPE_109M; break; + case 24: type = LLM_TYPE_335M; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_pockettts::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output_norm_b = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "bias"), {n_embd}, 0); + // no output head, the logits are unused; reuse the embedding table so a sampler can still run + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias", i), {n_embd}, 0); + + create_tensor_qkv(layer, i, n_embd, n_embd, n_embd_gqa, n_embd_gqa, TENSOR_NOT_REQUIRED); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", i), {n_embd}, 0); + + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } +} + +std::unique_ptr llama_model_pockettts::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_pockettts::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + GGML_ASSERT(n_embd_head == n_rot); + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + ggml_tensor * inp_pos = build_inp_pos(); + + auto * inp_attn = build_attn_inp_kv(); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + for (int il = 0; il < n_layer; ++il) { + cur = build_norm(inpL, + model.layers[il].attn_norm, + model.layers[il].attn_norm_b, + LLM_NORM, il); + cb(cur, "attn_norm", il); + + // self-attention + { + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, n_head_kv, il); + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + cur = build_attn(inp_attn, + model.layers[il].wo, NULL, model.layers[il].wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il); + } + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpL = ggml_get_rows(ctx0, inpL, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); + cb(ffn_inp, "ffn_inp", il); + + // FF + { + cur = build_norm(ffn_inp, + model.layers[il].ffn_norm, + model.layers[il].ffn_norm_b, + LLM_NORM, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + model.layers[il].ffn_up, NULL, NULL, + NULL, NULL, NULL, + model.layers[il].ffn_down, NULL, NULL, + NULL, + LLM_FFN_GELU, LLM_FFN_SEQ, il); + cb(cur, "ffn_out", il); + } + + cur = ggml_add(ctx0, cur, ffn_inp); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + // input for next layer + inpL = cur; + } + + cur = build_norm(inpL, + model.output_norm, + model.output_norm_b, + LLM_NORM, -1); + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur, model.output_s); + + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} From 915b3a1da7628ba5cd7cf33c1823f70e5c679087 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Wed, 5 Aug 2026 15:06:40 +0200 Subject: [PATCH 03/17] working impl, need verify and clean up --- conversion/pockettts.py | 23 +- gguf-py/gguf/constants.py | 5 +- tools/mtmd/CMakeLists.txt | 3 + tools/mtmd/clip-impl.h | 36 +++ tools/mtmd/clip-model.h | 83 ++++++ tools/mtmd/clip.cpp | 255 +++++++++++++++++-- tools/mtmd/models/models.h | 56 +++++ tools/mtmd/models/pockettts-gen.cpp | 286 +++++++++++++++++++++ tools/mtmd/models/pockettts-seanet.cpp | 158 ++++++++++++ tools/mtmd/models/pockettts-spkenc.cpp | 77 ++++++ tools/mtmd/models/qwen3tts-gen.cpp | 4 + tools/mtmd/mtmd-audio.cpp | 30 +++ tools/mtmd/mtmd-audio.h | 7 + tools/mtmd/mtmd-helper-gen.cpp | 335 +++++++++++++++++++++++++ tools/mtmd/mtmd.cpp | 11 + tools/mtmd/mtmd.h | 2 + 16 files changed, 1339 insertions(+), 32 deletions(-) create mode 100644 tools/mtmd/models/pockettts-gen.cpp create mode 100644 tools/mtmd/models/pockettts-seanet.cpp create mode 100644 tools/mtmd/models/pockettts-spkenc.cpp diff --git a/conversion/pockettts.py b/conversion/pockettts.py index 4aa41f23010..a62243bb226 100644 --- a/conversion/pockettts.py +++ b/conversion/pockettts.py @@ -50,21 +50,33 @@ class PocketTTSModel(TextModel): } def set_vocab(self): + # this is a unigram sentencepiece model; llama.cpp's SPM tokenizer greedily merges + # bigrams and cannot reproduce unigram segmentation, so use the UGM tokenizer instead + from sentencepiece import sentencepiece_model_pb2 as model + + proto = model.ModelProto() # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute] + proto.ParseFromString(open(self.dir_model / "tokenizer.model", "rb").read()) + assert proto.trainer_spec.model_type == 1, "expected a unigram tokenizer" + tokens, scores, toktypes = self._create_vocab_sentencepiece() # the last 3 rows of the embedding table are not sentencepiece pieces: the conditioner's - # padding row, then the two learned vectors appended by generate_extra_tensors() + # padding row, then the two learned vectors appended by _embd_table() extra = ["<|pad|>", "<|bos_before_voice|>", "<|audio_bos|>"] for i, name in enumerate(extra): tokens[len(tokens) - len(extra) + i] = name.encode("utf-8") toktypes[len(tokens) - len(extra) + i] = SentencePieceTokenTypes.CONTROL scores[len(tokens) - len(extra) + i] = -1000.0 - self.gguf_writer.add_tokenizer_model("llama") + self.gguf_writer.add_tokenizer_model("t5") self.gguf_writer.add_tokenizer_pre("default") self.gguf_writer.add_token_list(tokens) self.gguf_writer.add_token_scores(scores) self.gguf_writer.add_token_types(toktypes) + self.gguf_writer.add_add_space_prefix(proto.normalizer_spec.add_dummy_prefix) + self.gguf_writer.add_remove_extra_whitespaces(proto.normalizer_spec.remove_extra_whitespaces) + if proto.normalizer_spec.precompiled_charsmap: + self.gguf_writer.add_precompiled_charsmap(proto.normalizer_spec.precompiled_charsmap) self.gguf_writer.add_add_bos_token(False) self.gguf_writer.add_add_eos_token(False) @@ -122,7 +134,7 @@ class PocketTTSMmprojModel(MmprojModel): "linear1": (gguf.MODEL_TENSOR.A_ENC_FFN_UP, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP), "linear2": (gguf.MODEL_TENSOR.A_ENC_FFN_DOWN, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN), "layer_scale_1.scale": (gguf.MODEL_TENSOR.A_ENC_ATTN_SCALE, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE), - "layer_scale_2.scale": (gguf.MODEL_TENSOR.A_ENC_FFN_SCALE, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE), + "layer_scale_2.scale": (gguf.MODEL_TENSOR.A_ENC_FFN_SCALE_LS, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE), } _MIMI_TFM_QKV = ( (gguf.MODEL_TENSOR.A_ENC_ATTN_Q, gguf.MODEL_TENSOR.A_ENC_ATTN_K, gguf.MODEL_TENSOR.A_ENC_ATTN_V), @@ -143,6 +155,8 @@ def set_gguf_parameters(self): self.gguf_writer.add_audio_feed_forward_length(self.hparams_audio["intermediate_size"]) self.gguf_writer.add_audio_head_count(self.hparams_audio["num_attention_heads"]) self.gguf_writer.add_audio_attention_layernorm_eps(1e-5) + # mimi convolves the waveform directly, it is passed around as a 1-row "mel" + self.gguf_writer.add_audio_num_mel_bins(1) # generation: flow-matching decoder + mimi decoder # note: the SEANet and flow net hparams are hardcoded on the clip.cpp side for now @@ -273,8 +287,7 @@ def _mimi_tfm_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, if entry is None: return tensor = entry[1 if is_decoder else 0] - # layer_scale is stored without a .weight/.bias suffix - suffix = "" if key_with_suffix.endswith(".scale") else "." + suffix + suffix = ".weight" if key_with_suffix.endswith(".scale") else "." + suffix yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch) def _seanet_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 445144e13f5..991d930691f 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -1039,6 +1039,7 @@ class MODEL_TENSOR(IntEnum): A_ENC_SEANET_RES_CONV2 = auto() # residual unit, pointwise conv A_ENC_SEANET_SCALE_CONV = auto() # strided downsample conv A_ENC_ATTN_SCALE = auto() # layer scale (gamma) on the attn output + A_ENC_FFN_SCALE_LS = auto() # layer scale (gamma) on the FFN output A_ENC_SPEAKER_PROJ = auto() # voice latent -> backbone embd A_GEN_FLOW_INPUT_PROJ = auto() A_GEN_FLOW_COND_EMBD = auto() @@ -1736,7 +1737,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_ENC_SEANET_RES_CONV1: "a.seanet.blk.{bid}.res_conv1", MODEL_TENSOR.A_ENC_SEANET_RES_CONV2: "a.seanet.blk.{bid}.res_conv2", MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV: "a.seanet.blk.{bid}.scale_conv", - MODEL_TENSOR.A_ENC_ATTN_SCALE: "a.blk.{bid}.attn_scale", + MODEL_TENSOR.A_ENC_ATTN_SCALE: "a.blk.{bid}.ls1", + MODEL_TENSOR.A_ENC_FFN_SCALE_LS: "a.blk.{bid}.ls2", MODEL_TENSOR.A_ENC_SPEAKER_PROJ: "a.speaker_proj", MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ: "a.gen.flow.input_proj", MODEL_TENSOR.A_GEN_FLOW_COND_EMBD: "a.gen.flow.cond_embd", @@ -2078,6 +2080,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_ENC_SEANET_RES_CONV2, MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV, MODEL_TENSOR.A_ENC_ATTN_SCALE, + MODEL_TENSOR.A_ENC_FFN_SCALE_LS, MODEL_TENSOR.A_ENC_SPEAKER_PROJ, MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ, MODEL_TENSOR.A_GEN_FLOW_COND_EMBD, diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 4675fb9a97b..a1ff01788a8 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -56,6 +56,9 @@ add_library(mtmd models/mimo-audio.cpp models/qwen3tts-spkenc.cpp models/qwen3tts-gen.cpp + models/pockettts-seanet.cpp + models/pockettts-spkenc.cpp + models/pockettts-gen.cpp models/step3vl.cpp models/siglip.cpp models/whisper-enc.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index e1567ee5baf..0a31a7aec1a 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -246,6 +246,38 @@ #define TN_A_GEN_WAV_DAC_POST_SNAKE "a.gen.wav.dac.post_snake.%s" #define TN_A_GEN_WAV_DAC_POST_CONV "a.gen.wav.dac.post_conv.%s" +// pocket-tts +#define TN_A_SEANET_CONV_IN "a.seanet.conv_in.%s" +#define TN_A_SEANET_CONV_OUT "a.seanet.conv_out.%s" +#define TN_A_SEANET_RES_CONV1 "a.seanet.blk.%d.res_conv1.%s" +#define TN_A_SEANET_RES_CONV2 "a.seanet.blk.%d.res_conv2.%s" +#define TN_A_SEANET_SCALE_CONV "a.seanet.blk.%d.scale_conv.%s" +#define TN_A_SPEAKER_PROJ "a.speaker_proj.%s" +#define TN_A_DOWNSAMPLE_CONV "a.downsample.conv.%s" +#define TN_A_GEN_FLOW_INPUT_PROJ "a.gen.flow.input_proj.%s" +#define TN_A_GEN_FLOW_COND_EMBD "a.gen.flow.cond_embd.%s" +#define TN_A_GEN_FLOW_TIME_FREQS "a.gen.flow.time.%d.freqs" +#define TN_A_GEN_FLOW_TIME_UP "a.gen.flow.time.%d.up.%s" +#define TN_A_GEN_FLOW_TIME_DOWN "a.gen.flow.time.%d.down.%s" +#define TN_A_GEN_FLOW_TIME_NORM "a.gen.flow.time.%d.norm" +#define TN_A_GEN_FLOW_BLK_NORM "a.gen.flow.blk.%d.norm.%s" +#define TN_A_GEN_FLOW_BLK_UP "a.gen.flow.blk.%d.up.%s" +#define TN_A_GEN_FLOW_BLK_DOWN "a.gen.flow.blk.%d.down.%s" +#define TN_A_GEN_FLOW_BLK_ADA "a.gen.flow.blk.%d.ada.%s" +#define TN_A_GEN_FLOW_FINAL_ADA "a.gen.flow.final.ada.%s" +#define TN_A_GEN_FLOW_FINAL_PROJ "a.gen.flow.final.proj.%s" +#define TN_A_GEN_OUT_EOS "a.gen.out_eos.%s" +#define TN_A_GEN_INPUT_LINEAR "a.gen.input_linear.%s" +#define TN_A_GEN_EMB_MEAN "a.gen.emb_mean" +#define TN_A_GEN_EMB_STD "a.gen.emb_std" +#define TN_A_GEN_WAV_QUANT_OUT "a.gen.wav.quant_out.%s" +#define TN_A_GEN_WAV_UPSAMPLE "a.gen.wav.upsample.%s" +#define TN_A_GEN_WAV_SEANET_CONV_IN "a.gen.wav.seanet.conv_in.%s" +#define TN_A_GEN_WAV_SEANET_CONV_OUT "a.gen.wav.seanet.conv_out.%s" +#define TN_A_GEN_WAV_SEANET_RES_CONV1 "a.gen.wav.seanet.blk.%d.res_conv1.%s" +#define TN_A_GEN_WAV_SEANET_RES_CONV2 "a.gen.wav.seanet.blk.%d.res_conv2.%s" +#define TN_A_GEN_WAV_SEANET_SCALE_CONV "a.gen.wav.seanet.blk.%d.scale_conv.%s" + // cogvlm #define TN_MM_POST_FC_NORM "mm.post_fc_norm.%s" #define TN_MM_H_TO_4H "mm.up.%s" @@ -455,6 +487,8 @@ enum projector_type { PROJECTOR_TYPE_MIMO_AUDIO, PROJECTOR_TYPE_QWEN3TTS_SPKENC, PROJECTOR_TYPE_QWEN3TTS_GEN, + PROJECTOR_TYPE_POCKETTTS_SPKENC, + PROJECTOR_TYPE_POCKETTTS_GEN, PROJECTOR_TYPE_UNKNOWN, }; @@ -514,6 +548,8 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_PARAKEET, "parakeet"}, { PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"}, { PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"}, + { PROJECTOR_TYPE_POCKETTTS_SPKENC, "pockettts_spkenc"}, + { PROJECTOR_TYPE_POCKETTTS_GEN, "pockettts_gen"}, }; static projector_type clip_projector_type_from_string(const std::string & str) { diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index f3533f7eb49..9b0b2c5e2fd 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -139,6 +139,14 @@ struct clip_hparams { // threshold for the "out_eos_score" graph output float gen_eos_threshold = 0.0f; + // pocket-tts + int32_t seanet_n_stage = 0; + std::vector seanet_ratios; // encoder order (reversed compared to the config) + int32_t mimi_downsample = 0; // encoder frame rate / model frame rate + int32_t mimi_tfm_context = 0; // attention window of the mimi transformers, in frames + int32_t flow_n_step = 1; // lsd_decode steps + float flow_temp = 0.0f; // noise std is sqrt(temp) + // qwen3tts code2wav int32_t wav_tfm_n_layer = 0; int32_t wav_tfm_n_embd = 0; @@ -389,6 +397,63 @@ struct qf_block { std::vector qf_proj_layers; }; +// pocket-tts SEANet stack, used in both directions: +// encoder = conv_in -> per stage (residual unit, strided conv) -> conv_out +// decoder = conv_in -> per stage (strided convtr, residual unit) -> conv_out +struct clip_seanet { + // one residual unit: ELU -> dilated conv -> ELU -> pointwise conv, added to the input + struct stage { + ggml_tensor * res_conv1_w = nullptr; + ggml_tensor * res_conv1_b = nullptr; + ggml_tensor * res_conv2_w = nullptr; + ggml_tensor * res_conv2_b = nullptr; + ggml_tensor * scale_conv_w = nullptr; // strided conv (encoder) or convtr (decoder) + ggml_tensor * scale_conv_b = nullptr; + }; + + ggml_tensor * conv_in_w = nullptr; + ggml_tensor * conv_in_b = nullptr; + ggml_tensor * conv_out_w = nullptr; + ggml_tensor * conv_out_b = nullptr; + std::vector stages; +}; + +// pocket-tts flow-matching decoder (SimpleMLPAdaLN) +struct clip_flow_net { + // AdaLN res block: in_ln -> modulate -> Linear -> SiLU -> Linear, gated residual + struct block { + ggml_tensor * norm_w = nullptr; + ggml_tensor * norm_b = nullptr; + ggml_tensor * up_w = nullptr; + ggml_tensor * up_b = nullptr; + ggml_tensor * down_w = nullptr; + ggml_tensor * down_b = nullptr; + ggml_tensor * ada_w = nullptr; // -> shift, scale, gate + ggml_tensor * ada_b = nullptr; + }; + + // timestep embedder: cos/sin(t * freqs) -> Linear -> SiLU -> Linear -> RMSNorm + struct time_embd { + ggml_tensor * freqs = nullptr; + ggml_tensor * up_w = nullptr; + ggml_tensor * up_b = nullptr; + ggml_tensor * down_w = nullptr; + ggml_tensor * down_b = nullptr; + ggml_tensor * norm = nullptr; // RMSNorm alpha + }; + + ggml_tensor * input_proj_w = nullptr; + ggml_tensor * input_proj_b = nullptr; + ggml_tensor * cond_embd_w = nullptr; + ggml_tensor * cond_embd_b = nullptr; + ggml_tensor * final_ada_w = nullptr; // -> shift, scale + ggml_tensor * final_ada_b = nullptr; + ggml_tensor * final_proj_w = nullptr; + ggml_tensor * final_proj_b = nullptr; + std::vector time; + std::vector blocks; +}; + // qwen3tts code2wav: RVQ codes -> raw PCM struct clip_code2wav { // "upsample" stage: one ConvNeXt block plus the causal ConvTranspose1d before it @@ -686,6 +751,24 @@ struct clip_model { // qwen3tts code2wav: RVQ codes -> raw PCM clip_code2wav c2w; + // pocket-tts: SEANet stack, shared by the encoder (speaker path) and the decoder (gen path) + clip_seanet seanet; + + // pocket-tts: voice latent -> backbone embd (speaker path) + ggml_tensor * spk_proj_w = nullptr; + ggml_tensor * downsample_w = nullptr; + + // pocket-tts: flow-matching decoder, backbone hidden state -> next latent + clip_flow_net flow; + ggml_tensor * gen_out_eos_w = nullptr; + ggml_tensor * gen_out_eos_b = nullptr; + ggml_tensor * gen_input_lin_w = nullptr; // latent -> backbone embd + ggml_tensor * gen_emb_mean = nullptr; + ggml_tensor * gen_emb_std = nullptr; + ggml_tensor * gen_quant_out_w = nullptr; // latent -> decoder dim + ggml_tensor * gen_upsample_w = nullptr; // depthwise convtr, frame rate -> encoder frame rate + std::vector gen_tfm_layers; // mimi decoder_transformer + // cogvlm ggml_tensor * mm_post_fc_norm_w = nullptr; ggml_tensor * mm_post_fc_norm_b = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index fcf022b8d28..bcb7cb083ea 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1058,6 +1058,18 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + builder = std::make_unique(ctx, img); + } break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE; + const int n_step = params && params->n_steps > 0 ? params->n_steps : ctx->model.hparams.flow_n_step; + const int64_t n_latent = ctx->model.gen_input_lin_w->ne[0]; + const int n_frames = params && params->feats ? (int) (params->feats->size() / n_latent) : 1; + builder = std::make_unique(ctx, img, gen_process, n_step, n_frames); + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE; @@ -1730,6 +1742,23 @@ struct clip_model_loader { // matches the reference decoder's sliding_window (speech_tokenizer/config.json) hparams.wav_tfm_swa = 72; } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + // mimi front-end takes the raw waveform, no mel + hparams.audio_sample_rate = 24000; + // seanet ratios are [6,5,4] in the config, the encoder reverses them + hparams.seanet_ratios = { 4, 5, 6 }; + hparams.seanet_n_stage = (int32_t) hparams.seanet_ratios.size(); + hparams.mimi_downsample = 16; + // matches the reference transformer's "context" + hparams.mimi_tfm_context = 250; + hparams.rope_theta = 10000.0f; + // flow_lm defaults, see pocket_tts/default_parameters.py and the language config + hparams.flow_n_step = 1; + hparams.flow_temp = 0.3f; + hparams.gen_eos_threshold = -4.0f; + } break; case PROJECTOR_TYPE_PADDLEOCR: { hparams.n_merge = 2; @@ -1925,7 +1954,9 @@ struct clip_model_loader { // GEMMA4UA is encoder-free: it uses n_mel_bins as a raw-waveform frame size (640) and has no FFT/filterbank, so the mel-range and FFT // checks below do not apply to it. - const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA; + // pocket-tts is encoder-free in the same sense: mimi convolves the raw waveform + const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA && + model.proj_type != PROJECTOR_TYPE_POCKETTTS_SPKENC; // Validate audio hparams loaded from GGUF metadata if (hparams.n_mel_bins <= 0 || (fft_based && hparams.n_mel_bins > 256)) { @@ -1998,6 +2029,31 @@ struct clip_model_loader { return cur; }; + // pocket-tts: the encoder and the decoder share the same layout, only the prefix differs + auto load_seanet = [&](clip_seanet & seanet, bool is_decoder) { + const char * conv_in = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_IN : TN_A_SEANET_CONV_IN; + const char * conv_out = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_OUT : TN_A_SEANET_CONV_OUT; + const char * res1 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV1 : TN_A_SEANET_RES_CONV1; + const char * res2 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV2 : TN_A_SEANET_RES_CONV2; + const char * scale = is_decoder ? TN_A_GEN_WAV_SEANET_SCALE_CONV : TN_A_SEANET_SCALE_CONV; + + seanet.conv_in_w = get_tensor(string_format(conv_in, "weight")); + seanet.conv_in_b = get_tensor(string_format(conv_in, "bias")); + seanet.conv_out_w = get_tensor(string_format(conv_out, "weight")); + seanet.conv_out_b = get_tensor(string_format(conv_out, "bias")); + + seanet.stages.resize(hparams.seanet_n_stage); + for (int i = 0; i < hparams.seanet_n_stage; i++) { + auto & stage = seanet.stages[i]; + stage.res_conv1_w = get_tensor(string_format(res1, i, "weight")); + stage.res_conv1_b = get_tensor(string_format(res1, i, "bias")); + stage.res_conv2_w = get_tensor(string_format(res2, i, "weight")); + stage.res_conv2_b = get_tensor(string_format(res2, i, "bias")); + stage.scale_conv_w = get_tensor(string_format(scale, i, "weight")); + stage.scale_conv_b = get_tensor(string_format(scale, i, "bias")); + } + }; + auto get_vector = [&](const std::string & name) { std::vector result; auto it = tensor_offset.find(name); @@ -2059,7 +2115,8 @@ struct clip_model_loader { const bool has_standard_layers = ( model.proj_type != PROJECTOR_TYPE_GEMMA3NV && - model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC); + model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC && + model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN); // layers const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0; @@ -2726,6 +2783,81 @@ struct clip_model_loader { model.mm_fc_w = get_tensor(string_format(TN_MM_AUDIO_FC, "weight")); model.mm_fc_b = get_tensor(string_format(TN_MM_AUDIO_FC, "bias")); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + load_seanet(model.seanet, false); + model.downsample_w = get_tensor(string_format(TN_A_DOWNSAMPLE_CONV, "weight")); + model.spk_proj_w = get_tensor(string_format(TN_A_SPEAKER_PROJ, "weight")); + } break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + auto & flow = model.flow; + flow.input_proj_w = get_tensor(string_format(TN_A_GEN_FLOW_INPUT_PROJ, "weight")); + flow.input_proj_b = get_tensor(string_format(TN_A_GEN_FLOW_INPUT_PROJ, "bias")); + flow.cond_embd_w = get_tensor(string_format(TN_A_GEN_FLOW_COND_EMBD, "weight")); + flow.cond_embd_b = get_tensor(string_format(TN_A_GEN_FLOW_COND_EMBD, "bias")); + flow.final_ada_w = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_ADA, "weight")); + flow.final_ada_b = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_ADA, "bias")); + flow.final_proj_w = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_PROJ, "weight")); + flow.final_proj_b = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_PROJ, "bias")); + + flow.time.resize(2); + for (size_t i = 0; i < flow.time.size(); i++) { + auto & t = flow.time[i]; + t.freqs = get_tensor(string_format(TN_A_GEN_FLOW_TIME_FREQS, (int) i)); + t.up_w = get_tensor(string_format(TN_A_GEN_FLOW_TIME_UP, (int) i, "weight")); + t.up_b = get_tensor(string_format(TN_A_GEN_FLOW_TIME_UP, (int) i, "bias")); + t.down_w = get_tensor(string_format(TN_A_GEN_FLOW_TIME_DOWN, (int) i, "weight")); + t.down_b = get_tensor(string_format(TN_A_GEN_FLOW_TIME_DOWN, (int) i, "bias")); + t.norm = get_tensor(string_format(TN_A_GEN_FLOW_TIME_NORM, (int) i)); + } + + // one AdaLN block per flow depth, the count is only known from the tensors + for (int il = 0; ; il++) { + ggml_tensor * probe = get_tensor(string_format(TN_A_GEN_FLOW_BLK_NORM, il, "weight"), false); + if (probe == nullptr) { + break; + } + clip_flow_net::block blk; + blk.norm_w = probe; + blk.norm_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_NORM, il, "bias")); + blk.up_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_UP, il, "weight")); + blk.up_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_UP, il, "bias")); + blk.down_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_DOWN, il, "weight")); + blk.down_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_DOWN, il, "bias")); + blk.ada_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_ADA, il, "weight")); + blk.ada_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_ADA, il, "bias")); + flow.blocks.push_back(blk); + } + + model.gen_out_eos_w = get_tensor(string_format(TN_A_GEN_OUT_EOS, "weight")); + model.gen_out_eos_b = get_tensor(string_format(TN_A_GEN_OUT_EOS, "bias")); + model.gen_input_lin_w = get_tensor(string_format(TN_A_GEN_INPUT_LINEAR, "weight")); + model.gen_emb_mean = get_tensor(TN_A_GEN_EMB_MEAN); + model.gen_emb_std = get_tensor(TN_A_GEN_EMB_STD); + + // mimi decoder + model.gen_quant_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_OUT, "weight")); + model.gen_upsample_w = get_tensor(string_format(TN_A_GEN_WAV_UPSAMPLE, "weight")); + load_seanet(model.seanet, true); + model.gen_tfm_layers.resize(hparams.n_layer); + for (int il = 0; il < hparams.n_layer; il++) { + auto & layer = model.gen_tfm_layers[il]; + const char * p = "a.gen.wav.tfm"; + layer.ln_1_w = get_tensor(string_format(TN_LN_1, p, il, "weight")); + layer.ln_1_b = get_tensor(string_format(TN_LN_1, p, il, "bias")); + layer.q_w = get_tensor(string_format(TN_ATTN_Q, p, il, "weight")); + layer.k_w = get_tensor(string_format(TN_ATTN_K, p, il, "weight")); + layer.v_w = get_tensor(string_format(TN_ATTN_V, p, il, "weight")); + layer.o_w = get_tensor(string_format(TN_ATTN_OUTPUT, p, il, "weight")); + layer.ls_1_w = get_tensor(string_format(TN_LS_1, p, il, "weight")); + layer.ln_2_w = get_tensor(string_format(TN_LN_2, p, il, "weight")); + layer.ln_2_b = get_tensor(string_format(TN_LN_2, p, il, "bias")); + layer.ff_up_w = get_tensor(string_format(TN_FFN_UP, p, il, "weight")); + layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN, p, il, "weight")); + layer.ls_2_w = get_tensor(string_format(TN_LS_2, p, il, "weight")); + } + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { // code_predictor @@ -4028,6 +4160,17 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { // one hidden-state vector fed back to the talker per call n_patches = 1; } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + // one conditioning row per 12.5Hz frame + const int hop = ctx->model.hparams.mimi_downsample * 120; + n_patches = img->nx() / hop; + } break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + // one latent per call for GEN_CODE, GEN_WAV sizes its input from the caller + n_patches = 1; + } break; case PROJECTOR_TYPE_GRANITE4_VISION: { // Per-tile output token count: each projector block outputs @@ -4069,6 +4212,15 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 return clip_encode(ctx, ¶ms); } +// persisted state slots of the gen-audio decoder, per pipeline +static std::vector list_gen_state_slots(const clip_hparams & hparams, const clip_model & model) { + switch (model.proj_type) { + case PROJECTOR_TYPE_QWEN3TTS_GEN: return list_c2w_state_slots(hparams, model); + case PROJECTOR_TYPE_POCKETTTS_GEN: return list_pockettts_state_slots(hparams, model); + default: return {}; + } +} + bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { const clip_image_f32_batch & imgs = *params->imgs; int n_batch_cur = imgs.entries.size(); @@ -4133,6 +4285,45 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { ggml_backend_tensor_set(cur, values.data(), 0, ggml_nbytes(cur)); }; + // upload the decoder state from the previous call, or zero-fill on a cold start + auto set_gen_state_in = [&]() { + size_t offset = 0; + for (const auto & slot : list_gen_state_slots(hparams, model)) { + ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str()); + const size_t nb = ggml_nbytes(t); + if (params->state_in && params->state_in->size() >= offset + nb) { + ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb); + } else { + std::vector zeros(nb, 0); + ggml_backend_tensor_set(t, zeros.data(), 0, nb); + } + offset += nb; + } + }; + + // rope positions and attention mask of the mimi transformers (pocket-tts). + // the mask is causal with a sliding window, see _build_attention_mask() in the reference + auto set_pockettts_tfm_inputs = [&]() { + const int64_t n_pos = ggml_nelements(get_inp_tensor("inp_pos")); + std::vector positions((size_t) n_pos); + for (int64_t i = 0; i < n_pos; i++) { + positions[(size_t) i] = (int32_t) i; + } + set_input_i32("inp_pos", positions); + + const int64_t context = hparams.mimi_tfm_context; + std::vector mask((size_t) n_pos * n_pos, -INFINITY); + for (int64_t q = 0; q < n_pos; q++) { + for (int64_t k = 0; k < n_pos; k++) { + const int64_t delta = q - k; + if (delta >= 0 && delta < context) { + mask[(size_t) q * n_pos + k] = 0.0f; + } + } + } + set_input_f32("kq_mask", mask); + }; + // set input pixel values if (!imgs.is_audio) { size_t nelem = 0; @@ -4176,8 +4367,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } set_input_f32("inp_raw", inp_raw); - } else if (!(ctx->proj_type() == PROJECTOR_TYPE_QWEN3TTS_GEN && params->gen_process == CLIP_GEN_PROCESS_GEN_WAV)) { - // audio input, code2wav is not here: its only input is "inp_codes", set in the switch below + } else if (params->gen_process != CLIP_GEN_PROCESS_GEN_WAV) { + // audio input. GEN_WAV is not here: it takes codes or feats, set in the switch below GGML_ASSERT(imgs.entries.size() == 1); const auto & mel_inp = imgs.entries[0]; @@ -4646,6 +4837,28 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } set_input_i32("patches", patches); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + set_pockettts_tfm_inputs(); + } break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) { + GGML_ASSERT(params->feats != nullptr); + set_input_f32("inp_feats", *params->feats); + // positions and mask are derived in-graph from the persisted counter + set_gen_state_in(); + } else { + // flow matching starts from gaussian noise, std = sqrt(temp) + ggml_tensor * t = get_inp_tensor("inp_noise"); + std::normal_distribution dist(0.0f, std::sqrt(hparams.flow_temp)); + std::vector noise(ggml_nelements(t)); + for (auto & v : noise) { + v = dist(ctx->rng); + } + set_input_f32("inp_noise", noise); + } + } break; case PROJECTOR_TYPE_GEMMA4V: case PROJECTOR_TYPE_GEMMA4UV: { @@ -4770,20 +4983,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } } set_input_i32("inp_codes", codes); - - // upload the state from the previous call, or zero-fill on a cold start - size_t offset = 0; - for (const auto & slot : list_c2w_state_slots(hparams, model)) { - ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str()); - const size_t nb = ggml_nbytes(t); - if (params->state_in && params->state_in->size() >= offset + nb) { - ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb); - } else { - std::vector zeros(nb, 0); - ggml_backend_tensor_set(t, zeros.data(), 0, nb); - } - offset += nb; - } + set_gen_state_in(); } else { // code0 indexes gen_code_out_embd_w via ggml_get_rows; bound it const int64_t vocab0 = model.gen_code_out_embd_w->ne[1]; @@ -5251,16 +5451,15 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { // for audio gen models // + // optional outputs: a pipeline yields codes or feats, and not all have an eos head if (params->out_codes != nullptr) { ggml_tensor * codes = ggml_graph_get_tensor(gf, "out_codes"); - if (codes == nullptr) { - GGML_ABORT("out_codes requested but graph has no \"out_codes\" tensor"); + if (codes != nullptr) { + auto & out_codes = *params->out_codes; + out_codes.resize(ggml_nelements(codes)); + ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes)); } - auto & out_codes = *params->out_codes; - out_codes.resize(ggml_nelements(codes)); - ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes)); } - // optional outputs, a missing tensor is not an error if (params->out_feats != nullptr) { ggml_tensor * feats = ggml_graph_get_tensor(gf, "out_feats"); if (feats != nullptr) { @@ -5299,12 +5498,12 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { if (params->state_out != nullptr) { auto & state_out = *params->state_out; size_t total = 0; - for (const auto & slot : list_c2w_state_slots(hparams, model)) { + for (const auto & slot : list_gen_state_slots(hparams, model)) { total += (size_t) (slot.ne0 * slot.ne1) * sizeof(float); } state_out.resize(total); size_t offset = 0; - for (const auto & slot : list_c2w_state_slots(hparams, model)) { + for (const auto & slot : list_gen_state_slots(hparams, model)) { ggml_tensor * t = ggml_graph_get_tensor(gf, ("state_out_" + slot.name).c_str()); if (t == nullptr) { GGML_ABORT("state_out requested but graph has no \"state_out_%s\" tensor", slot.name.c_str()); @@ -5450,6 +5649,10 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_fc_w->ne[2]; case PROJECTOR_TYPE_QWEN3TTS_GEN: return ctx->model.gen_code_out_embd_w->ne[0]; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + return ctx->model.spk_proj_w->ne[1]; + case PROJECTOR_TYPE_POCKETTTS_GEN: + return ctx->model.gen_input_lin_w->ne[1]; case PROJECTOR_TYPE_PARAKEET: return ctx->model.mm_1_w->ne[1]; default: diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index eb924972bf5..1663c6620e8 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -317,6 +317,59 @@ struct clip_graph_qwen3tts_gen : clip_graph { }; }; +// +// pocket-tts: SEANet convolution stack, shared by the voice encoder and the mimi decoder. +// stateless unless state_in is populated: convs then pad instead of carrying left-context. +// +struct clip_graph_pockettts_seanet : clip_graph { + clip_graph_pockettts_seanet(const clip_graph & parent) : clip_graph(parent) {} + ggml_cgraph * build() override { GGML_ABORT("call encode()/decode() instead"); } + + // per-call streaming state, keyed by slot name (see list_pockettts_state_slots) + std::map state_in; + mutable std::vector> state_out; + + ggml_tensor * conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation, + bool pad_replicate = false, const std::string & state_name = "") const; + ggml_tensor * conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, + const std::string & state_name = "") const; + ggml_tensor * res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation, + const std::string & state_prefix = "") const; + + // x: [T, C] -> [T / hop, dim] + ggml_tensor * encode(ggml_tensor * x) const; + // x: [T, dim] -> [T * hop, 1], streams when state_in is populated + ggml_tensor * decode(ggml_tensor * x) const; +}; + +// mimi encoder + speaker_proj: reference waveform -> voice conditioning rows +struct clip_graph_pockettts_spkenc : clip_graph { + clip_graph_pockettts_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; + + ggml_tensor * tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const; +}; + +// +// pocket-tts generation: +// GEN_CODE = flow-matching decoder + end-of-speech head, one latent per call +// GEN_WAV = mimi decoder, a window of latents -> PCM +// +struct clip_graph_pockettts_gen : clip_graph { + clip_graph_pockettts_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int n_step, int n_frames) + : clip_graph(ctx, img), gen_process(gen_process), n_step(n_step), n_frames(n_frames) {} + ggml_cgraph * build() override; + + clip_gen_process_type gen_process; + int n_step; // lsd_decode steps, fixed at graph-build time + int n_frames; // GEN_WAV only: number of latents to decode + + // AdaLN modulation: x * (1 + scale) + shift + ggml_tensor * modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const; + ggml_tensor * time_embed(const clip_flow_net::time_embd & te, float t) const; + ggml_tensor * flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const; +}; + // one persisted state buffer used by code2wav, see qwen3tts-gen.cpp struct c2w_state_slot { std::string name; @@ -325,6 +378,9 @@ struct c2w_state_slot { }; std::vector list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model); +// same, for the streaming mimi decoder (pocket-tts GEN_WAV) +std::vector list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model); + struct clip_graph_kimik25 : clip_graph { clip_graph_kimik25(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/models/pockettts-gen.cpp b/tools/mtmd/models/pockettts-gen.cpp new file mode 100644 index 00000000000..b76d44da387 --- /dev/null +++ b/tools/mtmd/models/pockettts-gen.cpp @@ -0,0 +1,286 @@ +#include "models.h" + +#include + +// pocket-tts generation stages +// +// GEN_CODE: backbone hidden state -> next 32-d latent (flow matching) + end-of-speech score +// GEN_WAV : a window of latents -> PCM, through the mimi decoder +// +// there is no codebook anywhere, "codes" in the mtmd API are continuous features here + +// x * (1 + scale) + shift, all [D, 1] +ggml_tensor * clip_graph_pockettts_gen::modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const { + ggml_tensor * cur = ggml_mul(ctx0, x, ggml_scale_bias(ctx0, scale, 1.0f, 1.0f)); + return ggml_add(ctx0, cur, shift); +} + +// cos/sin(t * freqs) -> Linear -> SiLU -> Linear -> RMSNorm, see TimestepEmbedder +ggml_tensor * clip_graph_pockettts_gen::time_embed(const clip_flow_net::time_embd & te, float t) const { + // t is a graph-build constant, so the cos/sin table can be folded into a scaled copy + ggml_tensor * args = ggml_scale(ctx0, te.freqs, t); + ggml_tensor * emb = ggml_concat(ctx0, ggml_cos(ctx0, args), ggml_sin(ctx0, args), 0); + + ggml_tensor * cur = build_mm(te.up_w, emb); + cur = ggml_add(ctx0, cur, te.up_b); + cur = ggml_silu(ctx0, cur); + cur = build_mm(te.down_w, cur); + cur = ggml_add(ctx0, cur, te.down_b); + + // this "RMSNorm" divides by the unbiased variance, not the mean square, and it rescales + // the input rather than the centered value, see _rms_norm() in mlp.py + { + const int64_t n = cur->ne[0]; + ggml_tensor * mean = ggml_mean(ctx0, cur); + ggml_tensor * dev = ggml_sub(ctx0, cur, mean); + ggml_tensor * var = ggml_mean(ctx0, ggml_sqr(ctx0, dev)); + var = ggml_scale_bias(ctx0, var, (float) n / (float) (n - 1), 1e-5f); + cur = ggml_div(ctx0, cur, ggml_sqrt(ctx0, var)); + cur = ggml_mul(ctx0, cur, te.norm); + } + + return cur; +} + +// one velocity evaluation: v(cond, s, t, x) +ggml_tensor * clip_graph_pockettts_gen::flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const { + const auto & flow = model.flow; + + ggml_tensor * cur = build_mm(flow.input_proj_w, x); + cur = ggml_add(ctx0, cur, flow.input_proj_b); + + // the two time conditions are averaged, then added to the projected backbone state + ggml_tensor * ts = ggml_add(ctx0, time_embed(flow.time[0], s), time_embed(flow.time[1], t)); + ts = ggml_scale(ctx0, ts, 1.0f / (float) flow.time.size()); + + ggml_tensor * c = build_mm(flow.cond_embd_w, cond); + c = ggml_add(ctx0, c, flow.cond_embd_b); + + ggml_tensor * y = ggml_add(ctx0, ts, c); + cb(y, "flow_cond", -1); + + const int64_t n_ch = flow.blocks.empty() ? 0 : flow.blocks[0].norm_w->ne[0]; + + for (size_t il = 0; il < flow.blocks.size(); il++) { + const auto & blk = flow.blocks[il]; + + ggml_tensor * mod = build_mm(blk.ada_w, ggml_silu(ctx0, y)); + mod = ggml_add(ctx0, mod, blk.ada_b); + + ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0); + ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]); + ggml_tensor * gate = ggml_view_1d(ctx0, mod, n_ch, (size_t) 2 * n_ch * mod->nb[0]); + + ggml_tensor * h = build_norm(cur, blk.norm_w, blk.norm_b, NORM_TYPE_NORMAL, 1e-6f, (int) il); + h = modulate(h, shift, scale); + h = build_mm(blk.up_w, h); + h = ggml_add(ctx0, h, blk.up_b); + h = ggml_silu(ctx0, h); + h = build_mm(blk.down_w, h); + h = ggml_add(ctx0, h, blk.down_b); + + cur = ggml_add(ctx0, cur, ggml_mul(ctx0, gate, h)); + cb(cur, "flow_blk", (int) il); + } + + // final layer: the norm has no weights, only the AdaLN modulation + ggml_tensor * mod = build_mm(flow.final_ada_w, ggml_silu(ctx0, y)); + mod = ggml_add(ctx0, mod, flow.final_ada_b); + + ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0); + ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]); + + cur = build_norm(cur, nullptr, nullptr, NORM_TYPE_NORMAL, 1e-6f, -1); + cur = modulate(cur, shift, scale); + cur = build_mm(flow.final_proj_w, cur); + cur = ggml_add(ctx0, cur, flow.final_proj_b); + + return cur; +} + +// state carried between GEN_WAV calls: rope offset, per-layer KV window, conv left context +// and the transposed-conv overlap tails. shape lookup only, no graph needed +std::vector list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model) { + std::vector slots; + if (model.gen_upsample_w == nullptr) { + return slots; // not a pocket-tts decoder + } + const auto & seanet = model.seanet; + + slots.push_back({"tfm_pos", 1, 1}); + + const int64_t n_embd_a = model.gen_tfm_layers[0].q_w->ne[1]; + const int64_t prefix = hparams.mimi_tfm_context - 1; + for (size_t il = 0; il < model.gen_tfm_layers.size(); il++) { + slots.push_back({"tfm_k_" + std::to_string(il), n_embd_a, prefix}); + slots.push_back({"tfm_v_" + std::to_string(il), n_embd_a, prefix}); + } + + // upsample is depthwise, its output channel count is the input one + slots.push_back({"up", model.gen_upsample_w->ne[0] - hparams.mimi_downsample, model.gen_upsample_w->ne[2]}); + + slots.push_back({"dec_in", seanet.conv_in_w->ne[0] - 1, seanet.conv_in_w->ne[1]}); + for (int i = 0; i < hparams.seanet_n_stage; i++) { + const auto & stage = seanet.stages[i]; + const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i]; + slots.push_back({"dec_up_" + std::to_string(i), stage.scale_conv_w->ne[0] - stride, stage.scale_conv_w->ne[1]}); + slots.push_back({"dec_res_" + std::to_string(i), stage.res_conv1_w->ne[0] - 1, stage.res_conv1_w->ne[1]}); + } + slots.push_back({"dec_out", seanet.conv_out_w->ne[0] - 1, seanet.conv_out_w->ne[1]}); + + return slots; +} + +ggml_cgraph * clip_graph_pockettts_gen::build() { + if (gen_process == CLIP_GEN_PROCESS_GEN_CODE) { + // the backbone hidden state arrives as the single batch entry + ggml_tensor * h_state = build_inp_raw(1); + h_state = ggml_reshape_2d(ctx0, h_state, n_mmproj_embd, 1); + + // end-of-speech probe, thresholded on the host side + ggml_tensor * eos = build_mm(model.gen_out_eos_w, h_state); + eos = ggml_add(ctx0, eos, model.gen_out_eos_b); + ggml_set_name(eos, "out_eos_score"); + ggml_set_output(eos); + ggml_build_forward_expand(gf, eos); + + const int64_t n_latent = model.gen_input_lin_w->ne[0]; + + ggml_tensor * noise = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_latent, 1); + ggml_set_name(noise, "inp_noise"); + ggml_set_input(noise); + + // lsd_decode: integrate the velocity field from the noise sample + ggml_tensor * cur = noise; + for (int i = 0; i < n_step; i++) { + const float s = (float) i / (float) n_step; + const float t = (float) (i + 1) / (float) n_step; + ggml_tensor * v = flow_forward(h_state, cur, s, t); + cur = ggml_add(ctx0, cur, ggml_scale(ctx0, v, 1.0f / (float) n_step)); + } + cb(cur, "flow_latent", -1); + + ggml_set_name(cur, "out_feats"); + ggml_set_output(cur); + ggml_build_forward_expand(gf, cur); + + // the same latent, projected into the backbone's input space for the next step + ggml_tensor * embd = build_mm(model.gen_input_lin_w, cur); + cb(embd, "gen_embd", -1); + ggml_build_forward_expand(gf, embd); + + return gf; + } + + // GEN_WAV: [32, n_frames] latents -> PCM + ggml_tensor * feats = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, + model.gen_input_lin_w->ne[0], n_frames); + ggml_set_name(feats, "inp_feats"); + ggml_set_input(feats); + + // denormalize, then the DummyQuantizer up-projection + ggml_tensor * cur = ggml_add(ctx0, ggml_mul(ctx0, feats, model.gen_emb_std), model.gen_emb_mean); + cur = build_mm(model.gen_quant_out_w, cur); + cb(cur, "quant_out", -1); + + clip_graph_pockettts_seanet seanet(*this); + for (const auto & slot : list_pockettts_state_slots(hparams, model)) { + ggml_tensor * t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1); + ggml_set_name(t, ("state_in_" + slot.name).c_str()); + ggml_set_input(t); + seanet.state_in[slot.name] = t; + } + + // model frame rate -> encoder frame rate, depthwise transposed conv + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = seanet.conv_transpose1d(cur, model.gen_upsample_w, nullptr, hparams.mimi_downsample, "up"); + cb(cur, "mimi_upsample", -1); + + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + + // positions continue across calls, the counter lives in the state + const int64_t n_pos = cur->ne[1]; + const int64_t prefix = hparams.mimi_tfm_context - 1; + const int64_t n_kv = prefix + n_pos; + + ggml_tensor * base = ggml_reshape_1d(ctx0, seanet.state_in.at("tfm_pos"), 1); + ggml_tensor * inp_pos = ggml_cast(ctx0, ggml_add(ctx0, ggml_arange(ctx0, 0.0f, (float) n_pos, 1.0f), base), + GGML_TYPE_I32); + seanet.state_out.push_back({"tfm_pos", ggml_scale_bias(ctx0, seanet.state_in.at("tfm_pos"), 1.0f, (float) n_pos)}); + + // banded causal mask over [cached prefix | this chunk], and a cold-start mask for the + // cache rows that hold no real frame yet + ggml_tensor * pos_k = ggml_reshape_2d(ctx0, ggml_arange(ctx0, 0.0f, (float) n_kv, 1.0f), n_kv, 1); + ggml_tensor * pos_q = ggml_reshape_2d(ctx0, ggml_arange(ctx0, (float) prefix, (float) (prefix + n_pos), 1.0f), 1, n_pos); + ggml_tensor * diff = ggml_sub(ctx0, ggml_repeat_4d(ctx0, pos_q, n_kv, n_pos, 1, 1), pos_k); + + ggml_tensor * keep = ggml_mul(ctx0, + ggml_step(ctx0, ggml_scale_bias(ctx0, diff, 1.0f, 0.5f)), // delta >= 0 + ggml_step(ctx0, ggml_scale_bias(ctx0, diff, -1.0f, (float) hparams.mimi_tfm_context - 0.5f))); // delta < context + keep = ggml_mul(ctx0, keep, + ggml_step(ctx0, ggml_scale_bias(ctx0, ggml_add(ctx0, pos_k, base), 1.0f, 0.5f - (float) prefix))); + ggml_tensor * kq_mask = ggml_reshape_4d(ctx0, ggml_log(ctx0, keep), n_kv, n_pos, 1, 1); + + for (int il = 0; il < n_layer; il++) { + const auto & layer = model.gen_tfm_layers[il]; + ggml_tensor * inp = cur; + + cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il); + + ggml_tensor * Qcur = build_mm(layer.q_w, cur); + ggml_tensor * Kcur = build_mm(layer.k_w, cur); + ggml_tensor * Vcur = build_mm(layer.v_w, cur); + + Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos); + Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos); + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, + hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, + hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + // prepend the cached window, then keep this chunk's tail for the next call + const std::string k_name = "tfm_k_" + std::to_string(il); + const std::string v_name = "tfm_v_" + std::to_string(il); + ggml_tensor * k_full = ggml_concat(ctx0, seanet.state_in.at(k_name), + ggml_reshape_2d(ctx0, Kcur, d_head * n_head, n_pos), 1); + ggml_tensor * v_full = ggml_concat(ctx0, seanet.state_in.at(v_name), Vcur, 1); + seanet.state_out.push_back({k_name, ggml_cont(ctx0, ggml_view_2d(ctx0, k_full, k_full->ne[0], prefix, + k_full->nb[1], (size_t) n_pos * k_full->nb[1]))}); + seanet.state_out.push_back({v_name, ggml_cont(ctx0, ggml_view_2d(ctx0, v_full, v_full->ne[0], prefix, + v_full->nb[1], (size_t) n_pos * v_full->nb[1]))}); + + ggml_tensor * q_cur = ggml_reshape_4d(ctx0, Qcur, d_head, n_head, n_pos, 1); + ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k_full, d_head, n_head, n_kv, 1); + ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v_full, d_head, n_head, n_kv, 1); + + cur = build_attn(layer.o_w, nullptr, q_cur, k_cur, v_cur, kq_mask, kq_scale, il); + cur = ggml_mul(ctx0, cur, layer.ls_1_w); + cur = ggml_add(ctx0, cur, inp); + + inp = cur; + cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il); + cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il); + cur = ggml_mul(ctx0, cur, layer.ls_2_w); + cur = ggml_add(ctx0, cur, inp); + } + cb(cur, "mimi_dec_tfm", -1); + + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = seanet.decode(cur); + + for (const auto & s : seanet.state_out) { + ggml_set_name(s.second, ("state_out_" + s.first).c_str()); + ggml_set_output(s.second); + ggml_build_forward_expand(gf, s.second); + } + + // [n_samples, 1] -> [n_samples], clamped like the reference output + cur = ggml_reshape_1d(ctx0, cur, cur->ne[0]); + cur = ggml_clamp(ctx0, cur, -1.0f, 1.0f); + ggml_set_name(cur, "out_audio"); + ggml_set_output(cur); + ggml_build_forward_expand(gf, cur); + + return gf; +} diff --git a/tools/mtmd/models/pockettts-seanet.cpp b/tools/mtmd/models/pockettts-seanet.cpp new file mode 100644 index 00000000000..02c84c14342 --- /dev/null +++ b/tools/mtmd/models/pockettts-seanet.cpp @@ -0,0 +1,158 @@ +#include "models.h" + +// SEANet convolution stack of the mimi codec, see pocket_tts/modules/seanet.py +// +// tensors are T-first here: [T, C]. the convs are causal: they take left context from a +// state slot when given, otherwise they pad (cold start / one-shot encode) + +static int64_t div_ceil(int64_t a, int64_t b) { + return a / b + (a % b ? 1 : 0); +} + +// x: [T, IC], w: [K, IC, OC] -> [T / stride, OC] +// the convs are causal, so the whole K - stride padding goes on the left +ggml_tensor * clip_graph_pockettts_seanet::conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation, + bool pad_replicate, const std::string & state_name) const { + const int64_t k_size = (w->ne[0] - 1) * dilation + 1; + const int64_t p_total = k_size - stride; + + // trailing padding so the last frame is not dropped, see pad_for_conv1d() in conv.py + const int64_t n_frames = div_ceil(x->ne[0] - k_size + p_total, stride); + const int64_t ideal_len = n_frames * stride + k_size - p_total; + const int64_t p_extra = ideal_len - x->ne[0]; + + if (!state_name.empty() && p_total > 0) { + // streaming: the left context is the tail of the previous call + ggml_tensor * left = state_in.at(state_name); // [p_total, IC] + x = ggml_concat(ctx0, left, x, 0); + state_out.push_back({state_name, + ggml_cont(ctx0, ggml_view_2d(ctx0, x, p_total, x->ne[1], x->nb[1], + (size_t) (x->ne[0] - p_total) * x->nb[0]))}); + } else if (pad_replicate && p_total > 0) { + // the resamplers repeat the first frame instead of zero-padding + ggml_tensor * first = ggml_view_2d(ctx0, x, 1, x->ne[1], x->nb[1], 0); + ggml_tensor * left = ggml_repeat_4d(ctx0, first, p_total, x->ne[1], 1, 1); + x = ggml_concat(ctx0, left, x, 0); + x = ggml_pad_ext(ctx0, x, 0, p_extra, 0, 0, 0, 0, 0, 0); + } else { + x = ggml_pad_ext(ctx0, x, p_total, p_extra, 0, 0, 0, 0, 0, 0); + } + + ggml_tensor * y = ggml_conv_1d(ctx0, w, x, stride, 0, dilation); + y = ggml_reshape_2d(ctx0, y, y->ne[0], y->ne[1]); + if (b) { + y = ggml_add(ctx0, y, ggml_reshape_2d(ctx0, b, 1, b->ne[0])); + } + return y; +} + +// x: [T, IC], w: [K, OC/groups, IC] -> [T * stride, OC] +// the K - stride overlap tail belongs to the next call: it is added to the head of the next +// output when streaming, and simply dropped otherwise +ggml_tensor * clip_graph_pockettts_seanet::conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, + const std::string & state_name) const { + const int64_t p_total = w->ne[0] - stride; + const bool depthwise = w->ne[1] == 1 && w->ne[2] > 1; + const int64_t emit_len = x->ne[0] * stride; + + ggml_tensor * full = nullptr; + if (depthwise) { + // one group per channel, ggml_conv_transpose_1d has no grouped mode + for (int64_t ir = 0; ir < x->ne[1]; ir++) { + ggml_tensor * row = ggml_view_1d(ctx0, x, x->ne[0], ir * x->ne[0] * ggml_element_size(x)); + ggml_tensor * krn = ggml_view_1d(ctx0, w, w->ne[0], ir * w->ne[0] * ggml_element_size(w)); + row = ggml_conv_transpose_1d(ctx0, krn, row, stride, 0, 1); + full = full ? ggml_concat(ctx0, full, row, 1) : row; + } + } else { + full = ggml_conv_transpose_1d(ctx0, w, x, stride, 0, 1); + } + full = ggml_cont(ctx0, full); // [emit_len + p_total, OC] + + ggml_tensor * out; + if (state_name.empty() || p_total == 0) { + out = ggml_cont(ctx0, ggml_view_2d(ctx0, full, emit_len, full->ne[1], full->nb[1], 0)); + } else { + // overlap-add the tail the previous call held back + ggml_tensor * prev = state_in.at(state_name); // [p_total, OC] + ggml_tensor * head = ggml_add(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1], 0), prev); + if (emit_len > p_total) { + ggml_tensor * rest = ggml_view_2d(ctx0, full, emit_len - p_total, full->ne[1], full->nb[1], + (size_t) p_total * full->nb[0]); + out = ggml_concat(ctx0, head, rest, 0); + } else { + out = head; + } + state_out.push_back({state_name, + ggml_cont(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1], + (size_t) emit_len * full->nb[0]))}); + } + + if (b) { + out = ggml_add(ctx0, out, ggml_reshape_2d(ctx0, b, 1, b->ne[0])); + } + return out; +} + +// ELU -> dilated conv -> ELU -> pointwise conv, added back to the input +ggml_tensor * clip_graph_pockettts_seanet::res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation, + const std::string & state_prefix) const { + ggml_tensor * h = ggml_elu(ctx0, x); + h = conv1d(h, stage.res_conv1_w, stage.res_conv1_b, 1, dilation, false, state_prefix); + h = ggml_elu(ctx0, h); + // the second conv is pointwise, it needs no left context + h = conv1d(h, stage.res_conv2_w, stage.res_conv2_b, 1, 1); + return ggml_add(ctx0, x, h); +} + +ggml_tensor * clip_graph_pockettts_seanet::encode(ggml_tensor * x) const { + const auto & seanet = model.seanet; + + ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1); + cb(cur, "seanet_enc_in", -1); + + for (int i = 0; i < hparams.seanet_n_stage; i++) { + const auto & stage = seanet.stages[i]; + const int stride = hparams.seanet_ratios[i]; + + cur = res_unit(cur, stage, 1); + cur = ggml_elu(ctx0, cur); + cur = conv1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride, 1); + cb(cur, "seanet_enc_stage", i); + } + + cur = ggml_elu(ctx0, cur); + cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1); + cb(cur, "seanet_enc_out", -1); + + return cur; +} + +ggml_tensor * clip_graph_pockettts_seanet::decode(ggml_tensor * x) const { + const auto & seanet = model.seanet; + const bool stream = !state_in.empty(); + + ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1, false, + stream ? "dec_in" : ""); + cb(cur, "seanet_dec_in", -1); + + for (int i = 0; i < hparams.seanet_n_stage; i++) { + const auto & stage = seanet.stages[i]; + // the decoder mirrors the encoder, so the ratios are walked backwards + const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i]; + const std::string id = std::to_string(i); + + cur = ggml_elu(ctx0, cur); + cur = conv_transpose1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride, + stream ? "dec_up_" + id : ""); + cur = res_unit(cur, stage, 1, stream ? "dec_res_" + id : ""); + cb(cur, "seanet_dec_stage", i); + } + + cur = ggml_elu(ctx0, cur); + cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1, false, + stream ? "dec_out" : ""); + cb(cur, "seanet_dec_out", -1); + + return cur; +} diff --git a/tools/mtmd/models/pockettts-spkenc.cpp b/tools/mtmd/models/pockettts-spkenc.cpp new file mode 100644 index 00000000000..f802d90687d --- /dev/null +++ b/tools/mtmd/models/pockettts-spkenc.cpp @@ -0,0 +1,77 @@ +#include "models.h" + +// voice-prompt encoder: raw 24kHz waveform -> one conditioning row per 12.5Hz frame +// mimi encoder (SEANet + transformer + downsample), then flow_lm.speaker_proj_weight + +// pre-norm block with layer scale on both residual paths, see mimi_transformer.py +ggml_tensor * clip_graph_pockettts_spkenc::tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const { + ggml_tensor * inp = cur; + + cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il); + + ggml_tensor * Qcur = build_mm(layer.q_w, cur); + ggml_tensor * Kcur = build_mm(layer.k_w, cur); + ggml_tensor * Vcur = build_mm(layer.v_w, cur); + + const int64_t n_pos = cur->ne[1]; + Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos); + Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos); + Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos); + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, + hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, + hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + cur = build_attn(layer.o_w, nullptr, Qcur, Kcur, Vcur, kq_mask, kq_scale, il); + cur = ggml_mul(ctx0, cur, layer.ls_1_w); + cur = ggml_add(ctx0, cur, inp); + + inp = cur; + cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il); + cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il); + cur = ggml_mul(ctx0, cur, layer.ls_2_w); + cur = ggml_add(ctx0, cur, inp); + + return cur; +} + +ggml_cgraph * clip_graph_pockettts_spkenc::build() { + // the preprocessor hands over the waveform as a single-row "mel", already [n_samples, 1] + ggml_tensor * inp_raw = build_inp_raw(1); + ggml_tensor * cur = ggml_reshape_2d(ctx0, inp_raw, inp_raw->ne[0], inp_raw->ne[1]); + + clip_graph_pockettts_seanet seanet(*this); + cur = seanet.encode(cur); + cb(cur, "mimi_enc", -1); + + // [T, 512] -> transformer works on [512, T] + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + + ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, cur->ne[1]); + ggml_set_name(inp_pos, "inp_pos"); + ggml_set_input(inp_pos); + + // the mimi transformer is causal with a sliding window, see _build_attention_mask() + ggml_tensor * kq_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, cur->ne[1], cur->ne[1]); + ggml_set_name(kq_mask, "kq_mask"); + ggml_set_input(kq_mask); + + for (int il = 0; il < n_layer; il++) { + cur = tfm_layer_forward(cur, model.layers[il], inp_pos, kq_mask, il); + } + cb(cur, "mimi_enc_tfm", -1); + + // downsample to the model frame rate, [512, T] -> [T, 512] -> [T / 16, 32] + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = seanet.conv1d(cur, model.downsample_w, nullptr, hparams.mimi_downsample, 1, true); + cb(cur, "mimi_downsample", -1); + + // voice latent -> backbone embd + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = build_mm(model.spk_proj_w, cur); + cb(cur, "spk_proj", -1); + + ggml_build_forward_expand(gf, cur); + return gf; +} diff --git a/tools/mtmd/models/qwen3tts-gen.cpp b/tools/mtmd/models/qwen3tts-gen.cpp index b6c95efa941..84c77f4fad1 100644 --- a/tools/mtmd/models/qwen3tts-gen.cpp +++ b/tools/mtmd/models/qwen3tts-gen.cpp @@ -610,6 +610,10 @@ std::vector list_c2w_state_slots(const clip_hparams & hparams, c const auto & c2w = model.c2w; std::vector slots; + if (c2w.pre_conv_w == nullptr) { + return slots; // not a code2wav model, it keeps no state between calls + } + slots.push_back({"tfm_pos", 1, 1}); // prefix is (W-1) frames, the batch itself gives the other N=W frames (see tfm_layer_forward) diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index ca4b64efa4a..56f84f43c3a 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1423,3 +1423,33 @@ std::vector mtmd_audio_streaming_istft::flush() { return output; } + +// +// mtmd_audio_preprocessor_pockettts +// +// mimi takes the raw 24kHz waveform, there is no mel front-end. the samples are handed over +// as a single-row "mel" so they travel through the normal chunk path +// + +bool mtmd_audio_preprocessor_pockettts::preprocess(const float * samples, + size_t n_samples, + std::vector & output) { + // the encoder needs whole frames, see pad_for_conv1d() in the reference + const int64_t frame_size = (int64_t) hparams.mimi_downsample * 120; + if (n_samples == 0 || frame_size <= 0) { + return false; + } + + const int64_t n_frames = (int64_t) (n_samples + frame_size - 1) / frame_size; + const int64_t n_padded = n_frames * frame_size; + + mtmd_audio_mel out; + out.n_mel = 1; + out.n_len = n_padded; + out.n_len_org = (int64_t) n_samples; + out.data.assign((size_t) n_padded, 0.0f); + std::copy(samples, samples + n_samples, out.data.begin()); + + output.push_back(std::move(out)); + return true; +} diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index b4d6f725980..44ad098ae63 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -129,6 +129,13 @@ struct mtmd_audio_preprocessor_qwen3tts_spk : mtmd_audio_preprocessor { mtmd_audio_cache cache; }; +// mimi convolves the waveform directly, so this only pads it to a whole number of frames +struct mtmd_audio_preprocessor_pockettts : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_pockettts(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} + void initialize() override {} + bool preprocess(const float * samples, size_t n_samples, std::vector & output) override; +}; + struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor { mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { } void initialize() override; diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index bc83f5ba230..a0106b68958 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -5,6 +5,7 @@ #include "../src/llama-ext.h" #include +#include #include #include #include @@ -460,10 +461,344 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::vector out_buf; }; +// Pocket-TTS: the backbone emits no token at all, each step's hidden state is turned into one +// continuous latent by the flow net, and the end-of-speech head 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; + } + + 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 voice; + if (inp->speaker_ref) { + if (!encode_speaker(inp->speaker_ref, voice)) { + return 1; + } + } + + const std::string text = prepare_text(std::string(inp->prompt, inp->prompt_len)); + if (text.empty()) { + LOG_ERR("mtmd_helper_gen_audio: empty prompt\n"); + return 1; + } + frames_after_eos = count_words(text) <= 4 ? 5 : 3; + + std::vector 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); + + const int n_e = n_embd; + auto push_row = [&](llama_token t) { + prompt_embd_buf.insert(prompt_embd_buf.end(), + tok_embd.begin() + (size_t) t * n_e, + tok_embd.begin() + (size_t) (t + 1) * n_e); + }; + + // sequence order is voice, then text, then the audio BOS that starts generation + if (!voice.empty()) { + push_row(bos_before_voice); + prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end()); + } + for (llama_token t : ids) { + push_row(t); + } + push_row(audio_bos); + + 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{}; + inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; + inp.embd = const_cast(h_state_in); + // the same seed every step: clip only reseeds when it changes, so the noise + // stream keeps running instead of restarting on each frame + inp.seed = seed; + inp.n_steps = -1; + 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() + if (eos_step >= 0 && step_idx >= eos_step + frames_after_eos) { + *out_stop = true; + *h_state_out = nullptr; + return 0; + } + + 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(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 = find_special_token(vocab, "<|bos_before_voice|>"); + audio_bos = find_special_token(vocab, "<|audio_bos|>"); + for (llama_token t : { bos_before_voice, audio_bos }) { + 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; + } + + // same normalization as prepare_text_prompt() in the reference, it affects quality + static std::string prepare_text(const std::string & in) { + std::string s; + s.reserve(in.size() + 1); + for (char c : in) { + s += (c == '\n' || c == '\r') ? ' ' : 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 += '.'; + } + 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 & 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, carrying the mimi decoder state across calls so a window + // can be emitted as soon as it is full + bool flush_gen_wav() { + if (feats_buf.empty()) { + return true; + } + mtmd_gen_inp inp{}; + 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; + } + + bool specials_ok = false; + llama_token bos_before_voice = LLAMA_TOKEN_NULL; + llama_token audio_bos = LLAMA_TOKEN_NULL; + std::vector tok_embd; + + llama_seq_id seq_id = 0; + int pos = 0; + std::vector prompt_embd_buf; + std::unique_ptr 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; + // latents are decoded a window at a time, the decoder state bridges the windows + size_t window_frames = 8; + std::vector feats_buf; + std::vector dec_state; + std::vector audio_pcm; + std::vector h_state_buf; + mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; + std::vector out_buf; +}; + static std::unique_ptr 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(new qwen3tts_gen_audio_pipeline(lctx, mctx)); + case MTMD_GEN_AUDIO_TYPE_POCKETTTS: + return std::unique_ptr(new pockettts_gen_audio_pipeline(lctx, mctx)); default: return nullptr; } diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 6a5b3f5a606..efa66a972bb 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -762,6 +762,10 @@ struct mtmd_context { { audio_preproc = std::make_unique(ctx_a); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + audio_preproc = std::make_unique(ctx_a); + } break; default: throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj)); } @@ -1591,6 +1595,10 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) { 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; @@ -1688,6 +1696,9 @@ 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; + // gen_wav draws no randomness, but the seed must still match so it does not reseed + // the rng in the middle of a 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; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 32d40dca327..6502d73def6 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -334,6 +334,7 @@ MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname); enum mtmd_gen_audio_type { MTMD_GEN_AUDIO_TYPE_NONE, // not supported MTMD_GEN_AUDIO_TYPE_QWEN3TTS, + MTMD_GEN_AUDIO_TYPE_POCKETTTS, }; struct mtmd_gen_audio_info { enum mtmd_gen_audio_type type; @@ -345,6 +346,7 @@ enum mtmd_gen_process_type { MTMD_GEN_PROCESS_TYPE_GEN_CODE, // h_state to semantic (codes, mel-spectrogram, etc.) MTMD_GEN_PROCESS_TYPE_GEN_WAV, // convert semantic to PCM audio // for qwen3tts, this is code2wav + // for pocket-tts, this is mimi decoder }; struct mtmd_gen_inp { enum mtmd_gen_process_type type; From 9ff97e31578bd3b5cc144c22846e96ecb9c9dee2 Mon Sep 17 00:00:00 2001 From: Pascal Date: Wed, 5 Aug 2026 21:05:50 +0200 Subject: [PATCH 04/17] 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. --- tools/mtmd/models/pockettts-seanet.cpp | 30 +++++++++++++++----------- 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/tools/mtmd/models/pockettts-seanet.cpp b/tools/mtmd/models/pockettts-seanet.cpp index 02c84c14342..bf8fff14a7b 100644 --- a/tools/mtmd/models/pockettts-seanet.cpp +++ b/tools/mtmd/models/pockettts-seanet.cpp @@ -51,23 +51,29 @@ ggml_tensor * clip_graph_pockettts_seanet::conv1d(ggml_tensor * x, ggml_tensor * // output when streaming, and simply dropped otherwise ggml_tensor * clip_graph_pockettts_seanet::conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, const std::string & state_name) const { - const int64_t p_total = w->ne[0] - stride; + const int64_t K = w->ne[0]; + const int64_t T = x->ne[0]; + const int64_t p_total = K - stride; const bool depthwise = w->ne[1] == 1 && w->ne[2] > 1; - const int64_t emit_len = x->ne[0] * stride; + const int64_t OC = depthwise ? w->ne[2] : w->ne[1]; + const int64_t emit_len = T * stride; - ggml_tensor * full = nullptr; + // one column per input step, holding the [K, OC] window that col2im scatter-adds at t * stride + ggml_tensor * col; if (depthwise) { - // one group per channel, ggml_conv_transpose_1d has no grouped mode - for (int64_t ir = 0; ir < x->ne[1]; ir++) { - ggml_tensor * row = ggml_view_1d(ctx0, x, x->ne[0], ir * x->ne[0] * ggml_element_size(x)); - ggml_tensor * krn = ggml_view_1d(ctx0, w, w->ne[0], ir * w->ne[0] * ggml_element_size(w)); - row = ggml_conv_transpose_1d(ctx0, krn, row, stride, 0, 1); - full = full ? ggml_concat(ctx0, full, row, 1) : row; - } + // one group per channel: a batched matmul over the channels scales the kernel by each step + ggml_tensor * krn = ggml_reshape_3d(ctx0, w, 1, K, OC); // [1, K, OC] + ggml_tensor * xs = ggml_reshape_3d(ctx0, x, 1, T, OC); // [1, T, OC] + col = ggml_mul_mat(ctx0, krn, xs); // [K, T, OC] + col = ggml_cont(ctx0, ggml_permute(ctx0, col, 0, 2, 1, 3)); // [K, OC, T] + col = ggml_reshape_2d(ctx0, col, K * OC, T); } else { - full = ggml_conv_transpose_1d(ctx0, w, x, stride, 0, 1); + ggml_tensor * w2 = ggml_reshape_2d(ctx0, w, K * OC, w->ne[2]); + w2 = ggml_cont(ctx0, ggml_transpose(ctx0, w2)); // [IC, K * OC] + ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [IC, T] + col = ggml_mul_mat(ctx0, w2, xt); } - full = ggml_cont(ctx0, full); // [emit_len + p_total, OC] + ggml_tensor * full = ggml_col2im_1d(ctx0, col, stride, OC, 0); // [emit_len + p_total, OC] ggml_tensor * out; if (state_name.empty() || p_total == 0) { From b28295072584ee217d2c03c0cdd4e477360fa1eb Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Thu, 6 Aug 2026 00:13:26 +0200 Subject: [PATCH 05/17] flow_temp + frames_after_eos --- conversion/base.py | 5 ++-- conversion/pockettts.py | 47 +++++++++++++++++++++++++++++----- gguf-py/gguf/constants.py | 2 ++ gguf-py/gguf/gguf_writer.py | 4 +++ tools/mtmd/clip-impl.h | 4 ++- tools/mtmd/clip-model.h | 2 +- tools/mtmd/clip.cpp | 9 ++++--- tools/mtmd/mtmd-helper-gen.cpp | 20 ++++++++++----- 8 files changed, 73 insertions(+), 20 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index 69f0fd9e657..99bb00a1db0 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -2652,8 +2652,9 @@ def _load_hparams_pockettts(shapes: dict[str, tuple[int, ...]]) -> dict[str, Any "max_position_embeddings": 4096, # not stored anywhere in the checkpoint, but every released variant uses head_dim 64 "num_attention_heads": n_embd // 64, - # 2 learned vectors are appended to the embedding table as extra tokens, see pockettts.py - "vocab_size": n_vocab + 2, + # learned input vectors are appended to the embedding table as extra tokens, see + # pockettts.py. bos_before_voice only exists when the pack inserts it + "vocab_size": n_vocab + (2 if "flow_lm.bos_before_voice" in shapes else 1), "rope_theta": 10000.0, "layer_norm_eps": 1e-5, "audio_config": { diff --git a/conversion/pockettts.py b/conversion/pockettts.py index a62243bb226..b6608cf884b 100644 --- a/conversion/pockettts.py +++ b/conversion/pockettts.py @@ -7,7 +7,7 @@ if TYPE_CHECKING: from torch import Tensor -from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf +from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf, logger # Pocket TTS is a CALM: an autoregressive backbone conditions a flow-matching decoder that # generates one continuous 32-d latent per frame. There is no codebook anywhere in this model. @@ -36,6 +36,26 @@ _N_SEANET_STAGES = 3 _SAMPLE_RATE = 24000 +# The flow decoder's noise scale is tuned per language pack and is not derivable from the +# checkpoint: the english packs are byte-identical in shape and tokenizer yet disagree on it. +# It lives only in the pip package's pocket_tts/config/.yaml, so it is keyed on the +# model directory name here. 0.7 is the reference default (Config.default_temperature). +# +# The packs also tune pad_with_spaces_for_short_inputs and remove_semicolons, which only +# affect text normalization for one or two packs each. Those are not carried over. +_DEFAULT_TEMP = 0.7 +_PACK_TEMP = { + "english": 0.3, + "english_2026-04": 0.3, +} + + +def _pack_temp(name: str) -> float: + if name not in _PACK_TEMP: + logger.warning("pocket-tts: no tuned temperature for language pack %r, using %.1f", + name, _DEFAULT_TEMP) + return _PACK_TEMP.get(name, _DEFAULT_TEMP) + @ModelBase.register("PocketTTSModel") class PocketTTSModel(TextModel): @@ -60,9 +80,8 @@ def set_vocab(self): tokens, scores, toktypes = self._create_vocab_sentencepiece() - # the last 3 rows of the embedding table are not sentencepiece pieces: the conditioner's - # padding row, then the two learned vectors appended by _embd_table() - extra = ["<|pad|>", "<|bos_before_voice|>", "<|audio_bos|>"] + # the last rows of the embedding table are not sentencepiece pieces + extra = self._extra_tokens() for i, name in enumerate(extra): tokens[len(tokens) - len(extra) + i] = name.encode("utf-8") toktypes[len(tokens) - len(extra) + i] = SentencePieceTokenTypes.CONTROL @@ -112,14 +131,27 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter return + def _extra_tokens(self) -> list[str]: + # the conditioner's padding row, then the learned vectors appended by _embd_table(). + # bos_before_voice only exists when the pack sets insert_bos_before_voice + names = ["<|pad|>"] + if "flow_lm.bos_before_voice" in self.model_tensors: + names.append("<|bos_before_voice|>") + names.append("<|audio_bos|>") + return names + def _embd_table(self, embed: Tensor) -> Tensor: - bos_before_voice = self.model_tensors["flow_lm.bos_before_voice"]().reshape(1, -1) + rows = [embed] + if "flow_lm.bos_before_voice" in self.model_tensors: + rows.append(self.model_tensors["flow_lm.bos_before_voice"]().reshape(1, -1).to(embed.dtype)) + # bos_emb is a latent, it only enters the backbone through input_linear bos_emb = self.model_tensors["flow_lm.bos_emb"]() input_linear = self.model_tensors["flow_lm.input_linear.weight"]() audio_bos = torch.nn.functional.linear(bos_emb.float(), input_linear.float()).reshape(1, -1) + rows.append(audio_bos.to(embed.dtype)) - return torch.cat([embed, bos_before_voice.to(embed.dtype), audio_bos.to(embed.dtype)], dim=0) + return torch.cat(rows, dim=0) @ModelBase.register("PocketTTSModel") @@ -169,6 +201,9 @@ def set_gguf_parameters(self): self.gguf_writer.add_gen_audio_head_count(self.hparams_audio["num_attention_heads"]) self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5) + # the flow decoder draws its noise at this scale, see lsd_decode() in the reference + self.gguf_writer.add_gen_audio_flow_temperature(_pack_temp(self.dir_model.name)) + def tensor_force_quant(self, name, new_name, bid, n_dims): del name, bid, n_dims # conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 991d930691f..14deedf51c8 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -400,6 +400,8 @@ class Projector: class ClipGenAudio: PROJECTOR_TYPE = "clip.gen.audio.projector_type" # for mixed modality models + # noise scale of the flow decoder, differs between pocket-tts language packs + FLOW_TEMPERATURE = "clip.gen.audio.flow_temperature" EMBEDDING_LENGTH = "clip.gen.audio.embedding_length" FEED_FORWARD_LENGTH = "clip.gen.audio.feed_forward_length" BLOCK_COUNT = "clip.gen.audio.block_count" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 39da9f2c05f..9eacb946504 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1438,6 +1438,10 @@ def add_gen_audio_head_count_kv(self, value: int) -> None: def add_gen_audio_attention_layernorm_eps(self, value: float) -> None: self.add_float32(Keys.ClipGenAudio.Attention.LAYERNORM_EPS, value) + def add_gen_audio_flow_temperature(self, value: float) -> None: + self.add_float32(Keys.ClipGenAudio.FLOW_TEMPERATURE, value) + + def add_xielu_alpha_p(self, values: Sequence[float]): self.add_array(Keys.xIELU.ALPHA_P, values) diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 0a31a7aec1a..c3c35d6c3e9 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -92,7 +92,9 @@ #define KEY_A_LOCAL_GROUP_SIZE "clip.audio.local_group_size" // mimo-v2.5: input_local_transformer grouping size // audio generation (gen-audio)-specific #define KEY_GEN_AUDIO_PROJ_TYPE "clip.gen.audio.projector_type" // for models with mixed modalities -#define KEY_AUDIO_SUBSAMPLING_FACTOR "clip.audio.subsampling_factor" +// noise scale of the flow decoder, differs between pocket-tts language packs +#define KEY_GEN_AUDIO_FLOW_TEMP "clip.gen.audio.flow_temperature" +#define KEY_AUDIO_SUBSMPL_FACTOR "clip.audio.subsampling_factor" // // tensor name constants diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 9b0b2c5e2fd..1d6bea3952d 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -145,7 +145,7 @@ struct clip_hparams { int32_t mimi_downsample = 0; // encoder frame rate / model frame rate int32_t mimi_tfm_context = 0; // attention window of the mimi transformers, in frames int32_t flow_n_step = 1; // lsd_decode steps - float flow_temp = 0.0f; // noise std is sqrt(temp) + float flow_temp = 0.0f; // noise std is sqrt(temp), differs per language pack // qwen3tts code2wav int32_t wav_tfm_n_layer = 0; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index bcb7cb083ea..520e921bf0b 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1432,7 +1432,7 @@ struct clip_model_loader { } break; case PROJECTOR_TYPE_PARAKEET: { - get_u32(KEY_AUDIO_SUBSAMPLING_FACTOR, hparams.subsampling_factor); + get_u32(KEY_AUDIO_SUBSMPL_FACTOR, hparams.subsampling_factor); GGML_ASSERT(hparams.subsampling_factor == 8 && "subsampling_factor must match the conv strides in clip_graph_parakeet::build()"); get_u32(KEY_A_CONV_KERNEL_SIZE, hparams.audio_conv_kernel_size); @@ -1754,10 +1754,13 @@ struct clip_model_loader { // matches the reference transformer's "context" hparams.mimi_tfm_context = 250; hparams.rope_theta = 10000.0f; - // flow_lm defaults, see pocket_tts/default_parameters.py and the language config + // flow_lm defaults, see pocket_tts/default_parameters.py hparams.flow_n_step = 1; - hparams.flow_temp = 0.3f; hparams.gen_eos_threshold = -4.0f; + // differs per language pack, the converter writes it out. + // the fallback is the reference's own default + hparams.flow_temp = 0.7f; + get_f32(KEY_GEN_AUDIO_FLOW_TEMP, hparams.flow_temp, false); } break; case PROJECTOR_TYPE_PADDLEOCR: { diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index a0106b68958..f1ecb4f3172 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -503,6 +503,7 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { LOG_ERR("mtmd_helper_gen_audio: empty prompt\n"); return 1; } + // the model may pin the tail length, otherwise guess it from the text like the reference frames_after_eos = count_words(text) <= 4 ? 5 : 3; std::vector ids(text.size() + 16); @@ -523,7 +524,9 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { // sequence order is voice, then text, then the audio BOS that starts generation if (!voice.empty()) { - push_row(bos_before_voice); + if (bos_before_voice != LLAMA_TOKEN_NULL) { + push_row(bos_before_voice); + } prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end()); } for (llama_token t : ids) { @@ -651,13 +654,12 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { 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|>"); - for (llama_token t : { bos_before_voice, audio_bos }) { - if (t == LLAMA_TOKEN_NULL) { - LOG_ERR("mtmd_helper_gen_audio: missing a required special token in vocab\n"); - return false; - } + 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) { @@ -678,7 +680,11 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::string s; s.reserve(in.size() + 1); for (char c : in) { - s += (c == '\n' || c == '\r') ? ' ' : c; + if (c == '\n' || c == '\r') { + s += ' '; + } else { + s += c; + } } const size_t b = s.find_first_not_of(' '); const size_t e = s.find_last_not_of(' '); From c898cb90930e034743cc7175b6e4eb190c201764 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Thu, 6 Aug 2026 00:27:24 +0200 Subject: [PATCH 06/17] chunking --- tools/mtmd/mtmd-helper-gen.cpp | 180 +++++++++++++++++++++++++++++++-- 1 file changed, 172 insertions(+), 8 deletions(-) diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index f1ecb4f3172..884b11680e0 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -92,6 +93,7 @@ class mtmd_gen_audio_pipeline { 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; @@ -481,6 +483,10 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { 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 { @@ -503,8 +509,6 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { LOG_ERR("mtmd_helper_gen_audio: empty prompt\n"); return 1; } - // the model may pin the tail length, otherwise guess it from the text like the reference - frames_after_eos = count_words(text) <= 4 ? 5 : 3; std::vector ids(text.size() + 16); int n_ids = llama_tokenize(vocab, text.c_str(), (int32_t) text.size(), ids.data(), @@ -515,6 +519,14 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } ids.resize((size_t) n_ids); + // long inputs degrade badly, the reference splits them and restarts each piece from + // the voice conditioning, see split_into_best_sentences() + 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; auto push_row = [&](llama_token t) { prompt_embd_buf.insert(prompt_embd_buf.end(), @@ -529,10 +541,14 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end()); } - for (llama_token t : ids) { + // 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_row(t); } push_row(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)); @@ -591,11 +607,15 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { if (out.is_eos && eos_step < 0) { eos_step = step_idx; } - // the frame of the stopping step is discarded, matching _autoregressive_generation() - if (eos_step >= 0 && step_idx >= eos_step + frames_after_eos) { - *out_stop = true; - *h_state_out = nullptr; - return 0; + // 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); @@ -675,6 +695,142 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { return true; } + // token ids of the pieces the reference splits on, see split_into_best_sentences(). + // the leading token is dropped, it is the tokenizer's dummy prefix + std::vector punct_ids(const char * s) const { + std::vector 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(ids.begin() + 1, ids.begin() + n); + } + + // cut after runs of boundary tokens, so punctuation stays with the sentence it ends + static std::vector> split_on(const std::vector & ids, + const std::vector & boundary) { + std::vector> 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> split_chunks(const std::vector & ids) const { + if ((int) ids.size() <= max_chunk_tokens) { + return { ids }; + } + const std::vector eos_punct = punct_ids(".!...?"); + const std::vector mid_punct = punct_ids(",;:"); + + // oversized sentences are split again on weaker punctuation, else words get skipped + std::vector> 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> 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 reference guesses the tail from the word count, approximated here by tokens + 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]) { + prompt_embd_buf.insert(prompt_embd_buf.end(), + tok_embd.begin() + (size_t) t * n_e, + tok_embd.begin() + (size_t) (t + 1) * n_e); + } + prompt_embd_buf.insert(prompt_embd_buf.end(), + tok_embd.begin() + (size_t) audio_bos * n_e, + tok_embd.begin() + (size_t) (audio_bos + 1) * n_e); + arm_chunk_budget(chunk_idx); + + // bounded by max_chunk_tokens + 1, so one decode is enough + const int n_rows = (int) (prompt_embd_buf.size() / (size_t) n_e); + 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) { std::string s; @@ -789,6 +945,14 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { int step_idx = 0; int eos_step = -1; int frames_after_eos = 3; + // long inputs are split, each chunk restarts from the voice conditioning + static constexpr int max_chunk_tokens = 50; // MAX_TOKEN_PER_CHUNK in the reference + static constexpr double frame_rate = 12.5; + std::vector> 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 feats_buf; From d14bb68e5f60ffba4b4bf9541924d8a783fab30e Mon Sep 17 00:00:00 2001 From: Pascal Date: Thu, 6 Aug 2026 13:50:09 +0200 Subject: [PATCH 07/17] 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. --- conversion/pockettts.py | 16 ++++++++++++++-- gguf-py/gguf/constants.py | 2 ++ gguf-py/gguf/gguf_writer.py | 6 ++++++ tools/mtmd/clip-impl.h | 2 ++ tools/mtmd/clip-model.h | 2 ++ tools/mtmd/clip.cpp | 4 +++- tools/mtmd/mtmd-helper-gen.cpp | 15 +++++++++++---- tools/mtmd/mtmd.cpp | 10 +++++++--- tools/mtmd/mtmd.h | 2 ++ 9 files changed, 49 insertions(+), 10 deletions(-) diff --git a/conversion/pockettts.py b/conversion/pockettts.py index b6608cf884b..05934acab40 100644 --- a/conversion/pockettts.py +++ b/conversion/pockettts.py @@ -41,13 +41,21 @@ # It lives only in the pip package's pocket_tts/config/.yaml, so it is keyed on the # model directory name here. 0.7 is the reference default (Config.default_temperature). # -# The packs also tune pad_with_spaces_for_short_inputs and remove_semicolons, which only -# affect text normalization for one or two packs each. Those are not carried over. +# model_recommended_frames_after_eos and pad_with_spaces_for_short_inputs come from the same +# per-pack yaml. remove_semicolons does too, but it maps ";" to "," and is applied to every +# pack on the cpp side instead of being carried here. _DEFAULT_TEMP = 0.7 _PACK_TEMP = { "english": 0.3, "english_2026-04": 0.3, } +# 0 leaves the tail length to the caller, which guesses it from the text +_PACK_FRAMES_AFTER_EOS = { + "french_24l": 8, +} +_PACK_PAD_SHORT_TEXT = { + "english_2026-01": True, +} def _pack_temp(name: str) -> float: @@ -203,6 +211,10 @@ def set_gguf_parameters(self): # the flow decoder draws its noise at this scale, see lsd_decode() in the reference self.gguf_writer.add_gen_audio_flow_temperature(_pack_temp(self.dir_model.name)) + self.gguf_writer.add_gen_audio_frames_after_eos( + _PACK_FRAMES_AFTER_EOS.get(self.dir_model.name, 0)) + self.gguf_writer.add_gen_audio_pad_short_text( + _PACK_PAD_SHORT_TEXT.get(self.dir_model.name, False)) def tensor_force_quant(self, name, new_name, bid, n_dims): del name, bid, n_dims diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 14deedf51c8..3e6a9e17a8f 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -402,6 +402,8 @@ class ClipGenAudio: PROJECTOR_TYPE = "clip.gen.audio.projector_type" # for mixed modality models # noise scale of the flow decoder, differs between pocket-tts language packs FLOW_TEMPERATURE = "clip.gen.audio.flow_temperature" + FRAMES_AFTER_EOS = "clip.gen.audio.frames_after_eos" + PAD_SHORT_TEXT = "clip.gen.audio.pad_short_text" EMBEDDING_LENGTH = "clip.gen.audio.embedding_length" FEED_FORWARD_LENGTH = "clip.gen.audio.feed_forward_length" BLOCK_COUNT = "clip.gen.audio.block_count" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 9eacb946504..e67e2b6fef4 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1441,6 +1441,12 @@ def add_gen_audio_attention_layernorm_eps(self, value: float) -> None: def add_gen_audio_flow_temperature(self, value: float) -> None: self.add_float32(Keys.ClipGenAudio.FLOW_TEMPERATURE, value) + def add_gen_audio_frames_after_eos(self, value: int) -> None: + self.add_uint32(Keys.ClipGenAudio.FRAMES_AFTER_EOS, value) + + def add_gen_audio_pad_short_text(self, value: bool) -> None: + self.add_bool(Keys.ClipGenAudio.PAD_SHORT_TEXT, value) + def add_xielu_alpha_p(self, values: Sequence[float]): self.add_array(Keys.xIELU.ALPHA_P, values) diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index c3c35d6c3e9..f6657e0bd5a 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -94,6 +94,8 @@ #define KEY_GEN_AUDIO_PROJ_TYPE "clip.gen.audio.projector_type" // for models with mixed modalities // noise scale of the flow decoder, differs between pocket-tts language packs #define KEY_GEN_AUDIO_FLOW_TEMP "clip.gen.audio.flow_temperature" +#define KEY_GEN_AUDIO_FRAMES_EOS "clip.gen.audio.frames_after_eos" +#define KEY_GEN_AUDIO_PAD_SHORT "clip.gen.audio.pad_short_text" #define KEY_AUDIO_SUBSMPL_FACTOR "clip.audio.subsampling_factor" // diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 1d6bea3952d..9302f65a417 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -146,6 +146,8 @@ struct clip_hparams { int32_t mimi_tfm_context = 0; // attention window of the mimi transformers, in frames int32_t flow_n_step = 1; // lsd_decode steps float flow_temp = 0.0f; // noise std is sqrt(temp), differs per language pack + int32_t gen_frames_after_eos = 0; // tail the pack asks for, 0 leaves the guess to the caller + bool gen_pad_short_text = false; // qwen3tts code2wav int32_t wav_tfm_n_layer = 0; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 520e921bf0b..bf5e386a511 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1760,7 +1760,9 @@ struct clip_model_loader { // differs per language pack, the converter writes it out. // the fallback is the reference's own default hparams.flow_temp = 0.7f; - get_f32(KEY_GEN_AUDIO_FLOW_TEMP, hparams.flow_temp, false); + get_f32 (KEY_GEN_AUDIO_FLOW_TEMP, hparams.flow_temp, false); + get_u32 (KEY_GEN_AUDIO_FRAMES_EOS, hparams.gen_frames_after_eos, false); + get_bool(KEY_GEN_AUDIO_PAD_SHORT, hparams.gen_pad_short_text, false); } break; case PROJECTOR_TYPE_PADDLEOCR: { diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 884b11680e0..f2045bffd23 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -504,7 +504,8 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } } - const std::string text = prepare_text(std::string(inp->prompt, inp->prompt_len)); + const std::string text = prepare_text(std::string(inp->prompt, inp->prompt_len), + info.pad_short_text); if (text.empty()) { LOG_ERR("mtmd_helper_gen_audio: empty prompt\n"); return 1; @@ -775,8 +776,9 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { 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 reference guesses the tail from the word count, approximated here by tokens - frames_after_eos = n_tok <= 6 ? 5 : 3; + // the pack may pin the tail, otherwise the reference guesses it from the word count, + // approximated here by tokens + frames_after_eos = info.frames_after_eos > 0 ? info.frames_after_eos : (n_tok <= 6 ? 5 : 3); step_idx = 0; eos_step = -1; } @@ -832,12 +834,14 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } // same normalization as prepare_text_prompt() in the reference, it affects quality - static std::string prepare_text(const std::string & in) { + 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; } @@ -855,6 +859,9 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { if (std::isalnum(last)) { s += '.'; } + if (pad_short && count_words(s) < 5) { + s = std::string(8, ' ') + s; + } return s; } diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index efa66a972bb..a3d72314cb1 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1596,9 +1596,13 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) { info.sample_rate = 24000; break; case PROJECTOR_TYPE_POCKETTTS_GEN: - info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS; - info.sample_rate = 24000; - break; + { + const clip_hparams * hp = clip_get_hparams(ctx->ctx_gen_a); + info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS; + info.sample_rate = 24000; + info.frames_after_eos = hp->gen_frames_after_eos; + info.pad_short_text = hp->gen_pad_short_text; + } break; default: info.type = MTMD_GEN_AUDIO_TYPE_NONE; break; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 6502d73def6..e977ecd8a3b 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -339,6 +339,8 @@ enum mtmd_gen_audio_type { struct mtmd_gen_audio_info { enum mtmd_gen_audio_type type; int32_t sample_rate; // in Hz, for example 24000 for qwen3tts + int32_t frames_after_eos; // tail the model asks for, 0 to guess it from the text + bool pad_short_text; // the model wants short prompts padded with spaces }; MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx); From 4b9a4df9b913f4ddb1c8188f26649176c9ad07d0 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Thu, 6 Aug 2026 16:19:09 +0200 Subject: [PATCH 08/17] clip.gen.audio.model_variant --- conversion/pockettts.py | 37 ++------------------------------ gguf-py/gguf/constants.py | 6 ++---- gguf-py/gguf/gguf_writer.py | 10 ++------- tools/mtmd/clip-impl.h | 6 ++---- tools/mtmd/clip-model.h | 7 +++--- tools/mtmd/clip.cpp | 10 +++------ tools/mtmd/clip.h | 1 + tools/mtmd/mtmd-helper-gen.cpp | 39 ++++++++++++++++++++++++++++++---- tools/mtmd/mtmd.cpp | 13 ++++++------ tools/mtmd/mtmd.h | 5 +++-- 10 files changed, 60 insertions(+), 74 deletions(-) diff --git a/conversion/pockettts.py b/conversion/pockettts.py index 05934acab40..63e41e944ec 100644 --- a/conversion/pockettts.py +++ b/conversion/pockettts.py @@ -7,7 +7,7 @@ if TYPE_CHECKING: from torch import Tensor -from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf, logger +from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf # Pocket TTS is a CALM: an autoregressive backbone conditions a flow-matching decoder that # generates one continuous 32-d latent per frame. There is no codebook anywhere in this model. @@ -36,34 +36,6 @@ _N_SEANET_STAGES = 3 _SAMPLE_RATE = 24000 -# The flow decoder's noise scale is tuned per language pack and is not derivable from the -# checkpoint: the english packs are byte-identical in shape and tokenizer yet disagree on it. -# It lives only in the pip package's pocket_tts/config/.yaml, so it is keyed on the -# model directory name here. 0.7 is the reference default (Config.default_temperature). -# -# model_recommended_frames_after_eos and pad_with_spaces_for_short_inputs come from the same -# per-pack yaml. remove_semicolons does too, but it maps ";" to "," and is applied to every -# pack on the cpp side instead of being carried here. -_DEFAULT_TEMP = 0.7 -_PACK_TEMP = { - "english": 0.3, - "english_2026-04": 0.3, -} -# 0 leaves the tail length to the caller, which guesses it from the text -_PACK_FRAMES_AFTER_EOS = { - "french_24l": 8, -} -_PACK_PAD_SHORT_TEXT = { - "english_2026-01": True, -} - - -def _pack_temp(name: str) -> float: - if name not in _PACK_TEMP: - logger.warning("pocket-tts: no tuned temperature for language pack %r, using %.1f", - name, _DEFAULT_TEMP) - return _PACK_TEMP.get(name, _DEFAULT_TEMP) - @ModelBase.register("PocketTTSModel") class PocketTTSModel(TextModel): @@ -209,12 +181,7 @@ def set_gguf_parameters(self): self.gguf_writer.add_gen_audio_head_count(self.hparams_audio["num_attention_heads"]) self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5) - # the flow decoder draws its noise at this scale, see lsd_decode() in the reference - self.gguf_writer.add_gen_audio_flow_temperature(_pack_temp(self.dir_model.name)) - self.gguf_writer.add_gen_audio_frames_after_eos( - _PACK_FRAMES_AFTER_EOS.get(self.dir_model.name, 0)) - self.gguf_writer.add_gen_audio_pad_short_text( - _PACK_PAD_SHORT_TEXT.get(self.dir_model.name, False)) + self.gguf_writer.add_gen_audio_model_variant(self.dir_model.name) def tensor_force_quant(self, name, new_name, bid, n_dims): del name, bid, n_dims diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 3e6a9e17a8f..a2c1c53fd84 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -400,10 +400,8 @@ class Projector: class ClipGenAudio: PROJECTOR_TYPE = "clip.gen.audio.projector_type" # for mixed modality models - # noise scale of the flow decoder, differs between pocket-tts language packs - FLOW_TEMPERATURE = "clip.gen.audio.flow_temperature" - FRAMES_AFTER_EOS = "clip.gen.audio.frames_after_eos" - PAD_SHORT_TEXT = "clip.gen.audio.pad_short_text" + # name of the weight variant, for settings that are not in the checkpoint + MODEL_VARIANT = "clip.gen.audio.model_variant" EMBEDDING_LENGTH = "clip.gen.audio.embedding_length" FEED_FORWARD_LENGTH = "clip.gen.audio.feed_forward_length" BLOCK_COUNT = "clip.gen.audio.block_count" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index e67e2b6fef4..417d8dee232 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1438,14 +1438,8 @@ def add_gen_audio_head_count_kv(self, value: int) -> None: def add_gen_audio_attention_layernorm_eps(self, value: float) -> None: self.add_float32(Keys.ClipGenAudio.Attention.LAYERNORM_EPS, value) - def add_gen_audio_flow_temperature(self, value: float) -> None: - self.add_float32(Keys.ClipGenAudio.FLOW_TEMPERATURE, value) - - def add_gen_audio_frames_after_eos(self, value: int) -> None: - self.add_uint32(Keys.ClipGenAudio.FRAMES_AFTER_EOS, value) - - def add_gen_audio_pad_short_text(self, value: bool) -> None: - self.add_bool(Keys.ClipGenAudio.PAD_SHORT_TEXT, value) + def add_gen_audio_model_variant(self, value: str) -> None: + self.add_string(Keys.ClipGenAudio.MODEL_VARIANT, value) def add_xielu_alpha_p(self, values: Sequence[float]): diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index f6657e0bd5a..d6930720892 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -92,10 +92,8 @@ #define KEY_A_LOCAL_GROUP_SIZE "clip.audio.local_group_size" // mimo-v2.5: input_local_transformer grouping size // audio generation (gen-audio)-specific #define KEY_GEN_AUDIO_PROJ_TYPE "clip.gen.audio.projector_type" // for models with mixed modalities -// noise scale of the flow decoder, differs between pocket-tts language packs -#define KEY_GEN_AUDIO_FLOW_TEMP "clip.gen.audio.flow_temperature" -#define KEY_GEN_AUDIO_FRAMES_EOS "clip.gen.audio.frames_after_eos" -#define KEY_GEN_AUDIO_PAD_SHORT "clip.gen.audio.pad_short_text" +// name of the weight variant, for settings that are not in the checkpoint +#define KEY_GEN_AUDIO_VARIANT "clip.gen.audio.model_variant" #define KEY_AUDIO_SUBSMPL_FACTOR "clip.audio.subsampling_factor" // diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 9302f65a417..e7912c4dbb8 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -139,15 +139,16 @@ struct clip_hparams { // threshold for the "out_eos_score" graph output float gen_eos_threshold = 0.0f; + // name of the weight variant, some pipelines tune themselves on it + std::string gen_model_variant; + // pocket-tts int32_t seanet_n_stage = 0; std::vector seanet_ratios; // encoder order (reversed compared to the config) int32_t mimi_downsample = 0; // encoder frame rate / model frame rate int32_t mimi_tfm_context = 0; // attention window of the mimi transformers, in frames int32_t flow_n_step = 1; // lsd_decode steps - float flow_temp = 0.0f; // noise std is sqrt(temp), differs per language pack - int32_t gen_frames_after_eos = 0; // tail the pack asks for, 0 leaves the guess to the caller - bool gen_pad_short_text = false; + float flow_temp = 0.7f; // noise std is sqrt(temp), the caller can override it // qwen3tts code2wav int32_t wav_tfm_n_layer = 0; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index bf5e386a511..bb39d8b0905 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1293,6 +1293,7 @@ struct clip_model_loader { // these are unused, but still need to be set to avoid issues hparams.image_size = 0; hparams.patch_size = 1; + get_string(KEY_GEN_AUDIO_VARIANT, hparams.gen_model_variant, false); } else { GGML_ASSERT(false && "unknown modality"); @@ -1757,12 +1758,6 @@ struct clip_model_loader { // flow_lm defaults, see pocket_tts/default_parameters.py hparams.flow_n_step = 1; hparams.gen_eos_threshold = -4.0f; - // differs per language pack, the converter writes it out. - // the fallback is the reference's own default - hparams.flow_temp = 0.7f; - get_f32 (KEY_GEN_AUDIO_FLOW_TEMP, hparams.flow_temp, false); - get_u32 (KEY_GEN_AUDIO_FRAMES_EOS, hparams.gen_frames_after_eos, false); - get_bool(KEY_GEN_AUDIO_PAD_SHORT, hparams.gen_pad_short_text, false); } break; case PROJECTOR_TYPE_PADDLEOCR: { @@ -4856,7 +4851,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } else { // flow matching starts from gaussian noise, std = sqrt(temp) ggml_tensor * t = get_inp_tensor("inp_noise"); - std::normal_distribution dist(0.0f, std::sqrt(hparams.flow_temp)); + const float temp = params->flow_temp > 0.0f ? params->flow_temp : hparams.flow_temp; + std::normal_distribution dist(0.0f, std::sqrt(temp)); std::vector noise(ggml_nelements(t)); for (auto & v : noise) { v = dist(ctx->rng); diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index 97de0606dbf..a237a9466a3 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -107,6 +107,7 @@ struct clip_encode_params { std::vector * out_feats = nullptr; // continuous counterpart of out_codes uint32_t seed = UINT32_MAX; // UINT32_MAX for random int32_t n_steps = -1; // integration steps, for flow-matching decoders + float flow_temp = 0.0f; // noise scale of the flow decoder, 0 for default bool * out_is_eos = nullptr; // GEN_WAV diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index f2045bffd23..e8dc9c276c5 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -463,6 +463,33 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::vector out_buf; }; +// Settings that live only in the reference's per-pack yaml and are not derivable from the +// checkpoint: the english packs are identical in shape and tokenizer yet disagree on them. +// They are keyed on the weight variant name that the mmproj carries. +// remove_semicolons belongs here too, but it maps ";" to "," and is applied to every pack. +struct pockettts_pack_settings { + float temp = 0.7f; // Config.default_temperature + int frames_after_eos = 0; // 0 leaves the tail length to the caller + bool pad_short_text = false; +}; + +static pockettts_pack_settings pockettts_pack(const char * variant) { + static const std::unordered_map 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()) { + pockettts_pack_settings def; + LOG_WRN("mtmd_helper_gen_audio: no tuned settings for pocket-tts variant \"%s\", " + "using temperature %.1f\n", variant ? variant : "", def.temp); + return def; + } + return it->second; +} + // Pocket-TTS: the backbone emits no token at all, each step's hidden state is turned into one // continuous latent by the flow net, and the end-of-speech head lives in the mmproj class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { @@ -504,8 +531,10 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } } + pack = pockettts_pack(info.model_variant); + const std::string text = prepare_text(std::string(inp->prompt, inp->prompt_len), - info.pad_short_text); + pack.pad_short_text); if (text.empty()) { LOG_ERR("mtmd_helper_gen_audio: empty prompt\n"); return 1; @@ -598,8 +627,9 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { inp.embd = const_cast(h_state_in); // the same seed every step: clip only reseeds when it changes, so the noise // stream keeps running instead of restarting on each frame - inp.seed = seed; - inp.n_steps = -1; + inp.seed = seed; + inp.n_steps = -1; + inp.flow_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"); @@ -778,7 +808,7 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { chunk_budget = (int) std::ceil((n_tok / 3.0 + 2.0) * frame_rate); // the pack may pin the tail, otherwise the reference guesses it from the word count, // approximated here by tokens - frames_after_eos = info.frames_after_eos > 0 ? info.frames_after_eos : (n_tok <= 6 ? 5 : 3); + frames_after_eos = pack.frames_after_eos > 0 ? pack.frames_after_eos : (n_tok <= 6 ? 5 : 3); step_idx = 0; eos_step = -1; } @@ -936,6 +966,7 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { 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; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index a3d72314cb1..5d6f6c4ac97 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1586,23 +1586,21 @@ 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{}; + 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: - { - const clip_hparams * hp = clip_get_hparams(ctx->ctx_gen_a); - info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS; - info.sample_rate = 24000; - info.frames_after_eos = hp->gen_frames_after_eos; - info.pad_short_text = hp->gen_pad_short_text; - } break; + info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS; + info.sample_rate = 24000; + break; default: info.type = MTMD_GEN_AUDIO_TYPE_NONE; break; @@ -1647,6 +1645,7 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in params.top_p = inp->top_p; params.seed = inp->seed; params.n_steps = inp->n_steps; + params.flow_temp = inp->flow_temp; params.out_is_eos = &is_eos; if (!clip_encode(ctx_clip, ¶ms)) { diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index e977ecd8a3b..ee93c4bfdb5 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -339,8 +339,8 @@ enum mtmd_gen_audio_type { struct mtmd_gen_audio_info { enum mtmd_gen_audio_type type; int32_t sample_rate; // in Hz, for example 24000 for qwen3tts - int32_t frames_after_eos; // tail the model asks for, 0 to guess it from the text - bool pad_short_text; // the model wants short prompts padded with spaces + const char * model_variant; // name of the weight variant, empty if the mmproj has none + // some pipelines have settings that only exist per-variant }; MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx); @@ -360,6 +360,7 @@ struct mtmd_gen_inp { float top_p; uint32_t seed; // UINT32_MAX for random int32_t n_steps; // integration steps, for flow-matching decoders (-1 for default) + float flow_temp; // noise scale, for flow-matching decoders (0 for default) // for MTMD_GEN_PROCESS_TYPE_GEN_WAV // pass either codes (discrete) or feats (continuous), depending on the pipeline From a73f458f3b5cd9a1e4593069cbcbff3e66eed220 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Thu, 6 Aug 2026 16:42:25 +0200 Subject: [PATCH 09/17] clean up code comments --- conversion/base.py | 4 ++-- conversion/pockettts.py | 18 ++++++++---------- tools/mtmd/models/pockettts-gen.cpp | 13 ++++++------- tools/mtmd/models/pockettts-seanet.cpp | 8 +++----- tools/mtmd/mtmd-audio.cpp | 4 ++-- tools/mtmd/mtmd-helper-gen.cpp | 25 +++++++++---------------- tools/mtmd/mtmd.cpp | 3 +-- tools/mtmd/mtmd.h | 1 - 8 files changed, 31 insertions(+), 45 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index 99bb00a1db0..eb5a1d32a2e 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -2652,8 +2652,8 @@ def _load_hparams_pockettts(shapes: dict[str, tuple[int, ...]]) -> dict[str, Any "max_position_embeddings": 4096, # not stored anywhere in the checkpoint, but every released variant uses head_dim 64 "num_attention_heads": n_embd // 64, - # learned input vectors are appended to the embedding table as extra tokens, see - # pockettts.py. bos_before_voice only exists when the pack inserts it + # extra rows for the learned input vectors, see pockettts.py + # bos_before_voice only exists when the pack inserts it "vocab_size": n_vocab + (2 if "flow_lm.bos_before_voice" in shapes else 1), "rope_theta": 10000.0, "layer_norm_eps": 1e-5, diff --git a/conversion/pockettts.py b/conversion/pockettts.py index 63e41e944ec..04539e4681c 100644 --- a/conversion/pockettts.py +++ b/conversion/pockettts.py @@ -9,17 +9,15 @@ from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf -# Pocket TTS is a CALM: an autoregressive backbone conditions a flow-matching decoder that -# generates one continuous 32-d latent per frame. There is no codebook anywhere in this model. -# +# Pocket TTS is a CALM: the backbone conditions a flow-matching decoder that generates one +# continuous 32-d latent per frame. There is no codebook in this model. # The checkpoint ships no config.json, hparams are derived in base.load_hparams_non_hf(). # # Tricks being used to support this model via existing llama.cpp code paths: -# - bos_before_voice and bos_emb are learned input vectors, not tokens. they are appended to -# the embedding table as extra tokens so the helper can look them up like any other row. -# bos_emb lives in latent space, so input_linear is folded into it here -# - the backbone has no lm_head, the embedding table is reused as output so that a sampler -# can run over the (unused) logits +# - bos_before_voice and bos_emb are learned input vectors, not tokens +# they are appended to the embedding table as extra tokens, to be looked up like any other row +# - bos_emb lives in latent space, so input_linear is folded into it here +# - the backbone has no lm_head, the embedding table is reused as output for the unused logits # # pipeline stage mapping: # mimi encoder + speaker_proj --> mapped to normal mtmd audio encoder @@ -50,8 +48,8 @@ class PocketTTSModel(TextModel): } def set_vocab(self): - # this is a unigram sentencepiece model; llama.cpp's SPM tokenizer greedily merges - # bigrams and cannot reproduce unigram segmentation, so use the UGM tokenizer instead + # this is a unigram sentencepiece model, llama.cpp's SPM tokenizer cannot do + # unigram segmentation, so use the UGM tokenizer instead from sentencepiece import sentencepiece_model_pb2 as model proto = model.ModelProto() # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute] diff --git a/tools/mtmd/models/pockettts-gen.cpp b/tools/mtmd/models/pockettts-gen.cpp index b76d44da387..3266742150c 100644 --- a/tools/mtmd/models/pockettts-gen.cpp +++ b/tools/mtmd/models/pockettts-gen.cpp @@ -9,13 +9,12 @@ // // there is no codebook anywhere, "codes" in the mtmd API are continuous features here -// x * (1 + scale) + shift, all [D, 1] ggml_tensor * clip_graph_pockettts_gen::modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const { ggml_tensor * cur = ggml_mul(ctx0, x, ggml_scale_bias(ctx0, scale, 1.0f, 1.0f)); return ggml_add(ctx0, cur, shift); } -// cos/sin(t * freqs) -> Linear -> SiLU -> Linear -> RMSNorm, see TimestepEmbedder +// see TimestepEmbedder in the reference ggml_tensor * clip_graph_pockettts_gen::time_embed(const clip_flow_net::time_embd & te, float t) const { // t is a graph-build constant, so the cos/sin table can be folded into a scaled copy ggml_tensor * args = ggml_scale(ctx0, te.freqs, t); @@ -27,8 +26,8 @@ ggml_tensor * clip_graph_pockettts_gen::time_embed(const clip_flow_net::time_emb cur = build_mm(te.down_w, cur); cur = ggml_add(ctx0, cur, te.down_b); - // this "RMSNorm" divides by the unbiased variance, not the mean square, and it rescales - // the input rather than the centered value, see _rms_norm() in mlp.py + // this "RMSNorm" divides by the unbiased variance, not the mean square + // it also rescales the input, not the centered value, see _rms_norm() in mlp.py { const int64_t n = cur->ne[0]; ggml_tensor * mean = ggml_mean(ctx0, cur); @@ -99,7 +98,7 @@ ggml_tensor * clip_graph_pockettts_gen::flow_forward(ggml_tensor * cond, ggml_te } // state carried between GEN_WAV calls: rope offset, per-layer KV window, conv left context -// and the transposed-conv overlap tails. shape lookup only, no graph needed +// and the transposed-conv overlap tails std::vector list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model) { std::vector slots; if (model.gen_upsample_w == nullptr) { @@ -208,8 +207,8 @@ ggml_cgraph * clip_graph_pockettts_gen::build() { GGML_TYPE_I32); seanet.state_out.push_back({"tfm_pos", ggml_scale_bias(ctx0, seanet.state_in.at("tfm_pos"), 1.0f, (float) n_pos)}); - // banded causal mask over [cached prefix | this chunk], and a cold-start mask for the - // cache rows that hold no real frame yet + // banded causal mask over [cached prefix | this chunk] + // the last factor masks out cache rows that hold no real frame yet ggml_tensor * pos_k = ggml_reshape_2d(ctx0, ggml_arange(ctx0, 0.0f, (float) n_kv, 1.0f), n_kv, 1); ggml_tensor * pos_q = ggml_reshape_2d(ctx0, ggml_arange(ctx0, (float) prefix, (float) (prefix + n_pos), 1.0f), 1, n_pos); ggml_tensor * diff = ggml_sub(ctx0, ggml_repeat_4d(ctx0, pos_q, n_kv, n_pos, 1, 1), pos_k); diff --git a/tools/mtmd/models/pockettts-seanet.cpp b/tools/mtmd/models/pockettts-seanet.cpp index bf8fff14a7b..c47207f569b 100644 --- a/tools/mtmd/models/pockettts-seanet.cpp +++ b/tools/mtmd/models/pockettts-seanet.cpp @@ -2,8 +2,8 @@ // SEANet convolution stack of the mimi codec, see pocket_tts/modules/seanet.py // -// tensors are T-first here: [T, C]. the convs are causal: they take left context from a -// state slot when given, otherwise they pad (cold start / one-shot encode) +// tensors are T-first here: [T, C] +// the convs are causal: left context comes from a state slot, or from padding on a cold start static int64_t div_ceil(int64_t a, int64_t b) { return a / b + (a % b ? 1 : 0); @@ -47,8 +47,7 @@ ggml_tensor * clip_graph_pockettts_seanet::conv1d(ggml_tensor * x, ggml_tensor * } // x: [T, IC], w: [K, OC/groups, IC] -> [T * stride, OC] -// the K - stride overlap tail belongs to the next call: it is added to the head of the next -// output when streaming, and simply dropped otherwise +// the K - stride overlap tail belongs to the next call: added to its head when streaming, else dropped ggml_tensor * clip_graph_pockettts_seanet::conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, const std::string & state_name) const { const int64_t K = w->ne[0]; @@ -100,7 +99,6 @@ ggml_tensor * clip_graph_pockettts_seanet::conv_transpose1d(ggml_tensor * x, ggm return out; } -// ELU -> dilated conv -> ELU -> pointwise conv, added back to the input ggml_tensor * clip_graph_pockettts_seanet::res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation, const std::string & state_prefix) const { ggml_tensor * h = ggml_elu(ctx0, x); diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index 56f84f43c3a..67c533f40a7 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1427,8 +1427,8 @@ std::vector mtmd_audio_streaming_istft::flush() { // // mtmd_audio_preprocessor_pockettts // -// mimi takes the raw 24kHz waveform, there is no mel front-end. the samples are handed over -// as a single-row "mel" so they travel through the normal chunk path +// mimi takes the raw 24kHz waveform, there is no mel front-end +// the samples are handed over as a single-row "mel", to reuse the normal chunk path // bool mtmd_audio_preprocessor_pockettts::preprocess(const float * samples, diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index e8dc9c276c5..17eb5b3a828 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -93,7 +93,6 @@ class mtmd_gen_audio_pipeline { 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; @@ -463,10 +462,8 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::vector out_buf; }; -// Settings that live only in the reference's per-pack yaml and are not derivable from the -// checkpoint: the english packs are identical in shape and tokenizer yet disagree on them. -// They are keyed on the weight variant name that the mmproj carries. -// remove_semicolons belongs here too, but it maps ";" to "," and is applied to every pack. +// 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 struct pockettts_pack_settings { float temp = 0.7f; // Config.default_temperature int frames_after_eos = 0; // 0 leaves the tail length to the caller @@ -490,8 +487,8 @@ static pockettts_pack_settings pockettts_pack(const char * variant) { return it->second; } -// Pocket-TTS: the backbone emits no token at all, each step's hidden state is turned into one -// continuous latent by the flow net, and the end-of-speech head lives in the mmproj +// 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; @@ -549,8 +546,8 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } ids.resize((size_t) n_ids); - // long inputs degrade badly, the reference splits them and restarts each piece from - // the voice conditioning, see split_into_best_sentences() + // 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) { @@ -625,8 +622,7 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { mtmd_gen_inp inp{}; inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; inp.embd = const_cast(h_state_in); - // the same seed every step: clip only reseeds when it changes, so the noise - // stream keeps running instead of restarting on each frame + // clip only reseeds when the seed changes, so pass the same one on every step inp.seed = seed; inp.n_steps = -1; inp.flow_temp = pack.temp; @@ -806,8 +802,7 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { 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, otherwise the reference guesses it from the word count, - // approximated here by tokens + // 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; @@ -942,8 +937,7 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { return ok; } - // decodes the buffered latents, carrying the mimi decoder state across calls so a window - // can be emitted as soon as it is full + // decodes the buffered latents, the mimi decoder state carries over between calls bool flush_gen_wav() { if (feats_buf.empty()) { return true; @@ -983,7 +977,6 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { int step_idx = 0; int eos_step = -1; int frames_after_eos = 3; - // long inputs are split, each chunk restarts from the voice conditioning static constexpr int max_chunk_tokens = 50; // MAX_TOKEN_PER_CHUNK in the reference static constexpr double frame_rate = 12.5; std::vector> chunks; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 5d6f6c4ac97..906f0d0d94f 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1699,8 +1699,7 @@ 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; - // gen_wav draws no randomness, but the seed must still match so it does not reseed - // the rng in the middle of a generation + // 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; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index ee93c4bfdb5..2921c6e13a3 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -340,7 +340,6 @@ struct mtmd_gen_audio_info { enum mtmd_gen_audio_type type; int32_t sample_rate; // in Hz, for example 24000 for qwen3tts const char * model_variant; // name of the weight variant, empty if the mmproj has none - // some pipelines have settings that only exist per-variant }; MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx); From 4eef4072b29afc0d3f82d2c3f7d203accc488c0f Mon Sep 17 00:00:00 2001 From: Pascal Date: Thu, 6 Aug 2026 17:06:05 +0200 Subject: [PATCH 10/17] nit: drop the dead flow_temp hparam, the pack table holds the default --- conversion/pockettts.py | 2 +- tools/mtmd/clip-model.h | 1 - tools/mtmd/clip.cpp | 3 ++- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/conversion/pockettts.py b/conversion/pockettts.py index 04539e4681c..90e16cd2b49 100644 --- a/conversion/pockettts.py +++ b/conversion/pockettts.py @@ -169,7 +169,7 @@ def set_gguf_parameters(self): self.gguf_writer.add_audio_num_mel_bins(1) # generation: flow-matching decoder + mimi decoder - # note: the SEANet and flow net hparams are hardcoded on the clip.cpp side for now + # the SEANet and flow net hparams are constant across the family, clip.cpp holds them self.gguf_writer.add_clip_has_gen_audio_encoder(True) self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_GEN) self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text) diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index e7912c4dbb8..6e9ede71e69 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -148,7 +148,6 @@ struct clip_hparams { int32_t mimi_downsample = 0; // encoder frame rate / model frame rate int32_t mimi_tfm_context = 0; // attention window of the mimi transformers, in frames int32_t flow_n_step = 1; // lsd_decode steps - float flow_temp = 0.7f; // noise std is sqrt(temp), the caller can override it // qwen3tts code2wav int32_t wav_tfm_n_layer = 0; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index bb39d8b0905..f2d9cd839aa 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -4851,7 +4851,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } else { // flow matching starts from gaussian noise, std = sqrt(temp) ggml_tensor * t = get_inp_tensor("inp_noise"); - const float temp = params->flow_temp > 0.0f ? params->flow_temp : hparams.flow_temp; + // Config.default_temperature, a caller that knows its variant overrides it + const float temp = params->flow_temp > 0.0f ? params->flow_temp : 0.7f; std::normal_distribution dist(0.0f, std::sqrt(temp)); std::vector noise(ggml_nelements(t)); for (auto & v : noise) { From 57dc1116888b09f2fbb65accb364df531354d195 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 10 Aug 2026 19:39:22 +0200 Subject: [PATCH 11/17] update docs --- tools/mtmd/README-dev.md | 15 +++++++++++---- tools/mtmd/mtmd.h | 2 +- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/tools/mtmd/README-dev.md b/tools/mtmd/README-dev.md index 3cddd085ec6..ac43e1b81b1 100644 --- a/tools/mtmd/README-dev.md +++ b/tools/mtmd/README-dev.md @@ -59,8 +59,10 @@ Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system i ### Checklist for porting new audio generation models to mtmd -1. Establish a list of reusable and missing components from the current mtmd implementation. -2. For GGUF conversion: +1. Make sure to consult merged PRs about adding new TTS models, especially reviewer comments + - Example: https://github.com/ggml-org/llama.cpp/pulls?q=is%3Apr+mtmd+tts+is%3Amerged +2. Establish a list of reusable and missing components from the current mtmd implementation. +3. For GGUF conversion: - Backbone model should be converted to a normal text model (loadable via `libllama`) - If model used hard-coded embedding row ID, append them to token embeddings and assign token name for them (see `qwen3tts.py`) - If model have a specific output logits head for audio codes (usually semantic code), keep the head as-is and pad the logits at inference time (see `src/models/qwen3vl.cpp`) @@ -70,12 +72,17 @@ Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system i - For tensor naming: - Prefixed with `a.*` for tensors used by speaker encoder pipeline - Prefixed with `a.gen.*` for generation stages (code / mel-spectrogram / PCM generation) -3. Make sure most of the changes happen inside `mtmd-helper-gen.cpp`. A good PR looks like this: + - For GGUF metadata: + - Reuse as many existing keys as possible + - In most cases, you can hard-code model configs in the model graph class, or in `clip_hparams` + - If some values need to be exposed to the `mtmd_helper` layer, hard-code them in `mtmd_helper` and distinguish by pipeline and `mtmd_gen_audio_info::model_variant` if necessary + - Do NOT add new GGUF metadata or new fields to `mtmd_gen_audio_info` unless you can prove that you absolutely need them +4. 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 -4. Update usage documentation in `tools/tts/README.md` +5. Update usage documentation in `tools/tts/README.md` IMPORTANT: If your model needs changes that don't fit the existing infrastructure, **open an issue first for discussion**. diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index eb4f2246ea0..5e8c9a35e3c 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -349,7 +349,7 @@ enum mtmd_gen_audio_type { struct mtmd_gen_audio_info { enum mtmd_gen_audio_type type; int32_t sample_rate; // in Hz, for example 24000 for qwen3tts - const char * model_variant; // name of the weight variant, empty if the mmproj has none + const char * model_variant; // name of the weight variant, can be nullptr if not applicable }; MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx); From 6c197a77ea77ab77baaf5c3238beec924f373e64 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 10 Aug 2026 23:54:13 +0200 Subject: [PATCH 12/17] address security problems --- tools/mtmd/clip-model.h | 1 + tools/mtmd/clip.cpp | 12 +++++++++++ tools/mtmd/models/pockettts-gen.cpp | 6 ++++++ tools/mtmd/mtmd-audio.cpp | 8 ++++++++ tools/mtmd/mtmd-helper-gen.cpp | 32 +++++++++++++++-------------- 5 files changed, 44 insertions(+), 15 deletions(-) diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 206fc4be677..ad25c008e73 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -148,6 +148,7 @@ struct clip_hparams { std::string gen_model_variant; // pocket-tts + static constexpr int32_t pockettts_max_spk_seconds = 30; int32_t seanet_n_stage = 0; std::vector seanet_ratios; // encoder order (reversed compared to the config) int32_t mimi_downsample = 0; // encoder frame rate / model frame rate diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index cb03baabb1a..a0c1c169ba3 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1072,6 +1072,13 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE; const int n_step = params && params->n_steps > 0 ? params->n_steps : ctx->model.hparams.flow_n_step; const int64_t n_latent = ctx->model.gen_input_lin_w->ne[0]; + GGML_ASSERT(n_step > 0); + GGML_ASSERT(n_latent > 0); + // "inp_feats" takes the caller's buffer as-is, the graph must consume all of it + if (params && params->feats) { + GGML_ASSERT(params->feats->size() % (size_t) n_latent == 0); + GGML_ASSERT(params->feats->size() >= (size_t) n_latent); + } const int n_frames = params && params->feats ? (int) (params->feats->size() / n_latent) : 1; builder = std::make_unique(ctx, img, gen_process, n_step, n_frames); } break; @@ -4341,12 +4348,17 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { // the mask is causal with a sliding window, see _build_attention_mask() in the reference auto set_pockettts_tfm_inputs = [&]() { const int64_t n_pos = ggml_nelements(get_inp_tensor("inp_pos")); + GGML_ASSERT(n_pos > 0); std::vector positions((size_t) n_pos); for (int64_t i = 0; i < n_pos; i++) { positions[(size_t) i] = (int32_t) i; } set_input_i32("inp_pos", positions); + // the preprocessor truncates the waveform to keep this mask bounded + const int64_t max_pos = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate / 120; + GGML_ASSERT(n_pos <= max_pos && "pocket-tts speaker reference too long for a dense mask"); + const int64_t context = hparams.mimi_tfm_context; std::vector mask((size_t) n_pos * n_pos, -INFINITY); for (int64_t q = 0; q < n_pos; q++) { diff --git a/tools/mtmd/models/pockettts-gen.cpp b/tools/mtmd/models/pockettts-gen.cpp index 3266742150c..3fd613e5f7f 100644 --- a/tools/mtmd/models/pockettts-gen.cpp +++ b/tools/mtmd/models/pockettts-gen.cpp @@ -106,6 +106,12 @@ std::vector list_pockettts_state_slots(const clip_hparams & hpar } const auto & seanet = model.seanet; + // the slots below are sized from these + GGML_ASSERT(!model.gen_tfm_layers.empty()); + GGML_ASSERT((int) seanet.stages.size() >= hparams.seanet_n_stage); + GGML_ASSERT((int) hparams.seanet_ratios.size() >= hparams.seanet_n_stage); + GGML_ASSERT(hparams.mimi_tfm_context > 1 && hparams.mimi_downsample > 0); + slots.push_back({"tfm_pos", 1, 1}); const int64_t n_embd_a = model.gen_tfm_layers[0].q_w->ne[1]; diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index 67c533f40a7..98a8c11ee91 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1440,6 +1440,14 @@ bool mtmd_audio_preprocessor_pockettts::preprocess(const float * return false; } + // the mimi transformer mask is dense, so cost is quadratic in the reference length + const int64_t max_samples = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate; + if ((int64_t) n_samples > max_samples) { + LOG_WRN("%s: speaker reference is %.1f s, truncating to the first %d s\n", __func__, + (double) n_samples / hparams.audio_sample_rate, clip_hparams::pockettts_max_spk_seconds); + n_samples = (size_t) max_samples; + } + const int64_t n_frames = (int64_t) (n_samples + frame_size - 1) / frame_size; const int64_t n_padded = n_frames * frame_size; diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 40b0ef74724..72106aefb2d 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -549,16 +549,12 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } const int n_e = n_embd; - auto push_row = [&](llama_token t) { - prompt_embd_buf.insert(prompt_embd_buf.end(), - tok_embd.begin() + (size_t) t * n_e, - tok_embd.begin() + (size_t) (t + 1) * n_e); - }; // sequence order is voice, then text, then the audio BOS that starts generation if (!voice.empty()) { + GGML_ASSERT(voice.size() % (size_t) n_e == 0); if (bos_before_voice != LLAMA_TOKEN_NULL) { - push_row(bos_before_voice); + push_embd_row(prompt_embd_buf, bos_before_voice); } prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end()); } @@ -566,9 +562,9 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { n_voice_pos = (int) (prompt_embd_buf.size() / (size_t) n_e); for (llama_token t : chunks[0]) { - push_row(t); + push_embd_row(prompt_embd_buf, t); } - push_row(audio_bos); + push_embd_row(prompt_embd_buf, audio_bos); arm_chunk_budget(0); n_prompt = (int) (prompt_embd_buf.size() / (size_t) n_e); @@ -712,10 +708,20 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { LOG_ERR("mtmd_helper_gen_audio: token embedding copy failed\n"); return false; } + GGML_ASSERT(n_embd > 0 && n_tok_embd % (uint32_t) n_embd == 0); specials_ok = true; return true; } + // the table can be shorter than the vocab, so bound the row lookup + void push_embd_row(std::vector & dst, llama_token t) const { + const size_t n_rows = tok_embd.size() / (size_t) n_embd; + GGML_ASSERT(t >= 0 && (size_t) t < n_rows); + dst.insert(dst.end(), + tok_embd.begin() + (size_t) t * n_embd, + tok_embd.begin() + (size_t) (t + 1) * n_embd); + } + // token ids of the pieces the reference splits on, see split_into_best_sentences(). // the leading token is dropped, it is the tokenizer's dummy prefix std::vector punct_ids(const char * s) const { @@ -824,17 +830,13 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { const int n_e = n_embd; prompt_embd_buf.clear(); for (llama_token t : chunks[chunk_idx]) { - prompt_embd_buf.insert(prompt_embd_buf.end(), - tok_embd.begin() + (size_t) t * n_e, - tok_embd.begin() + (size_t) (t + 1) * n_e); + push_embd_row(prompt_embd_buf, t); } - prompt_embd_buf.insert(prompt_embd_buf.end(), - tok_embd.begin() + (size_t) audio_bos * n_e, - tok_embd.begin() + (size_t) (audio_bos + 1) * n_e); + push_embd_row(prompt_embd_buf, audio_bos); arm_chunk_budget(chunk_idx); - // bounded by max_chunk_tokens + 1, so one decode is enough const int n_rows = (int) (prompt_embd_buf.size() / (size_t) n_e); + GGML_ASSERT(n_rows > 0); decode_embd_batch batch(prompt_embd_buf.data(), n_rows, 1, n_e); batch.set_position_normal(pos, seq_id); batch.batch.logits[n_rows - 1] = 1; From 93055059cc4dba7af3cb7b803a3f645ea3da5764 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Tue, 11 Aug 2026 00:19:14 +0200 Subject: [PATCH 13/17] less invasive base.py --- conversion/base.py | 70 +++++++++++++++-------------------------- conversion/pockettts.py | 49 ++++++++++++++++++++++++++--- 2 files changed, 70 insertions(+), 49 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index eb5a1d32a2e..3572b77c21e 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -58,6 +58,11 @@ AnyModel = TypeVar("AnyModel", bound="type[ModelBase]") +# for checkpoints that ship no config.json, we will try to provide a synthetic one +HparamsMatcher = Callable[[Path], bool] +HparamsLoader = Callable[[Path], dict[str, Any]] + + class SentencePieceTokenTypes(IntEnum): NORMAL = 1 UNKNOWN = 2 @@ -77,6 +82,7 @@ class ModelBase: ModelType.TEXT: {}, ModelType.MMPROJ: {}, } + _hparams_loaders: list[tuple[HparamsMatcher, HparamsLoader]] = [] dir_model: Path ftype: gguf.LlamaFileType @@ -1040,6 +1046,24 @@ def get_model_part_names(dir_model: Path, prefix: str, suffix: str) -> list[str] return part_names + @staticmethod + def load_hparams_guess(dir_model: Path) -> dict[str, Any] | None: + # some models ship no config.json, will try to guess them + from conversion import load_all_models + load_all_models() + + for matcher, loader in ModelBase._hparams_loaders: + if matcher(dir_model): + return loader(dir_model) + return None + + @classmethod + def register_hparams_loader(cls, matcher: HparamsMatcher) -> Callable[[HparamsLoader], HparamsLoader]: + def inner(loader: HparamsLoader) -> HparamsLoader: + cls._hparams_loaders.append((matcher, loader)) + return loader + return inner + @staticmethod def load_hparams(dir_model: Path, is_mistral_format: bool): if is_mistral_format: @@ -1054,7 +1078,7 @@ def load_hparams(dir_model: Path, is_mistral_format: bool): except Exception as e: logger.warning(f"Failed to load model config from {dir_model}: {e}") if not (dir_model / "config.json").is_file(): - config = load_hparams_non_hf(dir_model) + config = ModelBase.load_hparams_guess(dir_model) if config is not None: return config logger.warning("Trying to load config.json instead") @@ -2622,50 +2646,6 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): LazyTorchTensor._dtype_str_map["F8_E8M0"] = torch.uint8 -def load_hparams_non_hf(dir_model: Path) -> dict[str, Any] | None: - # some models ship no config.json at all, their hparams are derived from the checkpoint - part_names = ModelBase.get_model_part_names(dir_model, "model", ".safetensors") - if len(part_names) != 1: - return None - with gguf.utility.SafetensorsLocal(dir_model / part_names[0]) as part: - shapes = {name: tuple(part[name].shape) for name in part.keys()} - - if "flow_lm.bos_emb" in shapes: - return _load_hparams_pockettts(shapes) - - return None - - -def _load_hparams_pockettts(shapes: dict[str, tuple[int, ...]]) -> dict[str, Any]: - logger.info("gguf: detected pocket-tts checkpoint, deriving hparams from tensor shapes") - n_vocab, n_embd = shapes["flow_lm.conditioner.embed.weight"] - n_layer = sum(1 for name in shapes if re.fullmatch(r"flow_lm\.transformer\.layers\.\d+\.norm1\.weight", name)) - n_layer_a = sum(1 for name in shapes if re.fullmatch(r"mimi\.encoder_transformer\.transformer\.layers\.\d+\.norm1\.weight", name)) - n_embd_a = shapes["mimi.encoder_transformer.transformer.layers.0.norm1.weight"][0] - return { - "architectures": ["PocketTTSModel"], - "model_type": "pockettts", - "num_hidden_layers": n_layer, - "hidden_size": n_embd, - "intermediate_size": shapes["flow_lm.transformer.layers.0.linear1.weight"][0], - # the transformer is fully causal with no context limit, this only bounds the KV cache - "max_position_embeddings": 4096, - # not stored anywhere in the checkpoint, but every released variant uses head_dim 64 - "num_attention_heads": n_embd // 64, - # extra rows for the learned input vectors, see pockettts.py - # bos_before_voice only exists when the pack inserts it - "vocab_size": n_vocab + (2 if "flow_lm.bos_before_voice" in shapes else 1), - "rope_theta": 10000.0, - "layer_norm_eps": 1e-5, - "audio_config": { - "num_hidden_layers": n_layer_a, - "hidden_size": n_embd_a, - "intermediate_size": shapes["mimi.encoder_transformer.transformer.layers.0.linear1.weight"][0], - "num_attention_heads": n_embd_a // 64, - }, - } - - def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> str: # TODO @ngxson : this won't work correctly if the model has both audio & vision encoders # maybe we should fallback to text model's arch in that case, since not many models have both diff --git a/conversion/pockettts.py b/conversion/pockettts.py index 90e16cd2b49..62ecb5acde7 100644 --- a/conversion/pockettts.py +++ b/conversion/pockettts.py @@ -1,17 +1,19 @@ from __future__ import annotations -from typing import Iterable, TYPE_CHECKING +import re +from pathlib import Path +from typing import Any, Iterable, TYPE_CHECKING import torch if TYPE_CHECKING: from torch import Tensor -from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf +from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf, logger # Pocket TTS is a CALM: the backbone conditions a flow-matching decoder that generates one # continuous 32-d latent per frame. There is no codebook in this model. -# The checkpoint ships no config.json, hparams are derived in base.load_hparams_non_hf(). +# The checkpoint ships no config.json, hparams come from _load_hparams() below. # # Tricks being used to support this model via existing llama.cpp code paths: # - bos_before_voice and bos_emb are learned input vectors, not tokens @@ -35,6 +37,45 @@ _SAMPLE_RATE = 24000 +def _tensor_shapes(dir_model: Path) -> dict[str, tuple[int, ...]]: + part_names = ModelBase.get_model_part_names(dir_model, "model", ".safetensors") + if len(part_names) != 1: + return {} + with gguf.utility.SafetensorsLocal(dir_model / part_names[0]) as part: + return {name: tuple(part[name].shape) for name in part.keys()} + + +@ModelBase.register_hparams_loader(lambda dir_model: "flow_lm.bos_emb" in _tensor_shapes(dir_model)) +def _load_hparams(dir_model: Path) -> dict[str, Any]: + logger.info("gguf: detected pocket-tts checkpoint, deriving hparams from tensor shapes") + shapes = _tensor_shapes(dir_model) + n_vocab, n_embd = shapes["flow_lm.conditioner.embed.weight"] + n_layer = sum(1 for name in shapes if re.fullmatch(r"flow_lm\.transformer\.layers\.\d+\.norm1\.weight", name)) + n_layer_a = sum(1 for name in shapes if re.fullmatch(r"mimi\.encoder_transformer\.transformer\.layers\.\d+\.norm1\.weight", name)) + n_embd_a = shapes["mimi.encoder_transformer.transformer.layers.0.norm1.weight"][0] + return { + "architectures": ["PocketTTSModel"], + "model_type": "pockettts", + "num_hidden_layers": n_layer, + "hidden_size": n_embd, + "intermediate_size": shapes["flow_lm.transformer.layers.0.linear1.weight"][0], + # the transformer is fully causal with no context limit, this only bounds the KV cache + "max_position_embeddings": 4096, + # not in the checkpoint, but every released variant uses head_dim 64 + "num_attention_heads": n_embd // 64, + # extra rows for the learned input vectors, see _embd_table() + "vocab_size": n_vocab + (2 if "flow_lm.bos_before_voice" in shapes else 1), + "rope_theta": 10000.0, + "layer_norm_eps": 1e-5, + "audio_config": { + "num_hidden_layers": n_layer_a, + "hidden_size": n_embd_a, + "intermediate_size": shapes["mimi.encoder_transformer.transformer.layers.0.linear1.weight"][0], + "num_attention_heads": n_embd_a // 64, + }, + } + + @ModelBase.register("PocketTTSModel") class PocketTTSModel(TextModel): model_arch = gguf.MODEL_ARCH.POCKETTTS @@ -324,7 +365,7 @@ def _seanet_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, T return for stage in range(_N_SEANET_STAGES): - res_idx = _DEC_RES_IDX(stage) if is_decoder else _ENC_RES_IDX(stage) + res_idx = _DEC_RES_IDX(stage) if is_decoder else _ENC_RES_IDX(stage) scale_idx = _DEC_SCALE_IDX(stage) if is_decoder else _ENC_SCALE_IDX(stage) if idx == scale_idx: yield (self.format_tensor_name(scale, stage, suffix=suffix), data_torch) From 3e47aaec49a9c54f35b6763586d9558b2a2b78bf Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Tue, 11 Aug 2026 00:26:53 +0200 Subject: [PATCH 14/17] lint --- gguf-py/gguf/gguf_writer.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index ba8c083a868..05f86396dc0 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1456,7 +1456,6 @@ def add_gen_audio_attention_layernorm_eps(self, value: float) -> None: def add_gen_audio_model_variant(self, value: str) -> None: self.add_string(Keys.ClipGenAudio.MODEL_VARIANT, value) - def add_xielu_alpha_p(self, values: Sequence[float]): self.add_array(Keys.xIELU.ALPHA_P, values) From 1d591cc2b7efbfd155676ef75e7d89bd9fcab92f Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Tue, 11 Aug 2026 00:51:10 +0200 Subject: [PATCH 15/17] add mtmd_gen_inp_default --- tools/mtmd/clip-model.h | 3 +++ tools/mtmd/clip.cpp | 6 +++--- tools/mtmd/clip.h | 3 +-- tools/mtmd/mtmd-helper-gen.cpp | 37 ++++++++++++++++++---------------- tools/mtmd/mtmd.cpp | 28 +++++++++++++++++++++++-- tools/mtmd/mtmd.h | 14 ++++++++++--- 6 files changed, 64 insertions(+), 27 deletions(-) diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index ad25c008e73..4465d9e0ef0 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -144,6 +144,9 @@ struct clip_hparams { // threshold for the "out_eos_score" graph output float gen_eos_threshold = 0.0f; + // default noise scale of a flow-matching decoder, see mtmd_gen_inp_default() + float gen_flow_temp = 0.0f; + // name of the weight variant, some pipelines tune themselves on it std::string gen_model_variant; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index a0c1c169ba3..469da5f4aae 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1070,7 +1070,7 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const case PROJECTOR_TYPE_POCKETTTS_GEN: { const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE; - const int n_step = params && params->n_steps > 0 ? params->n_steps : ctx->model.hparams.flow_n_step; + const int n_step = ctx->model.hparams.flow_n_step; const int64_t n_latent = ctx->model.gen_input_lin_w->ne[0]; GGML_ASSERT(n_step > 0); GGML_ASSERT(n_latent > 0); @@ -1784,6 +1784,7 @@ struct clip_model_loader { // flow_lm defaults, see pocket_tts/default_parameters.py hparams.flow_n_step = 1; hparams.gen_eos_threshold = -4.0f; + hparams.gen_flow_temp = 0.7f; // Config.default_temperature } break; case PROJECTOR_TYPE_PADDLEOCR: { @@ -4963,8 +4964,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } else { // flow matching starts from gaussian noise, std = sqrt(temp) ggml_tensor * t = get_inp_tensor("inp_noise"); - // Config.default_temperature, a caller that knows its variant overrides it - const float temp = params->flow_temp > 0.0f ? params->flow_temp : 0.7f; + const float temp = params->temp > 0.0f ? params->temp : hparams.gen_flow_temp; std::normal_distribution dist(0.0f, std::sqrt(temp)); std::vector noise(ggml_nelements(t)); for (auto & v : noise) { diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index a237a9466a3..a5b71377523 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -106,8 +106,7 @@ struct clip_encode_params { std::vector * out_codes = nullptr; // this frame's 16 sampled codes std::vector * out_feats = nullptr; // continuous counterpart of out_codes uint32_t seed = UINT32_MAX; // UINT32_MAX for random - int32_t n_steps = -1; // integration steps, for flow-matching decoders - float flow_temp = 0.0f; // noise scale of the flow decoder, 0 for default + float temp = 0.0f; // sampling temperature, noise scale for flow-matching decoders bool * out_is_eos = nullptr; // GEN_WAV diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 72106aefb2d..1c58d3ae195 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -203,8 +203,9 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { 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; @@ -258,7 +259,7 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { return 0; } - mtmd_gen_inp inp{}; + 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(h_state_in); @@ -401,10 +402,11 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { 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{}; @@ -458,9 +460,10 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { // 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.7f; // Config.default_temperature - int frames_after_eos = 0; // 0 leaves the tail length to the caller + float temp = 0.0f; + int frames_after_eos = 0; bool pad_short_text = false; }; @@ -473,10 +476,9 @@ static pockettts_pack_settings pockettts_pack(const char * variant) { }; auto it = packs.find(variant ? variant : ""); if (it == packs.end()) { - pockettts_pack_settings def; - LOG_WRN("mtmd_helper_gen_audio: no tuned settings for pocket-tts variant \"%s\", " - "using temperature %.1f\n", variant ? variant : "", def.temp); - return def; + LOG_WRN("mtmd_helper_gen_audio: no tuned settings for pocket-tts variant \"%s\"\n", + variant ? variant : ""); + return {}; } return it->second; } @@ -609,13 +611,14 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { 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{}; - inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; - inp.embd = const_cast(h_state_in); + mtmd_gen_inp inp = mtmd_gen_inp_default(mctx); + inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; + inp.embd = const_cast(h_state_in); // clip only reseeds when the seed changes, so pass the same one on every step - inp.seed = seed; - inp.n_steps = -1; - inp.flow_temp = pack.temp; + 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"); @@ -938,7 +941,7 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { if (feats_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.feats = feats_buf.data(); inp.n_feats = feats_buf.size(); diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 12497755450..1cf297829a3 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1826,6 +1826,31 @@ 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; + } + + const clip_hparams * hparams = clip_get_hparams(ctx->ctx_gen_a); + 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: + inp.temp = hparams->gen_flow_temp; + 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) { @@ -1862,8 +1887,7 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in params.top_k = inp->top_k; params.top_p = inp->top_p; params.seed = inp->seed; - params.n_steps = inp->n_steps; - params.flow_temp = inp->flow_temp; + params.temp = inp->temp; params.out_is_eos = &is_eos; if (!clip_encode(ctx_clip, ¶ms)) { diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 5e8c9a35e3c..c1a5921db2f 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -346,19 +346,23 @@ enum mtmd_gen_audio_type { MTMD_GEN_AUDIO_TYPE_QWEN3TTS, MTMD_GEN_AUDIO_TYPE_POCKETTTS, }; + struct mtmd_gen_audio_info { enum mtmd_gen_audio_type type; int32_t sample_rate; // in Hz, for example 24000 for qwen3tts const char * model_variant; // name of the weight variant, can be nullptr if not applicable }; + 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 semantic (codes, mel-spectrogram, etc.) MTMD_GEN_PROCESS_TYPE_GEN_WAV, // convert semantic to PCM audio // for qwen3tts, this is code2wav // for pocket-tts, this is mimi decoder }; + struct mtmd_gen_inp { enum mtmd_gen_process_type type; @@ -367,9 +371,8 @@ struct mtmd_gen_inp { float * embd; // the hidden state from backbone, must have n_text_embd elements int32_t top_k; float top_p; - uint32_t seed; // UINT32_MAX for random - int32_t n_steps; // integration steps, for flow-matching decoders (-1 for default) - float flow_temp; // noise scale, for flow-matching decoders (0 for default) + uint32_t seed; // UINT32_MAX for random + float temp; // sampling temperature, or noise scale for flow-matching decoders // for MTMD_GEN_PROCESS_TYPE_GEN_WAV // pass either codes (discrete) or feats (continuous), depending on the pipeline @@ -380,6 +383,7 @@ struct mtmd_gen_inp { const char * state_data; size_t state_size; }; + struct mtmd_gen_out { // note: output memory is allocated by the context, valid until next process() call @@ -398,6 +402,10 @@ struct mtmd_gen_out { const char * state_data; size_t state_size; }; + +// defaults tuned for the loaded pipeline, callers override only what they care about +MTMD_API struct mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx); + // 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, From 606375f8c3c6804bbe13883fcf5a09c66f0f6487 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Tue, 11 Aug 2026 00:56:18 +0200 Subject: [PATCH 16/17] add docs --- tools/tts/README.md | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tools/tts/README.md b/tools/tts/README.md index dd84336c399..1b08d5ef321 100644 --- a/tools/tts/README.md +++ b/tools/tts/README.md @@ -32,3 +32,28 @@ llama-tts -hf ggml-org/Qwen3-TTS-12Hz-1.7B-Base-GGUF \ --tts-speaker-file speaker.mp3 \ --output out.wav ``` + +## Pocket TTS + +Available params: +- `--tts-speaker-file` should point to a speaker reference audio file (wav, mp3). It is required, the model produces almost no audio without it +- Note: `lang` is not used, the language is a property of the weights + +Example usage: + +```sh +llama-tts -m pocket-tts.gguf \ + -mm mmproj-pocket-tts.gguf \ + -p "Hello world" \ + --tts-speaker-file speaker.mp3 \ + --output out.wav +``` + +**Note for GGUF conversion:** + +The [upstream repository](https://huggingface.co/kyutai/pocket-tts) holds one complete model per language under `languages/`, next to a set of shared files at the root. Convert one of the `languages/` directories, **not** the root directory: + +```sh +python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --outfile pocket-tts.gguf +python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --mmproj --outfile mmproj-pocket-tts.gguf +``` From f236e272f0053d20a60bd0925cf181c583bb221b Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Tue, 11 Aug 2026 01:01:16 +0200 Subject: [PATCH 17/17] rm gen_flow_temp --- tools/mtmd/clip-model.h | 3 --- tools/mtmd/clip.cpp | 4 ++-- tools/mtmd/mtmd.cpp | 6 ++++-- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 4465d9e0ef0..ad25c008e73 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -144,9 +144,6 @@ struct clip_hparams { // threshold for the "out_eos_score" graph output float gen_eos_threshold = 0.0f; - // default noise scale of a flow-matching decoder, see mtmd_gen_inp_default() - float gen_flow_temp = 0.0f; - // name of the weight variant, some pipelines tune themselves on it std::string gen_model_variant; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 469da5f4aae..2fb2b5041dc 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1784,7 +1784,6 @@ struct clip_model_loader { // flow_lm defaults, see pocket_tts/default_parameters.py hparams.flow_n_step = 1; hparams.gen_eos_threshold = -4.0f; - hparams.gen_flow_temp = 0.7f; // Config.default_temperature } break; case PROJECTOR_TYPE_PADDLEOCR: { @@ -4964,7 +4963,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } else { // flow matching starts from gaussian noise, std = sqrt(temp) ggml_tensor * t = get_inp_tensor("inp_noise"); - const float temp = params->temp > 0.0f ? params->temp : hparams.gen_flow_temp; + // Config.default_temperature, for a caller that does not set one + const float temp = params->temp > 0.0f ? params->temp : 0.7f; std::normal_distribution dist(0.0f, std::sqrt(temp)); std::vector noise(ggml_nelements(t)); for (auto & v : noise) { diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 1cf297829a3..4b9c45d6267 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1834,7 +1834,6 @@ mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) { return inp; } - const clip_hparams * hparams = clip_get_hparams(ctx->ctx_gen_a); 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 @@ -1843,7 +1842,10 @@ mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) { inp.temp = 0.9f; // TODO: handle this on graph break; case PROJECTOR_TYPE_POCKETTTS_GEN: - inp.temp = hparams->gen_flow_temp; + // 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;