clean up code comments

This commit is contained in:
Xuan Son Nguyen
2026-08-06 16:42:25 +02:00
parent 4b9a4df9b9
commit a73f458f3b
8 changed files with 31 additions and 45 deletions
+2 -2
View File
@@ -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,
+8 -10
View File
@@ -9,17 +9,15 @@ if TYPE_CHECKING:
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]
+6 -7
View File
@@ -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<c2w_state_slot> list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model) {
std::vector<c2w_state_slot> 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);
+3 -5
View File
@@ -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);
+2 -2
View File
@@ -1427,8 +1427,8 @@ std::vector<float> 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,
+9 -16
View File
@@ -93,7 +93,6 @@ public:
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 @@ private:
std::vector<char> 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 @@ public:
}
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 @@ public:
mtmd_gen_inp inp{};
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
inp.embd = const_cast<float *>(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 @@ private:
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 @@ private:
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 @@ private:
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<std::vector<llama_token>> chunks;
+1 -2
View File
@@ -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;
-1
View File
@@ -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);