LLM-Inference/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-08-12 22:31:11 +04:00
Code Issues Packages Projects Releases Wiki Activity

mtmd: support pocket-tts (#26871)

Browse Source 
* adapt the api

* text model ok

* working impl, need verify and clean up

* mtmd: build the pocket-tts transposed convolutions as GEMM + col2im

ggml_conv_transpose_1d has no grouped mode, so the depthwise upsample
was built as one convolution and one concat per channel, which floods
the graph with small nodes and makes kernel launches dominate the
decoder.

Fold both cases into the column form the seanet decoder already needs:
the general case reshapes the kernel to [IC, K * OC] and matmuls it
with the input, the depthwise case batches a matmul over the channels
so a step scales its own kernel. A single col2im_1d then scatter-adds
the columns back to the signal, with the same shape as before, so the
overlap-add tail, the streaming state and the bias are untouched.

Generation time per frame drops by 80% on CUDA and by 50% on CPU. The
output matches the previous implementation sample for sample, with a
correlation of 0.999994 and identical frame counts.

* flow_temp +  frames_after_eos

* chunking

* mtmd: carry the remaining pocket-tts per-pack settings

The language packs also tune the end-of-speech padding and the padding
of short prompts, next to the temperature already carried in the
mmproj: french_24l asks for 8 tail frames instead of the guessed 3,
english_2026-01 asks for short prompts to be padded with spaces.

Write both in the mmproj as clip.gen.audio.frames_after_eos and
clip.gen.audio.pad_short_text, keyed on the pack in the conversion
script like the temperature. The loader keeps them optional, so a
mmproj without them behaves as before. Map semicolons to commas for
every pack instead, the reference only asks for it on three of them and
it costs nothing elsewhere.

Existing mmproj files must be converted again to carry the two keys.

On a long french text the port now lands within 2% of the reference:
22.96s against 23.44s, with the same peak level and the same amount of
silence.

* clip.gen.audio.model_variant

* clean up code comments

* nit: drop the dead flow_temp hparam, the pack table holds the default

* update docs

* address security problems

* less invasive base.py

* lint

* add mtmd_gen_inp_default

* add docs

* rm gen_flow_temp

---------

Co-authored-by: Pascal <admin@serveurperso.com>
This commit is contained in:
Xuan-Son Nguyen
2026-08-11 14:18:30 +02:00
committed by GitHub
co-authored by Pascal
parent 8d274dd7c6
commit 6e62ba5384
29 changed files with 2458 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
conversion/__init__.py
+2
View File
@@ -214,6 +214,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
"Qwen3MoeForCausalLM": "qwen",
"Qwen3NextForCausalLM": "qwen",
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
"PocketTTSModel": "pockettts",
"Qwen3TTSForConditionalGeneration": "qwen3tts",
"Qwen3VLForConditionalGeneration": "qwen3vl",
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
@@ -310,6 +311,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
"Qwen2_5_VLForConditionalGeneration": "qwenvl",
"Qwen3ASRForConditionalGeneration": "qwen3vl",
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
"PocketTTSModel": "pockettts",
"Qwen3TTSForConditionalGeneration": "qwen3tts",
"Qwen3VLForConditionalGeneration": "qwen3vl",
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
conversion/base.py
+28
View File
@@ -58,6 +58,11 @@ logger = logging.getLogger("hf-to-gguf")
AnyModel = TypeVar("AnyModel", bound="type[ModelBase]")
# for checkpoints that ship no config.json, we will try to provide a synthetic one
HparamsMatcher = Callable[[Path], bool]
HparamsLoader = Callable[[Path], dict[str, Any]]
class SentencePieceTokenTypes(IntEnum):
NORMAL = 1
UNKNOWN = 2
@@ -77,6 +82,7 @@ class ModelBase:
ModelType.TEXT: {},
ModelType.MMPROJ: {},
}
_hparams_loaders: list[tuple[HparamsMatcher, HparamsLoader]] = []
dir_model: Path
ftype: gguf.LlamaFileType
@@ -1040,6 +1046,24 @@ class ModelBase:
return part_names
@staticmethod
def load_hparams_guess(dir_model: Path) -> dict[str, Any] | None:
# some models ship no config.json, will try to guess them
from conversion import load_all_models
load_all_models()
for matcher, loader in ModelBase._hparams_loaders:
if matcher(dir_model):
return loader(dir_model)
return None
@classmethod
def register_hparams_loader(cls, matcher: HparamsMatcher) -> Callable[[HparamsLoader], HparamsLoader]:
def inner(loader: HparamsLoader) -> HparamsLoader:
cls._hparams_loaders.append((matcher, loader))
return loader
return inner
@staticmethod
def load_hparams(dir_model: Path, is_mistral_format: bool):
if is_mistral_format:
@@ -1053,6 +1077,10 @@ class ModelBase:
config = AutoConfig.from_pretrained(dir_model, trust_remote_code=False).to_dict()
except Exception as e:
logger.warning(f"Failed to load model config from {dir_model}: {e}")
if not (dir_model / "config.json").is_file():
config = ModelBase.load_hparams_guess(dir_model)
if config is not None:
return config
logger.warning("Trying to load config.json instead")
with open(dir_model / "config.json", "r", encoding="utf-8") as f:
config = json.load(f)
conversion/pockettts.py
+378
View File
@@ -0,0 +1,378 @@
from __future__ import annotations
import re
from pathlib import Path
from typing import Any, Iterable, TYPE_CHECKING
import torch
if TYPE_CHECKING:
from torch import Tensor
from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf, logger
# 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 come from _load_hparams() below.
#
# 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, 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
# flow_lm.transformer --> mapped to normal libllama text model (autoregressive)
# flow_lm.flow_net + out_eos --> MTMD_GEN_PROCESS_TYPE_GEN_CODE
# mimi decoder --> MTMD_GEN_PROCESS_TYPE_GEN_WAV
# indices into mimi.encoder.model / mimi.decoder.model for stage i, see SEANetEncoder/SEANetDecoder
_ENC_RES_IDX = lambda i: 1 + 3 * i # noqa: E731
_ENC_SCALE_IDX = lambda i: 3 + 3 * i # noqa: E731
_DEC_SCALE_IDX = lambda i: 2 + 3 * i # noqa: E731
_DEC_RES_IDX = lambda i: 3 + 3 * i # noqa: E731
_N_SEANET_STAGES = 3
_SAMPLE_RATE = 24000
def _tensor_shapes(dir_model: Path) -> dict[str, tuple[int, ...]]:
part_names = ModelBase.get_model_part_names(dir_model, "model", ".safetensors")
if len(part_names) != 1:
return {}
with gguf.utility.SafetensorsLocal(dir_model / part_names[0]) as part:
return {name: tuple(part[name].shape) for name in part.keys()}
@ModelBase.register_hparams_loader(lambda dir_model: "flow_lm.bos_emb" in _tensor_shapes(dir_model))
def _load_hparams(dir_model: Path) -> dict[str, Any]:
logger.info("gguf: detected pocket-tts checkpoint, deriving hparams from tensor shapes")
shapes = _tensor_shapes(dir_model)
n_vocab, n_embd = shapes["flow_lm.conditioner.embed.weight"]
n_layer = sum(1 for name in shapes if re.fullmatch(r"flow_lm\.transformer\.layers\.\d+\.norm1\.weight", name))
n_layer_a = sum(1 for name in shapes if re.fullmatch(r"mimi\.encoder_transformer\.transformer\.layers\.\d+\.norm1\.weight", name))
n_embd_a = shapes["mimi.encoder_transformer.transformer.layers.0.norm1.weight"][0]
return {
"architectures": ["PocketTTSModel"],
"model_type": "pockettts",
"num_hidden_layers": n_layer,
"hidden_size": n_embd,
"intermediate_size": shapes["flow_lm.transformer.layers.0.linear1.weight"][0],
# the transformer is fully causal with no context limit, this only bounds the KV cache
"max_position_embeddings": 4096,
# not in the checkpoint, but every released variant uses head_dim 64
"num_attention_heads": n_embd // 64,
# extra rows for the learned input vectors, see _embd_table()
"vocab_size": n_vocab + (2 if "flow_lm.bos_before_voice" in shapes else 1),
"rope_theta": 10000.0,
"layer_norm_eps": 1e-5,
"audio_config": {
"num_hidden_layers": n_layer_a,
"hidden_size": n_embd_a,
"intermediate_size": shapes["mimi.encoder_transformer.transformer.layers.0.linear1.weight"][0],
"num_attention_heads": n_embd_a // 64,
},
}
@ModelBase.register("PocketTTSModel")
class PocketTTSModel(TextModel):
model_arch = gguf.MODEL_ARCH.POCKETTTS
_LAYER_TENSOR_MAP = {
"norm1": gguf.MODEL_TENSOR.ATTN_NORM,
"norm2": gguf.MODEL_TENSOR.FFN_NORM,
"self_attn.out_proj": gguf.MODEL_TENSOR.ATTN_OUT,
"linear1": gguf.MODEL_TENSOR.FFN_UP,
"linear2": gguf.MODEL_TENSOR.FFN_DOWN,
}
def set_vocab(self):
# 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]
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 rows of the embedding table are not sentencepiece pieces
extra = self._extra_tokens()
for i, name in enumerate(extra):
tokens[len(tokens) - len(extra) + i] = name.encode("utf-8")
toktypes[len(tokens) - len(extra) + i] = SentencePieceTokenTypes.CONTROL
scores[len(tokens) - len(extra) + i] = -1000.0
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)
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if not name.startswith("flow_lm."):
return # mimi and the flow net go to the mmproj
if name == "flow_lm.conditioner.embed.weight":
yield (self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD), self._embd_table(data_torch))
return
if name.startswith("flow_lm.out_norm."):
suffix = "." + name.rsplit(".", 1)[1]
yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT_NORM, suffix=suffix), data_torch)
return
if name.startswith("flow_lm.transformer.layers."):
assert bid is not None
key_with_suffix = name.split(f"layers.{bid}.", 1)[1]
key, suffix = key_with_suffix.rsplit(".", 1)
if key == "self_attn.in_proj":
q, k, v = data_torch.chunk(3, dim=0)
yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), q)
yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), k)
yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), v)
return
tensor = self._LAYER_TENSOR_MAP.get(key)
if tensor is not None:
yield (self.format_tensor_name(tensor, bid, suffix="." + suffix), data_torch)
return
return
def _extra_tokens(self) -> list[str]:
# the conditioner's padding row, then the learned vectors appended by _embd_table().
# bos_before_voice only exists when the pack sets insert_bos_before_voice
names = ["<|pad|>"]
if "flow_lm.bos_before_voice" in self.model_tensors:
names.append("<|bos_before_voice|>")
names.append("<|audio_bos|>")
return names
def _embd_table(self, embed: Tensor) -> Tensor:
rows = [embed]
if "flow_lm.bos_before_voice" in self.model_tensors:
rows.append(self.model_tensors["flow_lm.bos_before_voice"]().reshape(1, -1).to(embed.dtype))
# bos_emb is a latent, it only enters the backbone through input_linear
bos_emb = self.model_tensors["flow_lm.bos_emb"]()
input_linear = self.model_tensors["flow_lm.input_linear.weight"]()
audio_bos = torch.nn.functional.linear(bos_emb.float(), input_linear.float()).reshape(1, -1)
rows.append(audio_bos.to(embed.dtype))
return torch.cat(rows, dim=0)
@ModelBase.register("PocketTTSModel")
class PocketTTSMmprojModel(MmprojModel):
has_audio_encoder = True
has_vision_encoder = False
_MIMI_TFM_MAP = {
"norm1": (gguf.MODEL_TENSOR.A_ENC_INPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM),
"norm2": (gguf.MODEL_TENSOR.A_ENC_OUTPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM),
"self_attn.out_proj": (gguf.MODEL_TENSOR.A_ENC_OUTPUT, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT),
"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_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),
(gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V),
)
def set_gguf_parameters(self):
self.gguf_writer.add_file_type(self.ftype)
assert self.hparams_audio is not None
# voice-prompt encoder: mimi encoder + speaker_proj
self.gguf_writer.add_clip_has_audio_encoder(True)
# note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models
self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC)
self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"])
self.gguf_writer.add_audio_embedding_length(self.hparams_audio["hidden_size"])
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
# the SEANet and flow net hparams are constant across the family, clip.cpp holds them
self.gguf_writer.add_clip_has_gen_audio_encoder(True)
self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_GEN)
self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text)
self.gguf_writer.add_gen_audio_embedding_length(self.hparams_audio["hidden_size"])
self.gguf_writer.add_gen_audio_feed_forward_length(self.hparams_audio["intermediate_size"])
self.gguf_writer.add_gen_audio_block_count(self.hparams_audio["num_hidden_layers"])
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)
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
# conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path
if ".seanet." in new_name or new_name in ("a.downsample.conv.weight", "a.gen.wav.upsample.weight"):
return gguf.GGMLQuantizationType.F16
return False
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
del bid # the block index of the mimi transformers is parsed here, not by the base class
T = gguf.MODEL_TENSOR
if name in ("flow_lm.bos_emb", "flow_lm.bos_before_voice", "flow_lm.conditioner.embed.weight"):
return # folded into the backbone embedding table
if name.startswith("flow_lm.transformer.") or name.startswith("flow_lm.out_norm."):
return # backbone
if name == "flow_lm.speaker_proj_weight":
yield (self.format_tensor_name(T.A_ENC_SPEAKER_PROJ), data_torch)
return
if name == "flow_lm.input_linear.weight":
yield (self.format_tensor_name(T.A_GEN_INPUT_LINEAR), data_torch)
return
if name == "flow_lm.emb_mean":
yield (self.format_tensor_name(T.A_GEN_EMB_MEAN, suffix=""), data_torch)
return
if name == "flow_lm.emb_std":
yield (self.format_tensor_name(T.A_GEN_EMB_STD, suffix=""), data_torch)
return
if name.startswith("flow_lm.out_eos."):
suffix = "." + name.rsplit(".", 1)[1]
yield (self.format_tensor_name(T.A_GEN_OUT_EOS, suffix=suffix), data_torch)
return
if name.startswith("flow_lm.flow_net."):
yield from self._flow_net_tensor(name, data_torch)
return
if name == "mimi.downsample.conv.conv.weight":
yield (self.format_tensor_name(T.A_ENC_DOWNSAMPLE_CONV), data_torch)
return
if name == "mimi.upsample.convtr.convtr.weight":
yield (self.format_tensor_name(T.A_GEN_WAV_UPSAMPLE), data_torch)
return
if name == "mimi.quantizer.output_proj.weight":
yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_OUT), data_torch.squeeze(-1))
return
if "_transformer.transformer.layers." in name:
yield from self._mimi_tfm_tensor(name, data_torch)
return
if name.startswith("mimi.encoder.model.") or name.startswith("mimi.decoder.model."):
yield from self._seanet_tensor(name, data_torch)
return
return
def _flow_net_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]:
T = gguf.MODEL_TENSOR
key = name.split("flow_lm.flow_net.", 1)[1]
suffix = "." + key.rsplit(".", 1)[1]
simple = {
"input_proj": T.A_GEN_FLOW_INPUT_PROJ,
"cond_embed": T.A_GEN_FLOW_COND_EMBD,
"final_layer.linear": T.A_GEN_FLOW_FINAL_PROJ,
"final_layer.adaLN_modulation.1": T.A_GEN_FLOW_FINAL_ADA,
}
tensor = simple.get(key.rsplit(".", 1)[0])
if tensor is not None:
yield (self.format_tensor_name(tensor, suffix=suffix), data_torch)
return
if key.startswith("time_embed."):
bid = int(key.split(".")[1])
rest = key.split(f"time_embed.{bid}.", 1)[1]
time_map = {
"freqs": (T.A_GEN_FLOW_TIME_FREQS, ""),
"mlp.0": (T.A_GEN_FLOW_TIME_UP, suffix),
"mlp.2": (T.A_GEN_FLOW_TIME_DOWN, suffix),
"mlp.3.alpha": (T.A_GEN_FLOW_TIME_NORM, ""),
}
entry = time_map.get(rest) or time_map.get(rest.rsplit(".", 1)[0])
if entry is not None:
yield (self.format_tensor_name(entry[0], bid, suffix=entry[1]), data_torch)
return
if key.startswith("res_blocks."):
bid = int(key.split(".")[1])
rest = key.split(f"res_blocks.{bid}.", 1)[1].rsplit(".", 1)[0]
blk_map = {
"in_ln": T.A_GEN_FLOW_BLK_NORM,
"mlp.0": T.A_GEN_FLOW_BLK_UP,
"mlp.2": T.A_GEN_FLOW_BLK_DOWN,
"adaLN_modulation.1": T.A_GEN_FLOW_BLK_ADA,
}
tensor = blk_map.get(rest)
if tensor is not None:
yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch)
return
def _mimi_tfm_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]:
is_decoder = name.startswith("mimi.decoder_transformer.")
bid = int(name.split("_transformer.transformer.layers.", 1)[1].split(".")[0])
key_with_suffix = name.split(f".layers.{bid}.", 1)[1]
if key_with_suffix == "self_attn.in_proj.weight":
q, k, v = data_torch.chunk(3, dim=0)
names = self._MIMI_TFM_QKV[1 if is_decoder else 0]
for tensor, part in zip(names, (q, k, v)):
yield (self.format_tensor_name(tensor, bid), part)
return
key, suffix = key_with_suffix.rsplit(".", 1)
entry = self._MIMI_TFM_MAP.get(key) or self._MIMI_TFM_MAP.get(key_with_suffix)
if entry is None:
return
tensor = entry[1 if is_decoder else 0]
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]]:
T = gguf.MODEL_TENSOR
is_decoder = name.startswith("mimi.decoder.")
idx = int(name.split(".model.", 1)[1].split(".")[0])
suffix = "." + name.rsplit(".", 1)[1]
conv_in, conv_out, res1, res2, scale = (
(T.A_GEN_WAV_SEANET_CONV_IN, T.A_GEN_WAV_SEANET_CONV_OUT, T.A_GEN_WAV_SEANET_RES_CONV1,
T.A_GEN_WAV_SEANET_RES_CONV2, T.A_GEN_WAV_SEANET_SCALE_CONV)
if is_decoder else
(T.A_ENC_SEANET_CONV_IN, T.A_ENC_SEANET_CONV_OUT, T.A_ENC_SEANET_RES_CONV1,
T.A_ENC_SEANET_RES_CONV2, T.A_ENC_SEANET_SCALE_CONV)
)
if idx == 0:
yield (self.format_tensor_name(conv_in, suffix=suffix), data_torch)
return
if idx == 3 * _N_SEANET_STAGES + 2:
yield (self.format_tensor_name(conv_out, suffix=suffix), data_torch)
return
for stage in range(_N_SEANET_STAGES):
res_idx = _DEC_RES_IDX(stage) if is_decoder else _ENC_RES_IDX(stage)
scale_idx = _DEC_SCALE_IDX(stage) if is_decoder else _ENC_SCALE_IDX(stage)
if idx == scale_idx:
yield (self.format_tensor_name(scale, stage, suffix=suffix), data_torch)
return
if idx == res_idx:
# block.1 is the dilated conv, block.3 the pointwise one (0 and 2 are ELU)
inner = int(name.split(".block.", 1)[1].split(".")[0])
tensor = res1 if inner == 1 else res2
yield (self.format_tensor_name(tensor, stage, suffix=suffix), data_torch)
return
gguf-py/gguf/constants.py
+112
View File
@@ -407,6 +407,8 @@ class Keys:
class ClipGenAudio:
PROJECTOR_TYPE = "clip.gen.audio.projector_type" # for mixed modality models
# 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"
@@ -581,6 +583,7 @@ class MODEL_ARCH(IntEnum):
MELLUM = auto()
NANBEIGE = auto()
QWEN3TTS = auto()
POCKETTTS = auto()
class VISION_PROJECTOR_TYPE(IntEnum):
@@ -1040,6 +1043,38 @@ class MODEL_TENSOR(IntEnum):
A_GEN_WAV_DAC_RES_CONV2 = auto() # DAC residual unit, pointwise causal conv
A_GEN_WAV_DAC_POST_SNAKE = auto() # DAC final SnakeBeta
A_GEN_WAV_DAC_POST_CONV = auto() # DAC conv_post -> 1-channel PCM
# pocket-tts: SEANet encoder (speaker path) and decoder (a.gen.wav path)
A_ENC_SEANET_CONV_IN = auto()
A_ENC_SEANET_CONV_OUT = auto()
A_ENC_SEANET_RES_CONV1 = auto() # residual unit, dilated conv
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()
A_GEN_FLOW_TIME_FREQS = auto() # timestep embedder, stored cos/sin frequencies
A_GEN_FLOW_TIME_UP = auto()
A_GEN_FLOW_TIME_DOWN = auto()
A_GEN_FLOW_TIME_NORM = auto() # RMSNorm alpha
A_GEN_FLOW_BLK_NORM = auto() # AdaLN res block, in_ln
A_GEN_FLOW_BLK_UP = auto()
A_GEN_FLOW_BLK_DOWN = auto()
A_GEN_FLOW_BLK_ADA = auto() # AdaLN modulation, -> shift/scale/gate
A_GEN_FLOW_FINAL_ADA = auto() # final layer AdaLN modulation, -> shift/scale
A_GEN_FLOW_FINAL_PROJ = auto()
A_GEN_OUT_EOS = auto() # end-of-speech head on the backbone hidden state
A_GEN_INPUT_LINEAR = auto() # generated latent -> backbone embd
A_GEN_EMB_MEAN = auto() # latent denormalization stats
A_GEN_EMB_STD = auto()
A_GEN_WAV_QUANT_OUT = auto() # DummyQuantizer output_proj, latent -> decoder dim
A_GEN_WAV_UPSAMPLE = auto() # frame rate -> encoder frame rate, depthwise convtr
A_GEN_WAV_SEANET_CONV_IN = auto()
A_GEN_WAV_SEANET_CONV_OUT = auto() # -> 1-channel PCM
A_GEN_WAV_SEANET_RES_CONV1 = auto()
A_GEN_WAV_SEANET_RES_CONV2 = auto()
A_GEN_WAV_SEANET_SCALE_CONV = auto() # strided upsample convtr
A_MMPROJ = auto()
A_MMPROJ_FC = auto()
A_MM_NORM_PRE = auto()
@@ -1255,6 +1290,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
MODEL_ARCH.MELLUM: "mellum",
MODEL_ARCH.NANBEIGE: "nanbeige",
MODEL_ARCH.QWEN3TTS: "qwen3tts",
MODEL_ARCH.POCKETTTS: "pockettts",
}
VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = {
@@ -1709,6 +1745,37 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2: "a.gen.wav.dac.blk.{bid}.res.{xid}.conv2",
MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE: "a.gen.wav.dac.post_snake",
MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV: "a.gen.wav.dac.post_conv",
MODEL_TENSOR.A_ENC_SEANET_CONV_IN: "a.seanet.conv_in",
MODEL_TENSOR.A_ENC_SEANET_CONV_OUT: "a.seanet.conv_out",
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}.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",
MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS: "a.gen.flow.time.{bid}.freqs",
MODEL_TENSOR.A_GEN_FLOW_TIME_UP: "a.gen.flow.time.{bid}.up",
MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN: "a.gen.flow.time.{bid}.down",
MODEL_TENSOR.A_GEN_FLOW_TIME_NORM: "a.gen.flow.time.{bid}.norm",
MODEL_TENSOR.A_GEN_FLOW_BLK_NORM: "a.gen.flow.blk.{bid}.norm",
MODEL_TENSOR.A_GEN_FLOW_BLK_UP: "a.gen.flow.blk.{bid}.up",
MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN: "a.gen.flow.blk.{bid}.down",
MODEL_TENSOR.A_GEN_FLOW_BLK_ADA: "a.gen.flow.blk.{bid}.ada",
MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA: "a.gen.flow.final.ada",
MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ: "a.gen.flow.final.proj",
MODEL_TENSOR.A_GEN_OUT_EOS: "a.gen.out_eos",
MODEL_TENSOR.A_GEN_INPUT_LINEAR: "a.gen.input_linear",
MODEL_TENSOR.A_GEN_EMB_MEAN: "a.gen.emb_mean",
MODEL_TENSOR.A_GEN_EMB_STD: "a.gen.emb_std",
MODEL_TENSOR.A_GEN_WAV_QUANT_OUT: "a.gen.wav.quant_out",
MODEL_TENSOR.A_GEN_WAV_UPSAMPLE: "a.gen.wav.upsample",
MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN: "a.gen.wav.seanet.conv_in",
MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT: "a.gen.wav.seanet.conv_out",
MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1: "a.gen.wav.seanet.blk.{bid}.res_conv1",
MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2: "a.gen.wav.seanet.blk.{bid}.res_conv2",
MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV: "a.gen.wav.seanet.blk.{bid}.scale_conv",
MODEL_TENSOR.A_MMPROJ: "mm.a.mlp.{bid}",
MODEL_TENSOR.A_MMPROJ_FC: "mm.a.fc",
MODEL_TENSOR.A_MM_NORM_PRE: "mm.a.norm_pre",
@@ -2020,6 +2087,37 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2,
MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE,
MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV,
MODEL_TENSOR.A_ENC_SEANET_CONV_IN,
MODEL_TENSOR.A_ENC_SEANET_CONV_OUT,
MODEL_TENSOR.A_ENC_SEANET_RES_CONV1,
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,
MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS,
MODEL_TENSOR.A_GEN_FLOW_TIME_UP,
MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN,
MODEL_TENSOR.A_GEN_FLOW_TIME_NORM,
MODEL_TENSOR.A_GEN_FLOW_BLK_NORM,
MODEL_TENSOR.A_GEN_FLOW_BLK_UP,
MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN,
MODEL_TENSOR.A_GEN_FLOW_BLK_ADA,
MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA,
MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ,
MODEL_TENSOR.A_GEN_OUT_EOS,
MODEL_TENSOR.A_GEN_INPUT_LINEAR,
MODEL_TENSOR.A_GEN_EMB_MEAN,
MODEL_TENSOR.A_GEN_EMB_STD,
MODEL_TENSOR.A_GEN_WAV_QUANT_OUT,
MODEL_TENSOR.A_GEN_WAV_UPSAMPLE,
MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN,
MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT,
MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1,
MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2,
MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV,
MODEL_TENSOR.A_ENC_CONV_NORM_MEAN,
MODEL_TENSOR.A_ENC_CONV_NORM_VAR,
MODEL_TENSOR.A_ENC_MEL_FILTERS,
@@ -4903,6 +5001,18 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
],
MODEL_ARCH.POCKETTTS: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
MODEL_TENSOR.ATTN_NORM,
MODEL_TENSOR.ATTN_Q,
MODEL_TENSOR.ATTN_K,
MODEL_TENSOR.ATTN_V,
MODEL_TENSOR.ATTN_OUT,
MODEL_TENSOR.FFN_NORM,
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
],
}
# tensors that will not be serialized
@@ -5179,6 +5289,8 @@ class VisionProjectorType:
NEMOTRON_V2_VL = "nemotron_v2_vl"
QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder
QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor
POCKETTTS_SPKENC = "pockettts_spkenc" # audio: mimi encoder as voice-prompt encoder
POCKETTTS_GEN = "pockettts_gen" # audio generation: flow-matching decoder + mimi decoder
HUNYUANVL = "hunyuanvl"
PARAKEET = "parakeet" # audio
MINIMAXM3 = "minimax_m3"
gguf-py/gguf/gguf_writer.py
+3
View File
@@ -1453,6 +1453,9 @@ 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_model_variant(self, value: str) -> None:
self.add_string(Keys.ClipGenAudio.MODEL_VARIANT, value)
def add_xielu_alpha_p(self, values: Sequence[float]):
self.add_array(Keys.xIELU.ALPHA_P, values)
src/llama-arch.cpp
+1
View File
@@ -147,6 +147,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_MELLUM, "mellum" },
{ LLM_ARCH_NANBEIGE, "nanbeige" },
{ LLM_ARCH_QWEN3TTS, "qwen3tts" },
{ LLM_ARCH_POCKETTTS, "pockettts" },
{ LLM_ARCH_UNKNOWN, "(unknown)" },
};
src/llama-arch.h
+1
View File
@@ -152,6 +152,7 @@ enum llm_arch {
LLM_ARCH_DFLASH,
LLM_ARCH_NANBEIGE,
LLM_ARCH_QWEN3TTS,
LLM_ARCH_POCKETTTS,
LLM_ARCH_UNKNOWN,
};
src/llama-model.cpp
+3
View File
@@ -116,6 +116,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_qwen3vlmoe(params);
case LLM_ARCH_QWEN3TTS:
return new llama_model_qwen3tts(params);
case LLM_ARCH_POCKETTTS:
return new llama_model_pockettts(params);
case LLM_ARCH_PHI2:
return new llama_model_phi2(params);
case LLM_ARCH_PHI3:
@@ -2637,6 +2639,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_MAINCODER:
case LLM_ARCH_GLM_DSA:
case LLM_ARCH_NANBEIGE:
case LLM_ARCH_POCKETTTS:
return LLAMA_ROPE_TYPE_NORM;
// the pairs of head values are offset by n_rot/2
src/models/models.h
+13
View File
@@ -713,6 +713,19 @@ struct llama_model_gpt2 : public llama_model_base {
};
struct llama_model_pockettts : public llama_model_base {
llama_model_pockettts(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_graph_context {
graph(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_codeshell : public llama_model_base {
llama_model_codeshell(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
src/models/pockettts.cpp
+146
View File
@@ -0,0 +1,146 @@
#include "models.h"
// backbone of the pocket-tts CALM pipeline: the "text" side of a flow language model.
// it has no lm_head, the audio latents are produced by the flow net inside the mmproj
void llama_model_pockettts::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
switch (hparams.n_layer()) {
case 6: type = LLM_TYPE_109M; break;
case 24: type = LLM_TYPE_335M; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_pockettts::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output_norm_b = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "bias"), {n_embd}, 0);
// no output head, the logits are unused; reuse the embedding table so a sampler can still run
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias", i), {n_embd}, 0);
create_tensor_qkv(layer, i, n_embd, n_embd, n_embd_gqa, n_embd_gqa, TENSOR_NOT_REQUIRED);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", i), {n_embd}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
}
}
std::unique_ptr<llm_graph_context> llama_model_pockettts::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
llama_model_pockettts::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
GGML_ASSERT(n_embd_head == n_rot);
ggml_tensor * cur;
ggml_tensor * inpL;
inpL = build_inp_embd(model.tok_embd);
ggml_tensor * inp_pos = build_inp_pos();
auto * inp_attn = build_attn_inp_kv();
ggml_tensor * inp_out_ids = build_inp_out_ids();
for (int il = 0; il < n_layer; ++il) {
cur = build_norm(inpL,
model.layers[il].attn_norm,
model.layers[il].attn_norm_b,
LLM_NORM, il);
cb(cur, "attn_norm", il);
// self-attention
{
auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
n_embd_head, n_head, n_head_kv, il);
Qcur = ggml_rope_ext(
ctx0, Qcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow
);
Kcur = ggml_rope_ext(
ctx0, Kcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow
);
cb(Qcur, "Qcur", il);
cb(Kcur, "Kcur", il);
cb(Vcur, "Vcur", il);
cur = build_attn(inp_attn,
model.layers[il].wo, NULL, model.layers[il].wo_s,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
}
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
cb(ffn_inp, "ffn_inp", il);
// FF
{
cur = build_norm(ffn_inp,
model.layers[il].ffn_norm,
model.layers[il].ffn_norm_b,
LLM_NORM, il);
cb(cur, "ffn_norm", il);
cur = build_ffn(cur,
model.layers[il].ffn_up, NULL, NULL,
NULL, NULL, NULL,
model.layers[il].ffn_down, NULL, NULL,
NULL,
LLM_FFN_GELU, LLM_FFN_SEQ, il);
cb(cur, "ffn_out", il);
}
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
// input for next layer
inpL = cur;
}
cur = build_norm(inpL,
model.output_norm,
model.output_norm_b,
LLM_NORM, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
cur = build_lora_mm(model.output, cur, model.output_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
tools/mtmd/CMakeLists.txt
+3
View File
@@ -57,6 +57,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/README-dev.md
+11 -4
View File
@@ -59,8 +59,10 @@ Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system i
### Checklist for porting new audio generation models to mtmd
1. Establish a list of reusable and missing components from the current mtmd implementation.
2. For GGUF conversion:
1. Make sure to consult merged PRs about adding new TTS models, especially reviewer comments
- Example: https://github.com/ggml-org/llama.cpp/pulls?q=is%3Apr+mtmd+tts+is%3Amerged
2. Establish a list of reusable and missing components from the current mtmd implementation.
3. For GGUF conversion:
- Backbone model should be converted to a normal text model (loadable via `libllama`)
- If model used hard-coded embedding row ID, append them to token embeddings and assign token name for them (see `qwen3tts.py`)
- If model have a specific output logits head for audio codes (usually semantic code), keep the head as-is and pad the logits at inference time (see `src/models/qwen3vl.cpp`)
@@ -70,12 +72,17 @@ Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system i
- For tensor naming:
- Prefixed with `a.*` for tensors used by speaker encoder pipeline
- Prefixed with `a.gen.*` for generation stages (code / mel-spectrogram / PCM generation)
3. Make sure most of the changes happen inside `mtmd-helper-gen.cpp`. A good PR looks like this:
- For GGUF metadata:
- Reuse as many existing keys as possible
- In most cases, you can hard-code model configs in the model graph class, or in `clip_hparams`
- If some values need to be exposed to the `mtmd_helper` layer, hard-code them in `mtmd_helper` and distinguish by pipeline and `mtmd_gen_audio_info::model_variant` if necessary
- Do NOT add new GGUF metadata or new fields to `mtmd_gen_audio_info` unless you can prove that you absolutely need them
4. Make sure most of the changes happen inside `mtmd-helper-gen.cpp`. A good PR looks like this:
- 10-20% changes is to add new backbone (text) model and conversion
- 60% changes inside `mtmd-helper-gen.cpp`
- 10% changes inside `libmtmd` and `clip.cpp` systems
- The rest downstream code (CLI, server) should have no changes at all
4. Update usage documentation in `tools/tts/README.md`
5. Update usage documentation in `tools/tts/README.md`
IMPORTANT: If your model needs changes that don't fit the existing infrastructure, **open an issue first for discussion**.
tools/mtmd/clip-impl.h
+39 -1
View File
@@ -92,7 +92,9 @@
#define KEY_A_LOCAL_GROUP_SIZE "clip.audio.local_group_size" // mimo-v2.5: input_local_transformer grouping size
// audio generation (gen-audio)-specific
#define KEY_GEN_AUDIO_PROJ_TYPE "clip.gen.audio.projector_type" // for models with mixed modalities
#define KEY_AUDIO_SUBSAMPLING_FACTOR "clip.audio.subsampling_factor"
// 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"
//
// tensor name constants
@@ -246,6 +248,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 +489,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_MUSE_GLIMMER,
PROJECTOR_TYPE_UNKNOWN,
};
@@ -515,6 +551,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"},
{ PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"},
};
tools/mtmd/clip-model.h
+89
View File
@@ -141,6 +141,20 @@ struct clip_hparams {
int32_t rvq_num_quantizers = 0;
std::vector<int32_t> rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17)
// 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
static constexpr int32_t pockettts_max_spk_seconds = 30;
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
// qwen3tts code2wav
int32_t wav_tfm_n_layer = 0;
int32_t wav_tfm_n_embd = 0;
@@ -402,6 +416,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
@@ -699,6 +770,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
+273 -30
View File
@@ -174,6 +174,10 @@ struct clip_ctx {
bool support_batch = false;
// for audio gen, reseeded only when the caller asks for another seed
std::mt19937 rng{std::random_device{}()};
uint32_t rng_seed = UINT32_MAX;
clip_ctx(clip_context_params & ctx_params) {
flash_attn_type = ctx_params.flash_attn_type;
no_alloc = ctx_params.no_alloc;
@@ -1059,6 +1063,25 @@ 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 = ctx->model.hparams.flow_n_step;
const int64_t n_latent = ctx->model.gen_input_lin_w->ne[0];
GGML_ASSERT(n_step > 0);
GGML_ASSERT(n_latent > 0);
// "inp_feats" takes the caller's buffer as-is, the graph must consume all of it
if (params && params->feats) {
GGML_ASSERT(params->feats->size() % (size_t) n_latent == 0);
GGML_ASSERT(params->feats->size() >= (size_t) n_latent);
}
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;
@@ -1282,6 +1305,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");
@@ -1421,7 +1445,7 @@ struct clip_model_loader {
} break;
case PROJECTOR_TYPE_PARAKEET:
{
get_u32(KEY_AUDIO_SUBSAMPLING_FACTOR, hparams.subsampling_factor);
get_u32(KEY_AUDIO_SUBSMPL_FACTOR, hparams.subsampling_factor);
GGML_ASSERT(hparams.subsampling_factor == 8 &&
"subsampling_factor must match the conv strides in clip_graph_parakeet::build()");
get_u32(KEY_A_CONV_KERNEL_SIZE, hparams.audio_conv_kernel_size);
@@ -1745,6 +1769,22 @@ 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
hparams.flow_n_step = 1;
hparams.gen_eos_threshold = -4.0f;
} break;
case PROJECTOR_TYPE_PADDLEOCR:
{
hparams.n_merge = 2;
@@ -1947,7 +1987,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)) {
@@ -2020,6 +2062,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);
@@ -2081,7 +2148,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;
@@ -2755,6 +2823,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
@@ -4060,6 +4203,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
@@ -4101,6 +4255,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();
@@ -4116,6 +4279,11 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
clip_model_loader::warmup(*ctx, *params->imgs);
}
if (params->seed != ctx->rng_seed) {
ctx->rng_seed = params->seed;
ctx->rng.seed(params->seed == UINT32_MAX ? std::random_device{}() : params->seed);
}
// build the inference graph
ggml_backend_sched_reset(ctx->sched.get());
ggml_cgraph * gf = clip_get_graph_builder(ctx, imgs, params)->build();
@@ -4160,6 +4328,50 @@ 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"));
GGML_ASSERT(n_pos > 0);
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);
// the preprocessor truncates the waveform to keep this mask bounded
const int64_t max_pos = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate / 120;
GGML_ASSERT(n_pos <= max_pos && "pocket-tts speaker reference too long for a dense mask");
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;
@@ -4203,8 +4415,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];
@@ -4737,6 +4949,30 @@ 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");
// Config.default_temperature, for a caller that does not set one
const float temp = params->temp > 0.0f ? params->temp : 0.7f;
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);
}
set_input_f32("inp_noise", noise);
}
} break;
case PROJECTOR_TYPE_GEMMA4V:
case PROJECTOR_TYPE_GEMMA4UV:
{
@@ -4861,20 +5097,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];
@@ -4886,11 +5109,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
set_input_i32("inp_code0", code0);
// one uniform(0,1) draw per codebook, used by do_sampling()
static std::mt19937 rng{ std::random_device{}() };
std::uniform_real_distribution<float> dist(0.0f, 1.0f);
const int64_t n_acoustic = model.gen_code_head_w->ne[2];
for (int64_t g = 0; g < n_acoustic; g++) {
std::vector<float> r = { dist(rng) };
std::vector<float> r = { dist(ctx->rng) };
set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r);
}
}
@@ -5343,14 +5565,31 @@ 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));
}
}
if (params->out_feats != nullptr) {
ggml_tensor * feats = ggml_graph_get_tensor(gf, "out_feats");
if (feats != nullptr) {
auto & out_feats = *params->out_feats;
out_feats.resize(ggml_nelements(feats));
ggml_backend_tensor_get(feats, out_feats.data(), 0, ggml_nbytes(feats));
}
}
if (params->out_is_eos != nullptr) {
ggml_tensor * eos = ggml_graph_get_tensor(gf, "out_eos_score");
if (eos != nullptr) {
GGML_ASSERT(ggml_nelements(eos) == 1);
float score = 0.0f;
ggml_backend_tensor_get(eos, &score, 0, sizeof(float));
*params->out_is_eos = score > hparams.gen_eos_threshold;
}
auto & out_codes = *params->out_codes;
out_codes.resize(ggml_nelements(codes));
ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes));
}
if (params->out_audio != nullptr) {
ggml_tensor * audio = ggml_graph_get_tensor(gf, "out_audio");
@@ -5362,9 +5601,9 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
ggml_backend_tensor_get(audio, out_audio.data(), 0, ggml_nbytes(audio));
// drop the tail audio that comes from the code-0 rear padding
const int64_t n_codes = model.gen_code_head_w->ne[2] + 1;
const int64_t n_codes = params->codes ? model.gen_code_head_w->ne[2] + 1 : 0;
const int64_t n_frames_w = hparams.wav_tfm_swa;
const int64_t n_frames = (int64_t) params->codes->size() / n_codes;
const int64_t n_frames = params->codes ? (int64_t) params->codes->size() / n_codes : n_frames_w;
if (n_frames < n_frames_w) {
const size_t hop = out_audio.size() / n_frames_w;
out_audio.resize((size_t) n_frames * hop);
@@ -5373,12 +5612,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());
@@ -5526,6 +5765,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/clip.h
+5
View File
@@ -104,9 +104,14 @@ struct clip_encode_params {
int32_t top_k = 50;
float top_p = 1.0f;
std::vector<int32_t> * out_codes = nullptr; // this frame's 16 sampled codes
std::vector<float> * out_feats = nullptr; // continuous counterpart of out_codes
uint32_t seed = UINT32_MAX; // UINT32_MAX for random
float temp = 0.0f; // sampling temperature, noise scale for flow-matching decoders
bool * out_is_eos = nullptr;
// GEN_WAV
const std::vector<int32_t> * codes = nullptr; // this frame's 16 RVQ codes
const std::vector<float> * feats = nullptr; // continuous counterpart of codes
std::vector<float> * out_audio = nullptr; // decoded PCM samples, F32
const std::vector<uint8_t> * state_in = nullptr; // state from previous call, null or wrong size means cold start
std::vector<uint8_t> * state_out = nullptr; // state for the next call
tools/mtmd/models/models.h
+56
View File
@@ -318,6 +318,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;
@@ -326,6 +379,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
+291
View File
@@ -0,0 +1,291 @@
#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
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);
}
// 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);
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
// 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);
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
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;
// the slots below are sized from these
GGML_ASSERT(!model.gen_tfm_layers.empty());
GGML_ASSERT((int) seanet.stages.size() >= hparams.seanet_n_stage);
GGML_ASSERT((int) hparams.seanet_ratios.size() >= hparams.seanet_n_stage);
GGML_ASSERT(hparams.mimi_tfm_context > 1 && hparams.mimi_downsample > 0);
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]
// 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);
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
+162
View File
@@ -0,0 +1,162 @@
#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: 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);
}
// 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: 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];
const int64_t T = x->ne[0];
const int64_t p_total = K - stride;
const bool depthwise = w->ne[1] == 1 && w->ne[2] > 1;
const int64_t OC = depthwise ? w->ne[2] : w->ne[1];
const int64_t emit_len = T * stride;
// one column per input step, holding the [K, OC] window that col2im scatter-adds at t * stride
ggml_tensor * col;
if (depthwise) {
// one group per channel: a batched matmul over the channels scales the kernel by each step
ggml_tensor * krn = ggml_reshape_3d(ctx0, w, 1, K, OC); // [1, K, OC]
ggml_tensor * xs = ggml_reshape_3d(ctx0, x, 1, T, OC); // [1, T, OC]
col = ggml_mul_mat(ctx0, krn, xs); // [K, T, OC]
col = ggml_cont(ctx0, ggml_permute(ctx0, col, 0, 2, 1, 3)); // [K, OC, T]
col = ggml_reshape_2d(ctx0, col, K * OC, T);
} else {
ggml_tensor * w2 = ggml_reshape_2d(ctx0, w, K * OC, w->ne[2]);
w2 = ggml_cont(ctx0, ggml_transpose(ctx0, w2)); // [IC, K * OC]
ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [IC, T]
col = ggml_mul_mat(ctx0, w2, xt);
}
ggml_tensor * full = ggml_col2im_1d(ctx0, col, stride, OC, 0); // [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;
}
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
+38
View File
@@ -1423,3 +1423,41 @@ 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", to reuse 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;
}
// the mimi transformer mask is dense, so cost is quadratic in the reference length
const int64_t max_samples = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate;
if ((int64_t) n_samples > max_samples) {
LOG_WRN("%s: speaker reference is %.1f s, truncating to the first %d s\n", __func__,
(double) n_samples / hparams.audio_sample_rate, clip_hparams::pockettts_max_spk_seconds);
n_samples = (size_t) max_samples;
}
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
+573 -10
View File
@@ -5,6 +5,8 @@
#include "../src/llama-ext.h"
#include <algorithm>
#include <cctype>
#include <cmath>
#include <cstring>
#include <memory>
#include <string>
@@ -87,7 +89,8 @@ public:
virtual int32_t step_prompt(int32_t n_batch) = 0;
// sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token,
// those read what they need from h_state_in instead
virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) = 0;
// set out_stop on end-of-speech, h_state_out must be null if no frame is generated
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:
@@ -200,8 +203,10 @@ public:
prompt_pos = 0;
pos = 0;
top_k = inp->top_k > 0 ? inp->top_k : 50;
top_p = inp->top_p > 0 ? inp->top_p : 1.0f;
const mtmd_gen_inp def = mtmd_gen_inp_default(mctx);
top_k = inp->top_k > 0 ? inp->top_k : def.top_k;
top_p = inp->top_p > 0 ? inp->top_p : def.top_p;
seed = inp->seed;
out_type = inp->out_type;
// the prompt above holds the whole text stream up to tts_eos, so every generated
@@ -241,13 +246,26 @@ public:
return n_prompt - prompt_pos;
}
int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) override {
mtmd_gen_inp inp{};
int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override {
if (sampled == LLAMA_TOKEN_NULL) {
LOG_ERR("mtmd_helper_gen_audio: qwen3tts requires a token sampled from the backbone\n");
return 1;
}
// backbone signals end-of-speech with a token, no frame for this step
if (sampled == codec_eos || llama_vocab_is_eog(vocab, sampled)) {
*out_stop = true;
*h_state_out = nullptr;
return 0;
}
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
inp.code0 = sampled - codec_0;
inp.embd = const_cast<float *>(h_state_in);
inp.top_k = top_k;
inp.top_p = top_p;
inp.seed = seed;
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: gen_code process failed\n");
@@ -384,10 +402,11 @@ private:
if (codes_buf.empty()) {
return true;
}
mtmd_gen_inp inp{};
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
inp.codes = codes_buf.data();
inp.n_codes = codes_buf.size();
inp.seed = seed; // same seed as gen_code, else clip reseeds mid-generation
inp.state_data = c2w_state.empty() ? nullptr : (const char *) c2w_state.data();
inp.state_size = c2w_state.size();
mtmd_gen_out out{};
@@ -427,8 +446,9 @@ private:
std::unique_ptr<decode_embd_batch> prompt_batch;
int n_prompt = 0;
int prompt_pos = 0;
int32_t top_k = 50;
float top_p = 1.0f;
int32_t top_k = 50;
float top_p = 1.0f;
uint32_t seed = UINT32_MAX;
std::vector<int32_t> codes_buf;
std::vector<uint8_t> c2w_state;
std::vector<float> audio_pcm;
@@ -438,10 +458,547 @@ private:
std::vector<char> out_buf;
};
// 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
// all three are 0 / false when the pack does not tune them, the model default is then used
struct pockettts_pack_settings {
float temp = 0.0f;
int frames_after_eos = 0;
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()) {
LOG_WRN("mtmd_helper_gen_audio: no tuned settings for pocket-tts variant \"%s\"\n",
variant ? variant : "");
return {};
}
return it->second;
}
// 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;
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;
chunks.clear();
chunk_idx = 0;
n_voice_pos = 0;
chunk_budget = 0;
}
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;
}
}
pack = pockettts_pack(info.model_variant);
const std::string text = prepare_text(std::string(inp->prompt, inp->prompt_len),
pack.pad_short_text);
if (text.empty()) {
LOG_ERR("mtmd_helper_gen_audio: empty prompt\n");
return 1;
}
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);
// 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) {
LOG_INF("mtmd_helper_gen_audio: %d tokens split into %zu chunks\n", n_ids, chunks.size());
}
const int n_e = n_embd;
// sequence order is voice, then text, then the audio BOS that starts generation
if (!voice.empty()) {
GGML_ASSERT(voice.size() % (size_t) n_e == 0);
if (bos_before_voice != LLAMA_TOKEN_NULL) {
push_embd_row(prompt_embd_buf, bos_before_voice);
}
prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end());
}
// every later chunk rewinds to here and re-prompts, so the voice stays primed
n_voice_pos = (int) (prompt_embd_buf.size() / (size_t) n_e);
for (llama_token t : chunks[0]) {
push_embd_row(prompt_embd_buf, t);
}
push_embd_row(prompt_embd_buf, audio_bos);
arm_chunk_budget(0);
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 = mtmd_gen_inp_default(mctx);
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
inp.embd = const_cast<float *>(h_state_in);
// clip only reseeds when the seed changes, so pass the same one on every step
inp.seed = seed;
if (pack.temp > 0.0f) {
inp.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");
return 1;
}
if (out.is_eos && eos_step < 0) {
eos_step = step_idx;
}
// the frame of the stopping step is discarded, matching _autoregressive_generation().
// the budget is the reference's fallback for a chunk whose eos head never fires
const bool chunk_done = (eos_step >= 0 && step_idx >= eos_step + frames_after_eos) ||
step_idx >= chunk_budget;
if (chunk_done) {
if (eos_step < 0) {
LOG_WRN("mtmd_helper_gen_audio: chunk %zu hit its budget without end-of-speech\n", chunk_idx);
}
return finish_chunk(h_state_out, out_stop);
}
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 is optional, some packs do not insert it
bos_before_voice = find_special_token(vocab, "<|bos_before_voice|>");
audio_bos = find_special_token(vocab, "<|audio_bos|>");
if (audio_bos == LLAMA_TOKEN_NULL) {
LOG_ERR("mtmd_helper_gen_audio: missing <|audio_bos|> in vocab\n");
return false;
}
const uint32_t n_tok_embd = llama_model_get_tok_embd(model, nullptr);
if (n_tok_embd == 0) {
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;
}
GGML_ASSERT(n_embd > 0 && n_tok_embd % (uint32_t) n_embd == 0);
specials_ok = true;
return true;
}
// the table can be shorter than the vocab, so bound the row lookup
void push_embd_row(std::vector<float> & dst, llama_token t) const {
const size_t n_rows = tok_embd.size() / (size_t) n_embd;
GGML_ASSERT(t >= 0 && (size_t) t < n_rows);
dst.insert(dst.end(),
tok_embd.begin() + (size_t) t * n_embd,
tok_embd.begin() + (size_t) (t + 1) * n_embd);
}
// token ids of the pieces the reference splits on, see split_into_best_sentences().
// the leading token is dropped, it is the tokenizer's dummy prefix
std::vector<llama_token> punct_ids(const char * s) const {
std::vector<llama_token> ids(16);
const int n = llama_tokenize(vocab, s, (int32_t) strlen(s), ids.data(), (int32_t) ids.size(), false, false);
if (n <= 1) {
return {};
}
return std::vector<llama_token>(ids.begin() + 1, ids.begin() + n);
}
// cut after runs of boundary tokens, so punctuation stays with the sentence it ends
static std::vector<std::vector<llama_token>> split_on(const std::vector<llama_token> & ids,
const std::vector<llama_token> & boundary) {
std::vector<std::vector<llama_token>> out;
size_t start = 0;
bool prev_was_boundary = false;
for (size_t i = 0; i < ids.size(); i++) {
const bool is_boundary = std::find(boundary.begin(), boundary.end(), ids[i]) != boundary.end();
if (!is_boundary && prev_was_boundary) {
out.emplace_back(ids.begin() + start, ids.begin() + i);
start = i;
}
prev_was_boundary = is_boundary;
}
out.emplace_back(ids.begin() + start, ids.end());
return out;
}
std::vector<std::vector<llama_token>> split_chunks(const std::vector<llama_token> & ids) const {
if ((int) ids.size() <= max_chunk_tokens) {
return { ids };
}
const std::vector<llama_token> eos_punct = punct_ids(".!...?");
const std::vector<llama_token> mid_punct = punct_ids(",;:");
// oversized sentences are split again on weaker punctuation, else words get skipped
std::vector<std::vector<llama_token>> segments;
for (auto & seg : split_on(ids, eos_punct)) {
if ((int) seg.size() <= max_chunk_tokens) {
segments.push_back(std::move(seg));
continue;
}
auto sub = split_on(seg, mid_punct);
if (sub.size() > 1) {
for (auto & s : sub) {
segments.push_back(std::move(s));
}
} else {
segments.push_back(std::move(seg));
}
}
std::vector<std::vector<llama_token>> out;
for (auto & seg : segments) {
if (seg.empty()) {
continue;
}
if (!out.empty() && (int) (out.back().size() + seg.size()) <= max_chunk_tokens) {
out.back().insert(out.back().end(), seg.begin(), seg.end());
} else {
out.push_back(std::move(seg));
}
}
if (out.empty()) {
out.push_back(ids);
}
for (const auto & c : out) {
if ((int) c.size() > max_chunk_tokens) {
LOG_WRN("mtmd_helper_gen_audio: chunk of %zu tokens exceeds the %d token budget, "
"generation may skip words\n", c.size(), max_chunk_tokens);
}
}
return out;
}
// _estimate_max_gen_len() plus the per-chunk tail guess, both in frames
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, 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;
}
// ends the current chunk and, if there is another, re-prompts it on top of the voice
int32_t finish_chunk(const float ** h_state_out, bool * out_stop) {
if (!flush_gen_wav()) {
return 1;
}
// the decoder restarts too, the next chunk's audio is not continuous with this one
dec_state.clear();
if (chunk_idx + 1 >= chunks.size()) {
*out_stop = true;
*h_state_out = nullptr;
return 0;
}
chunk_idx++;
// drop this chunk's text and audio, keep the voice conditioning
llama_memory_seq_rm(llama_get_memory(lctx), seq_id, n_voice_pos, -1);
pos = n_voice_pos;
const int n_e = n_embd;
prompt_embd_buf.clear();
for (llama_token t : chunks[chunk_idx]) {
push_embd_row(prompt_embd_buf, t);
}
push_embd_row(prompt_embd_buf, audio_bos);
arm_chunk_budget(chunk_idx);
const int n_rows = (int) (prompt_embd_buf.size() / (size_t) n_e);
GGML_ASSERT(n_rows > 0);
decode_embd_batch batch(prompt_embd_buf.data(), n_rows, 1, n_e);
batch.set_position_normal(pos, seq_id);
batch.batch.logits[n_rows - 1] = 1;
if (llama_decode(lctx, batch.batch) != 0) {
LOG_ERR("mtmd_helper_gen_audio: chunk prompt decode failed\n");
return 1;
}
pos += n_rows;
prompt_embd_buf.clear();
const float * he = llama_get_embeddings_ith(lctx, -1);
h_state_buf.assign(he, he + n_embd);
*h_state_out = h_state_buf.data();
*out_stop = false;
return 0;
}
// same normalization as prepare_text_prompt() in the reference, it affects quality
static std::string prepare_text(const std::string & in, bool pad_short) {
std::string s;
s.reserve(in.size() + 1);
for (char c : in) {
if (c == '\n' || c == '\r') {
s += ' ';
} else if (c == ';') {
s += ',';
} else {
s += 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 += '.';
}
if (pad_short && count_words(s) < 5) {
s = std::string(8, ' ') + 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, the mimi decoder state carries over between calls
bool flush_gen_wav() {
if (feats_buf.empty()) {
return true;
}
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
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;
}
pockettts_pack_settings pack;
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;
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;
size_t chunk_idx = 0;
int n_voice_pos = 0; // KV positions held by the voice conditioning
int chunk_budget = 0;
// 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;
}
@@ -483,11 +1040,17 @@ int32_t mtmd_helper_gen_audio_step_prompt(mtmd_helper_gen_audio * ctx, int32_t n
}
int32_t mtmd_helper_gen_audio_step_gen(mtmd_helper_gen_audio * ctx, llama_token sampled,
const float * h_state_in, const float ** h_state_out) {
const float * h_state_in, const float ** h_state_out,
bool * out_stop) {
if (!ctx->pipeline) {
return 1;
}
return ctx->pipeline->step_gen(sampled, h_state_in, h_state_out);
bool stop = false;
const int32_t ret = ctx->pipeline->step_gen(sampled, h_state_in, h_state_out, &stop);
if (out_stop) {
*out_stop = stop;
}
return ret;
}
int32_t mtmd_helper_gen_audio_get_output(mtmd_helper_gen_audio * ctx, int32_t * out_sample_rate,
tools/mtmd/mtmd-helper.h
+10 -6
View File
@@ -183,8 +183,9 @@ struct mtmd_helper_gen_audio_inp {
mtmd_bitmap * speaker_ref; // optional, can be NULL
const char * lang; // optional, can be NULL
int32_t top_k;
float top_p;
int32_t top_k;
float top_p;
uint32_t seed; // UINT32_MAX for random (default: random)
enum mtmd_helper_gen_audio_outtype out_type;
};
@@ -208,12 +209,15 @@ MTMD_API int32_t mtmd_helper_gen_audio_step_prompt(
int32_t n_batch);
// generates one frame; must only be called after step_prompt() has returned 0
// h_state_out is valid until next step_gen() or reset() call
// sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token
// out_stop (optional) is set on end-of-speech, the caller must then stop the loop
// h_state_out is valid until next step_gen() or reset() call, null if no frame is generated
MTMD_API int32_t mtmd_helper_gen_audio_step_gen(
mtmd_helper_gen_audio * ctx,
llama_token sampled,
const float * h_state_in,
const float ** h_state_out);
const float ** h_state_out,
bool * out_stop);
// out_data valid until next get_output() or reset() call
// out_n_samples (optional, can be NULL) receives the number of generated PCM samples
@@ -261,8 +265,8 @@ struct gen_audio {
int32_t step_prompt(int32_t n_batch) {
return mtmd_helper_gen_audio_step_prompt(ctx.get(), n_batch);
}
int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out) {
return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out);
int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out, bool * out_stop = nullptr) {
return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out, out_stop);
}
int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples = nullptr) {
return mtmd_helper_gen_audio_get_output(ctx.get(), out_sample_rate, out_data, out_data_len, out_n_samples);
tools/mtmd/mtmd.cpp
+77 -16
View File
@@ -477,6 +477,7 @@ struct mtmd_context {
// generation context
struct clip_ctx * ctx_gen_a; // audio
std::vector<int32_t> gen_out_codes; // this frame's 16 sampled codes (GEN_CODE)
std::vector<float> gen_out_feats; // this frame's continuous features, if any (GEN_CODE)
std::vector<float> gen_out_embd; // next-step hidden state fed back to backbone (GEN_CODE)
std::vector<float> gen_out_audio; // decoded PCM samples for the current frame (GEN_WAV)
std::vector<uint8_t> gen_out_state; // state to feed into the next GEN_WAV call
@@ -979,6 +980,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));
}
@@ -1798,16 +1803,22 @@ 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;
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:
info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS;
info.sample_rate = 24000;
break;
default:
info.type = MTMD_GEN_AUDIO_TYPE_NONE;
break;
@@ -1815,6 +1826,33 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
return info;
}
mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) {
mtmd_gen_inp inp{};
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
inp.seed = UINT32_MAX;
if (!ctx->ctx_gen_a) {
return inp;
}
switch (clip_get_projector_type(ctx->ctx_gen_a)) {
case PROJECTOR_TYPE_QWEN3TTS_GEN:
// https://huggingface.co/Qwen/Qwen3-TTS-12Hz-1.7B-Base/blob/main/generation_config.json
inp.top_k = 50;
inp.top_p = 1.0f;
inp.temp = 0.9f; // TODO: handle this on graph
break;
case PROJECTOR_TYPE_POCKETTTS_GEN:
// https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/default_parameters.py
inp.top_k = 50;
inp.top_p = 1.0f;
inp.temp = 0.7f;
break;
default:
break;
}
return inp;
}
static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) {
clip_ctx * ctx_clip = ctx->ctx_gen_a;
if (!ctx_clip) {
@@ -1822,6 +1860,8 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
return 1;
}
*out = {};
if (inp->type == MTMD_GEN_PROCESS_TYPE_GEN_CODE) {
const size_t n_embd = (size_t) clip_n_mmproj_embd(ctx_clip);
@@ -1835,16 +1875,22 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
std::vector<float> out_embd(n_embd);
std::vector<int32_t> out_codes;
std::vector<float> out_feats;
bool is_eos = false;
clip_encode_params params;
params.imgs = &batch;
params.n_threads = ctx->n_threads;
params.gen_process = CLIP_GEN_PROCESS_GEN_CODE;
params.out_embd = &out_embd;
params.out_codes = &out_codes;
params.code0 = inp->code0;
params.top_k = inp->top_k;
params.top_p = inp->top_p;
params.imgs = &batch;
params.n_threads = ctx->n_threads;
params.gen_process = CLIP_GEN_PROCESS_GEN_CODE;
params.out_embd = &out_embd;
params.out_codes = &out_codes;
params.out_feats = &out_feats;
params.code0 = inp->code0;
params.top_k = inp->top_k;
params.top_p = inp->top_p;
params.seed = inp->seed;
params.temp = inp->temp;
params.out_is_eos = &is_eos;
if (!clip_encode(ctx_clip, &params)) {
LOG_ERR("%s: clip_encode failed (gen_code)\n", __func__);
@@ -1853,19 +1899,31 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
ctx->gen_out_embd = std::move(out_embd);
ctx->gen_out_codes = std::move(out_codes);
ctx->gen_out_feats = std::move(out_feats);
out->embd = ctx->gen_out_embd.data();
out->codes = ctx->gen_out_codes.data();
out->n_codes = ctx->gen_out_codes.size();
out->embd = ctx->gen_out_embd.data();
out->codes = ctx->gen_out_codes.data();
out->n_codes = ctx->gen_out_codes.size();
out->feats = ctx->gen_out_feats.data();
out->n_feats = ctx->gen_out_feats.size();
out->is_eos = is_eos;
return 0;
}
// MTMD_GEN_PROCESS_TYPE_GEN_WAV
if (!inp->codes || inp->n_codes == 0) {
LOG_ERR("%s: codes required for gen_wav\n", __func__);
const bool has_codes = inp->codes && inp->n_codes > 0;
const bool has_feats = inp->feats && inp->n_feats > 0;
if (has_codes == has_feats) {
LOG_ERR("%s: gen_wav requires exactly one of codes or feats\n", __func__);
return 1;
}
std::vector<int32_t> in_codes(inp->codes, inp->codes + inp->n_codes);
std::vector<int32_t> in_codes;
std::vector<float> in_feats;
if (has_codes) {
in_codes.assign(inp->codes, inp->codes + inp->n_codes);
} else {
in_feats.assign(inp->feats, inp->feats + inp->n_feats);
}
std::vector<uint8_t> in_state;
if (inp->state_data) {
in_state.assign(inp->state_data, inp->state_data + inp->state_size);
@@ -1885,7 +1943,10 @@ 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;
params.codes = &in_codes;
// 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;
params.out_audio = &ctx->gen_out_audio;
params.state_in = inp->state_data ? &in_state : nullptr;
params.state_out = &ctx->gen_out_state;
tools/mtmd/mtmd.h
+21 -1
View File
@@ -344,18 +344,25 @@ 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;
int32_t sample_rate; // in Hz, for example 24000 for qwen3tts
const char * model_variant; // name of the weight variant, can be nullptr if not applicable
};
MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx);
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;
@@ -364,21 +371,30 @@ struct mtmd_gen_inp {
float * embd; // the hidden state from backbone, must have n_text_embd elements
int32_t top_k;
float top_p;
uint32_t seed; // UINT32_MAX for random
float temp; // sampling temperature, or noise scale for flow-matching decoders
// for MTMD_GEN_PROCESS_TYPE_GEN_WAV
// pass either codes (discrete) or feats (continuous), depending on the pipeline
int32_t * codes;
size_t n_codes;
const float * feats;
size_t n_feats;
const char * state_data;
size_t state_size;
};
struct mtmd_gen_out {
// note: output memory is allocated by the context, valid until next process() call
// for MTMD_GEN_PROCESS_TYPE_GEN_CODE
const int32_t * codes;
size_t n_codes;
size_t n_codes;
const float * feats; // continuous counterpart of codes
size_t n_feats;
const float * embd; // the generated hidden state, to be fed back to backbone
// it must have n_text_embd elements
bool is_eos; // only set by pipelines having the EOS head inside mmproj
// for MTMD_GEN_PROCESS_TYPE_GEN_WAV
const float * audio;
@@ -386,6 +402,10 @@ struct mtmd_gen_out {
const char * state_data;
size_t state_size;
};
// defaults tuned for the loaded pipeline, callers override only what they care about
MTMD_API struct mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx);
// note: this API is stateless, caller must handle state management and audio frame accumulation
MTMD_API int32_t mtmd_gen_audio_process(mtmd_context * ctx,
const struct mtmd_gen_inp * inp,
tools/tts/README.md
+25
View File
@@ -32,3 +32,28 @@ llama-tts -hf ggml-org/Qwen3-TTS-12Hz-1.7B-Base-GGUF \
--tts-speaker-file speaker.mp3 \
--output out.wav
```
## Pocket TTS
Available params:
- `--tts-speaker-file` should point to a speaker reference audio file (wav, mp3). It is required, the model produces almost no audio without it
- Note: `lang` is not used, the language is a property of the weights
Example usage:
```sh
llama-tts -m pocket-tts.gguf \
-mm mmproj-pocket-tts.gguf \
-p "Hello world" \
--tts-speaker-file speaker.mp3 \
--output out.wav
```
**Note for GGUF conversion:**
The [upstream repository](https://huggingface.co/kyutai/pocket-tts) holds one complete model per language under `languages/`, next to a set of shared files at the root. Convert one of the `languages/<name>` directories, **not** the root directory:
```sh
python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --outfile pocket-tts.gguf
python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --mmproj --outfile mmproj-pocket-tts.gguf
```
tools/tts/tts.cpp
+10 -5
View File
@@ -119,6 +119,7 @@ int main(int argc, char ** argv) {
inp.lang = params.tts_lang.c_str();
inp.top_k = params.sampling.top_k;
inp.top_p = params.sampling.top_p;
inp.seed = params.sampling.seed;
inp.out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
//
@@ -143,8 +144,7 @@ int main(int argc, char ** argv) {
}
}
const llama_vocab * vocab = llama_model_get_vocab(model);
// note: some pipelines ignore this token and use the hidden state instead
auto sample_semantic_code = [&]() -> llama_token {
llama_token t = common_sampler_sample(smpl, lctx, -1);
common_sampler_accept(smpl, t, true);
@@ -159,19 +159,24 @@ int main(int argc, char ** argv) {
tts_timings timings;
const int64_t t_gen_start_us = ggml_time_us();
for (; n_frames < max_new && !llama_vocab_is_eog(vocab, sampled); n_frames++) {
bool stop = false;
while (!stop && n_frames < max_new) {
const float * h_next = nullptr;
// stage 2+3: semantic --> acoustic details --> audio waveform
// step_gen() runs both stages and returns new h_state for next step
if (gen.step_gen(sampled, h_state, &h_next) != 0) {
if (gen.step_gen(sampled, h_state, &h_next, &stop) != 0) {
LOG_ERR("step_gen failed at frame %d\n", n_frames);
return 1;
}
if (!h_next) {
break; // stopped without generating a frame
}
n_frames++;
h_state = h_next;
sampled = sample_semantic_code();
timings.report(n_frames + 1);
timings.report(n_frames);
}
const double t_gen_s = (ggml_time_us() - t_gen_start_us) / 1e6;