upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-10-02 02:47:26 -05:00
Code Issues Packages Projects Releases Wiki Activity

working impl, need verify and clean up

Browse Source
This commit is contained in:
Xuan Son Nguyen
2026-08-05 15:06:40 +02:00
parent a0c869f043
commit 915b3a1da7
16 changed files with 1339 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
conversion/pockettts.py
+18 -5
View File
@@ -50,21 +50,33 @@ 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
from sentencepiece import sentencepiece_model_pb2 as model
proto = model.ModelProto() # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
proto.ParseFromString(open(self.dir_model / "tokenizer.model", "rb").read())
assert proto.trainer_spec.model_type == 1, "expected a unigram tokenizer"
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 generate_extra_tensors()
# padding row, then the two learned vectors appended by _embd_table()
extra = ["<|pad|>", "<|bos_before_voice|>", "<|audio_bos|>"]
for i, name in enumerate(extra):
tokens[len(tokens) - len(extra) + i] = name.encode("utf-8")
toktypes[len(tokens) - len(extra) + i] = SentencePieceTokenTypes.CONTROL
scores[len(tokens) - len(extra) + i] = -1000.0
self.gguf_writer.add_tokenizer_model("llama")
self.gguf_writer.add_tokenizer_model("t5")
self.gguf_writer.add_tokenizer_pre("default")
self.gguf_writer.add_token_list(tokens)
self.gguf_writer.add_token_scores(scores)
self.gguf_writer.add_token_types(toktypes)
self.gguf_writer.add_add_space_prefix(proto.normalizer_spec.add_dummy_prefix)
self.gguf_writer.add_remove_extra_whitespaces(proto.normalizer_spec.remove_extra_whitespaces)
if proto.normalizer_spec.precompiled_charsmap:
self.gguf_writer.add_precompiled_charsmap(proto.normalizer_spec.precompiled_charsmap)
self.gguf_writer.add_add_bos_token(False)
self.gguf_writer.add_add_eos_token(False)
@@ -122,7 +134,7 @@ class PocketTTSMmprojModel(MmprojModel):
"linear1": (gguf.MODEL_TENSOR.A_ENC_FFN_UP, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP),
"linear2": (gguf.MODEL_TENSOR.A_ENC_FFN_DOWN, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN),
"layer_scale_1.scale": (gguf.MODEL_TENSOR.A_ENC_ATTN_SCALE, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE),
"layer_scale_2.scale": (gguf.MODEL_TENSOR.A_ENC_FFN_SCALE, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE),
"layer_scale_2.scale": (gguf.MODEL_TENSOR.A_ENC_FFN_SCALE_LS, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE),
}
_MIMI_TFM_QKV = (
(gguf.MODEL_TENSOR.A_ENC_ATTN_Q, gguf.MODEL_TENSOR.A_ENC_ATTN_K, gguf.MODEL_TENSOR.A_ENC_ATTN_V),
@@ -143,6 +155,8 @@ class PocketTTSMmprojModel(MmprojModel):
self.gguf_writer.add_audio_feed_forward_length(self.hparams_audio["intermediate_size"])
self.gguf_writer.add_audio_head_count(self.hparams_audio["num_attention_heads"])
self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
# mimi convolves the waveform directly, it is passed around as a 1-row "mel"
self.gguf_writer.add_audio_num_mel_bins(1)
# generation: flow-matching decoder + mimi decoder
# note: the SEANet and flow net hparams are hardcoded on the clip.cpp side for now
@@ -273,8 +287,7 @@ class PocketTTSMmprojModel(MmprojModel):
if entry is None:
return
tensor = entry[1 if is_decoder else 0]
# layer_scale is stored without a .weight/.bias suffix
suffix = "" if key_with_suffix.endswith(".scale") else "." + suffix
suffix = ".weight" if key_with_suffix.endswith(".scale") else "." + suffix
yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch)
def _seanet_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]:
gguf-py/gguf/constants.py
+4 -1
View File
@@ -1039,6 +1039,7 @@ class MODEL_TENSOR(IntEnum):
A_ENC_SEANET_RES_CONV2 = auto() # residual unit, pointwise conv
A_ENC_SEANET_SCALE_CONV = auto() # strided downsample conv
A_ENC_ATTN_SCALE = auto() # layer scale (gamma) on the attn output
A_ENC_FFN_SCALE_LS = auto() # layer scale (gamma) on the FFN output
A_ENC_SPEAKER_PROJ = auto() # voice latent -> backbone embd
A_GEN_FLOW_INPUT_PROJ = auto()
A_GEN_FLOW_COND_EMBD = auto()
@@ -1736,7 +1737,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
MODEL_TENSOR.A_ENC_SEANET_RES_CONV1: "a.seanet.blk.{bid}.res_conv1",
MODEL_TENSOR.A_ENC_SEANET_RES_CONV2: "a.seanet.blk.{bid}.res_conv2",
MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV: "a.seanet.blk.{bid}.scale_conv",
MODEL_TENSOR.A_ENC_ATTN_SCALE: "a.blk.{bid}.attn_scale",
MODEL_TENSOR.A_ENC_ATTN_SCALE: "a.blk.{bid}.ls1",
MODEL_TENSOR.A_ENC_FFN_SCALE_LS: "a.blk.{bid}.ls2",
MODEL_TENSOR.A_ENC_SPEAKER_PROJ: "a.speaker_proj",
MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ: "a.gen.flow.input_proj",
MODEL_TENSOR.A_GEN_FLOW_COND_EMBD: "a.gen.flow.cond_embd",
@@ -2078,6 +2080,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.A_ENC_SEANET_RES_CONV2,
MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV,
MODEL_TENSOR.A_ENC_ATTN_SCALE,
MODEL_TENSOR.A_ENC_FFN_SCALE_LS,
MODEL_TENSOR.A_ENC_SPEAKER_PROJ,
MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ,
MODEL_TENSOR.A_GEN_FLOW_COND_EMBD,
tools/mtmd/CMakeLists.txt
+3
View File
@@ -56,6 +56,9 @@ add_library(mtmd
models/mimo-audio.cpp
models/qwen3tts-spkenc.cpp
models/qwen3tts-gen.cpp
models/pockettts-seanet.cpp
models/pockettts-spkenc.cpp
models/pockettts-gen.cpp
models/step3vl.cpp
models/siglip.cpp
models/whisper-enc.cpp
tools/mtmd/clip-impl.h
+36
View File
@@ -246,6 +246,38 @@
#define TN_A_GEN_WAV_DAC_POST_SNAKE "a.gen.wav.dac.post_snake.%s"
#define TN_A_GEN_WAV_DAC_POST_CONV "a.gen.wav.dac.post_conv.%s"
// pocket-tts
#define TN_A_SEANET_CONV_IN "a.seanet.conv_in.%s"
#define TN_A_SEANET_CONV_OUT "a.seanet.conv_out.%s"
#define TN_A_SEANET_RES_CONV1 "a.seanet.blk.%d.res_conv1.%s"
#define TN_A_SEANET_RES_CONV2 "a.seanet.blk.%d.res_conv2.%s"
#define TN_A_SEANET_SCALE_CONV "a.seanet.blk.%d.scale_conv.%s"
#define TN_A_SPEAKER_PROJ "a.speaker_proj.%s"
#define TN_A_DOWNSAMPLE_CONV "a.downsample.conv.%s"
#define TN_A_GEN_FLOW_INPUT_PROJ "a.gen.flow.input_proj.%s"
#define TN_A_GEN_FLOW_COND_EMBD "a.gen.flow.cond_embd.%s"
#define TN_A_GEN_FLOW_TIME_FREQS "a.gen.flow.time.%d.freqs"
#define TN_A_GEN_FLOW_TIME_UP "a.gen.flow.time.%d.up.%s"
#define TN_A_GEN_FLOW_TIME_DOWN "a.gen.flow.time.%d.down.%s"
#define TN_A_GEN_FLOW_TIME_NORM "a.gen.flow.time.%d.norm"
#define TN_A_GEN_FLOW_BLK_NORM "a.gen.flow.blk.%d.norm.%s"
#define TN_A_GEN_FLOW_BLK_UP "a.gen.flow.blk.%d.up.%s"
#define TN_A_GEN_FLOW_BLK_DOWN "a.gen.flow.blk.%d.down.%s"
#define TN_A_GEN_FLOW_BLK_ADA "a.gen.flow.blk.%d.ada.%s"
#define TN_A_GEN_FLOW_FINAL_ADA "a.gen.flow.final.ada.%s"
#define TN_A_GEN_FLOW_FINAL_PROJ "a.gen.flow.final.proj.%s"
#define TN_A_GEN_OUT_EOS "a.gen.out_eos.%s"
#define TN_A_GEN_INPUT_LINEAR "a.gen.input_linear.%s"
#define TN_A_GEN_EMB_MEAN "a.gen.emb_mean"
#define TN_A_GEN_EMB_STD "a.gen.emb_std"
#define TN_A_GEN_WAV_QUANT_OUT "a.gen.wav.quant_out.%s"
#define TN_A_GEN_WAV_UPSAMPLE "a.gen.wav.upsample.%s"
#define TN_A_GEN_WAV_SEANET_CONV_IN "a.gen.wav.seanet.conv_in.%s"
#define TN_A_GEN_WAV_SEANET_CONV_OUT "a.gen.wav.seanet.conv_out.%s"
#define TN_A_GEN_WAV_SEANET_RES_CONV1 "a.gen.wav.seanet.blk.%d.res_conv1.%s"
#define TN_A_GEN_WAV_SEANET_RES_CONV2 "a.gen.wav.seanet.blk.%d.res_conv2.%s"
#define TN_A_GEN_WAV_SEANET_SCALE_CONV "a.gen.wav.seanet.blk.%d.scale_conv.%s"
// cogvlm
#define TN_MM_POST_FC_NORM "mm.post_fc_norm.%s"
#define TN_MM_H_TO_4H "mm.up.%s"
@@ -455,6 +487,8 @@ enum projector_type {
PROJECTOR_TYPE_MIMO_AUDIO,
PROJECTOR_TYPE_QWEN3TTS_SPKENC,
PROJECTOR_TYPE_QWEN3TTS_GEN,
PROJECTOR_TYPE_POCKETTTS_SPKENC,
PROJECTOR_TYPE_POCKETTTS_GEN,
PROJECTOR_TYPE_UNKNOWN,
};
@@ -514,6 +548,8 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
{ PROJECTOR_TYPE_PARAKEET, "parakeet"},
{ PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"},
{ PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"},
{ PROJECTOR_TYPE_POCKETTTS_SPKENC, "pockettts_spkenc"},
{ PROJECTOR_TYPE_POCKETTTS_GEN, "pockettts_gen"},
};
static projector_type clip_projector_type_from_string(const std::string & str) {
tools/mtmd/clip-model.h
+83
View File
@@ -139,6 +139,14 @@ struct clip_hparams {
// threshold for the "out_eos_score" graph output
float gen_eos_threshold = 0.0f;
// 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)
// qwen3tts code2wav
int32_t wav_tfm_n_layer = 0;
int32_t wav_tfm_n_embd = 0;
@@ -389,6 +397,63 @@ struct qf_block {
std::vector<clip_layer> qf_proj_layers;
};
// pocket-tts SEANet stack, used in both directions:
// encoder = conv_in -> per stage (residual unit, strided conv) -> conv_out
// decoder = conv_in -> per stage (strided convtr, residual unit) -> conv_out
struct clip_seanet {
// one residual unit: ELU -> dilated conv -> ELU -> pointwise conv, added to the input
struct stage {
ggml_tensor * res_conv1_w = nullptr;
ggml_tensor * res_conv1_b = nullptr;
ggml_tensor * res_conv2_w = nullptr;
ggml_tensor * res_conv2_b = nullptr;
ggml_tensor * scale_conv_w = nullptr; // strided conv (encoder) or convtr (decoder)
ggml_tensor * scale_conv_b = nullptr;
};
ggml_tensor * conv_in_w = nullptr;
ggml_tensor * conv_in_b = nullptr;
ggml_tensor * conv_out_w = nullptr;
ggml_tensor * conv_out_b = nullptr;
std::vector<stage> stages;
};
// pocket-tts flow-matching decoder (SimpleMLPAdaLN)
struct clip_flow_net {
// AdaLN res block: in_ln -> modulate -> Linear -> SiLU -> Linear, gated residual
struct block {
ggml_tensor * norm_w = nullptr;
ggml_tensor * norm_b = nullptr;
ggml_tensor * up_w = nullptr;
ggml_tensor * up_b = nullptr;
ggml_tensor * down_w = nullptr;
ggml_tensor * down_b = nullptr;
ggml_tensor * ada_w = nullptr; // -> shift, scale, gate
ggml_tensor * ada_b = nullptr;
};
// timestep embedder: cos/sin(t * freqs) -> Linear -> SiLU -> Linear -> RMSNorm
struct time_embd {
ggml_tensor * freqs = nullptr;
ggml_tensor * up_w = nullptr;
ggml_tensor * up_b = nullptr;
ggml_tensor * down_w = nullptr;
ggml_tensor * down_b = nullptr;
ggml_tensor * norm = nullptr; // RMSNorm alpha
};
ggml_tensor * input_proj_w = nullptr;
ggml_tensor * input_proj_b = nullptr;
ggml_tensor * cond_embd_w = nullptr;
ggml_tensor * cond_embd_b = nullptr;
ggml_tensor * final_ada_w = nullptr; // -> shift, scale
ggml_tensor * final_ada_b = nullptr;
ggml_tensor * final_proj_w = nullptr;
ggml_tensor * final_proj_b = nullptr;
std::vector<time_embd> time;
std::vector<block> blocks;
};
// qwen3tts code2wav: RVQ codes -> raw PCM
struct clip_code2wav {
// "upsample" stage: one ConvNeXt block plus the causal ConvTranspose1d before it
@@ -686,6 +751,24 @@ struct clip_model {
// qwen3tts code2wav: RVQ codes -> raw PCM
clip_code2wav c2w;
// pocket-tts: SEANet stack, shared by the encoder (speaker path) and the decoder (gen path)
clip_seanet seanet;
// pocket-tts: voice latent -> backbone embd (speaker path)
ggml_tensor * spk_proj_w = nullptr;
ggml_tensor * downsample_w = nullptr;
// pocket-tts: flow-matching decoder, backbone hidden state -> next latent
clip_flow_net flow;
ggml_tensor * gen_out_eos_w = nullptr;
ggml_tensor * gen_out_eos_b = nullptr;
ggml_tensor * gen_input_lin_w = nullptr; // latent -> backbone embd
ggml_tensor * gen_emb_mean = nullptr;
ggml_tensor * gen_emb_std = nullptr;
ggml_tensor * gen_quant_out_w = nullptr; // latent -> decoder dim
ggml_tensor * gen_upsample_w = nullptr; // depthwise convtr, frame rate -> encoder frame rate
std::vector<clip_layer> gen_tfm_layers; // mimi decoder_transformer
// cogvlm
ggml_tensor * mm_post_fc_norm_w = nullptr;
ggml_tensor * mm_post_fc_norm_b = nullptr;
tools/mtmd/clip.cpp
+229 -26
View File
@@ -1058,6 +1058,18 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
{
builder = std::make_unique<clip_graph_qwen3tts_spkenc>(ctx, img);
} break;
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
{
builder = std::make_unique<clip_graph_pockettts_spkenc>(ctx, img);
} break;
case PROJECTOR_TYPE_POCKETTTS_GEN:
{
const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
const int n_step = params && params->n_steps > 0 ? params->n_steps : ctx->model.hparams.flow_n_step;
const int64_t n_latent = ctx->model.gen_input_lin_w->ne[0];
const int n_frames = params && params->feats ? (int) (params->feats->size() / n_latent) : 1;
builder = std::make_unique<clip_graph_pockettts_gen>(ctx, img, gen_process, n_step, n_frames);
} break;
case PROJECTOR_TYPE_QWEN3TTS_GEN:
{
const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
@@ -1730,6 +1742,23 @@ struct clip_model_loader {
// matches the reference decoder's sliding_window (speech_tokenizer/config.json)
hparams.wav_tfm_swa = 72;
} break;
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
case PROJECTOR_TYPE_POCKETTTS_GEN:
{
// mimi front-end takes the raw waveform, no mel
hparams.audio_sample_rate = 24000;
// seanet ratios are [6,5,4] in the config, the encoder reverses them
hparams.seanet_ratios = { 4, 5, 6 };
hparams.seanet_n_stage = (int32_t) hparams.seanet_ratios.size();
hparams.mimi_downsample = 16;
// 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
hparams.flow_n_step = 1;
hparams.flow_temp = 0.3f;
hparams.gen_eos_threshold = -4.0f;
} break;
case PROJECTOR_TYPE_PADDLEOCR:
{
hparams.n_merge = 2;
@@ -1925,7 +1954,9 @@ struct clip_model_loader {
// GEMMA4UA is encoder-free: it uses n_mel_bins as a raw-waveform frame size (640) and has no FFT/filterbank, so the mel-range and FFT
// checks below do not apply to it.
const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA;
// pocket-tts is encoder-free in the same sense: mimi convolves the raw waveform
const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA &&
model.proj_type != PROJECTOR_TYPE_POCKETTTS_SPKENC;
// Validate audio hparams loaded from GGUF metadata
if (hparams.n_mel_bins <= 0 || (fft_based && hparams.n_mel_bins > 256)) {
@@ -1998,6 +2029,31 @@ struct clip_model_loader {
return cur;
};
// pocket-tts: the encoder and the decoder share the same layout, only the prefix differs
auto load_seanet = [&](clip_seanet & seanet, bool is_decoder) {
const char * conv_in = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_IN : TN_A_SEANET_CONV_IN;
const char * conv_out = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_OUT : TN_A_SEANET_CONV_OUT;
const char * res1 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV1 : TN_A_SEANET_RES_CONV1;
const char * res2 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV2 : TN_A_SEANET_RES_CONV2;
const char * scale = is_decoder ? TN_A_GEN_WAV_SEANET_SCALE_CONV : TN_A_SEANET_SCALE_CONV;
seanet.conv_in_w = get_tensor(string_format(conv_in, "weight"));
seanet.conv_in_b = get_tensor(string_format(conv_in, "bias"));
seanet.conv_out_w = get_tensor(string_format(conv_out, "weight"));
seanet.conv_out_b = get_tensor(string_format(conv_out, "bias"));
seanet.stages.resize(hparams.seanet_n_stage);
for (int i = 0; i < hparams.seanet_n_stage; i++) {
auto & stage = seanet.stages[i];
stage.res_conv1_w = get_tensor(string_format(res1, i, "weight"));
stage.res_conv1_b = get_tensor(string_format(res1, i, "bias"));
stage.res_conv2_w = get_tensor(string_format(res2, i, "weight"));
stage.res_conv2_b = get_tensor(string_format(res2, i, "bias"));
stage.scale_conv_w = get_tensor(string_format(scale, i, "weight"));
stage.scale_conv_b = get_tensor(string_format(scale, i, "bias"));
}
};
auto get_vector = [&](const std::string & name) {
std::vector<float> result;
auto it = tensor_offset.find(name);
@@ -2059,7 +2115,8 @@ struct clip_model_loader {
const bool has_standard_layers = (
model.proj_type != PROJECTOR_TYPE_GEMMA3NV &&
model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC);
model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC &&
model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN);
// layers
const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0;
@@ -2726,6 +2783,81 @@ struct clip_model_loader {
model.mm_fc_w = get_tensor(string_format(TN_MM_AUDIO_FC, "weight"));
model.mm_fc_b = get_tensor(string_format(TN_MM_AUDIO_FC, "bias"));
} break;
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
{
load_seanet(model.seanet, false);
model.downsample_w = get_tensor(string_format(TN_A_DOWNSAMPLE_CONV, "weight"));
model.spk_proj_w = get_tensor(string_format(TN_A_SPEAKER_PROJ, "weight"));
} break;
case PROJECTOR_TYPE_POCKETTTS_GEN:
{
auto & flow = model.flow;
flow.input_proj_w = get_tensor(string_format(TN_A_GEN_FLOW_INPUT_PROJ, "weight"));
flow.input_proj_b = get_tensor(string_format(TN_A_GEN_FLOW_INPUT_PROJ, "bias"));
flow.cond_embd_w = get_tensor(string_format(TN_A_GEN_FLOW_COND_EMBD, "weight"));
flow.cond_embd_b = get_tensor(string_format(TN_A_GEN_FLOW_COND_EMBD, "bias"));
flow.final_ada_w = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_ADA, "weight"));
flow.final_ada_b = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_ADA, "bias"));
flow.final_proj_w = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_PROJ, "weight"));
flow.final_proj_b = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_PROJ, "bias"));
flow.time.resize(2);
for (size_t i = 0; i < flow.time.size(); i++) {
auto & t = flow.time[i];
t.freqs = get_tensor(string_format(TN_A_GEN_FLOW_TIME_FREQS, (int) i));
t.up_w = get_tensor(string_format(TN_A_GEN_FLOW_TIME_UP, (int) i, "weight"));
t.up_b = get_tensor(string_format(TN_A_GEN_FLOW_TIME_UP, (int) i, "bias"));
t.down_w = get_tensor(string_format(TN_A_GEN_FLOW_TIME_DOWN, (int) i, "weight"));
t.down_b = get_tensor(string_format(TN_A_GEN_FLOW_TIME_DOWN, (int) i, "bias"));
t.norm = get_tensor(string_format(TN_A_GEN_FLOW_TIME_NORM, (int) i));
}
// one AdaLN block per flow depth, the count is only known from the tensors
for (int il = 0; ; il++) {
ggml_tensor * probe = get_tensor(string_format(TN_A_GEN_FLOW_BLK_NORM, il, "weight"), false);
if (probe == nullptr) {
break;
}
clip_flow_net::block blk;
blk.norm_w = probe;
blk.norm_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_NORM, il, "bias"));
blk.up_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_UP, il, "weight"));
blk.up_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_UP, il, "bias"));
blk.down_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_DOWN, il, "weight"));
blk.down_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_DOWN, il, "bias"));
blk.ada_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_ADA, il, "weight"));
blk.ada_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_ADA, il, "bias"));
flow.blocks.push_back(blk);
}
model.gen_out_eos_w = get_tensor(string_format(TN_A_GEN_OUT_EOS, "weight"));
model.gen_out_eos_b = get_tensor(string_format(TN_A_GEN_OUT_EOS, "bias"));
model.gen_input_lin_w = get_tensor(string_format(TN_A_GEN_INPUT_LINEAR, "weight"));
model.gen_emb_mean = get_tensor(TN_A_GEN_EMB_MEAN);
model.gen_emb_std = get_tensor(TN_A_GEN_EMB_STD);
// mimi decoder
model.gen_quant_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_OUT, "weight"));
model.gen_upsample_w = get_tensor(string_format(TN_A_GEN_WAV_UPSAMPLE, "weight"));
load_seanet(model.seanet, true);
model.gen_tfm_layers.resize(hparams.n_layer);
for (int il = 0; il < hparams.n_layer; il++) {
auto & layer = model.gen_tfm_layers[il];
const char * p = "a.gen.wav.tfm";
layer.ln_1_w = get_tensor(string_format(TN_LN_1, p, il, "weight"));
layer.ln_1_b = get_tensor(string_format(TN_LN_1, p, il, "bias"));
layer.q_w = get_tensor(string_format(TN_ATTN_Q, p, il, "weight"));
layer.k_w = get_tensor(string_format(TN_ATTN_K, p, il, "weight"));
layer.v_w = get_tensor(string_format(TN_ATTN_V, p, il, "weight"));
layer.o_w = get_tensor(string_format(TN_ATTN_OUTPUT, p, il, "weight"));
layer.ls_1_w = get_tensor(string_format(TN_LS_1, p, il, "weight"));
layer.ln_2_w = get_tensor(string_format(TN_LN_2, p, il, "weight"));
layer.ln_2_b = get_tensor(string_format(TN_LN_2, p, il, "bias"));
layer.ff_up_w = get_tensor(string_format(TN_FFN_UP, p, il, "weight"));
layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN, p, il, "weight"));
layer.ls_2_w = get_tensor(string_format(TN_LS_2, p, il, "weight"));
}
} break;
case PROJECTOR_TYPE_QWEN3TTS_GEN:
{
// code_predictor
@@ -4028,6 +4160,17 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
// one hidden-state vector fed back to the talker per call
n_patches = 1;
} break;
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
{
// one conditioning row per 12.5Hz frame
const int hop = ctx->model.hparams.mimi_downsample * 120;
n_patches = img->nx() / hop;
} break;
case PROJECTOR_TYPE_POCKETTTS_GEN:
{
// one latent per call for GEN_CODE, GEN_WAV sizes its input from the caller
n_patches = 1;
} break;
case PROJECTOR_TYPE_GRANITE4_VISION:
{
// Per-tile output token count: each projector block outputs
@@ -4069,6 +4212,15 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
return clip_encode(ctx, &params);
}
// persisted state slots of the gen-audio decoder, per pipeline
static std::vector<c2w_state_slot> list_gen_state_slots(const clip_hparams & hparams, const clip_model & model) {
switch (model.proj_type) {
case PROJECTOR_TYPE_QWEN3TTS_GEN: return list_c2w_state_slots(hparams, model);
case PROJECTOR_TYPE_POCKETTTS_GEN: return list_pockettts_state_slots(hparams, model);
default: return {};
}
}
bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
const clip_image_f32_batch & imgs = *params->imgs;
int n_batch_cur = imgs.entries.size();
@@ -4133,6 +4285,45 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
ggml_backend_tensor_set(cur, values.data(), 0, ggml_nbytes(cur));
};
// upload the decoder state from the previous call, or zero-fill on a cold start
auto set_gen_state_in = [&]() {
size_t offset = 0;
for (const auto & slot : list_gen_state_slots(hparams, model)) {
ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str());
const size_t nb = ggml_nbytes(t);
if (params->state_in && params->state_in->size() >= offset + nb) {
ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb);
} else {
std::vector<uint8_t> zeros(nb, 0);
ggml_backend_tensor_set(t, zeros.data(), 0, nb);
}
offset += nb;
}
};
// rope positions and attention mask of the mimi transformers (pocket-tts).
// the mask is causal with a sliding window, see _build_attention_mask() in the reference
auto set_pockettts_tfm_inputs = [&]() {
const int64_t n_pos = ggml_nelements(get_inp_tensor("inp_pos"));
std::vector<int32_t> positions((size_t) n_pos);
for (int64_t i = 0; i < n_pos; i++) {
positions[(size_t) i] = (int32_t) i;
}
set_input_i32("inp_pos", positions);
const int64_t context = hparams.mimi_tfm_context;
std::vector<float> mask((size_t) n_pos * n_pos, -INFINITY);
for (int64_t q = 0; q < n_pos; q++) {
for (int64_t k = 0; k < n_pos; k++) {
const int64_t delta = q - k;
if (delta >= 0 && delta < context) {
mask[(size_t) q * n_pos + k] = 0.0f;
}
}
}
set_input_f32("kq_mask", mask);
};
// set input pixel values
if (!imgs.is_audio) {
size_t nelem = 0;
@@ -4176,8 +4367,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
}
set_input_f32("inp_raw", inp_raw);
} else if (!(ctx->proj_type() == PROJECTOR_TYPE_QWEN3TTS_GEN && params->gen_process == CLIP_GEN_PROCESS_GEN_WAV)) {
// audio input, code2wav is not here: its only input is "inp_codes", set in the switch below
} else if (params->gen_process != CLIP_GEN_PROCESS_GEN_WAV) {
// audio input. GEN_WAV is not here: it takes codes or feats, set in the switch below
GGML_ASSERT(imgs.entries.size() == 1);
const auto & mel_inp = imgs.entries[0];
@@ -4646,6 +4837,28 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
}
set_input_i32("patches", patches);
} break;
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
{
set_pockettts_tfm_inputs();
} break;
case PROJECTOR_TYPE_POCKETTTS_GEN:
{
if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) {
GGML_ASSERT(params->feats != nullptr);
set_input_f32("inp_feats", *params->feats);
// positions and mask are derived in-graph from the persisted counter
set_gen_state_in();
} 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));
std::vector<float> noise(ggml_nelements(t));
for (auto & v : noise) {
v = dist(ctx->rng);
}
set_input_f32("inp_noise", noise);
}
} break;
case PROJECTOR_TYPE_GEMMA4V:
case PROJECTOR_TYPE_GEMMA4UV:
{
@@ -4770,20 +4983,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
}
}
set_input_i32("inp_codes", codes);
// upload the state from the previous call, or zero-fill on a cold start
size_t offset = 0;
for (const auto & slot : list_c2w_state_slots(hparams, model)) {
ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str());
const size_t nb = ggml_nbytes(t);
if (params->state_in && params->state_in->size() >= offset + nb) {
ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb);
} else {
std::vector<uint8_t> zeros(nb, 0);
ggml_backend_tensor_set(t, zeros.data(), 0, nb);
}
offset += nb;
}
set_gen_state_in();
} else {
// code0 indexes gen_code_out_embd_w via ggml_get_rows; bound it
const int64_t vocab0 = model.gen_code_out_embd_w->ne[1];
@@ -5251,16 +5451,15 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
// for audio gen models
//
// optional outputs: a pipeline yields codes or feats, and not all have an eos head
if (params->out_codes != nullptr) {
ggml_tensor * codes = ggml_graph_get_tensor(gf, "out_codes");
if (codes == nullptr) {
GGML_ABORT("out_codes requested but graph has no \"out_codes\" tensor");
if (codes != nullptr) {
auto & out_codes = *params->out_codes;
out_codes.resize(ggml_nelements(codes));
ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes));
}
auto & out_codes = *params->out_codes;
out_codes.resize(ggml_nelements(codes));
ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes));
}
// optional outputs, a missing tensor is not an error
if (params->out_feats != nullptr) {
ggml_tensor * feats = ggml_graph_get_tensor(gf, "out_feats");
if (feats != nullptr) {
@@ -5299,12 +5498,12 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
if (params->state_out != nullptr) {
auto & state_out = *params->state_out;
size_t total = 0;
for (const auto & slot : list_c2w_state_slots(hparams, model)) {
for (const auto & slot : list_gen_state_slots(hparams, model)) {
total += (size_t) (slot.ne0 * slot.ne1) * sizeof(float);
}
state_out.resize(total);
size_t offset = 0;
for (const auto & slot : list_c2w_state_slots(hparams, model)) {
for (const auto & slot : list_gen_state_slots(hparams, model)) {
ggml_tensor * t = ggml_graph_get_tensor(gf, ("state_out_" + slot.name).c_str());
if (t == nullptr) {
GGML_ABORT("state_out requested but graph has no \"state_out_%s\" tensor", slot.name.c_str());
@@ -5450,6 +5649,10 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
return ctx->model.mm_fc_w->ne[2];
case PROJECTOR_TYPE_QWEN3TTS_GEN:
return ctx->model.gen_code_out_embd_w->ne[0];
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
return ctx->model.spk_proj_w->ne[1];
case PROJECTOR_TYPE_POCKETTTS_GEN:
return ctx->model.gen_input_lin_w->ne[1];
case PROJECTOR_TYPE_PARAKEET:
return ctx->model.mm_1_w->ne[1];
default:
tools/mtmd/models/models.h
+56
View File
@@ -317,6 +317,59 @@ struct clip_graph_qwen3tts_gen : clip_graph {
};
};
//
// pocket-tts: SEANet convolution stack, shared by the voice encoder and the mimi decoder.
// stateless unless state_in is populated: convs then pad instead of carrying left-context.
//
struct clip_graph_pockettts_seanet : clip_graph {
clip_graph_pockettts_seanet(const clip_graph & parent) : clip_graph(parent) {}
ggml_cgraph * build() override { GGML_ABORT("call encode()/decode() instead"); }
// per-call streaming state, keyed by slot name (see list_pockettts_state_slots)
std::map<std::string, ggml_tensor *> state_in;
mutable std::vector<std::pair<std::string, ggml_tensor *>> state_out;
ggml_tensor * conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation,
bool pad_replicate = false, const std::string & state_name = "") const;
ggml_tensor * conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride,
const std::string & state_name = "") const;
ggml_tensor * res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation,
const std::string & state_prefix = "") const;
// x: [T, C] -> [T / hop, dim]
ggml_tensor * encode(ggml_tensor * x) const;
// x: [T, dim] -> [T * hop, 1], streams when state_in is populated
ggml_tensor * decode(ggml_tensor * x) const;
};
// mimi encoder + speaker_proj: reference waveform -> voice conditioning rows
struct clip_graph_pockettts_spkenc : clip_graph {
clip_graph_pockettts_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
ggml_cgraph * build() override;
ggml_tensor * tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const;
};
//
// pocket-tts generation:
// GEN_CODE = flow-matching decoder + end-of-speech head, one latent per call
// GEN_WAV = mimi decoder, a window of latents -> PCM
//
struct clip_graph_pockettts_gen : clip_graph {
clip_graph_pockettts_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int n_step, int n_frames)
: clip_graph(ctx, img), gen_process(gen_process), n_step(n_step), n_frames(n_frames) {}
ggml_cgraph * build() override;
clip_gen_process_type gen_process;
int n_step; // lsd_decode steps, fixed at graph-build time
int n_frames; // GEN_WAV only: number of latents to decode
// AdaLN modulation: x * (1 + scale) + shift
ggml_tensor * modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const;
ggml_tensor * time_embed(const clip_flow_net::time_embd & te, float t) const;
ggml_tensor * flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const;
};
// one persisted state buffer used by code2wav, see qwen3tts-gen.cpp
struct c2w_state_slot {
std::string name;
@@ -325,6 +378,9 @@ struct c2w_state_slot {
};
std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model);
// same, for the streaming mimi decoder (pocket-tts GEN_WAV)
std::vector<c2w_state_slot> list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model);
struct clip_graph_kimik25 : clip_graph {
clip_graph_kimik25(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
ggml_cgraph * build() override;
tools/mtmd/models/pockettts-gen.cpp
+286
View File
@@ -0,0 +1,286 @@
#include "models.h"
#include <cmath>
// pocket-tts generation stages
//
// GEN_CODE: backbone hidden state -> next 32-d latent (flow matching) + end-of-speech score
// GEN_WAV : a window of latents -> PCM, through the mimi decoder
//
// 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
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);
ggml_tensor * emb = ggml_concat(ctx0, ggml_cos(ctx0, args), ggml_sin(ctx0, args), 0);
ggml_tensor * cur = build_mm(te.up_w, emb);
cur = ggml_add(ctx0, cur, te.up_b);
cur = ggml_silu(ctx0, cur);
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
{
const int64_t n = cur->ne[0];
ggml_tensor * mean = ggml_mean(ctx0, cur);
ggml_tensor * dev = ggml_sub(ctx0, cur, mean);
ggml_tensor * var = ggml_mean(ctx0, ggml_sqr(ctx0, dev));
var = ggml_scale_bias(ctx0, var, (float) n / (float) (n - 1), 1e-5f);
cur = ggml_div(ctx0, cur, ggml_sqrt(ctx0, var));
cur = ggml_mul(ctx0, cur, te.norm);
}
return cur;
}
// one velocity evaluation: v(cond, s, t, x)
ggml_tensor * clip_graph_pockettts_gen::flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const {
const auto & flow = model.flow;
ggml_tensor * cur = build_mm(flow.input_proj_w, x);
cur = ggml_add(ctx0, cur, flow.input_proj_b);
// the two time conditions are averaged, then added to the projected backbone state
ggml_tensor * ts = ggml_add(ctx0, time_embed(flow.time[0], s), time_embed(flow.time[1], t));
ts = ggml_scale(ctx0, ts, 1.0f / (float) flow.time.size());
ggml_tensor * c = build_mm(flow.cond_embd_w, cond);
c = ggml_add(ctx0, c, flow.cond_embd_b);
ggml_tensor * y = ggml_add(ctx0, ts, c);
cb(y, "flow_cond", -1);
const int64_t n_ch = flow.blocks.empty() ? 0 : flow.blocks[0].norm_w->ne[0];
for (size_t il = 0; il < flow.blocks.size(); il++) {
const auto & blk = flow.blocks[il];
ggml_tensor * mod = build_mm(blk.ada_w, ggml_silu(ctx0, y));
mod = ggml_add(ctx0, mod, blk.ada_b);
ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0);
ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]);
ggml_tensor * gate = ggml_view_1d(ctx0, mod, n_ch, (size_t) 2 * n_ch * mod->nb[0]);
ggml_tensor * h = build_norm(cur, blk.norm_w, blk.norm_b, NORM_TYPE_NORMAL, 1e-6f, (int) il);
h = modulate(h, shift, scale);
h = build_mm(blk.up_w, h);
h = ggml_add(ctx0, h, blk.up_b);
h = ggml_silu(ctx0, h);
h = build_mm(blk.down_w, h);
h = ggml_add(ctx0, h, blk.down_b);
cur = ggml_add(ctx0, cur, ggml_mul(ctx0, gate, h));
cb(cur, "flow_blk", (int) il);
}
// final layer: the norm has no weights, only the AdaLN modulation
ggml_tensor * mod = build_mm(flow.final_ada_w, ggml_silu(ctx0, y));
mod = ggml_add(ctx0, mod, flow.final_ada_b);
ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0);
ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]);
cur = build_norm(cur, nullptr, nullptr, NORM_TYPE_NORMAL, 1e-6f, -1);
cur = modulate(cur, shift, scale);
cur = build_mm(flow.final_proj_w, cur);
cur = ggml_add(ctx0, cur, flow.final_proj_b);
return cur;
}
// 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
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) {
return slots; // not a pocket-tts decoder
}
const auto & seanet = model.seanet;
slots.push_back({"tfm_pos", 1, 1});
const int64_t n_embd_a = model.gen_tfm_layers[0].q_w->ne[1];
const int64_t prefix = hparams.mimi_tfm_context - 1;
for (size_t il = 0; il < model.gen_tfm_layers.size(); il++) {
slots.push_back({"tfm_k_" + std::to_string(il), n_embd_a, prefix});
slots.push_back({"tfm_v_" + std::to_string(il), n_embd_a, prefix});
}
// upsample is depthwise, its output channel count is the input one
slots.push_back({"up", model.gen_upsample_w->ne[0] - hparams.mimi_downsample, model.gen_upsample_w->ne[2]});
slots.push_back({"dec_in", seanet.conv_in_w->ne[0] - 1, seanet.conv_in_w->ne[1]});
for (int i = 0; i < hparams.seanet_n_stage; i++) {
const auto & stage = seanet.stages[i];
const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i];
slots.push_back({"dec_up_" + std::to_string(i), stage.scale_conv_w->ne[0] - stride, stage.scale_conv_w->ne[1]});
slots.push_back({"dec_res_" + std::to_string(i), stage.res_conv1_w->ne[0] - 1, stage.res_conv1_w->ne[1]});
}
slots.push_back({"dec_out", seanet.conv_out_w->ne[0] - 1, seanet.conv_out_w->ne[1]});
return slots;
}
ggml_cgraph * clip_graph_pockettts_gen::build() {
if (gen_process == CLIP_GEN_PROCESS_GEN_CODE) {
// the backbone hidden state arrives as the single batch entry
ggml_tensor * h_state = build_inp_raw(1);
h_state = ggml_reshape_2d(ctx0, h_state, n_mmproj_embd, 1);
// end-of-speech probe, thresholded on the host side
ggml_tensor * eos = build_mm(model.gen_out_eos_w, h_state);
eos = ggml_add(ctx0, eos, model.gen_out_eos_b);
ggml_set_name(eos, "out_eos_score");
ggml_set_output(eos);
ggml_build_forward_expand(gf, eos);
const int64_t n_latent = model.gen_input_lin_w->ne[0];
ggml_tensor * noise = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_latent, 1);
ggml_set_name(noise, "inp_noise");
ggml_set_input(noise);
// lsd_decode: integrate the velocity field from the noise sample
ggml_tensor * cur = noise;
for (int i = 0; i < n_step; i++) {
const float s = (float) i / (float) n_step;
const float t = (float) (i + 1) / (float) n_step;
ggml_tensor * v = flow_forward(h_state, cur, s, t);
cur = ggml_add(ctx0, cur, ggml_scale(ctx0, v, 1.0f / (float) n_step));
}
cb(cur, "flow_latent", -1);
ggml_set_name(cur, "out_feats");
ggml_set_output(cur);
ggml_build_forward_expand(gf, cur);
// the same latent, projected into the backbone's input space for the next step
ggml_tensor * embd = build_mm(model.gen_input_lin_w, cur);
cb(embd, "gen_embd", -1);
ggml_build_forward_expand(gf, embd);
return gf;
}
// GEN_WAV: [32, n_frames] latents -> PCM
ggml_tensor * feats = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32,
model.gen_input_lin_w->ne[0], n_frames);
ggml_set_name(feats, "inp_feats");
ggml_set_input(feats);
// denormalize, then the DummyQuantizer up-projection
ggml_tensor * cur = ggml_add(ctx0, ggml_mul(ctx0, feats, model.gen_emb_std), model.gen_emb_mean);
cur = build_mm(model.gen_quant_out_w, cur);
cb(cur, "quant_out", -1);
clip_graph_pockettts_seanet seanet(*this);
for (const auto & slot : list_pockettts_state_slots(hparams, model)) {
ggml_tensor * t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1);
ggml_set_name(t, ("state_in_" + slot.name).c_str());
ggml_set_input(t);
seanet.state_in[slot.name] = t;
}
// model frame rate -> encoder frame rate, depthwise transposed conv
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
cur = seanet.conv_transpose1d(cur, model.gen_upsample_w, nullptr, hparams.mimi_downsample, "up");
cb(cur, "mimi_upsample", -1);
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
// positions continue across calls, the counter lives in the state
const int64_t n_pos = cur->ne[1];
const int64_t prefix = hparams.mimi_tfm_context - 1;
const int64_t n_kv = prefix + n_pos;
ggml_tensor * base = ggml_reshape_1d(ctx0, seanet.state_in.at("tfm_pos"), 1);
ggml_tensor * inp_pos = ggml_cast(ctx0, ggml_add(ctx0, ggml_arange(ctx0, 0.0f, (float) n_pos, 1.0f), base),
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
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);
ggml_tensor * keep = ggml_mul(ctx0,
ggml_step(ctx0, ggml_scale_bias(ctx0, diff, 1.0f, 0.5f)), // delta >= 0
ggml_step(ctx0, ggml_scale_bias(ctx0, diff, -1.0f, (float) hparams.mimi_tfm_context - 0.5f))); // delta < context
keep = ggml_mul(ctx0, keep,
ggml_step(ctx0, ggml_scale_bias(ctx0, ggml_add(ctx0, pos_k, base), 1.0f, 0.5f - (float) prefix)));
ggml_tensor * kq_mask = ggml_reshape_4d(ctx0, ggml_log(ctx0, keep), n_kv, n_pos, 1, 1);
for (int il = 0; il < n_layer; il++) {
const auto & layer = model.gen_tfm_layers[il];
ggml_tensor * inp = cur;
cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
ggml_tensor * Qcur = build_mm(layer.q_w, cur);
ggml_tensor * Kcur = build_mm(layer.k_w, cur);
ggml_tensor * Vcur = build_mm(layer.v_w, cur);
Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos);
Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos);
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
// prepend the cached window, then keep this chunk's tail for the next call
const std::string k_name = "tfm_k_" + std::to_string(il);
const std::string v_name = "tfm_v_" + std::to_string(il);
ggml_tensor * k_full = ggml_concat(ctx0, seanet.state_in.at(k_name),
ggml_reshape_2d(ctx0, Kcur, d_head * n_head, n_pos), 1);
ggml_tensor * v_full = ggml_concat(ctx0, seanet.state_in.at(v_name), Vcur, 1);
seanet.state_out.push_back({k_name, ggml_cont(ctx0, ggml_view_2d(ctx0, k_full, k_full->ne[0], prefix,
k_full->nb[1], (size_t) n_pos * k_full->nb[1]))});
seanet.state_out.push_back({v_name, ggml_cont(ctx0, ggml_view_2d(ctx0, v_full, v_full->ne[0], prefix,
v_full->nb[1], (size_t) n_pos * v_full->nb[1]))});
ggml_tensor * q_cur = ggml_reshape_4d(ctx0, Qcur, d_head, n_head, n_pos, 1);
ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k_full, d_head, n_head, n_kv, 1);
ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v_full, d_head, n_head, n_kv, 1);
cur = build_attn(layer.o_w, nullptr, q_cur, k_cur, v_cur, kq_mask, kq_scale, il);
cur = ggml_mul(ctx0, cur, layer.ls_1_w);
cur = ggml_add(ctx0, cur, inp);
inp = cur;
cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il);
cur = ggml_mul(ctx0, cur, layer.ls_2_w);
cur = ggml_add(ctx0, cur, inp);
}
cb(cur, "mimi_dec_tfm", -1);
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
cur = seanet.decode(cur);
for (const auto & s : seanet.state_out) {
ggml_set_name(s.second, ("state_out_" + s.first).c_str());
ggml_set_output(s.second);
ggml_build_forward_expand(gf, s.second);
}
// [n_samples, 1] -> [n_samples], clamped like the reference output
cur = ggml_reshape_1d(ctx0, cur, cur->ne[0]);
cur = ggml_clamp(ctx0, cur, -1.0f, 1.0f);
ggml_set_name(cur, "out_audio");
ggml_set_output(cur);
ggml_build_forward_expand(gf, cur);
return gf;
}
tools/mtmd/models/pockettts-seanet.cpp
+158
View File
@@ -0,0 +1,158 @@
#include "models.h"
// 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)
static int64_t div_ceil(int64_t a, int64_t b) {
return a / b + (a % b ? 1 : 0);
}
// x: [T, IC], w: [K, IC, OC] -> [T / stride, OC]
// the convs are causal, so the whole K - stride padding goes on the left
ggml_tensor * clip_graph_pockettts_seanet::conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation,
bool pad_replicate, const std::string & state_name) const {
const int64_t k_size = (w->ne[0] - 1) * dilation + 1;
const int64_t p_total = k_size - stride;
// trailing padding so the last frame is not dropped, see pad_for_conv1d() in conv.py
const int64_t n_frames = div_ceil(x->ne[0] - k_size + p_total, stride);
const int64_t ideal_len = n_frames * stride + k_size - p_total;
const int64_t p_extra = ideal_len - x->ne[0];
if (!state_name.empty() && p_total > 0) {
// streaming: the left context is the tail of the previous call
ggml_tensor * left = state_in.at(state_name); // [p_total, IC]
x = ggml_concat(ctx0, left, x, 0);
state_out.push_back({state_name,
ggml_cont(ctx0, ggml_view_2d(ctx0, x, p_total, x->ne[1], x->nb[1],
(size_t) (x->ne[0] - p_total) * x->nb[0]))});
} else if (pad_replicate && p_total > 0) {
// the resamplers repeat the first frame instead of zero-padding
ggml_tensor * first = ggml_view_2d(ctx0, x, 1, x->ne[1], x->nb[1], 0);
ggml_tensor * left = ggml_repeat_4d(ctx0, first, p_total, x->ne[1], 1, 1);
x = ggml_concat(ctx0, left, x, 0);
x = ggml_pad_ext(ctx0, x, 0, p_extra, 0, 0, 0, 0, 0, 0);
} else {
x = ggml_pad_ext(ctx0, x, p_total, p_extra, 0, 0, 0, 0, 0, 0);
}
ggml_tensor * y = ggml_conv_1d(ctx0, w, x, stride, 0, dilation);
y = ggml_reshape_2d(ctx0, y, y->ne[0], y->ne[1]);
if (b) {
y = ggml_add(ctx0, y, ggml_reshape_2d(ctx0, b, 1, b->ne[0]));
}
return y;
}
// 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
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 p_total = w->ne[0] - stride;
const bool depthwise = w->ne[1] == 1 && w->ne[2] > 1;
const int64_t emit_len = x->ne[0] * stride;
ggml_tensor * full = nullptr;
if (depthwise) {
// one group per channel, ggml_conv_transpose_1d has no grouped mode
for (int64_t ir = 0; ir < x->ne[1]; ir++) {
ggml_tensor * row = ggml_view_1d(ctx0, x, x->ne[0], ir * x->ne[0] * ggml_element_size(x));
ggml_tensor * krn = ggml_view_1d(ctx0, w, w->ne[0], ir * w->ne[0] * ggml_element_size(w));
row = ggml_conv_transpose_1d(ctx0, krn, row, stride, 0, 1);
full = full ? ggml_concat(ctx0, full, row, 1) : row;
}
} else {
full = ggml_conv_transpose_1d(ctx0, w, x, stride, 0, 1);
}
full = ggml_cont(ctx0, full); // [emit_len + p_total, OC]
ggml_tensor * out;
if (state_name.empty() || p_total == 0) {
out = ggml_cont(ctx0, ggml_view_2d(ctx0, full, emit_len, full->ne[1], full->nb[1], 0));
} else {
// overlap-add the tail the previous call held back
ggml_tensor * prev = state_in.at(state_name); // [p_total, OC]
ggml_tensor * head = ggml_add(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1], 0), prev);
if (emit_len > p_total) {
ggml_tensor * rest = ggml_view_2d(ctx0, full, emit_len - p_total, full->ne[1], full->nb[1],
(size_t) p_total * full->nb[0]);
out = ggml_concat(ctx0, head, rest, 0);
} else {
out = head;
}
state_out.push_back({state_name,
ggml_cont(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1],
(size_t) emit_len * full->nb[0]))});
}
if (b) {
out = ggml_add(ctx0, out, ggml_reshape_2d(ctx0, b, 1, b->ne[0]));
}
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);
h = conv1d(h, stage.res_conv1_w, stage.res_conv1_b, 1, dilation, false, state_prefix);
h = ggml_elu(ctx0, h);
// the second conv is pointwise, it needs no left context
h = conv1d(h, stage.res_conv2_w, stage.res_conv2_b, 1, 1);
return ggml_add(ctx0, x, h);
}
ggml_tensor * clip_graph_pockettts_seanet::encode(ggml_tensor * x) const {
const auto & seanet = model.seanet;
ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1);
cb(cur, "seanet_enc_in", -1);
for (int i = 0; i < hparams.seanet_n_stage; i++) {
const auto & stage = seanet.stages[i];
const int stride = hparams.seanet_ratios[i];
cur = res_unit(cur, stage, 1);
cur = ggml_elu(ctx0, cur);
cur = conv1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride, 1);
cb(cur, "seanet_enc_stage", i);
}
cur = ggml_elu(ctx0, cur);
cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1);
cb(cur, "seanet_enc_out", -1);
return cur;
}
ggml_tensor * clip_graph_pockettts_seanet::decode(ggml_tensor * x) const {
const auto & seanet = model.seanet;
const bool stream = !state_in.empty();
ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1, false,
stream ? "dec_in" : "");
cb(cur, "seanet_dec_in", -1);
for (int i = 0; i < hparams.seanet_n_stage; i++) {
const auto & stage = seanet.stages[i];
// the decoder mirrors the encoder, so the ratios are walked backwards
const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i];
const std::string id = std::to_string(i);
cur = ggml_elu(ctx0, cur);
cur = conv_transpose1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride,
stream ? "dec_up_" + id : "");
cur = res_unit(cur, stage, 1, stream ? "dec_res_" + id : "");
cb(cur, "seanet_dec_stage", i);
}
cur = ggml_elu(ctx0, cur);
cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1, false,
stream ? "dec_out" : "");
cb(cur, "seanet_dec_out", -1);
return cur;
}
tools/mtmd/models/pockettts-spkenc.cpp
+77
View File
@@ -0,0 +1,77 @@
#include "models.h"
// voice-prompt encoder: raw 24kHz waveform -> one conditioning row per 12.5Hz frame
// mimi encoder (SEANet + transformer + downsample), then flow_lm.speaker_proj_weight
// pre-norm block with layer scale on both residual paths, see mimi_transformer.py
ggml_tensor * clip_graph_pockettts_spkenc::tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const {
ggml_tensor * inp = cur;
cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
ggml_tensor * Qcur = build_mm(layer.q_w, cur);
ggml_tensor * Kcur = build_mm(layer.k_w, cur);
ggml_tensor * Vcur = build_mm(layer.v_w, cur);
const int64_t n_pos = cur->ne[1];
Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos);
Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos);
Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos);
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
cur = build_attn(layer.o_w, nullptr, Qcur, Kcur, Vcur, kq_mask, kq_scale, il);
cur = ggml_mul(ctx0, cur, layer.ls_1_w);
cur = ggml_add(ctx0, cur, inp);
inp = cur;
cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il);
cur = ggml_mul(ctx0, cur, layer.ls_2_w);
cur = ggml_add(ctx0, cur, inp);
return cur;
}
ggml_cgraph * clip_graph_pockettts_spkenc::build() {
// the preprocessor hands over the waveform as a single-row "mel", already [n_samples, 1]
ggml_tensor * inp_raw = build_inp_raw(1);
ggml_tensor * cur = ggml_reshape_2d(ctx0, inp_raw, inp_raw->ne[0], inp_raw->ne[1]);
clip_graph_pockettts_seanet seanet(*this);
cur = seanet.encode(cur);
cb(cur, "mimi_enc", -1);
// [T, 512] -> transformer works on [512, T]
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, cur->ne[1]);
ggml_set_name(inp_pos, "inp_pos");
ggml_set_input(inp_pos);
// the mimi transformer is causal with a sliding window, see _build_attention_mask()
ggml_tensor * kq_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, cur->ne[1], cur->ne[1]);
ggml_set_name(kq_mask, "kq_mask");
ggml_set_input(kq_mask);
for (int il = 0; il < n_layer; il++) {
cur = tfm_layer_forward(cur, model.layers[il], inp_pos, kq_mask, il);
}
cb(cur, "mimi_enc_tfm", -1);
// downsample to the model frame rate, [512, T] -> [T, 512] -> [T / 16, 32]
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
cur = seanet.conv1d(cur, model.downsample_w, nullptr, hparams.mimi_downsample, 1, true);
cb(cur, "mimi_downsample", -1);
// voice latent -> backbone embd
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
cur = build_mm(model.spk_proj_w, cur);
cb(cur, "spk_proj", -1);
ggml_build_forward_expand(gf, cur);
return gf;
}
tools/mtmd/models/qwen3tts-gen.cpp
+4
View File
@@ -610,6 +610,10 @@ std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, c
const auto & c2w = model.c2w;
std::vector<c2w_state_slot> slots;
if (c2w.pre_conv_w == nullptr) {
return slots; // not a code2wav model, it keeps no state between calls
}
slots.push_back({"tfm_pos", 1, 1});
// prefix is (W-1) frames, the batch itself gives the other N=W frames (see tfm_layer_forward)
tools/mtmd/mtmd-audio.cpp
+30
View File
@@ -1423,3 +1423,33 @@ std::vector<float> mtmd_audio_streaming_istft::flush() {
return output;
}
//
// 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
//
bool mtmd_audio_preprocessor_pockettts::preprocess(const float * samples,
size_t n_samples,
std::vector<mtmd_audio_mel> & output) {
// the encoder needs whole frames, see pad_for_conv1d() in the reference
const int64_t frame_size = (int64_t) hparams.mimi_downsample * 120;
if (n_samples == 0 || frame_size <= 0) {
return false;
}
const int64_t n_frames = (int64_t) (n_samples + frame_size - 1) / frame_size;
const int64_t n_padded = n_frames * frame_size;
mtmd_audio_mel out;
out.n_mel = 1;
out.n_len = n_padded;
out.n_len_org = (int64_t) n_samples;
out.data.assign((size_t) n_padded, 0.0f);
std::copy(samples, samples + n_samples, out.data.begin());
output.push_back(std::move(out));
return true;
}
tools/mtmd/mtmd-audio.h
+7
View File
@@ -129,6 +129,13 @@ struct mtmd_audio_preprocessor_qwen3tts_spk : mtmd_audio_preprocessor {
mtmd_audio_cache cache;
};
// mimi convolves the waveform directly, so this only pads it to a whole number of frames
struct mtmd_audio_preprocessor_pockettts : mtmd_audio_preprocessor {
mtmd_audio_preprocessor_pockettts(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {}
void initialize() override {}
bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) override;
};
struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor {
mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { }
void initialize() override;
tools/mtmd/mtmd-helper-gen.cpp
+335
View File
@@ -5,6 +5,7 @@
#include "../src/llama-ext.h"
#include <algorithm>
#include <cctype>
#include <cstring>
#include <memory>
#include <string>
@@ -460,10 +461,344 @@ private:
std::vector<char> out_buf;
};
// 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 {
public:
using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline;
void reset() override {
seq_id = 0;
pos = 0;
feats_buf.clear();
dec_state.clear();
audio_pcm.clear();
h_state_buf.clear();
out_buf.clear();
prompt_embd_buf.clear();
prompt_batch.reset();
n_prompt = 0;
prompt_pos = 0;
step_idx = 0;
eos_step = -1;
}
int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
reset();
seq_id = inp->seq_id;
if (!ensure_cache()) {
return 1;
}
std::vector<float> voice;
if (inp->speaker_ref) {
if (!encode_speaker(inp->speaker_ref, voice)) {
return 1;
}
}
const std::string text = prepare_text(std::string(inp->prompt, inp->prompt_len));
if (text.empty()) {
LOG_ERR("mtmd_helper_gen_audio: empty prompt\n");
return 1;
}
frames_after_eos = count_words(text) <= 4 ? 5 : 3;
std::vector<llama_token> ids(text.size() + 16);
int n_ids = llama_tokenize(vocab, text.c_str(), (int32_t) text.size(), ids.data(),
(int32_t) ids.size(), false, false);
if (n_ids <= 0) {
LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n");
return 1;
}
ids.resize((size_t) n_ids);
const int n_e = n_embd;
auto push_row = [&](llama_token t) {
prompt_embd_buf.insert(prompt_embd_buf.end(),
tok_embd.begin() + (size_t) t * n_e,
tok_embd.begin() + (size_t) (t + 1) * n_e);
};
// sequence order is voice, then text, then the audio BOS that starts generation
if (!voice.empty()) {
push_row(bos_before_voice);
prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end());
}
for (llama_token t : ids) {
push_row(t);
}
push_row(audio_bos);
n_prompt = (int) (prompt_embd_buf.size() / (size_t) n_e);
prompt_batch.reset(new decode_embd_batch(prompt_embd_buf.data(), n_prompt, 1, n_e));
prompt_batch->set_position_normal(0, seq_id);
prompt_pos = 0;
seed = inp->seed;
out_type = inp->out_type;
return 0;
}
int32_t step_prompt(int32_t n_batch) override {
GGML_ASSERT(n_batch > 0);
if (prompt_pos >= n_prompt) {
return 0;
}
const int32_t n_tokens_batch = std::min(n_batch, n_prompt - prompt_pos);
llama_batch batch_view = prompt_batch->get_view(prompt_pos, n_tokens_batch);
if ((prompt_pos + n_tokens_batch) == n_prompt) {
batch_view.logits[n_tokens_batch - 1] = 1;
}
if (llama_decode(lctx, batch_view) != 0) {
LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n");
return -1;
}
pos += n_tokens_batch;
prompt_pos += n_tokens_batch;
if (prompt_pos >= n_prompt) {
prompt_batch.reset();
prompt_embd_buf.clear();
return 0;
}
return n_prompt - prompt_pos;
}
int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override {
(void) sampled; // the backbone output is continuous, there is no token to consume
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
inp.seed = seed;
inp.n_steps = -1;
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: flow decode failed\n");
return 1;
}
if (out.is_eos && eos_step < 0) {
eos_step = step_idx;
}
// the frame of the stopping step is discarded, matching _autoregressive_generation()
if (eos_step >= 0 && step_idx >= eos_step + frames_after_eos) {
*out_stop = true;
*h_state_out = nullptr;
return 0;
}
feats_buf.insert(feats_buf.end(), out.feats, out.feats + out.n_feats);
step_idx++;
if (out.n_feats > 0 && feats_buf.size() / out.n_feats >= window_frames) {
if (!flush_gen_wav()) {
return 1;
}
}
decode_embd_batch batch_embd(const_cast<float *>(out.embd), 1, 1, n_embd);
batch_embd.set_position_normal(pos, seq_id);
batch_embd.batch.logits[0] = 1;
pos++;
if (llama_decode(lctx, batch_embd.batch) != 0) {
LOG_ERR("mtmd_helper_gen_audio: decode failed\n");
return 1;
}
const float * he = llama_get_embeddings_ith(lctx, -1);
h_state_buf.assign(he, he + n_embd);
*h_state_out = h_state_buf.data();
return 0;
}
int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override {
if (!flush_gen_wav()) {
return 1;
}
*out_sample_rate = info.sample_rate;
if (out_n_samples) {
*out_n_samples = (int64_t) audio_pcm.size();
}
if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
*out_data = (const char *) audio_pcm.data();
*out_data_len = audio_pcm.size() * sizeof(float);
return 0;
}
out_buf.clear();
if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
return 1;
}
*out_data = out_buf.data();
*out_data_len = out_buf.size();
return 0;
}
private:
bool ensure_cache() {
if (specials_ok) {
return true;
}
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;
}
}
const uint32_t n_tok_embd = llama_model_get_tok_embd(model, nullptr);
if (n_tok_embd == 0) {
LOG_ERR("mtmd_helper_gen_audio: model has no token embeddings\n");
return false;
}
tok_embd.resize(n_tok_embd);
if (llama_model_get_tok_embd(model, tok_embd.data()) != n_tok_embd) {
LOG_ERR("mtmd_helper_gen_audio: token embedding copy failed\n");
return false;
}
specials_ok = true;
return true;
}
// same normalization as prepare_text_prompt() in the reference, it affects quality
static std::string prepare_text(const std::string & in) {
std::string s;
s.reserve(in.size() + 1);
for (char c : in) {
s += (c == '\n' || c == '\r') ? ' ' : c;
}
const size_t b = s.find_first_not_of(' ');
const size_t e = s.find_last_not_of(' ');
if (b == std::string::npos) {
return "";
}
s = s.substr(b, e - b + 1);
if (s[0] >= 'a' && s[0] <= 'z') {
s[0] = (char) (s[0] - 'a' + 'A');
}
const unsigned char last = (unsigned char) s.back();
if (std::isalnum(last)) {
s += '.';
}
return s;
}
static int count_words(const std::string & s) {
int n = 0;
bool in_word = false;
for (char c : s) {
if (c == ' ') {
in_word = false;
} else if (!in_word) {
in_word = true;
n++;
}
}
return n;
}
// runs the reference wav through the mimi encoder, returns one row per 12.5Hz frame
bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
if (!mtmd_support_audio(mctx)) {
LOG_ERR("mtmd_helper_gen_audio: mmproj has no voice encoder\n");
return false;
}
const std::string marker = mtmd_default_marker();
mtmd_input_text text{ marker.c_str(), marker.size(), false, true };
mtmd_input_chunks * chunks = mtmd_input_chunks_init();
const mtmd_bitmap * bptr = bitmap;
bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0;
if (ok) {
ok = false;
for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {
const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) {
continue;
}
if (mtmd_encode_chunk(mctx, chunk) != 0) {
LOG_ERR("mtmd_helper_gen_audio: voice encode failed\n");
break;
}
const float * embd = mtmd_get_output_embd(mctx);
const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk);
out.assign(embd, embd + n);
ok = true;
break;
}
}
mtmd_input_chunks_free(chunks);
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
bool flush_gen_wav() {
if (feats_buf.empty()) {
return true;
}
mtmd_gen_inp inp{};
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
inp.feats = feats_buf.data();
inp.n_feats = feats_buf.size();
inp.seed = seed;
inp.state_data = dec_state.empty() ? nullptr : (const char *) dec_state.data();
inp.state_size = dec_state.size();
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: mimi decode failed\n");
return false;
}
audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples);
dec_state.assign(out.state_data, out.state_data + out.state_size);
feats_buf.clear();
return true;
}
bool specials_ok = false;
llama_token bos_before_voice = LLAMA_TOKEN_NULL;
llama_token audio_bos = LLAMA_TOKEN_NULL;
std::vector<float> tok_embd;
llama_seq_id seq_id = 0;
int pos = 0;
std::vector<float> prompt_embd_buf;
std::unique_ptr<decode_embd_batch> prompt_batch;
int n_prompt = 0;
int prompt_pos = 0;
uint32_t seed = UINT32_MAX;
// end-of-speech is latched, then a few more frames are generated as tail padding
int step_idx = 0;
int eos_step = -1;
int frames_after_eos = 3;
// latents are decoded a window at a time, the decoder state bridges the windows
size_t window_frames = 8;
std::vector<float> feats_buf;
std::vector<uint8_t> dec_state;
std::vector<float> audio_pcm;
std::vector<float> h_state_buf;
mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
std::vector<char> out_buf;
};
static std::unique_ptr<mtmd_gen_audio_pipeline> make_pipeline(llama_context * lctx, mtmd_context * mctx) {
switch (mtmd_gen_audio_get_info(mctx).type) {
case MTMD_GEN_AUDIO_TYPE_QWEN3TTS:
return std::unique_ptr<mtmd_gen_audio_pipeline>(new qwen3tts_gen_audio_pipeline(lctx, mctx));
case MTMD_GEN_AUDIO_TYPE_POCKETTTS:
return std::unique_ptr<mtmd_gen_audio_pipeline>(new pockettts_gen_audio_pipeline(lctx, mctx));
default:
return nullptr;
}
tools/mtmd/mtmd.cpp
+11
View File
@@ -762,6 +762,10 @@ struct mtmd_context {
{
audio_preproc = std::make_unique<mtmd_audio_preprocessor_qwen3tts_spk>(ctx_a);
} break;
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
{
audio_preproc = std::make_unique<mtmd_audio_preprocessor_pockettts>(ctx_a);
} break;
default:
throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj));
}
@@ -1591,6 +1595,10 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
info.type = MTMD_GEN_AUDIO_TYPE_QWEN3TTS;
info.sample_rate = 24000;
break;
case PROJECTOR_TYPE_POCKETTTS_GEN:
info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS;
info.sample_rate = 24000;
break;
default:
info.type = MTMD_GEN_AUDIO_TYPE_NONE;
break;
@@ -1688,6 +1696,9 @@ 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
params.seed = inp->seed;
params.codes = has_codes ? &in_codes : nullptr;
params.feats = has_feats ? &in_feats : nullptr;
params.out_audio = &ctx->gen_out_audio;
tools/mtmd/mtmd.h
+2
View File
@@ -334,6 +334,7 @@ MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname);
enum mtmd_gen_audio_type {
MTMD_GEN_AUDIO_TYPE_NONE, // not supported
MTMD_GEN_AUDIO_TYPE_QWEN3TTS,
MTMD_GEN_AUDIO_TYPE_POCKETTTS,
};
struct mtmd_gen_audio_info {
enum mtmd_gen_audio_type type;
@@ -345,6 +346,7 @@ enum mtmd_gen_process_type {
MTMD_GEN_PROCESS_TYPE_GEN_CODE, // h_state to semantic (codes, mel-spectrogram, etc.)
MTMD_GEN_PROCESS_TYPE_GEN_WAV, // convert semantic to PCM audio
// for qwen3tts, this is code2wav
// for pocket-tts, this is mimi decoder
};
struct mtmd_gen_inp {
enum mtmd_gen_process_type type;