diff --git a/conversion/pockettts.py b/conversion/pockettts.py index 05934acab4..63e41e944e 100644 --- a/conversion/pockettts.py +++ b/conversion/pockettts.py @@ -7,7 +7,7 @@ import torch 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 @@ _DEC_RES_IDX = lambda i: 3 + 3 * i # noqa: E731 _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 @@ class PocketTTSMmprojModel(MmprojModel): 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 3e6a9e17a8..a2c1c53fd8 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -400,10 +400,8 @@ class Keys: 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 e67e2b6fef..417d8dee23 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1438,14 +1438,8 @@ class GGUFWriter: 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 f6657e0bd5..d693072089 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 9302f65a41..e7912c4dbb 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 bf5e386a51..bb39d8b090 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 97de0606db..a237a9466a 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 f2045bffd2..e8dc9c276c 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -463,6 +463,33 @@ private: 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 @@ public: } } + 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 @@ public: 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 @@ private: 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 @@ private: 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 a3d72314cb..5d6f6c4ac9 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 e977ecd8a3..ee93c4bfdb 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