"Qwen3MoeForCausalLM": "qwen",
"Qwen3NextForCausalLM": "qwen",
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
+ "PocketTTSModel": "pockettts",
"Qwen3TTSForConditionalGeneration": "qwen3tts",
"Qwen3VLForConditionalGeneration": "qwen3vl",
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
"Qwen2_5_VLForConditionalGeneration": "qwenvl",
"Qwen3ASRForConditionalGeneration": "qwen3vl",
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
+ "PocketTTSModel": "pockettts",
"Qwen3TTSForConditionalGeneration": "qwen3tts",
"Qwen3VLForConditionalGeneration": "qwen3vl",
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
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
ModelType.TEXT: {},
ModelType.MMPROJ: {},
}
+ _hparams_loaders: list[tuple[HparamsMatcher, HparamsLoader]] = []
dir_model: Path
ftype: gguf.LlamaFileType
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:
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)
--- /dev/null
+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
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"
MELLUM = auto()
NANBEIGE = auto()
QWEN3TTS = auto()
+ POCKETTTS = auto()
class VISION_PROJECTOR_TYPE(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()
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] = {
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",
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,
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
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"
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)
{ LLM_ARCH_MELLUM, "mellum" },
{ LLM_ARCH_NANBEIGE, "nanbeige" },
{ LLM_ARCH_QWEN3TTS, "qwen3tts" },
+ { LLM_ARCH_POCKETTTS, "pockettts" },
{ LLM_ARCH_UNKNOWN, "(unknown)" },
};
LLM_ARCH_DFLASH,
LLM_ARCH_NANBEIGE,
LLM_ARCH_QWEN3TTS,
+ LLM_ARCH_POCKETTTS,
LLM_ARCH_UNKNOWN,
};
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:
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
};
+struct llama_model_pockettts : public llama_model_base {
+ llama_model_pockettts(const struct llama_model_params & params) : llama_model_base(params) {}
+ void load_arch_hparams(llama_model_loader & ml) override;
+ void load_arch_tensors(llama_model_loader & ml) override;
+
+ struct graph : public llm_graph_context {
+ graph(const llama_model & model, const llm_graph_params & params);
+ };
+
+ std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
+
+
struct llama_model_codeshell : public llama_model_base {
llama_model_codeshell(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
--- /dev/null
+#include "models.h"
+
+// backbone of the pocket-tts CALM pipeline: the "text" side of a flow language model.
+// it has no lm_head, the audio latents are produced by the flow net inside the mmproj
+
+void llama_model_pockettts::load_arch_hparams(llama_model_loader & ml) {
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
+
+ switch (hparams.n_layer()) {
+ case 6: type = LLM_TYPE_109M; break;
+ case 24: type = LLM_TYPE_335M; break;
+ default: type = LLM_TYPE_UNKNOWN;
+ }
+}
+
+void llama_model_pockettts::load_arch_tensors(llama_model_loader &) {
+ LLAMA_LOAD_LOCALS;
+
+ tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+ output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+ output_norm_b = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "bias"), {n_embd}, 0);
+ // no output head, the logits are unused; reuse the embedding table so a sampler can still run
+ output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
+
+ for (int i = 0; i < n_layer; ++i) {
+ auto & layer = layers[i];
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+ layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias", i), {n_embd}, 0);
+
+ create_tensor_qkv(layer, i, n_embd, n_embd, n_embd_gqa, n_embd_gqa, TENSOR_NOT_REQUIRED);
+
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0);
+
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+ layer.ffn_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", i), {n_embd}, 0);
+
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
+ }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_pockettts::build_arch_graph(const llm_graph_params & params) const {
+ return std::make_unique<graph>(*this, params);
+}
+
+llama_model_pockettts::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
+ const int64_t n_embd_head = hparams.n_embd_head_v();
+
+ GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+ GGML_ASSERT(n_embd_head == n_rot);
+
+ ggml_tensor * cur;
+ ggml_tensor * inpL;
+
+ inpL = build_inp_embd(model.tok_embd);
+
+ ggml_tensor * inp_pos = build_inp_pos();
+
+ auto * inp_attn = build_attn_inp_kv();
+
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ for (int il = 0; il < n_layer; ++il) {
+ cur = build_norm(inpL,
+ model.layers[il].attn_norm,
+ model.layers[il].attn_norm_b,
+ LLM_NORM, il);
+ cb(cur, "attn_norm", il);
+
+ // self-attention
+ {
+ auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
+ n_embd_head, n_head, n_head_kv, il);
+
+ Qcur = ggml_rope_ext(
+ ctx0, Qcur, inp_pos, nullptr,
+ n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow
+ );
+
+ Kcur = ggml_rope_ext(
+ ctx0, Kcur, inp_pos, nullptr,
+ n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow
+ );
+
+ cb(Qcur, "Qcur", il);
+ cb(Kcur, "Kcur", il);
+ cb(Vcur, "Vcur", il);
+
+ cur = build_attn(inp_attn,
+ model.layers[il].wo, NULL, model.layers[il].wo_s,
+ Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
+ }
+
+ if (il == n_layer - 1 && inp_out_ids) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
+ }
+
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
+ cb(ffn_inp, "ffn_inp", il);
+
+ // FF
+ {
+ cur = build_norm(ffn_inp,
+ model.layers[il].ffn_norm,
+ model.layers[il].ffn_norm_b,
+ LLM_NORM, il);
+ cb(cur, "ffn_norm", il);
+
+ cur = build_ffn(cur,
+ model.layers[il].ffn_up, NULL, NULL,
+ NULL, NULL, NULL,
+ model.layers[il].ffn_down, NULL, NULL,
+ NULL,
+ LLM_FFN_GELU, LLM_FFN_SEQ, il);
+ cb(cur, "ffn_out", il);
+ }
+
+ cur = ggml_add(ctx0, cur, ffn_inp);
+
+ cur = build_cvec(cur, il);
+ cb(cur, "l_out", il);
+
+ // input for next layer
+ inpL = cur;
+ }
+
+ cur = build_norm(inpL,
+ model.output_norm,
+ model.output_norm_b,
+ LLM_NORM, -1);
+
+ cb(cur, "result_norm", -1);
+ res->t_embd = cur;
+
+ cur = build_lora_mm(model.output, cur, model.output_s);
+
+ cb(cur, "result_output", -1);
+ res->t_logits = cur;
+
+ ggml_build_forward_expand(gf, cur);
+}
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
### 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`)
- 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**.
#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
#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"
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,
};
{ 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"},
};
int32_t rvq_num_quantizers = 0;
std::vector<int32_t> rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17)
+ // threshold for the "out_eos_score" graph output
+ float gen_eos_threshold = 0.0f;
+
+ // name of the weight variant, some pipelines tune themselves on it
+ std::string gen_model_variant;
+
+ // pocket-tts
+ static constexpr int32_t pockettts_max_spk_seconds = 30;
+ int32_t seanet_n_stage = 0;
+ std::vector<int32_t> seanet_ratios; // encoder order (reversed compared to the config)
+ int32_t mimi_downsample = 0; // encoder frame rate / model frame rate
+ int32_t mimi_tfm_context = 0; // attention window of the mimi transformers, in frames
+ int32_t flow_n_step = 1; // lsd_decode steps
+
// qwen3tts code2wav
int32_t wav_tfm_n_layer = 0;
int32_t wav_tfm_n_embd = 0;
std::vector<clip_layer> qf_proj_layers;
};
+// pocket-tts SEANet stack, used in both directions:
+// encoder = conv_in -> per stage (residual unit, strided conv) -> conv_out
+// decoder = conv_in -> per stage (strided convtr, residual unit) -> conv_out
+struct clip_seanet {
+ // one residual unit: ELU -> dilated conv -> ELU -> pointwise conv, added to the input
+ struct stage {
+ ggml_tensor * res_conv1_w = nullptr;
+ ggml_tensor * res_conv1_b = nullptr;
+ ggml_tensor * res_conv2_w = nullptr;
+ ggml_tensor * res_conv2_b = nullptr;
+ ggml_tensor * scale_conv_w = nullptr; // strided conv (encoder) or convtr (decoder)
+ ggml_tensor * scale_conv_b = nullptr;
+ };
+
+ ggml_tensor * conv_in_w = nullptr;
+ ggml_tensor * conv_in_b = nullptr;
+ ggml_tensor * conv_out_w = nullptr;
+ ggml_tensor * conv_out_b = nullptr;
+ std::vector<stage> stages;
+};
+
+// pocket-tts flow-matching decoder (SimpleMLPAdaLN)
+struct clip_flow_net {
+ // AdaLN res block: in_ln -> modulate -> Linear -> SiLU -> Linear, gated residual
+ struct block {
+ ggml_tensor * norm_w = nullptr;
+ ggml_tensor * norm_b = nullptr;
+ ggml_tensor * up_w = nullptr;
+ ggml_tensor * up_b = nullptr;
+ ggml_tensor * down_w = nullptr;
+ ggml_tensor * down_b = nullptr;
+ ggml_tensor * ada_w = nullptr; // -> shift, scale, gate
+ ggml_tensor * ada_b = nullptr;
+ };
+
+ // timestep embedder: cos/sin(t * freqs) -> Linear -> SiLU -> Linear -> RMSNorm
+ struct time_embd {
+ ggml_tensor * freqs = nullptr;
+ ggml_tensor * up_w = nullptr;
+ ggml_tensor * up_b = nullptr;
+ ggml_tensor * down_w = nullptr;
+ ggml_tensor * down_b = nullptr;
+ ggml_tensor * norm = nullptr; // RMSNorm alpha
+ };
+
+ ggml_tensor * input_proj_w = nullptr;
+ ggml_tensor * input_proj_b = nullptr;
+ ggml_tensor * cond_embd_w = nullptr;
+ ggml_tensor * cond_embd_b = nullptr;
+ ggml_tensor * final_ada_w = nullptr; // -> shift, scale
+ ggml_tensor * final_ada_b = nullptr;
+ ggml_tensor * final_proj_w = nullptr;
+ ggml_tensor * final_proj_b = nullptr;
+ std::vector<time_embd> time;
+ std::vector<block> blocks;
+};
+
// qwen3tts code2wav: RVQ codes -> raw PCM
struct clip_code2wav {
// "upsample" stage: one ConvNeXt block plus the causal ConvTranspose1d before it
// qwen3tts code2wav: RVQ codes -> raw PCM
clip_code2wav c2w;
+ // pocket-tts: SEANet stack, shared by the encoder (speaker path) and the decoder (gen path)
+ clip_seanet seanet;
+
+ // pocket-tts: voice latent -> backbone embd (speaker path)
+ ggml_tensor * spk_proj_w = nullptr;
+ ggml_tensor * downsample_w = nullptr;
+
+ // pocket-tts: flow-matching decoder, backbone hidden state -> next latent
+ clip_flow_net flow;
+ ggml_tensor * gen_out_eos_w = nullptr;
+ ggml_tensor * gen_out_eos_b = nullptr;
+ ggml_tensor * gen_input_lin_w = nullptr; // latent -> backbone embd
+ ggml_tensor * gen_emb_mean = nullptr;
+ ggml_tensor * gen_emb_std = nullptr;
+ ggml_tensor * gen_quant_out_w = nullptr; // latent -> decoder dim
+ ggml_tensor * gen_upsample_w = nullptr; // depthwise convtr, frame rate -> encoder frame rate
+ std::vector<clip_layer> gen_tfm_layers; // mimi decoder_transformer
+
// cogvlm
ggml_tensor * mm_post_fc_norm_w = nullptr;
ggml_tensor * mm_post_fc_norm_b = nullptr;
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;
{
builder = std::make_unique<clip_graph_qwen3tts_spkenc>(ctx, img);
} break;
+ case PROJECTOR_TYPE_POCKETTTS_SPKENC:
+ {
+ builder = std::make_unique<clip_graph_pockettts_spkenc>(ctx, img);
+ } break;
+ case PROJECTOR_TYPE_POCKETTTS_GEN:
+ {
+ const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
+ const int n_step = ctx->model.hparams.flow_n_step;
+ const int64_t n_latent = ctx->model.gen_input_lin_w->ne[0];
+ GGML_ASSERT(n_step > 0);
+ GGML_ASSERT(n_latent > 0);
+ // "inp_feats" takes the caller's buffer as-is, the graph must consume all of it
+ if (params && params->feats) {
+ GGML_ASSERT(params->feats->size() % (size_t) n_latent == 0);
+ GGML_ASSERT(params->feats->size() >= (size_t) n_latent);
+ }
+ const int n_frames = params && params->feats ? (int) (params->feats->size() / n_latent) : 1;
+ builder = std::make_unique<clip_graph_pockettts_gen>(ctx, img, gen_process, n_step, n_frames);
+ } break;
case PROJECTOR_TYPE_QWEN3TTS_GEN:
{
const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
// 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");
} 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);
// 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;
// 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)) {
return cur;
};
+ // pocket-tts: the encoder and the decoder share the same layout, only the prefix differs
+ auto load_seanet = [&](clip_seanet & seanet, bool is_decoder) {
+ const char * conv_in = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_IN : TN_A_SEANET_CONV_IN;
+ const char * conv_out = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_OUT : TN_A_SEANET_CONV_OUT;
+ const char * res1 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV1 : TN_A_SEANET_RES_CONV1;
+ const char * res2 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV2 : TN_A_SEANET_RES_CONV2;
+ const char * scale = is_decoder ? TN_A_GEN_WAV_SEANET_SCALE_CONV : TN_A_SEANET_SCALE_CONV;
+
+ seanet.conv_in_w = get_tensor(string_format(conv_in, "weight"));
+ seanet.conv_in_b = get_tensor(string_format(conv_in, "bias"));
+ seanet.conv_out_w = get_tensor(string_format(conv_out, "weight"));
+ seanet.conv_out_b = get_tensor(string_format(conv_out, "bias"));
+
+ seanet.stages.resize(hparams.seanet_n_stage);
+ for (int i = 0; i < hparams.seanet_n_stage; i++) {
+ auto & stage = seanet.stages[i];
+ stage.res_conv1_w = get_tensor(string_format(res1, i, "weight"));
+ stage.res_conv1_b = get_tensor(string_format(res1, i, "bias"));
+ stage.res_conv2_w = get_tensor(string_format(res2, i, "weight"));
+ stage.res_conv2_b = get_tensor(string_format(res2, i, "bias"));
+ stage.scale_conv_w = get_tensor(string_format(scale, i, "weight"));
+ stage.scale_conv_b = get_tensor(string_format(scale, i, "bias"));
+ }
+ };
+
auto get_vector = [&](const std::string & name) {
std::vector<float> result;
auto it = tensor_offset.find(name);
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;
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
// 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
return clip_encode(ctx, ¶ms);
}
+// persisted state slots of the gen-audio decoder, per pipeline
+static std::vector<c2w_state_slot> list_gen_state_slots(const clip_hparams & hparams, const clip_model & model) {
+ switch (model.proj_type) {
+ case PROJECTOR_TYPE_QWEN3TTS_GEN: return list_c2w_state_slots(hparams, model);
+ case PROJECTOR_TYPE_POCKETTTS_GEN: return list_pockettts_state_slots(hparams, model);
+ default: return {};
+ }
+}
+
bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
const clip_image_f32_batch & imgs = *params->imgs;
int n_batch_cur = imgs.entries.size();
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();
ggml_backend_tensor_set(cur, values.data(), 0, ggml_nbytes(cur));
};
+ // upload the decoder state from the previous call, or zero-fill on a cold start
+ auto set_gen_state_in = [&]() {
+ size_t offset = 0;
+ for (const auto & slot : list_gen_state_slots(hparams, model)) {
+ ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str());
+ const size_t nb = ggml_nbytes(t);
+ if (params->state_in && params->state_in->size() >= offset + nb) {
+ ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb);
+ } else {
+ std::vector<uint8_t> zeros(nb, 0);
+ ggml_backend_tensor_set(t, zeros.data(), 0, nb);
+ }
+ offset += nb;
+ }
+ };
+
+ // rope positions and attention mask of the mimi transformers (pocket-tts).
+ // the mask is causal with a sliding window, see _build_attention_mask() in the reference
+ auto set_pockettts_tfm_inputs = [&]() {
+ const int64_t n_pos = ggml_nelements(get_inp_tensor("inp_pos"));
+ GGML_ASSERT(n_pos > 0);
+ std::vector<int32_t> positions((size_t) n_pos);
+ for (int64_t i = 0; i < n_pos; i++) {
+ positions[(size_t) i] = (int32_t) i;
+ }
+ set_input_i32("inp_pos", positions);
+
+ // the preprocessor truncates the waveform to keep this mask bounded
+ const int64_t max_pos = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate / 120;
+ GGML_ASSERT(n_pos <= max_pos && "pocket-tts speaker reference too long for a dense mask");
+
+ const int64_t context = hparams.mimi_tfm_context;
+ std::vector<float> mask((size_t) n_pos * n_pos, -INFINITY);
+ for (int64_t q = 0; q < n_pos; q++) {
+ for (int64_t k = 0; k < n_pos; k++) {
+ const int64_t delta = q - k;
+ if (delta >= 0 && delta < context) {
+ mask[(size_t) q * n_pos + k] = 0.0f;
+ }
+ }
+ }
+ set_input_f32("kq_mask", mask);
+ };
+
// set input pixel values
if (!imgs.is_audio) {
size_t nelem = 0;
}
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];
}
set_input_i32("patches", patches);
} break;
+ case PROJECTOR_TYPE_POCKETTTS_SPKENC:
+ {
+ set_pockettts_tfm_inputs();
+ } break;
+ case PROJECTOR_TYPE_POCKETTTS_GEN:
+ {
+ if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) {
+ GGML_ASSERT(params->feats != nullptr);
+ set_input_f32("inp_feats", *params->feats);
+ // positions and mask are derived in-graph from the persisted counter
+ set_gen_state_in();
+ } else {
+ // flow matching starts from gaussian noise, std = sqrt(temp)
+ ggml_tensor * t = get_inp_tensor("inp_noise");
+ // Config.default_temperature, for a caller that does not set one
+ const float temp = params->temp > 0.0f ? params->temp : 0.7f;
+ std::normal_distribution<float> dist(0.0f, std::sqrt(temp));
+ std::vector<float> noise(ggml_nelements(t));
+ for (auto & v : noise) {
+ v = dist(ctx->rng);
+ }
+ set_input_f32("inp_noise", noise);
+ }
+ } break;
case PROJECTOR_TYPE_GEMMA4V:
case PROJECTOR_TYPE_GEMMA4UV:
{
}
}
set_input_i32("inp_codes", codes);
-
- // upload the state from the previous call, or zero-fill on a cold start
- size_t offset = 0;
- for (const auto & slot : list_c2w_state_slots(hparams, model)) {
- ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str());
- const size_t nb = ggml_nbytes(t);
- if (params->state_in && params->state_in->size() >= offset + nb) {
- ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb);
- } else {
- std::vector<uint8_t> zeros(nb, 0);
- ggml_backend_tensor_set(t, zeros.data(), 0, nb);
- }
- offset += nb;
- }
+ set_gen_state_in();
} else {
// code0 indexes gen_code_out_embd_w via ggml_get_rows; bound it
const int64_t vocab0 = model.gen_code_out_embd_w->ne[1];
set_input_i32("inp_code0", code0);
// one uniform(0,1) draw per codebook, used by do_sampling()
- static std::mt19937 rng{ std::random_device{}() };
std::uniform_real_distribution<float> dist(0.0f, 1.0f);
const int64_t n_acoustic = model.gen_code_head_w->ne[2];
for (int64_t g = 0; g < n_acoustic; g++) {
- std::vector<float> r = { dist(rng) };
+ std::vector<float> r = { dist(ctx->rng) };
set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r);
}
}
// 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");
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);
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());
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:
int32_t top_k = 50;
float top_p = 1.0f;
std::vector<int32_t> * out_codes = nullptr; // this frame's 16 sampled codes
+ std::vector<float> * out_feats = nullptr; // continuous counterpart of out_codes
+ uint32_t seed = UINT32_MAX; // UINT32_MAX for random
+ float temp = 0.0f; // sampling temperature, noise scale for flow-matching decoders
+ bool * out_is_eos = nullptr;
// GEN_WAV
const std::vector<int32_t> * codes = nullptr; // this frame's 16 RVQ codes
+ const std::vector<float> * feats = nullptr; // continuous counterpart of codes
std::vector<float> * out_audio = nullptr; // decoded PCM samples, F32
const std::vector<uint8_t> * state_in = nullptr; // state from previous call, null or wrong size means cold start
std::vector<uint8_t> * state_out = nullptr; // state for the next call
};
};
+//
+// pocket-tts: SEANet convolution stack, shared by the voice encoder and the mimi decoder.
+// stateless unless state_in is populated: convs then pad instead of carrying left-context.
+//
+struct clip_graph_pockettts_seanet : clip_graph {
+ clip_graph_pockettts_seanet(const clip_graph & parent) : clip_graph(parent) {}
+ ggml_cgraph * build() override { GGML_ABORT("call encode()/decode() instead"); }
+
+ // per-call streaming state, keyed by slot name (see list_pockettts_state_slots)
+ std::map<std::string, ggml_tensor *> state_in;
+ mutable std::vector<std::pair<std::string, ggml_tensor *>> state_out;
+
+ ggml_tensor * conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation,
+ bool pad_replicate = false, const std::string & state_name = "") const;
+ ggml_tensor * conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride,
+ const std::string & state_name = "") const;
+ ggml_tensor * res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation,
+ const std::string & state_prefix = "") const;
+
+ // x: [T, C] -> [T / hop, dim]
+ ggml_tensor * encode(ggml_tensor * x) const;
+ // x: [T, dim] -> [T * hop, 1], streams when state_in is populated
+ ggml_tensor * decode(ggml_tensor * x) const;
+};
+
+// mimi encoder + speaker_proj: reference waveform -> voice conditioning rows
+struct clip_graph_pockettts_spkenc : clip_graph {
+ clip_graph_pockettts_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
+ ggml_cgraph * build() override;
+
+ ggml_tensor * tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const;
+};
+
+//
+// pocket-tts generation:
+// GEN_CODE = flow-matching decoder + end-of-speech head, one latent per call
+// GEN_WAV = mimi decoder, a window of latents -> PCM
+//
+struct clip_graph_pockettts_gen : clip_graph {
+ clip_graph_pockettts_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int n_step, int n_frames)
+ : clip_graph(ctx, img), gen_process(gen_process), n_step(n_step), n_frames(n_frames) {}
+ ggml_cgraph * build() override;
+
+ clip_gen_process_type gen_process;
+ int n_step; // lsd_decode steps, fixed at graph-build time
+ int n_frames; // GEN_WAV only: number of latents to decode
+
+ // AdaLN modulation: x * (1 + scale) + shift
+ ggml_tensor * modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const;
+ ggml_tensor * time_embed(const clip_flow_net::time_embd & te, float t) const;
+ ggml_tensor * flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const;
+};
+
// one persisted state buffer used by code2wav, see qwen3tts-gen.cpp
struct c2w_state_slot {
std::string name;
};
std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model);
+// same, for the streaming mimi decoder (pocket-tts GEN_WAV)
+std::vector<c2w_state_slot> list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model);
+
struct clip_graph_kimik25 : clip_graph {
clip_graph_kimik25(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
ggml_cgraph * build() override;
--- /dev/null
+#include "models.h"
+
+#include <cmath>
+
+// pocket-tts generation stages
+//
+// GEN_CODE: backbone hidden state -> next 32-d latent (flow matching) + end-of-speech score
+// GEN_WAV : a window of latents -> PCM, through the mimi decoder
+//
+// there is no codebook anywhere, "codes" in the mtmd API are continuous features here
+
+ggml_tensor * clip_graph_pockettts_gen::modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const {
+ ggml_tensor * cur = ggml_mul(ctx0, x, ggml_scale_bias(ctx0, scale, 1.0f, 1.0f));
+ return ggml_add(ctx0, cur, shift);
+}
+
+// see TimestepEmbedder in the reference
+ggml_tensor * clip_graph_pockettts_gen::time_embed(const clip_flow_net::time_embd & te, float t) const {
+ // t is a graph-build constant, so the cos/sin table can be folded into a scaled copy
+ ggml_tensor * args = ggml_scale(ctx0, te.freqs, t);
+ ggml_tensor * emb = ggml_concat(ctx0, ggml_cos(ctx0, args), ggml_sin(ctx0, args), 0);
+
+ ggml_tensor * cur = build_mm(te.up_w, emb);
+ cur = ggml_add(ctx0, cur, te.up_b);
+ cur = ggml_silu(ctx0, cur);
+ cur = build_mm(te.down_w, cur);
+ cur = ggml_add(ctx0, cur, te.down_b);
+
+ // this "RMSNorm" divides by the unbiased variance, not the mean square
+ // it also rescales the input, not the centered value, see _rms_norm() in mlp.py
+ {
+ const int64_t n = cur->ne[0];
+ ggml_tensor * mean = ggml_mean(ctx0, cur);
+ ggml_tensor * dev = ggml_sub(ctx0, cur, mean);
+ ggml_tensor * var = ggml_mean(ctx0, ggml_sqr(ctx0, dev));
+ var = ggml_scale_bias(ctx0, var, (float) n / (float) (n - 1), 1e-5f);
+ cur = ggml_div(ctx0, cur, ggml_sqrt(ctx0, var));
+ cur = ggml_mul(ctx0, cur, te.norm);
+ }
+
+ return cur;
+}
+
+// one velocity evaluation: v(cond, s, t, x)
+ggml_tensor * clip_graph_pockettts_gen::flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const {
+ const auto & flow = model.flow;
+
+ ggml_tensor * cur = build_mm(flow.input_proj_w, x);
+ cur = ggml_add(ctx0, cur, flow.input_proj_b);
+
+ // the two time conditions are averaged, then added to the projected backbone state
+ ggml_tensor * ts = ggml_add(ctx0, time_embed(flow.time[0], s), time_embed(flow.time[1], t));
+ ts = ggml_scale(ctx0, ts, 1.0f / (float) flow.time.size());
+
+ ggml_tensor * c = build_mm(flow.cond_embd_w, cond);
+ c = ggml_add(ctx0, c, flow.cond_embd_b);
+
+ ggml_tensor * y = ggml_add(ctx0, ts, c);
+ cb(y, "flow_cond", -1);
+
+ const int64_t n_ch = flow.blocks.empty() ? 0 : flow.blocks[0].norm_w->ne[0];
+
+ for (size_t il = 0; il < flow.blocks.size(); il++) {
+ const auto & blk = flow.blocks[il];
+
+ ggml_tensor * mod = build_mm(blk.ada_w, ggml_silu(ctx0, y));
+ mod = ggml_add(ctx0, mod, blk.ada_b);
+
+ ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0);
+ ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]);
+ ggml_tensor * gate = ggml_view_1d(ctx0, mod, n_ch, (size_t) 2 * n_ch * mod->nb[0]);
+
+ ggml_tensor * h = build_norm(cur, blk.norm_w, blk.norm_b, NORM_TYPE_NORMAL, 1e-6f, (int) il);
+ h = modulate(h, shift, scale);
+ h = build_mm(blk.up_w, h);
+ h = ggml_add(ctx0, h, blk.up_b);
+ h = ggml_silu(ctx0, h);
+ h = build_mm(blk.down_w, h);
+ h = ggml_add(ctx0, h, blk.down_b);
+
+ cur = ggml_add(ctx0, cur, ggml_mul(ctx0, gate, h));
+ cb(cur, "flow_blk", (int) il);
+ }
+
+ // final layer: the norm has no weights, only the AdaLN modulation
+ ggml_tensor * mod = build_mm(flow.final_ada_w, ggml_silu(ctx0, y));
+ mod = ggml_add(ctx0, mod, flow.final_ada_b);
+
+ ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0);
+ ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]);
+
+ cur = build_norm(cur, nullptr, nullptr, NORM_TYPE_NORMAL, 1e-6f, -1);
+ cur = modulate(cur, shift, scale);
+ cur = build_mm(flow.final_proj_w, cur);
+ cur = ggml_add(ctx0, cur, flow.final_proj_b);
+
+ return cur;
+}
+
+// state carried between GEN_WAV calls: rope offset, per-layer KV window, conv left context
+// and the transposed-conv overlap tails
+std::vector<c2w_state_slot> list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model) {
+ std::vector<c2w_state_slot> slots;
+ if (model.gen_upsample_w == nullptr) {
+ return slots; // not a pocket-tts decoder
+ }
+ const auto & seanet = model.seanet;
+
+ // the slots below are sized from these
+ GGML_ASSERT(!model.gen_tfm_layers.empty());
+ GGML_ASSERT((int) seanet.stages.size() >= hparams.seanet_n_stage);
+ GGML_ASSERT((int) hparams.seanet_ratios.size() >= hparams.seanet_n_stage);
+ GGML_ASSERT(hparams.mimi_tfm_context > 1 && hparams.mimi_downsample > 0);
+
+ slots.push_back({"tfm_pos", 1, 1});
+
+ const int64_t n_embd_a = model.gen_tfm_layers[0].q_w->ne[1];
+ const int64_t prefix = hparams.mimi_tfm_context - 1;
+ for (size_t il = 0; il < model.gen_tfm_layers.size(); il++) {
+ slots.push_back({"tfm_k_" + std::to_string(il), n_embd_a, prefix});
+ slots.push_back({"tfm_v_" + std::to_string(il), n_embd_a, prefix});
+ }
+
+ // upsample is depthwise, its output channel count is the input one
+ slots.push_back({"up", model.gen_upsample_w->ne[0] - hparams.mimi_downsample, model.gen_upsample_w->ne[2]});
+
+ slots.push_back({"dec_in", seanet.conv_in_w->ne[0] - 1, seanet.conv_in_w->ne[1]});
+ for (int i = 0; i < hparams.seanet_n_stage; i++) {
+ const auto & stage = seanet.stages[i];
+ const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i];
+ slots.push_back({"dec_up_" + std::to_string(i), stage.scale_conv_w->ne[0] - stride, stage.scale_conv_w->ne[1]});
+ slots.push_back({"dec_res_" + std::to_string(i), stage.res_conv1_w->ne[0] - 1, stage.res_conv1_w->ne[1]});
+ }
+ slots.push_back({"dec_out", seanet.conv_out_w->ne[0] - 1, seanet.conv_out_w->ne[1]});
+
+ return slots;
+}
+
+ggml_cgraph * clip_graph_pockettts_gen::build() {
+ if (gen_process == CLIP_GEN_PROCESS_GEN_CODE) {
+ // the backbone hidden state arrives as the single batch entry
+ ggml_tensor * h_state = build_inp_raw(1);
+ h_state = ggml_reshape_2d(ctx0, h_state, n_mmproj_embd, 1);
+
+ // end-of-speech probe, thresholded on the host side
+ ggml_tensor * eos = build_mm(model.gen_out_eos_w, h_state);
+ eos = ggml_add(ctx0, eos, model.gen_out_eos_b);
+ ggml_set_name(eos, "out_eos_score");
+ ggml_set_output(eos);
+ ggml_build_forward_expand(gf, eos);
+
+ const int64_t n_latent = model.gen_input_lin_w->ne[0];
+
+ ggml_tensor * noise = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_latent, 1);
+ ggml_set_name(noise, "inp_noise");
+ ggml_set_input(noise);
+
+ // lsd_decode: integrate the velocity field from the noise sample
+ ggml_tensor * cur = noise;
+ for (int i = 0; i < n_step; i++) {
+ const float s = (float) i / (float) n_step;
+ const float t = (float) (i + 1) / (float) n_step;
+ ggml_tensor * v = flow_forward(h_state, cur, s, t);
+ cur = ggml_add(ctx0, cur, ggml_scale(ctx0, v, 1.0f / (float) n_step));
+ }
+ cb(cur, "flow_latent", -1);
+
+ ggml_set_name(cur, "out_feats");
+ ggml_set_output(cur);
+ ggml_build_forward_expand(gf, cur);
+
+ // the same latent, projected into the backbone's input space for the next step
+ ggml_tensor * embd = build_mm(model.gen_input_lin_w, cur);
+ cb(embd, "gen_embd", -1);
+ ggml_build_forward_expand(gf, embd);
+
+ return gf;
+ }
+
+ // GEN_WAV: [32, n_frames] latents -> PCM
+ ggml_tensor * feats = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32,
+ model.gen_input_lin_w->ne[0], n_frames);
+ ggml_set_name(feats, "inp_feats");
+ ggml_set_input(feats);
+
+ // denormalize, then the DummyQuantizer up-projection
+ ggml_tensor * cur = ggml_add(ctx0, ggml_mul(ctx0, feats, model.gen_emb_std), model.gen_emb_mean);
+ cur = build_mm(model.gen_quant_out_w, cur);
+ cb(cur, "quant_out", -1);
+
+ clip_graph_pockettts_seanet seanet(*this);
+ for (const auto & slot : list_pockettts_state_slots(hparams, model)) {
+ ggml_tensor * t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1);
+ ggml_set_name(t, ("state_in_" + slot.name).c_str());
+ ggml_set_input(t);
+ seanet.state_in[slot.name] = t;
+ }
+
+ // model frame rate -> encoder frame rate, depthwise transposed conv
+ cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
+ cur = seanet.conv_transpose1d(cur, model.gen_upsample_w, nullptr, hparams.mimi_downsample, "up");
+ cb(cur, "mimi_upsample", -1);
+
+ cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
+
+ // positions continue across calls, the counter lives in the state
+ const int64_t n_pos = cur->ne[1];
+ const int64_t prefix = hparams.mimi_tfm_context - 1;
+ const int64_t n_kv = prefix + n_pos;
+
+ ggml_tensor * base = ggml_reshape_1d(ctx0, seanet.state_in.at("tfm_pos"), 1);
+ ggml_tensor * inp_pos = ggml_cast(ctx0, ggml_add(ctx0, ggml_arange(ctx0, 0.0f, (float) n_pos, 1.0f), base),
+ GGML_TYPE_I32);
+ seanet.state_out.push_back({"tfm_pos", ggml_scale_bias(ctx0, seanet.state_in.at("tfm_pos"), 1.0f, (float) n_pos)});
+
+ // banded causal mask over [cached prefix | this chunk]
+ // the last factor masks out cache rows that hold no real frame yet
+ ggml_tensor * pos_k = ggml_reshape_2d(ctx0, ggml_arange(ctx0, 0.0f, (float) n_kv, 1.0f), n_kv, 1);
+ ggml_tensor * pos_q = ggml_reshape_2d(ctx0, ggml_arange(ctx0, (float) prefix, (float) (prefix + n_pos), 1.0f), 1, n_pos);
+ ggml_tensor * diff = ggml_sub(ctx0, ggml_repeat_4d(ctx0, pos_q, n_kv, n_pos, 1, 1), pos_k);
+
+ ggml_tensor * keep = ggml_mul(ctx0,
+ ggml_step(ctx0, ggml_scale_bias(ctx0, diff, 1.0f, 0.5f)), // delta >= 0
+ ggml_step(ctx0, ggml_scale_bias(ctx0, diff, -1.0f, (float) hparams.mimi_tfm_context - 0.5f))); // delta < context
+ keep = ggml_mul(ctx0, keep,
+ ggml_step(ctx0, ggml_scale_bias(ctx0, ggml_add(ctx0, pos_k, base), 1.0f, 0.5f - (float) prefix)));
+ ggml_tensor * kq_mask = ggml_reshape_4d(ctx0, ggml_log(ctx0, keep), n_kv, n_pos, 1, 1);
+
+ for (int il = 0; il < n_layer; il++) {
+ const auto & layer = model.gen_tfm_layers[il];
+ ggml_tensor * inp = cur;
+
+ cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
+
+ ggml_tensor * Qcur = build_mm(layer.q_w, cur);
+ ggml_tensor * Kcur = build_mm(layer.k_w, cur);
+ ggml_tensor * Vcur = build_mm(layer.v_w, cur);
+
+ Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos);
+ Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos);
+
+ Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
+ hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+ Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
+ hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+
+ // prepend the cached window, then keep this chunk's tail for the next call
+ const std::string k_name = "tfm_k_" + std::to_string(il);
+ const std::string v_name = "tfm_v_" + std::to_string(il);
+ ggml_tensor * k_full = ggml_concat(ctx0, seanet.state_in.at(k_name),
+ ggml_reshape_2d(ctx0, Kcur, d_head * n_head, n_pos), 1);
+ ggml_tensor * v_full = ggml_concat(ctx0, seanet.state_in.at(v_name), Vcur, 1);
+ seanet.state_out.push_back({k_name, ggml_cont(ctx0, ggml_view_2d(ctx0, k_full, k_full->ne[0], prefix,
+ k_full->nb[1], (size_t) n_pos * k_full->nb[1]))});
+ seanet.state_out.push_back({v_name, ggml_cont(ctx0, ggml_view_2d(ctx0, v_full, v_full->ne[0], prefix,
+ v_full->nb[1], (size_t) n_pos * v_full->nb[1]))});
+
+ ggml_tensor * q_cur = ggml_reshape_4d(ctx0, Qcur, d_head, n_head, n_pos, 1);
+ ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k_full, d_head, n_head, n_kv, 1);
+ ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v_full, d_head, n_head, n_kv, 1);
+
+ cur = build_attn(layer.o_w, nullptr, q_cur, k_cur, v_cur, kq_mask, kq_scale, il);
+ cur = ggml_mul(ctx0, cur, layer.ls_1_w);
+ cur = ggml_add(ctx0, cur, inp);
+
+ inp = cur;
+ cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
+ cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il);
+ cur = ggml_mul(ctx0, cur, layer.ls_2_w);
+ cur = ggml_add(ctx0, cur, inp);
+ }
+ cb(cur, "mimi_dec_tfm", -1);
+
+ cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
+ cur = seanet.decode(cur);
+
+ for (const auto & s : seanet.state_out) {
+ ggml_set_name(s.second, ("state_out_" + s.first).c_str());
+ ggml_set_output(s.second);
+ ggml_build_forward_expand(gf, s.second);
+ }
+
+ // [n_samples, 1] -> [n_samples], clamped like the reference output
+ cur = ggml_reshape_1d(ctx0, cur, cur->ne[0]);
+ cur = ggml_clamp(ctx0, cur, -1.0f, 1.0f);
+ ggml_set_name(cur, "out_audio");
+ ggml_set_output(cur);
+ ggml_build_forward_expand(gf, cur);
+
+ return gf;
+}
--- /dev/null
+#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;
+}
--- /dev/null
+#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;
+}
const auto & c2w = model.c2w;
std::vector<c2w_state_slot> slots;
+ if (c2w.pre_conv_w == nullptr) {
+ return slots; // not a code2wav model, it keeps no state between calls
+ }
+
slots.push_back({"tfm_pos", 1, 1});
// prefix is (W-1) frames, the batch itself gives the other N=W frames (see tfm_layer_forward)
return output;
}
+
+//
+// mtmd_audio_preprocessor_pockettts
+//
+// mimi takes the raw 24kHz waveform, there is no mel front-end
+// the samples are handed over as a single-row "mel", to reuse the normal chunk path
+//
+
+bool mtmd_audio_preprocessor_pockettts::preprocess(const float * samples,
+ size_t n_samples,
+ std::vector<mtmd_audio_mel> & output) {
+ // the encoder needs whole frames, see pad_for_conv1d() in the reference
+ const int64_t frame_size = (int64_t) hparams.mimi_downsample * 120;
+ if (n_samples == 0 || frame_size <= 0) {
+ return false;
+ }
+
+ // the mimi transformer mask is dense, so cost is quadratic in the reference length
+ const int64_t max_samples = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate;
+ if ((int64_t) n_samples > max_samples) {
+ LOG_WRN("%s: speaker reference is %.1f s, truncating to the first %d s\n", __func__,
+ (double) n_samples / hparams.audio_sample_rate, clip_hparams::pockettts_max_spk_seconds);
+ n_samples = (size_t) max_samples;
+ }
+
+ const int64_t n_frames = (int64_t) (n_samples + frame_size - 1) / frame_size;
+ const int64_t n_padded = n_frames * frame_size;
+
+ mtmd_audio_mel out;
+ out.n_mel = 1;
+ out.n_len = n_padded;
+ out.n_len_org = (int64_t) n_samples;
+ out.data.assign((size_t) n_padded, 0.0f);
+ std::copy(samples, samples + n_samples, out.data.begin());
+
+ output.push_back(std::move(out));
+ return true;
+}
mtmd_audio_cache cache;
};
+// mimi convolves the waveform directly, so this only pads it to a whole number of frames
+struct mtmd_audio_preprocessor_pockettts : mtmd_audio_preprocessor {
+ mtmd_audio_preprocessor_pockettts(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {}
+ void initialize() override {}
+ bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) override;
+};
+
struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor {
mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { }
void initialize() override;
#include "../src/llama-ext.h"
#include <algorithm>
+#include <cctype>
+#include <cmath>
#include <cstring>
#include <memory>
#include <string>
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:
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
return n_prompt - prompt_pos;
}
- int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) override {
- mtmd_gen_inp inp{};
+ int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override {
+ if (sampled == LLAMA_TOKEN_NULL) {
+ LOG_ERR("mtmd_helper_gen_audio: qwen3tts requires a token sampled from the backbone\n");
+ return 1;
+ }
+
+ // backbone signals end-of-speech with a token, no frame for this step
+ if (sampled == codec_eos || llama_vocab_is_eog(vocab, sampled)) {
+ *out_stop = true;
+ *h_state_out = nullptr;
+ return 0;
+ }
+
+ mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
inp.code0 = sampled - codec_0;
inp.embd = const_cast<float *>(h_state_in);
inp.top_k = top_k;
inp.top_p = top_p;
+ inp.seed = seed;
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: gen_code process failed\n");
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{};
std::unique_ptr<decode_embd_batch> prompt_batch;
int n_prompt = 0;
int prompt_pos = 0;
- int32_t top_k = 50;
- float top_p = 1.0f;
+ int32_t top_k = 50;
+ float top_p = 1.0f;
+ uint32_t seed = UINT32_MAX;
std::vector<int32_t> codes_buf;
std::vector<uint8_t> c2w_state;
std::vector<float> audio_pcm;
std::vector<char> out_buf;
};
+// settings that only live in the reference's per-pack yaml, not in the checkpoint
+// the english packs share the same shapes and tokenizer, but disagree on these
+// all three are 0 / false when the pack does not tune them, the model default is then used
+struct pockettts_pack_settings {
+ float temp = 0.0f;
+ int frames_after_eos = 0;
+ bool pad_short_text = false;
+};
+
+static pockettts_pack_settings pockettts_pack(const char * variant) {
+ static const std::unordered_map<std::string, pockettts_pack_settings> packs = {
+ { "english", { 0.3f, 0, false } },
+ { "english_2026-01", { 0.7f, 0, true } },
+ { "english_2026-04", { 0.3f, 0, false } },
+ { "french_24l", { 0.7f, 8, false } },
+ };
+ auto it = packs.find(variant ? variant : "");
+ if (it == packs.end()) {
+ LOG_WRN("mtmd_helper_gen_audio: no tuned settings for pocket-tts variant \"%s\"\n",
+ variant ? variant : "");
+ return {};
+ }
+ return it->second;
+}
+
+// pocket-tts: the backbone emits no token, the flow net turns each hidden state into a latent
+// the end-of-speech head also lives in the mmproj
+class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
+public:
+ using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline;
+
+ void reset() override {
+ seq_id = 0;
+ pos = 0;
+ feats_buf.clear();
+ dec_state.clear();
+ audio_pcm.clear();
+ h_state_buf.clear();
+ out_buf.clear();
+ prompt_embd_buf.clear();
+ prompt_batch.reset();
+ n_prompt = 0;
+ prompt_pos = 0;
+ step_idx = 0;
+ eos_step = -1;
+ chunks.clear();
+ chunk_idx = 0;
+ n_voice_pos = 0;
+ chunk_budget = 0;
+ }
+
+ int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
+ reset();
+ seq_id = inp->seq_id;
+
+ if (!ensure_cache()) {
+ return 1;
+ }
+
+ std::vector<float> voice;
+ if (inp->speaker_ref) {
+ if (!encode_speaker(inp->speaker_ref, voice)) {
+ return 1;
+ }
+ }
+
+ pack = pockettts_pack(info.model_variant);
+
+ const std::string text = prepare_text(std::string(inp->prompt, inp->prompt_len),
+ pack.pad_short_text);
+ if (text.empty()) {
+ LOG_ERR("mtmd_helper_gen_audio: empty prompt\n");
+ return 1;
+ }
+
+ std::vector<llama_token> ids(text.size() + 16);
+ int n_ids = llama_tokenize(vocab, text.c_str(), (int32_t) text.size(), ids.data(),
+ (int32_t) ids.size(), false, false);
+ if (n_ids <= 0) {
+ LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n");
+ return 1;
+ }
+ ids.resize((size_t) n_ids);
+
+ // long inputs degrade badly, so each chunk restarts from the voice conditioning
+ // see split_into_best_sentences() in the reference
+ chunks = split_chunks(ids);
+ chunk_idx = 0;
+ if (chunks.size() > 1) {
+ LOG_INF("mtmd_helper_gen_audio: %d tokens split into %zu chunks\n", n_ids, chunks.size());
+ }
+
+ const int n_e = n_embd;
+
+ // sequence order is voice, then text, then the audio BOS that starts generation
+ if (!voice.empty()) {
+ GGML_ASSERT(voice.size() % (size_t) n_e == 0);
+ if (bos_before_voice != LLAMA_TOKEN_NULL) {
+ push_embd_row(prompt_embd_buf, bos_before_voice);
+ }
+ prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end());
+ }
+ // every later chunk rewinds to here and re-prompts, so the voice stays primed
+ n_voice_pos = (int) (prompt_embd_buf.size() / (size_t) n_e);
+
+ for (llama_token t : chunks[0]) {
+ push_embd_row(prompt_embd_buf, t);
+ }
+ push_embd_row(prompt_embd_buf, audio_bos);
+ arm_chunk_budget(0);
+
+ n_prompt = (int) (prompt_embd_buf.size() / (size_t) n_e);
+ prompt_batch.reset(new decode_embd_batch(prompt_embd_buf.data(), n_prompt, 1, n_e));
+ prompt_batch->set_position_normal(0, seq_id);
+ prompt_pos = 0;
+
+ seed = inp->seed;
+ out_type = inp->out_type;
+
+ return 0;
+ }
+
+ int32_t step_prompt(int32_t n_batch) override {
+ GGML_ASSERT(n_batch > 0);
+ if (prompt_pos >= n_prompt) {
+ return 0;
+ }
+ const int32_t n_tokens_batch = std::min(n_batch, n_prompt - prompt_pos);
+ llama_batch batch_view = prompt_batch->get_view(prompt_pos, n_tokens_batch);
+
+ if ((prompt_pos + n_tokens_batch) == n_prompt) {
+ batch_view.logits[n_tokens_batch - 1] = 1;
+ }
+
+ if (llama_decode(lctx, batch_view) != 0) {
+ LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n");
+ return -1;
+ }
+
+ pos += n_tokens_batch;
+ prompt_pos += n_tokens_batch;
+
+ if (prompt_pos >= n_prompt) {
+ prompt_batch.reset();
+ prompt_embd_buf.clear();
+ return 0;
+ }
+ return n_prompt - prompt_pos;
+ }
+
+ int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override {
+ (void) sampled; // the backbone output is continuous, there is no token to consume
+
+ mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
+ inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
+ inp.embd = const_cast<float *>(h_state_in);
+ // clip only reseeds when the seed changes, so pass the same one on every step
+ inp.seed = seed;
+ if (pack.temp > 0.0f) {
+ inp.temp = pack.temp;
+ }
+ mtmd_gen_out out{};
+ if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
+ LOG_ERR("mtmd_helper_gen_audio: flow decode failed\n");
+ return 1;
+ }
+ if (out.is_eos && eos_step < 0) {
+ eos_step = step_idx;
+ }
+ // the frame of the stopping step is discarded, matching _autoregressive_generation().
+ // the budget is the reference's fallback for a chunk whose eos head never fires
+ const bool chunk_done = (eos_step >= 0 && step_idx >= eos_step + frames_after_eos) ||
+ step_idx >= chunk_budget;
+ if (chunk_done) {
+ if (eos_step < 0) {
+ LOG_WRN("mtmd_helper_gen_audio: chunk %zu hit its budget without end-of-speech\n", chunk_idx);
+ }
+ return finish_chunk(h_state_out, out_stop);
+ }
+
+ feats_buf.insert(feats_buf.end(), out.feats, out.feats + out.n_feats);
+ step_idx++;
+ if (out.n_feats > 0 && feats_buf.size() / out.n_feats >= window_frames) {
+ if (!flush_gen_wav()) {
+ return 1;
+ }
+ }
+
+ decode_embd_batch batch_embd(const_cast<float *>(out.embd), 1, 1, n_embd);
+ batch_embd.set_position_normal(pos, seq_id);
+ batch_embd.batch.logits[0] = 1;
+ pos++;
+
+ if (llama_decode(lctx, batch_embd.batch) != 0) {
+ LOG_ERR("mtmd_helper_gen_audio: decode failed\n");
+ return 1;
+ }
+
+ const float * he = llama_get_embeddings_ith(lctx, -1);
+ h_state_buf.assign(he, he + n_embd);
+ *h_state_out = h_state_buf.data();
+
+ return 0;
+ }
+
+ int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override {
+ if (!flush_gen_wav()) {
+ return 1;
+ }
+
+ *out_sample_rate = info.sample_rate;
+ if (out_n_samples) {
+ *out_n_samples = (int64_t) audio_pcm.size();
+ }
+
+ if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
+ *out_data = (const char *) audio_pcm.data();
+ *out_data_len = audio_pcm.size() * sizeof(float);
+ return 0;
+ }
+
+ out_buf.clear();
+ if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
+ LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
+ return 1;
+ }
+ *out_data = out_buf.data();
+ *out_data_len = out_buf.size();
+ return 0;
+ }
+
+private:
+ bool ensure_cache() {
+ if (specials_ok) {
+ return true;
+ }
+ // bos_before_voice is optional, some packs do not insert it
+ bos_before_voice = find_special_token(vocab, "<|bos_before_voice|>");
+ audio_bos = find_special_token(vocab, "<|audio_bos|>");
+ if (audio_bos == LLAMA_TOKEN_NULL) {
+ LOG_ERR("mtmd_helper_gen_audio: missing <|audio_bos|> in vocab\n");
+ return false;
+ }
+ const uint32_t n_tok_embd = llama_model_get_tok_embd(model, nullptr);
+ if (n_tok_embd == 0) {
+ LOG_ERR("mtmd_helper_gen_audio: model has no token embeddings\n");
+ return false;
+ }
+ tok_embd.resize(n_tok_embd);
+ if (llama_model_get_tok_embd(model, tok_embd.data()) != n_tok_embd) {
+ LOG_ERR("mtmd_helper_gen_audio: token embedding copy failed\n");
+ return false;
+ }
+ GGML_ASSERT(n_embd > 0 && n_tok_embd % (uint32_t) n_embd == 0);
+ specials_ok = true;
+ return true;
+ }
+
+ // the table can be shorter than the vocab, so bound the row lookup
+ void push_embd_row(std::vector<float> & dst, llama_token t) const {
+ const size_t n_rows = tok_embd.size() / (size_t) n_embd;
+ GGML_ASSERT(t >= 0 && (size_t) t < n_rows);
+ dst.insert(dst.end(),
+ tok_embd.begin() + (size_t) t * n_embd,
+ tok_embd.begin() + (size_t) (t + 1) * n_embd);
+ }
+
+ // token ids of the pieces the reference splits on, see split_into_best_sentences().
+ // the leading token is dropped, it is the tokenizer's dummy prefix
+ std::vector<llama_token> punct_ids(const char * s) const {
+ std::vector<llama_token> ids(16);
+ const int n = llama_tokenize(vocab, s, (int32_t) strlen(s), ids.data(), (int32_t) ids.size(), false, false);
+ if (n <= 1) {
+ return {};
+ }
+ return std::vector<llama_token>(ids.begin() + 1, ids.begin() + n);
+ }
+
+ // cut after runs of boundary tokens, so punctuation stays with the sentence it ends
+ static std::vector<std::vector<llama_token>> split_on(const std::vector<llama_token> & ids,
+ const std::vector<llama_token> & boundary) {
+ std::vector<std::vector<llama_token>> out;
+ size_t start = 0;
+ bool prev_was_boundary = false;
+ for (size_t i = 0; i < ids.size(); i++) {
+ const bool is_boundary = std::find(boundary.begin(), boundary.end(), ids[i]) != boundary.end();
+ if (!is_boundary && prev_was_boundary) {
+ out.emplace_back(ids.begin() + start, ids.begin() + i);
+ start = i;
+ }
+ prev_was_boundary = is_boundary;
+ }
+ out.emplace_back(ids.begin() + start, ids.end());
+ return out;
+ }
+
+ std::vector<std::vector<llama_token>> split_chunks(const std::vector<llama_token> & ids) const {
+ if ((int) ids.size() <= max_chunk_tokens) {
+ return { ids };
+ }
+ const std::vector<llama_token> eos_punct = punct_ids(".!...?");
+ const std::vector<llama_token> mid_punct = punct_ids(",;:");
+
+ // oversized sentences are split again on weaker punctuation, else words get skipped
+ std::vector<std::vector<llama_token>> segments;
+ for (auto & seg : split_on(ids, eos_punct)) {
+ if ((int) seg.size() <= max_chunk_tokens) {
+ segments.push_back(std::move(seg));
+ continue;
+ }
+ auto sub = split_on(seg, mid_punct);
+ if (sub.size() > 1) {
+ for (auto & s : sub) {
+ segments.push_back(std::move(s));
+ }
+ } else {
+ segments.push_back(std::move(seg));
+ }
+ }
+
+ std::vector<std::vector<llama_token>> out;
+ for (auto & seg : segments) {
+ if (seg.empty()) {
+ continue;
+ }
+ if (!out.empty() && (int) (out.back().size() + seg.size()) <= max_chunk_tokens) {
+ out.back().insert(out.back().end(), seg.begin(), seg.end());
+ } else {
+ out.push_back(std::move(seg));
+ }
+ }
+ if (out.empty()) {
+ out.push_back(ids);
+ }
+ for (const auto & c : out) {
+ if ((int) c.size() > max_chunk_tokens) {
+ LOG_WRN("mtmd_helper_gen_audio: chunk of %zu tokens exceeds the %d token budget, "
+ "generation may skip words\n", c.size(), max_chunk_tokens);
+ }
+ }
+ return out;
+ }
+
+ // _estimate_max_gen_len() plus the per-chunk tail guess, both in frames
+ void arm_chunk_budget(size_t idx) {
+ const int n_tok = (int) chunks[idx].size();
+ chunk_budget = (int) std::ceil((n_tok / 3.0 + 2.0) * frame_rate);
+ // the pack may pin the tail, else the reference guesses it from the word count
+ frames_after_eos = pack.frames_after_eos > 0 ? pack.frames_after_eos : (n_tok <= 6 ? 5 : 3);
+ step_idx = 0;
+ eos_step = -1;
+ }
+
+ // ends the current chunk and, if there is another, re-prompts it on top of the voice
+ int32_t finish_chunk(const float ** h_state_out, bool * out_stop) {
+ if (!flush_gen_wav()) {
+ return 1;
+ }
+ // the decoder restarts too, the next chunk's audio is not continuous with this one
+ dec_state.clear();
+
+ if (chunk_idx + 1 >= chunks.size()) {
+ *out_stop = true;
+ *h_state_out = nullptr;
+ return 0;
+ }
+ chunk_idx++;
+
+ // drop this chunk's text and audio, keep the voice conditioning
+ llama_memory_seq_rm(llama_get_memory(lctx), seq_id, n_voice_pos, -1);
+ pos = n_voice_pos;
+
+ const int n_e = n_embd;
+ prompt_embd_buf.clear();
+ for (llama_token t : chunks[chunk_idx]) {
+ push_embd_row(prompt_embd_buf, t);
+ }
+ push_embd_row(prompt_embd_buf, audio_bos);
+ arm_chunk_budget(chunk_idx);
+
+ const int n_rows = (int) (prompt_embd_buf.size() / (size_t) n_e);
+ GGML_ASSERT(n_rows > 0);
+ decode_embd_batch batch(prompt_embd_buf.data(), n_rows, 1, n_e);
+ batch.set_position_normal(pos, seq_id);
+ batch.batch.logits[n_rows - 1] = 1;
+ if (llama_decode(lctx, batch.batch) != 0) {
+ LOG_ERR("mtmd_helper_gen_audio: chunk prompt decode failed\n");
+ return 1;
+ }
+ pos += n_rows;
+ prompt_embd_buf.clear();
+
+ const float * he = llama_get_embeddings_ith(lctx, -1);
+ h_state_buf.assign(he, he + n_embd);
+ *h_state_out = h_state_buf.data();
+ *out_stop = false;
+ return 0;
+ }
+
+ // same normalization as prepare_text_prompt() in the reference, it affects quality
+ static std::string prepare_text(const std::string & in, bool pad_short) {
+ std::string s;
+ s.reserve(in.size() + 1);
+ for (char c : in) {
+ if (c == '\n' || c == '\r') {
+ s += ' ';
+ } else if (c == ';') {
+ s += ',';
+ } else {
+ s += c;
+ }
+ }
+ const size_t b = s.find_first_not_of(' ');
+ const size_t e = s.find_last_not_of(' ');
+ if (b == std::string::npos) {
+ return "";
+ }
+ s = s.substr(b, e - b + 1);
+ if (s[0] >= 'a' && s[0] <= 'z') {
+ s[0] = (char) (s[0] - 'a' + 'A');
+ }
+ const unsigned char last = (unsigned char) s.back();
+ if (std::isalnum(last)) {
+ s += '.';
+ }
+ if (pad_short && count_words(s) < 5) {
+ s = std::string(8, ' ') + s;
+ }
+ return s;
+ }
+
+ static int count_words(const std::string & s) {
+ int n = 0;
+ bool in_word = false;
+ for (char c : s) {
+ if (c == ' ') {
+ in_word = false;
+ } else if (!in_word) {
+ in_word = true;
+ n++;
+ }
+ }
+ return n;
+ }
+
+ // runs the reference wav through the mimi encoder, returns one row per 12.5Hz frame
+ bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
+ if (!mtmd_support_audio(mctx)) {
+ LOG_ERR("mtmd_helper_gen_audio: mmproj has no voice encoder\n");
+ return false;
+ }
+ const std::string marker = mtmd_default_marker();
+ mtmd_input_text text{ marker.c_str(), marker.size(), false, true };
+ mtmd_input_chunks * chunks = mtmd_input_chunks_init();
+ const mtmd_bitmap * bptr = bitmap;
+ bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0;
+ if (ok) {
+ ok = false;
+ for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {
+ const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
+ if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) {
+ continue;
+ }
+ if (mtmd_encode_chunk(mctx, chunk) != 0) {
+ LOG_ERR("mtmd_helper_gen_audio: voice encode failed\n");
+ break;
+ }
+ const float * embd = mtmd_get_output_embd(mctx);
+ const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk);
+ out.assign(embd, embd + n);
+ ok = true;
+ break;
+ }
+ }
+ mtmd_input_chunks_free(chunks);
+ return ok;
+ }
+
+ // decodes the buffered latents, the mimi decoder state carries over between calls
+ bool flush_gen_wav() {
+ if (feats_buf.empty()) {
+ return true;
+ }
+ mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
+ inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
+ inp.feats = feats_buf.data();
+ inp.n_feats = feats_buf.size();
+ inp.seed = seed;
+ inp.state_data = dec_state.empty() ? nullptr : (const char *) dec_state.data();
+ inp.state_size = dec_state.size();
+ mtmd_gen_out out{};
+ if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
+ LOG_ERR("mtmd_helper_gen_audio: mimi decode failed\n");
+ return false;
+ }
+ audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples);
+ dec_state.assign(out.state_data, out.state_data + out.state_size);
+ feats_buf.clear();
+ return true;
+ }
+
+ pockettts_pack_settings pack;
+ bool specials_ok = false;
+ llama_token bos_before_voice = LLAMA_TOKEN_NULL;
+ llama_token audio_bos = LLAMA_TOKEN_NULL;
+ std::vector<float> tok_embd;
+
+ llama_seq_id seq_id = 0;
+ int pos = 0;
+ std::vector<float> prompt_embd_buf;
+ std::unique_ptr<decode_embd_batch> prompt_batch;
+ int n_prompt = 0;
+ int prompt_pos = 0;
+ uint32_t seed = UINT32_MAX;
+ // end-of-speech is latched, then a few more frames are generated as tail padding
+ int step_idx = 0;
+ int eos_step = -1;
+ int frames_after_eos = 3;
+ static constexpr int max_chunk_tokens = 50; // MAX_TOKEN_PER_CHUNK in the reference
+ static constexpr double frame_rate = 12.5;
+ std::vector<std::vector<llama_token>> chunks;
+ size_t chunk_idx = 0;
+ int n_voice_pos = 0; // KV positions held by the voice conditioning
+ int chunk_budget = 0;
+
+ // latents are decoded a window at a time, the decoder state bridges the windows
+ size_t window_frames = 8;
+ std::vector<float> feats_buf;
+ std::vector<uint8_t> dec_state;
+ std::vector<float> audio_pcm;
+ std::vector<float> h_state_buf;
+ mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
+ std::vector<char> out_buf;
+};
+
static std::unique_ptr<mtmd_gen_audio_pipeline> make_pipeline(llama_context * lctx, mtmd_context * mctx) {
switch (mtmd_gen_audio_get_info(mctx).type) {
case MTMD_GEN_AUDIO_TYPE_QWEN3TTS:
return std::unique_ptr<mtmd_gen_audio_pipeline>(new qwen3tts_gen_audio_pipeline(lctx, mctx));
+ case MTMD_GEN_AUDIO_TYPE_POCKETTTS:
+ return std::unique_ptr<mtmd_gen_audio_pipeline>(new pockettts_gen_audio_pipeline(lctx, mctx));
default:
return nullptr;
}
}
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,
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;
};
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
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);
// generation context
struct clip_ctx * ctx_gen_a; // audio
std::vector<int32_t> gen_out_codes; // this frame's 16 sampled codes (GEN_CODE)
+ std::vector<float> gen_out_feats; // this frame's continuous features, if any (GEN_CODE)
std::vector<float> gen_out_embd; // next-step hidden state fed back to backbone (GEN_CODE)
std::vector<float> gen_out_audio; // decoded PCM samples for the current frame (GEN_WAV)
std::vector<uint8_t> gen_out_state; // state to feed into the next GEN_WAV call
{
audio_preproc = std::make_unique<mtmd_audio_preprocessor_qwen3tts_spk>(ctx_a);
} break;
+ case PROJECTOR_TYPE_POCKETTTS_SPKENC:
+ {
+ audio_preproc = std::make_unique<mtmd_audio_preprocessor_pockettts>(ctx_a);
+ } break;
default:
throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj));
}
//
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;
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) {
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);
std::vector<float> out_embd(n_embd);
std::vector<int32_t> out_codes;
+ std::vector<float> out_feats;
+ bool is_eos = false;
clip_encode_params params;
- params.imgs = &batch;
- params.n_threads = ctx->n_threads;
- params.gen_process = CLIP_GEN_PROCESS_GEN_CODE;
- params.out_embd = &out_embd;
- params.out_codes = &out_codes;
- params.code0 = inp->code0;
- params.top_k = inp->top_k;
- params.top_p = inp->top_p;
+ params.imgs = &batch;
+ params.n_threads = ctx->n_threads;
+ params.gen_process = CLIP_GEN_PROCESS_GEN_CODE;
+ params.out_embd = &out_embd;
+ params.out_codes = &out_codes;
+ params.out_feats = &out_feats;
+ params.code0 = inp->code0;
+ params.top_k = inp->top_k;
+ params.top_p = inp->top_p;
+ params.seed = inp->seed;
+ params.temp = inp->temp;
+ params.out_is_eos = &is_eos;
if (!clip_encode(ctx_clip, ¶ms)) {
LOG_ERR("%s: clip_encode failed (gen_code)\n", __func__);
ctx->gen_out_embd = std::move(out_embd);
ctx->gen_out_codes = std::move(out_codes);
-
- out->embd = ctx->gen_out_embd.data();
- out->codes = ctx->gen_out_codes.data();
- out->n_codes = ctx->gen_out_codes.size();
+ 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->feats = ctx->gen_out_feats.data();
+ out->n_feats = ctx->gen_out_feats.size();
+ out->is_eos = is_eos;
return 0;
}
// MTMD_GEN_PROCESS_TYPE_GEN_WAV
- if (!inp->codes || inp->n_codes == 0) {
- LOG_ERR("%s: codes required for gen_wav\n", __func__);
+ const bool has_codes = inp->codes && inp->n_codes > 0;
+ const bool has_feats = inp->feats && inp->n_feats > 0;
+ if (has_codes == has_feats) {
+ LOG_ERR("%s: gen_wav requires exactly one of codes or feats\n", __func__);
return 1;
}
- std::vector<int32_t> in_codes(inp->codes, inp->codes + inp->n_codes);
+ std::vector<int32_t> in_codes;
+ std::vector<float> in_feats;
+ if (has_codes) {
+ in_codes.assign(inp->codes, inp->codes + inp->n_codes);
+ } else {
+ in_feats.assign(inp->feats, inp->feats + inp->n_feats);
+ }
std::vector<uint8_t> in_state;
if (inp->state_data) {
in_state.assign(inp->state_data, inp->state_data + inp->state_size);
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;
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;
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;
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,
--tts-speaker-file speaker.mp3 \
--output out.wav
```
+
+## Pocket TTS
+
+Available params:
+- `--tts-speaker-file` should point to a speaker reference audio file (wav, mp3). It is required, the model produces almost no audio without it
+- Note: `lang` is not used, the language is a property of the weights
+
+Example usage:
+
+```sh
+llama-tts -m pocket-tts.gguf \
+ -mm mmproj-pocket-tts.gguf \
+ -p "Hello world" \
+ --tts-speaker-file speaker.mp3 \
+ --output out.wav
+```
+
+**Note for GGUF conversion:**
+
+The [upstream repository](https://huggingface.co/kyutai/pocket-tts) holds one complete model per language under `languages/`, next to a set of shared files at the root. Convert one of the `languages/<name>` directories, **not** the root directory:
+
+```sh
+python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --outfile pocket-tts.gguf
+python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --mmproj --outfile mmproj-pocket-tts.gguf
+```
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;
//
}
}
- 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);
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;
overview / pkg / ggml / sources / llama.cpp / commitdiff