From 6e62ba538478202094edc6c100c782719e310aa3 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Tue, 11 Aug 2026 14:18:30 +0200 Subject: [PATCH] mtmd: support pocket-tts (#26871) * 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 --- conversion/__init__.py | 2 + conversion/base.py | 28 ++ conversion/pockettts.py | 378 ++++++++++++++++ gguf-py/gguf/constants.py | 112 +++++ gguf-py/gguf/gguf_writer.py | 3 + src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/llama-model.cpp | 3 + src/models/models.h | 13 + src/models/pockettts.cpp | 146 +++++++ tools/mtmd/CMakeLists.txt | 3 + tools/mtmd/README-dev.md | 15 +- tools/mtmd/clip-impl.h | 40 +- tools/mtmd/clip-model.h | 89 ++++ tools/mtmd/clip.cpp | 303 +++++++++++-- tools/mtmd/clip.h | 5 + tools/mtmd/models/models.h | 56 +++ tools/mtmd/models/pockettts-gen.cpp | 291 ++++++++++++ tools/mtmd/models/pockettts-seanet.cpp | 162 +++++++ tools/mtmd/models/pockettts-spkenc.cpp | 77 ++++ tools/mtmd/models/qwen3tts-gen.cpp | 4 + tools/mtmd/mtmd-audio.cpp | 38 ++ tools/mtmd/mtmd-audio.h | 7 + tools/mtmd/mtmd-helper-gen.cpp | 583 ++++++++++++++++++++++++- tools/mtmd/mtmd-helper.h | 16 +- tools/mtmd/mtmd.cpp | 93 +++- tools/mtmd/mtmd.h | 22 +- tools/tts/README.md | 25 ++ tools/tts/tts.cpp | 15 +- 29 files changed, 2458 insertions(+), 73 deletions(-) create mode 100644 conversion/pockettts.py create mode 100644 src/models/pockettts.cpp create mode 100644 tools/mtmd/models/pockettts-gen.cpp create mode 100644 tools/mtmd/models/pockettts-seanet.cpp create mode 100644 tools/mtmd/models/pockettts-spkenc.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index c7d8046c49..695289b73a 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -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", diff --git a/conversion/base.py b/conversion/base.py index a7cd3fd904..3572b77c21 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -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) diff --git a/conversion/pockettts.py b/conversion/pockettts.py new file mode 100644 index 0000000000..62ecb5acde --- /dev/null +++ b/conversion/pockettts.py @@ -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 diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 8d0cad59b2..a197cacd3f 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -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" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 81ae07c11a..05f86396dc 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -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) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 73fb8b9813..8ed9391d7c 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -147,6 +147,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_MELLUM, "mellum" }, { LLM_ARCH_NANBEIGE, "nanbeige" }, { LLM_ARCH_QWEN3TTS, "qwen3tts" }, + { LLM_ARCH_POCKETTTS, "pockettts" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; diff --git a/src/llama-arch.h b/src/llama-arch.h index 51dfd288ce..18d9de186f 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -152,6 +152,7 @@ enum llm_arch { LLM_ARCH_DFLASH, LLM_ARCH_NANBEIGE, LLM_ARCH_QWEN3TTS, + LLM_ARCH_POCKETTTS, LLM_ARCH_UNKNOWN, }; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index cc1917b6da..0e27cb4171 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -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 diff --git a/src/models/models.h b/src/models/models.h index 9230345208..ddb9ae2f12 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -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 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; diff --git a/src/models/pockettts.cpp b/src/models/pockettts.cpp new file mode 100644 index 0000000000..1b3bb6c648 --- /dev/null +++ b/src/models/pockettts.cpp @@ -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 llama_model_pockettts::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*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); +} diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index fe22cb1254..3f4a6c670d 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -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 diff --git a/tools/mtmd/README-dev.md b/tools/mtmd/README-dev.md index 3cddd085ec..ac43e1b81b 100644 --- a/tools/mtmd/README-dev.md +++ b/tools/mtmd/README-dev.md @@ -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**. diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index bf73222f9d..b2c8b40213 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -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_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"}, }; diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 761aabf64e..ad25c008e7 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -141,6 +141,20 @@ struct clip_hparams { int32_t rvq_num_quantizers = 0; std::vector 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 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 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 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; + std::vector 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 gen_tfm_layers; // mimi decoder_transformer + // cogvlm ggml_tensor * mm_post_fc_norm_w = nullptr; ggml_tensor * mm_post_fc_norm_b = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 1e53eddf81..2fb2b5041d 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -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_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + builder = std::make_unique(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(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 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, ¶ms); } +// persisted state slots of the gen-audio decoder, per pipeline +static std::vector 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 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 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 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 dist(0.0f, std::sqrt(temp)); + std::vector 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 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 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 r = { dist(rng) }; + std::vector 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: diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index 7f706d976e..a5b7137752 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -104,9 +104,14 @@ struct clip_encode_params { int32_t top_k = 50; float top_p = 1.0f; std::vector * out_codes = nullptr; // this frame's 16 sampled codes + std::vector * 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 * codes = nullptr; // this frame's 16 RVQ codes + const std::vector * feats = nullptr; // continuous counterpart of codes std::vector * out_audio = nullptr; // decoded PCM samples, F32 const std::vector * state_in = nullptr; // state from previous call, null or wrong size means cold start std::vector * state_out = nullptr; // state for the next call diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 519c0d019a..ed8c1ea518 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -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 state_in; + mutable std::vector> 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 list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model); +// same, for the streaming mimi decoder (pocket-tts GEN_WAV) +std::vector 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; diff --git a/tools/mtmd/models/pockettts-gen.cpp b/tools/mtmd/models/pockettts-gen.cpp new file mode 100644 index 0000000000..3fd613e5f7 --- /dev/null +++ b/tools/mtmd/models/pockettts-gen.cpp @@ -0,0 +1,291 @@ +#include "models.h" + +#include + +// 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 list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model) { + std::vector 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; +} diff --git a/tools/mtmd/models/pockettts-seanet.cpp b/tools/mtmd/models/pockettts-seanet.cpp new file mode 100644 index 0000000000..c47207f569 --- /dev/null +++ b/tools/mtmd/models/pockettts-seanet.cpp @@ -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; +} diff --git a/tools/mtmd/models/pockettts-spkenc.cpp b/tools/mtmd/models/pockettts-spkenc.cpp new file mode 100644 index 0000000000..f802d90687 --- /dev/null +++ b/tools/mtmd/models/pockettts-spkenc.cpp @@ -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; +} diff --git a/tools/mtmd/models/qwen3tts-gen.cpp b/tools/mtmd/models/qwen3tts-gen.cpp index b6c95efa94..84c77f4fad 100644 --- a/tools/mtmd/models/qwen3tts-gen.cpp +++ b/tools/mtmd/models/qwen3tts-gen.cpp @@ -610,6 +610,10 @@ std::vector list_c2w_state_slots(const clip_hparams & hparams, c const auto & c2w = model.c2w; std::vector 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) diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index ca4b64efa4..98a8c11ee9 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1423,3 +1423,41 @@ std::vector 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 & 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; +} diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index b4d6f72598..44ad098ae6 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -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 & output) override; +}; + struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor { mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { } void initialize() override; diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 85671d1a33..1c58d3ae19 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -5,6 +5,8 @@ #include "../src/llama-ext.h" #include +#include +#include #include #include #include @@ -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(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 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 codes_buf; std::vector c2w_state; std::vector audio_pcm; @@ -438,10 +458,547 @@ private: std::vector 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 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 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 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(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(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 & 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 punct_ids(const char * s) const { + std::vector 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(ids.begin() + 1, ids.begin() + n); + } + + // cut after runs of boundary tokens, so punctuation stays with the sentence it ends + static std::vector> split_on(const std::vector & ids, + const std::vector & boundary) { + std::vector> 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> split_chunks(const std::vector & ids) const { + if ((int) ids.size() <= max_chunk_tokens) { + return { ids }; + } + const std::vector eos_punct = punct_ids(".!...?"); + const std::vector mid_punct = punct_ids(",;:"); + + // oversized sentences are split again on weaker punctuation, else words get skipped + std::vector> 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> 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 & 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 tok_embd; + + llama_seq_id seq_id = 0; + int pos = 0; + std::vector prompt_embd_buf; + std::unique_ptr 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> 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 feats_buf; + std::vector dec_state; + std::vector audio_pcm; + std::vector h_state_buf; + mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; + std::vector out_buf; +}; + static std::unique_ptr 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(new qwen3tts_gen_audio_pipeline(lctx, mctx)); + case MTMD_GEN_AUDIO_TYPE_POCKETTTS: + return std::unique_ptr(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, diff --git a/tools/mtmd/mtmd-helper.h b/tools/mtmd/mtmd-helper.h index 7e5cf9b509..832f7171ac 100644 --- a/tools/mtmd/mtmd-helper.h +++ b/tools/mtmd/mtmd-helper.h @@ -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); diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 82b73d5cd1..4b9c45d626 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -477,6 +477,7 @@ struct mtmd_context { // generation context struct clip_ctx * ctx_gen_a; // audio std::vector gen_out_codes; // this frame's 16 sampled codes (GEN_CODE) + std::vector gen_out_feats; // this frame's continuous features, if any (GEN_CODE) std::vector gen_out_embd; // next-step hidden state fed back to backbone (GEN_CODE) std::vector gen_out_audio; // decoded PCM samples for the current frame (GEN_WAV) std::vector gen_out_state; // state to feed into the next GEN_WAV call @@ -979,6 +980,10 @@ struct mtmd_context { { audio_preproc = std::make_unique(ctx_a); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + audio_preproc = std::make_unique(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 out_embd(n_embd); std::vector out_codes; + std::vector 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, ¶ms)) { 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 in_codes(inp->codes, inp->codes + inp->n_codes); + std::vector in_codes; + std::vector 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 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; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index e5063d114c..c1a5921db2 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -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, diff --git a/tools/tts/README.md b/tools/tts/README.md index dd84336c39..1b08d5ef32 100644 --- a/tools/tts/README.md +++ b/tools/tts/README.md @@ -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/` 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 +``` diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index 49c405e147..fd7522f8de 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -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;