flow_temp + frames_after_eos

This commit is contained in:
Xuan Son Nguyen
2026-08-06 00:13:26 +02:00
parent 9ff97e3157
commit b282950725
8 changed files with 73 additions and 20 deletions
+3 -2
View File
@@ -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": {
+41 -6
View File
@@ -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/<name>.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
+2
View File
@@ -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"
+4
View File
@@ -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)
+3 -1
View File
@@ -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
+1 -1
View File
@@ -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;
+6 -3
View File
@@ -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:
{
+13 -7
View File
@@ -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<llama_token> 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(' ');