clip.gen.audio.model_variant

This commit is contained in:
Xuan Son Nguyen
2026-08-06 16:19:09 +02:00
parent d14bb68e5f
commit 4b9a4df9b9
10 changed files with 60 additions and 74 deletions
+2 -35
View File
@@ -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/<name>.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
+2 -4
View File
@@ -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"
+2 -8
View File
@@ -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]):
+2 -4
View File
@@ -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"
//
+4 -3
View File
@@ -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<int32_t> 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;
+3 -7
View File
@@ -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<float> 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<float> dist(0.0f, std::sqrt(temp));
std::vector<float> noise(ggml_nelements(t));
for (auto & v : noise) {
v = dist(ctx->rng);
+1
View File
@@ -107,6 +107,7 @@ struct clip_encode_params {
std::vector<float> * 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
+35 -4
View File
@@ -463,6 +463,33 @@ 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.
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<std::string, pockettts_pack_settings> 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<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
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;
+6 -7
View File
@@ -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, &params)) {
+3 -2
View File
@@ -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