From a73f458f3b5cd9a1e4593069cbcbff3e66eed220 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Thu, 6 Aug 2026 16:42:25 +0200 Subject: [PATCH] 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 99bb00a1db..eb5a1d32a2 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 63e41e944e..04539e4681 100644 --- a/conversion/pockettts.py +++ b/conversion/pockettts.py @@ -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] diff --git a/tools/mtmd/models/pockettts-gen.cpp b/tools/mtmd/models/pockettts-gen.cpp index b76d44da38..3266742150 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 bf8fff14a7..c47207f569 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 56f84f43c3..67c533f40a 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 e8dc9c276c..17eb5b3a82 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -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 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(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> chunks; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 5d6f6c4ac9..906f0d0d94 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 ee93c4bfdb..2921c6e13a 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);