diff --git a/conversion/base.py b/conversion/base.py index 69f0fd9e65..99bb00a1db 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 a62243bb22..b6608cf884 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 +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 @@ _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). +# +# 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 @@ class PocketTTSModel(TextModel): 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 @@ class PocketTTSModel(TextModel): 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 @@ 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)) + 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 991d930691..14deedf51c 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -400,6 +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" 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 39da9f2c05..9eacb94650 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1438,6 +1438,10 @@ 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_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 0a31a7aec1..c3c35d6c3e 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 9b0b2c5e2f..1d6bea3952 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 bcb7cb083e..520e921bf0 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 a0106b6895..f1ecb4f317 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -503,6 +503,7 @@ public: 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 @@ public: // 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 @@ private: 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 @@ private: 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(' ');