Commit b97627967 for llama.cpp
commit b97627967c89a5ef135c2da945a036abe2d22ed3
Author: Xuan-Son Nguyen <son@huggingface.co>
Date: Tue Nov 3 23:47:11 2026 +0100
model, mtmd: add chatterbox TTS support (#30993)
* add chatterbox TTS support
Co-authored-by: Pascal <admin@serveurperso.com>
* add built-in voice support
* rm redundant code
* clean up
* clean up 2
* Apply batched suggestions from code review
Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
* improve
---------
Co-authored-by: Pascal <admin@serveurperso.com>
Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
diff --git a/conversion/__init__.py b/conversion/__init__.py
index 6cea1ec78..1370dd3b1 100644
--- a/conversion/__init__.py
+++ b/conversion/__init__.py
@@ -45,6 +45,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
"ChameleonForConditionalGeneration": "chameleon",
"ChatGLMForConditionalGeneration": "chatglm",
"ChatGLMModel": "chatglm",
+ "ChatterboxModel": "chatterbox",
"ClefModel": "clef",
"CloverMoeModelLM": "clover",
"CodeShellForCausalLM": "codeshell",
@@ -330,6 +331,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
MMPROJ_MODEL_MAP: dict[str, str] = {
"AudioFlamingo3ForConditionalGeneration": "ultravox",
+ "ChatterboxModel": "chatterbox",
"ClefModel": "clef",
"CogVLMForCausalLM": "cogvlm",
"Cohere2VisionForConditionalGeneration": "command_r",
diff --git a/conversion/base.py b/conversion/base.py
index 9d9d6e875..28f2bf241 100644
--- a/conversion/base.py
+++ b/conversion/base.py
@@ -225,8 +225,6 @@ class ModelBase:
tensors: dict[str, Callable[[], Tensor]] = {}
if remote_hf_model_id is not None:
- is_safetensors = True
-
logger.info(f"Using remote model with HuggingFace id: {remote_hf_model_id}")
remote_tensors = gguf.utility.SafetensorRemote.get_list_tensors_hf_model(remote_hf_model_id)
for name, remote_tensor in remote_tensors.items():
@@ -237,9 +235,13 @@ class ModelBase:
return tensors
- prefix = "model" if not self.is_mistral_format else "consolidated"
+ return self.index_local_tensors("model" if not self.is_mistral_format else "consolidated")
+
+ # index the local "<prefix>*.safetensors" files (or pytorch_model*.bin), through filter_tensors
+ def index_local_tensors(self, prefix: str) -> dict[str, Callable[[], Tensor]]:
+ tensors: dict[str, Callable[[], Tensor]] = {}
part_names: list[str] = ModelBase.get_model_part_names(self.dir_model, prefix, ".safetensors")
- is_safetensors: bool = len(part_names) > 0 or (not self.is_mistral_format and (self.dir_model / "model.safetensors.index.json").is_file())
+ is_safetensors: bool = len(part_names) > 0 or (not self.is_mistral_format and (self.dir_model / f"{prefix}.safetensors.index.json").is_file())
if not is_safetensors:
part_names = ModelBase.get_model_part_names(self.dir_model, "pytorch_model", ".bin")
@@ -247,7 +249,7 @@ class ModelBase:
tensor_names_from_parts: set[str] = set()
if not self.is_mistral_format:
- index_name = "model.safetensors" if is_safetensors else "pytorch_model.bin"
+ index_name = f"{prefix}.safetensors" if is_safetensors else "pytorch_model.bin"
index_name += ".index.json"
index_file = self.dir_model / index_name
diff --git a/conversion/chatterbox.py b/conversion/chatterbox.py
new file mode 100644
index 000000000..4ab93860c
--- /dev/null
+++ b/conversion/chatterbox.py
@@ -0,0 +1,471 @@
+# Chatterbox (ResembleAI) turbo and multilingual v3, from the raw checkpoints of the official repos
+# the checkpoints ship no config.json, the variant is detected from the talker file and the hparams come from the reference code
+# talker: text_emb and speech_emb are fused into one [text | speech] vocab, speech token i is <|speech_i|>
+# mmproj: the s3gen sidecar, the voice and conditioning encoders, the position and speech tables, and the conds.pt voice
+
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Iterable
+
+import torch
+import torch.nn.functional as F
+
+if TYPE_CHECKING:
+ from torch import Tensor
+
+from .base import ModelBase, TextModel, MmprojModel, LazyTorchTensor, gguf, logger
+from .llama import LlamaModel
+
+# checkpoint file prefixes, the files are "<prefix>.safetensors"
+TURBO_TALKER = "t3_turbo_v1"
+MTL_TALKER = "t3_mtl23ls_v3"
+TURBO_S3GEN = "s3gen_meanflow"
+MTL_S3GEN = "s3gen_v3"
+VE = "ve"
+
+# relative speech ids shared by both variants (start/stop_speech_token in the reference)
+SPEECH_BOS = 6561
+SPEECH_EOS = 6562
+
+# sampling defaults of the reference, samplers it does not use are disabled so the common defaults do not apply
+TURBO_TOP_K = 1000
+MTL_TOP_K = 0
+TURBO_SAMPLING = {"min_p": 0.0, "top_p": 0.95, "temp": 0.8, "penalty_repeat": 1.2}
+MTL_SAMPLING = {"min_p": 0.05, "top_p": 1.0, "temp": 0.8, "penalty_repeat": 1.2}
+# the reference repetition penalty covers the whole generated history
+PENALTY_LAST_N = -1
+
+
+def _has_file(dir_model: Path, prefix: str) -> bool:
+ return (dir_model / f"{prefix}.safetensors").is_file()
+
+
+def _is_turbo(dir_model: Path) -> bool:
+ if _has_file(dir_model, TURBO_TALKER):
+ return True
+ if _has_file(dir_model, MTL_TALKER):
+ return False
+ raise FileNotFoundError(f"no chatterbox talker checkpoint in {dir_model}")
+
+
+# talker hparams of the reference (GPT2_MEDIUM_CONFIG / LLAMA_520M_CONFIG_DICT, T3Config)
+TURBO_TEXT_CONFIG = {
+ "n_layer": 24, "n_embd": 1024, "n_inner": 4096, "n_head": 16, "n_ctx": 8196, "layer_norm_epsilon": 1e-05,
+ "vocab_size": 50276, "speech_vocab_size": 6563,
+}
+MTL_TEXT_CONFIG = {
+ "num_hidden_layers": 30, "hidden_size": 1024, "intermediate_size": 4096, "num_attention_heads": 16,
+ "num_key_value_heads": 16, "head_dim": 64, "max_position_embeddings": 131072, "rms_norm_eps": 1e-05,
+ "rope_theta": 500000.0,
+ "rope_scaling": {"factor": 8.0, "high_freq_factor": 4.0, "low_freq_factor": 1.0, "original_max_position_embeddings": 8192},
+ "vocab_size": 2454, "speech_vocab_size": 8194,
+}
+
+
+def _is_chatterbox_checkpoint(dir_model: Path) -> bool:
+ return _has_file(dir_model, VE) and (_has_file(dir_model, TURBO_TALKER) or _has_file(dir_model, MTL_TALKER))
+
+
+@ModelBase.register_hparams_loader(_is_chatterbox_checkpoint)
+def _load_chatterbox_hparams(dir_model: Path) -> dict[str, Any]:
+ turbo = _is_turbo(dir_model)
+ logger.info(f"gguf: detected Chatterbox {'turbo' if turbo else 'multilingual v3'} checkpoint")
+ text_config = dict(TURBO_TEXT_CONFIG if turbo else MTL_TEXT_CONFIG)
+ # the mmproj tensors are named explicitly, they do not go through the block tensor map
+ return {"architectures": ["ChatterboxModel"], "text_config": text_config, "audio_config": {"num_hidden_layers": 0}}
+
+
+@ModelBase.register("ChatterboxModel")
+class ChatterboxTalkerModel(TextModel):
+ model_arch = gguf.MODEL_ARCH.LLAMA # multilingual; the turbo constructor switches to GPT2
+
+ def __init__(self, dir_model: Path, *args, **kwargs):
+ self.is_turbo = _is_turbo(dir_model)
+ if self.is_turbo:
+ self.model_arch = gguf.MODEL_ARCH.GPT2
+ super().__init__(dir_model, *args, **kwargs)
+ self._text_embd: Tensor | None = None
+ self._speech_embd: Tensor | None = None
+ self._text_head: Tensor | None = None
+ self._speech_head: Tensor | None = None
+
+ def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+ return self.index_local_tensors(TURBO_TALKER if self.is_turbo else MTL_TALKER)
+
+ @classmethod
+ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+ name, _ = item
+ # the reference feeds inputs_embeds, the transformer input embeddings are unused
+ if name in ("tfmr.wte.weight", "tfmr.embed_tokens.weight"):
+ return None
+ if name.startswith(("cond_enc.", "text_pos_emb.", "speech_pos_emb.")):
+ return None
+ return super().filter_tensors(item)
+
+ def set_vocab(self):
+ if self.is_turbo:
+ self._set_vocab_turbo()
+ else:
+ self._set_vocab_mtl()
+
+ def _speech_token_names(self, n_speech: int) -> list[str]:
+ return [f"<|speech_{i}|>" for i in range(n_speech)]
+
+ def _set_vocab_turbo(self):
+ # stock GPT-2 BPE from the checkpoint dir, extended with the speech tokens
+ tokens, toktypes, tokpre = self.get_vocab_base()
+ n_text = len(tokens)
+ n_speech = self.hparams["speech_vocab_size"]
+ speech = self._speech_token_names(n_speech)
+ tokens += speech
+ toktypes += [gguf.TokenType.CONTROL] * n_speech
+
+ with open(self.dir_model / "merges.txt", "r", encoding="utf-8") as f:
+ merges = [line.rstrip("\n") for line in f if line.strip() and not line.startswith("#version")]
+
+ self._write_vocab(tokens, toktypes, merges, tokpre, n_text)
+
+ def _write_vocab(self, tokens: list[str], toktypes: list[int], merges: list[str], tokpre: str, n_text: int):
+ self.gguf_writer.add_tokenizer_model("gpt2")
+ self.gguf_writer.add_tokenizer_pre(tokpre)
+ self.gguf_writer.add_token_list(tokens)
+ self.gguf_writer.add_token_types(toktypes)
+ self.gguf_writer.add_token_merges(merges)
+
+ self.gguf_writer.add_bos_token_id(n_text + SPEECH_BOS)
+ self.gguf_writer.add_eos_token_id(n_text + SPEECH_EOS)
+ self.gguf_writer.add_add_bos_token(False)
+ self.gguf_writer.add_add_eos_token(False)
+
+ # the reference samples the speech head only
+ self.gguf_writer.add_suppress_tokens(list(range(n_text)))
+
+ def _set_vocab_mtl(self):
+ # the reference BPE is char-level, it is re-encoded to byte-level with extra merges that rebuild each multi-byte char
+ # its lowercasing is folded into the embedding table, so the prompt needs no text preprocessing
+ with open(self.dir_model / "mtl_tokenizer.json", "r", encoding="utf-8") as f:
+ tok = json.load(f)
+
+ byte_map = gguf.vocab.bytes_to_unicode()
+
+ def enc(s: str) -> str:
+ return "".join(byte_map[b] for b in s.encode("utf-8"))
+
+ n_text = self.hparams["vocab_size"]
+ tokens: list[str] = [f"[unused_{i}]" for i in range(n_text)]
+ toktypes = [int(gguf.TokenType.UNUSED)] * n_text
+ char_merges: list[str] = []
+ for t, i in tok["model"]["vocab"].items():
+ if t == " ":
+ # the reference substitutes [SPACE] before encoding, the raw space token is never used
+ continue
+ tokens[i] = enc(t)
+ toktypes[i] = int(gguf.TokenType.NORMAL)
+ if len(t) == 1 and len(t.encode("utf-8")) > 1:
+ parts = [byte_map[b] for b in t.encode("utf-8")]
+ for k in range(1, len(parts)):
+ char_merges.append("".join(parts[:k]) + " " + parts[k])
+ for entry in tok.get("added_tokens", []):
+ tokens[entry["id"]] = entry["content"]
+ toktypes[entry["id"]] = int(gguf.TokenType.CONTROL)
+
+ # raw spaces tokenize to [SPACE] directly
+ space_id = tok["model"]["vocab"]["[SPACE]"]
+ tokens[space_id] = enc(" ")
+ toktypes[space_id] = int(gguf.TokenType.NORMAL)
+
+ n_speech = self.hparams["speech_vocab_size"]
+ tokens += self._speech_token_names(n_speech)
+ toktypes += [int(gguf.TokenType.CONTROL)] * n_speech
+
+ # chars are atomic in the reference, they must form before any of its merges
+ merges = char_merges
+ for m in tok["model"].get("merges", []):
+ a, b = m if isinstance(m, list) else m.split(" ")
+ merges.append(enc(a) + " " + enc(b))
+
+ self._write_vocab(tokens, toktypes, merges, "default", n_text)
+
+ def _mtl_case_fold_pairs(self) -> list[tuple[int, int]]:
+ # (upper_id, lower_id) of the single-char vocab entries, for every script in the vocab
+ with open(self.dir_model / "mtl_tokenizer.json", "r", encoding="utf-8") as f:
+ vocab = json.load(f)["model"]["vocab"]
+ pairs: list[tuple[int, int]] = []
+ for t, i in vocab.items():
+ if len(t) != 1:
+ continue
+ low = t.lower()
+ if low != t and len(low) == 1 and low in vocab:
+ pairs.append((i, vocab[low]))
+ return pairs
+
+ def set_gguf_parameters(self):
+ super().set_gguf_parameters()
+
+ sampling = TURBO_SAMPLING if self.is_turbo else MTL_SAMPLING
+ self.gguf_writer.add_sampling_top_k(TURBO_TOP_K if self.is_turbo else MTL_TOP_K)
+ self.gguf_writer.add_sampling_min_p(sampling["min_p"])
+ self.gguf_writer.add_sampling_top_p(sampling["top_p"])
+ self.gguf_writer.add_sampling_temp(sampling["temp"])
+ self.gguf_writer.add_sampling_penalty_repeat(sampling["penalty_repeat"])
+ self.gguf_writer.add_sampling_penalty_last_n(PENALTY_LAST_N)
+
+ def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+ if self.is_turbo:
+ return
+ # llama3 rope scaling baked into the rope_freqs factors tensor
+ rp = self.hparams["rope_scaling"]
+ dim = self.hparams["head_dim"]
+ base = self.hparams["rope_theta"]
+ freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+ factor = rp["factor"]
+ low_freq_wavelen = rp["original_max_position_embeddings"] / rp["low_freq_factor"]
+ high_freq_wavelen = rp["original_max_position_embeddings"] / rp["high_freq_factor"]
+ rope_factors = []
+ for freq in freqs:
+ wavelen = 2 * math.pi / freq
+ if wavelen < high_freq_wavelen:
+ rope_factors.append(1)
+ elif wavelen > low_freq_wavelen:
+ rope_factors.append(factor)
+ else:
+ smooth = (rp["original_max_position_embeddings"] / wavelen - rp["low_freq_factor"]) / (rp["high_freq_factor"] - rp["low_freq_factor"])
+ rope_factors.append(1 / ((1 - smooth) / factor + smooth))
+ yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))
+
+ def _maybe_emit_fused(self) -> Iterable[tuple[str, Tensor]]:
+ if self._text_embd is not None and self._speech_embd is not None:
+ yield (self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD),
+ torch.cat([self._text_embd, self._speech_embd], dim=0))
+ self._text_embd = None
+ self._speech_embd = None
+ if self._text_head is not None and self._speech_head is not None:
+ yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT),
+ torch.cat([self._text_head, self._speech_head], dim=0))
+ self._text_head = None
+ self._speech_head = None
+
+ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+ # generate_extra_tensors output comes back through here with its final name
+ if name.startswith("rope_freqs"):
+ yield (name, data_torch)
+ return
+ # fused [text | speech] vocab: embeddings and output head
+ if name == "text_emb.weight":
+ if not self.is_turbo:
+ # the reference lowercases its input, the uppercase rows take the lowercase ones
+ fold = list(range(data_torch.shape[0]))
+ for upper, lower in self._mtl_case_fold_pairs():
+ fold[upper] = lower
+ data_torch = data_torch[torch.tensor(fold)]
+ self._text_embd = data_torch
+ yield from self._maybe_emit_fused()
+ return
+ if name == "speech_emb.weight":
+ self._speech_embd = data_torch
+ yield from self._maybe_emit_fused()
+ return
+ if name == "text_head.weight":
+ self._text_head = data_torch
+ yield from self._maybe_emit_fused()
+ return
+ if name == "speech_head.weight":
+ self._speech_head = data_torch
+ yield from self._maybe_emit_fused()
+ return
+ if name == "speech_head.bias":
+ # the text head has no bias, its zone is suppressed at sampling anyway
+ yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT, suffix=".bias"), F.pad(data_torch, (self.hparams["vocab_size"], 0)))
+ return
+
+ assert name.startswith("tfmr.")
+ name = name[len("tfmr."):]
+
+ if self.is_turbo:
+ # HF GPT-2 layout: Conv1D style weights are stored transposed
+ if name.endswith((".c_attn.weight", ".c_proj.weight", ".c_fc.weight")):
+ data_torch = data_torch.transpose(1, 0)
+ yield (self.map_tensor_name(name), data_torch)
+ return
+
+ if name.endswith("q_proj.weight"):
+ data_torch = LlamaModel.permute(data_torch, self.hparams["num_attention_heads"], None)
+ if name.endswith("k_proj.weight"):
+ data_torch = LlamaModel.permute(data_torch, self.hparams["num_attention_heads"], self.hparams["num_key_value_heads"])
+ yield (self.map_tensor_name("model." + name), data_torch)
+
+
+@ModelBase.register("ChatterboxModel")
+class ChatterboxMmprojModel(MmprojModel):
+ has_vision_encoder = False
+ has_audio_encoder = True
+
+ def __init__(self, dir_model: Path, *args, **kwargs):
+ self.is_turbo = _is_turbo(dir_model)
+ super().__init__(dir_model, *args, **kwargs)
+ self._wnorm_g: dict[str, Tensor] = {}
+ self._wnorm_v: dict[str, Tensor] = {}
+
+ def get_audio_config(self) -> dict[str, Any] | None:
+ return self.global_config.get("audio_config")
+
+ def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+ tensors: dict[str, Callable[[], Tensor]] = {}
+ for prefix in (TURBO_S3GEN if self.is_turbo else MTL_S3GEN, VE, TURBO_TALKER if self.is_turbo else MTL_TALKER):
+ tensors.update(self.index_local_tensors(prefix))
+ return tensors
+
+ @classmethod
+ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+ # the s3gen, ve and talker checkpoints have disjoint top-level names
+ name, gen = item
+ new_name: str | None = None
+ # batchnorm bookkeeping (over the gguf name length cap) and the tokenizer dsp buffers are unused
+ if name.endswith("num_batches_tracked") or name in ("tokenizer.window", "tokenizer._mel_filters"):
+ return None
+ for src, dst in (
+ ("flow.encoder.", "a.gen.fenc."),
+ ("flow.decoder.estimator.", "a.gen.est."),
+ ("mel2wav.", "a.gen.hift."),
+ ("speaker_encoder.", "a.spk."),
+ ("tokenizer.", "a.s3tok."),
+ ("flow.", "a.gen.flow."), # input_embedding, encoder_proj, spk_embed_affine_layer
+ ("lstm.", "a.ve.lstm."),
+ ("proj.", "a.ve.proj."),
+ ("cond_enc.", "a.cenc."),
+ ):
+ if name.startswith(src):
+ new_name = dst + name[len(src):]
+ break
+ if name == "text_pos_emb.emb.weight":
+ new_name = "a.gen.t3.text_pos_emb"
+ elif name == "speech_pos_emb.emb.weight":
+ new_name = "a.gen.t3.speech_pos_emb"
+ elif name == "speech_emb.weight":
+ new_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.A_GEN_CODE_OUT_EMBD] + ".weight"
+ if new_name is None:
+ return None
+ return super().filter_tensors((new_name, gen))
+
+ def set_gguf_parameters(self):
+ self.gguf_writer.add_file_type(self.ftype)
+
+ # speaker encoder, the front-end hparams are hard-coded in clip.cpp
+ self.gguf_writer.add_clip_has_audio_encoder(True)
+ self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.CHATTERBOX_SPKENC)
+ self.gguf_writer.add_audio_projection_dim(80)
+ self.gguf_writer.add_audio_num_mel_bins(80)
+ self.gguf_writer.add_audio_block_count(0)
+ self.gguf_writer.add_audio_embedding_length(192)
+ self.gguf_writer.add_audio_head_count(1)
+ self.gguf_writer.add_audio_feed_forward_length(192)
+ self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
+
+ # audio generator (s3gen, chatterbox projector); the flow encoder shape
+ self.gguf_writer.add_clip_has_gen_audio_encoder(True)
+ self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.CHATTERBOX)
+ self.gguf_writer.add_gen_audio_projection_dim(80)
+ self.gguf_writer.add_gen_audio_embedding_length(512)
+ self.gguf_writer.add_gen_audio_feed_forward_length(2048)
+ self.gguf_writer.add_gen_audio_block_count(6)
+ self.gguf_writer.add_gen_audio_head_count(8)
+ self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5)
+
+ def _fuse_weight_norm(self, name: str, data_torch: Tensor) -> tuple[str, Tensor] | None:
+ # torch weight_norm parametrization: weight = g * v / |v| over dims 1..n
+ base = name.split(".parametrizations.weight.original")[0]
+ if name.endswith("original0"):
+ self._wnorm_g[base] = data_torch
+ else:
+ self._wnorm_v[base] = data_torch
+ if base in self._wnorm_g and base in self._wnorm_v:
+ g = self._wnorm_g.pop(base)
+ v = self._wnorm_v.pop(base)
+ norm = v.float().norm(dim=tuple(range(1, v.dim())), keepdim=True)
+ return (base + ".weight", g.float() * v.float() / norm)
+ return None
+
+ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+ if ".parametrizations.weight.original" in name:
+ fused = self._fuse_weight_norm(name, data_torch)
+ if fused is not None:
+ yield fused
+ return
+ yield (name, data_torch)
+
+ def tensor_force_quant(self, name, new_name, bid, n_dims):
+ del name, bid
+ # tensors read on the host stay F32
+ if new_name.startswith(("a.ve.", "a.cenc.", "voice.", "a.gen.t3.", "a.gen.hift.m_source.")):
+ return gguf.GGMLQuantizationType.F32
+ # the conv ops have no BF16 kernels
+ if n_dims >= 2 and new_name.endswith((".weight", ".weight_v")):
+ return gguf.GGMLQuantizationType.F16
+ return gguf.GGMLQuantizationType.F32
+
+ def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+ conds = torch.load(self.dir_model / "conds.pt", map_location="cpu", weights_only=True)
+ t3c = conds["t3"] if isinstance(conds, dict) else conds.t3
+ genc = conds["gen"] if isinstance(conds, dict) else conds.gen
+ t3c = vars(t3c) if not isinstance(t3c, dict) else t3c
+ genc = vars(genc) if not isinstance(genc, dict) else genc
+
+ def talker_tensor(name: str) -> Tensor:
+ return LazyTorchTensor.to_eager(self.model_tensors[name]())
+
+ # the conds.pt voice: voice.default holds the talker conditioning rows, voice.default.ref_* the flow decoder reference
+ # 1D tensors are written as F32, the speech ids stay exact
+ yield ("voice.default.ref_tokens", genc["prompt_token"][0].float())
+ yield ("voice.default.ref_feat", genc["prompt_feat"][0].float())
+ yield ("voice.default.ref_spk", genc["embedding"][0].float())
+
+ spkr_w = talker_tensor("a.cenc.spkr_enc.weight").float()
+ spkr_b = talker_tensor("a.cenc.spkr_enc.bias").float()
+ spkr_row = spkr_w @ t3c["speaker_emb"][0].float() + spkr_b
+ speech_emb = talker_tensor(self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_OUT_EMBD)).float()
+ cond_tokens = t3c["cond_prompt_speech_tokens"][0]
+
+ if self.is_turbo:
+ # turbo talker conditioning: [spkr, speech rows of the cond tokens]
+ yield ("voice.default", torch.cat([spkr_row[None], speech_emb[cond_tokens]], dim=0))
+ return
+
+ # multilingual talker conditioning [spkr, perceiver x32, emotion], from the reference perceiver over the cond speech tokens
+ with torch.no_grad():
+ speech_pos = talker_tensor("a.gen.t3.speech_pos_emb").float()
+ pse = speech_emb[cond_tokens] + speech_pos[: cond_tokens.shape[0]]
+
+ ln_w = talker_tensor("a.cenc.perceiver.attn.norm.weight").float()
+ ln_b = talker_tensor("a.cenc.perceiver.attn.norm.bias").float()
+ wq = talker_tensor("a.cenc.perceiver.attn.to_q.weight").float()
+ bq = talker_tensor("a.cenc.perceiver.attn.to_q.bias").float()
+ wk = talker_tensor("a.cenc.perceiver.attn.to_k.weight").float()
+ bk = talker_tensor("a.cenc.perceiver.attn.to_k.bias").float()
+ wv = talker_tensor("a.cenc.perceiver.attn.to_v.weight").float()
+ bv = talker_tensor("a.cenc.perceiver.attn.to_v.bias").float()
+ wo = talker_tensor("a.cenc.perceiver.attn.proj_out.weight").float()
+ bo = talker_tensor("a.cenc.perceiver.attn.proj_out.bias").float()
+ query = talker_tensor("a.cenc.perceiver.pre_attention_query")[0].float()
+
+ n_head = 4
+ n_e = query.shape[1]
+
+ def attn_block(x1: Tensor, x2: Tensor) -> Tensor:
+ nx1 = F.layer_norm(x1, (n_e,), ln_w, ln_b)
+ nx2 = F.layer_norm(x2, (n_e,), ln_w, ln_b)
+ q = (nx1 @ wq.T + bq).view(-1, n_head, n_e // n_head).transpose(0, 1)
+ k = (nx2 @ wk.T + bk).view(-1, n_head, n_e // n_head).transpose(0, 1)
+ v = (nx2 @ wv.T + bv).view(-1, n_head, n_e // n_head).transpose(0, 1)
+ ctx = F.scaled_dot_product_attention(q, k, v)
+ ctx = ctx.transpose(0, 1).reshape(-1, n_e)
+ return ctx @ wo.T + bo + x1
+
+ pre = attn_block(query, pse)
+ p32 = attn_block(pre, pre)
+
+ emo = talker_tensor("a.cenc.emotion_adv_fc.weight").float()
+ emo_row = emo[:, 0] * t3c["emotion_adv"].reshape(-1)[0]
+ yield ("voice.default", torch.cat([spkr_row[None], p32, emo_row[None]], dim=0))
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index c4c4a4c65..fcd17f4b5 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -6916,6 +6916,8 @@ class VisionProjectorType:
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
+ CHATTERBOX_SPKENC = "chatterbox_spkenc" # audio: CAMPPlus speaker encoder
+ CHATTERBOX = "chatterbox" # audio generation: s3gen flow matching + HiFT vocoder
SOPRANO = "soprano" # audio generation: Vocos decoder
HUNYUANVL = "hunyuanvl"
PARAKEET = "parakeet" # audio
diff --git a/src/models/gpt2.cpp b/src/models/gpt2.cpp
index d6ca64744..102dad97e 100644
--- a/src/models/gpt2.cpp
+++ b/src/models/gpt2.cpp
@@ -22,6 +22,7 @@ void llama_model_gpt2::load_arch_tensors(llama_model_loader &) {
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);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
+ output_b = create_tensor(tn(LLM_TENSOR_OUTPUT, "bias"), {n_vocab}, TENSOR_NOT_REQUIRED);
// if output is NULL, init from the input tok embed
if (output == NULL) {
@@ -140,6 +141,9 @@ llama_model_gpt2::graph::graph(const llama_model & model, const llm_graph_params
res->t_embd = cur;
cur = build_lora_mm(model.output, cur, model.output_s);
+ if (model.output_b) {
+ cur = ggml_add(ctx0, cur, model.output_b);
+ }
cb(cur, "result_output", -1);
res->t_logits = cur;
diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt
index e514de60b..1b178111d 100644
--- a/tools/mtmd/CMakeLists.txt
+++ b/tools/mtmd/CMakeLists.txt
@@ -28,6 +28,8 @@ add_library(mtmd
clip-model.h
clip-graph.h
models/models.h
+ models/chatterbox-gen.cpp
+ models/chatterbox-spkenc.cpp
models/cogvlm.cpp
models/conformer.cpp
models/deepseek4v.cpp
diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h
index 52d5fffe3..dcd6ec216 100644
--- a/tools/mtmd/clip-impl.h
+++ b/tools/mtmd/clip-impl.h
@@ -599,6 +599,8 @@ enum projector_type {
PROJECTOR_TYPE_SENSENOVA_U1,
PROJECTOR_TYPE_COHERE2V,
PROJECTOR_TYPE_MUSE_SPARK_AUDIO,
+ PROJECTOR_TYPE_CHATTERBOX,
+ PROJECTOR_TYPE_CHATTERBOX_SPKENC,
PROJECTOR_TYPE_UNKNOWN,
};
@@ -681,6 +683,8 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
{ PROJECTOR_TYPE_SENSENOVA_U1, "sensenova_u1"},
{ PROJECTOR_TYPE_COHERE2V, "cohere2v"},
{ PROJECTOR_TYPE_MUSE_SPARK_AUDIO, "muse-spark-audio"},
+ { PROJECTOR_TYPE_CHATTERBOX, "chatterbox"},
+ { PROJECTOR_TYPE_CHATTERBOX_SPKENC, "chatterbox_spkenc"},
};
static projector_type clip_projector_type_from_string(const std::string & str) {
diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h
index c42f417c8..b5a4a1225 100644
--- a/tools/mtmd/clip-model.h
+++ b/tools/mtmd/clip-model.h
@@ -6,6 +6,7 @@
#include <algorithm>
#include <array>
+#include <map>
#include <vector>
#include <unordered_set>
#include <cstdint>
@@ -611,6 +612,235 @@ struct clip_code2wav {
ggml_tensor * dac_post_conv_b = nullptr;
};
+struct clip_chatterbox {
+ // espnet rel-pos conformer layer of the flow encoder
+ struct enc_layer {
+ ggml_tensor * norm_mha_w = nullptr;
+ ggml_tensor * norm_mha_b = nullptr;
+ ggml_tensor * attn_q_w = nullptr;
+ ggml_tensor * attn_q_b = nullptr;
+ ggml_tensor * attn_k_w = nullptr;
+ ggml_tensor * attn_k_b = nullptr;
+ ggml_tensor * attn_v_w = nullptr;
+ ggml_tensor * attn_v_b = nullptr;
+ ggml_tensor * attn_pos_w = nullptr; // linear_pos, no bias
+ ggml_tensor * attn_pos_bias_u = nullptr;
+ ggml_tensor * attn_pos_bias_v = nullptr;
+ ggml_tensor * attn_out_w = nullptr;
+ ggml_tensor * attn_out_b = nullptr;
+ ggml_tensor * norm_ff_w = nullptr;
+ ggml_tensor * norm_ff_b = nullptr;
+ ggml_tensor * ffn_1_w = nullptr;
+ ggml_tensor * ffn_1_b = nullptr;
+ ggml_tensor * ffn_2_w = nullptr;
+ ggml_tensor * ffn_2_b = nullptr;
+ };
+
+ // causal conv block of the estimator (conv k3 left padded + layer norm)
+ struct causal_block {
+ ggml_tensor * conv_w = nullptr;
+ ggml_tensor * conv_b = nullptr;
+ ggml_tensor * norm_w = nullptr;
+ ggml_tensor * norm_b = nullptr;
+ };
+
+ // time conditioned resnet of the estimator
+ struct resnet {
+ causal_block block1;
+ causal_block block2;
+ ggml_tensor * mlp_w = nullptr; // time projection (mlp.1)
+ ggml_tensor * mlp_b = nullptr;
+ ggml_tensor * res_conv_w = nullptr;
+ ggml_tensor * res_conv_b = nullptr;
+ };
+
+ // diffusers style transformer block of the estimator
+ struct tfm_block {
+ ggml_tensor * norm1_w = nullptr;
+ ggml_tensor * norm1_b = nullptr;
+ ggml_tensor * attn_q_w = nullptr; // no bias on qkv
+ ggml_tensor * attn_k_w = nullptr;
+ ggml_tensor * attn_v_w = nullptr;
+ ggml_tensor * attn_out_w = nullptr; // to_out.0
+ ggml_tensor * attn_out_b = nullptr;
+ ggml_tensor * norm3_w = nullptr;
+ ggml_tensor * norm3_b = nullptr;
+ ggml_tensor * ff_in_w = nullptr; // ff.net.0.proj
+ ggml_tensor * ff_in_b = nullptr;
+ ggml_tensor * ff_out_w = nullptr; // ff.net.2
+ ggml_tensor * ff_out_b = nullptr;
+ };
+
+ // one estimator stage: resnet, transformer stack, boundary causal conv (null on mid stages)
+ struct est_stage {
+ resnet res;
+ std::vector<tfm_block> tfm;
+ ggml_tensor * conv_w = nullptr;
+ ggml_tensor * conv_b = nullptr;
+ };
+
+ // snake resblock unit of the hift vocoder
+ struct hift_res_unit {
+ ggml_tensor * act1_alpha = nullptr;
+ ggml_tensor * act2_alpha = nullptr;
+ ggml_tensor * conv1_w = nullptr;
+ ggml_tensor * conv1_b = nullptr;
+ ggml_tensor * conv2_w = nullptr;
+ ggml_tensor * conv2_b = nullptr;
+ };
+ struct hift_res {
+ std::vector<hift_res_unit> units; // dilations 1/3/5 on conv1
+ };
+
+ // one hift upsample stage: conv transpose, source injection, 3 resblocks
+ struct hift_up {
+ ggml_tensor * up_w = nullptr;
+ ggml_tensor * up_b = nullptr;
+ ggml_tensor * source_down_w = nullptr;
+ ggml_tensor * source_down_b = nullptr;
+ hift_res source_res;
+ hift_res res_0;
+ hift_res res_1;
+ hift_res res_2;
+ };
+
+ // f0 predictor conv of the hift vocoder
+ struct f0_conv {
+ ggml_tensor * w = nullptr;
+ ggml_tensor * b = nullptr;
+ };
+
+ // s3 tokenizer attention block (neox rope on q/k, fsmn memory on v)
+ struct s3tok_block {
+ ggml_tensor * attn_ln_w = nullptr;
+ ggml_tensor * attn_ln_b = nullptr;
+ ggml_tensor * attn_q_w = nullptr;
+ ggml_tensor * attn_q_b = nullptr;
+ ggml_tensor * attn_k_w = nullptr; // no bias
+ ggml_tensor * attn_v_w = nullptr;
+ ggml_tensor * attn_v_b = nullptr;
+ ggml_tensor * fsmn_w = nullptr;
+ ggml_tensor * attn_out_w = nullptr;
+ ggml_tensor * attn_out_b = nullptr;
+ ggml_tensor * mlp_ln_w = nullptr;
+ ggml_tensor * mlp_ln_b = nullptr;
+ ggml_tensor * mlp_in_w = nullptr;
+ ggml_tensor * mlp_in_b = nullptr;
+ ggml_tensor * mlp_out_w = nullptr;
+ ggml_tensor * mlp_out_b = nullptr;
+ };
+
+ // folded batchnorm (w/b stay null on the affine free variant)
+ struct bn {
+ ggml_tensor * mean = nullptr;
+ ggml_tensor * var = nullptr;
+ ggml_tensor * w = nullptr;
+ ggml_tensor * b = nullptr;
+ };
+
+ // fcm residual 2d block of the speaker encoder, the shortcut is null on stride 1 blocks
+ struct spk_res2d {
+ ggml_tensor * conv1_w = nullptr;
+ bn bn1;
+ ggml_tensor * conv2_w = nullptr;
+ bn bn2;
+ ggml_tensor * shortcut_w = nullptr;
+ bn shortcut_bn;
+ };
+
+ // cam dense tdnn layer of the speaker encoder
+ struct spk_cam_layer {
+ bn nl1_bn;
+ ggml_tensor * linear1_w = nullptr;
+ bn nl2_bn;
+ ggml_tensor * local_w = nullptr; // cam_layer.linear_local
+ ggml_tensor * ctx1_w = nullptr; // cam_layer.linear1
+ ggml_tensor * ctx1_b = nullptr;
+ ggml_tensor * ctx2_w = nullptr; // cam_layer.linear2
+ ggml_tensor * ctx2_b = nullptr;
+ };
+ struct spk_cam_block {
+ std::vector<spk_cam_layer> layers;
+ bn transit_bn;
+ ggml_tensor * transit_w = nullptr;
+ };
+
+ // flow encoder
+ ggml_tensor * input_embedding_w = nullptr; // flow.input_embedding
+ ggml_tensor * embed_linear_w = nullptr; // fenc.embed.out.0
+ ggml_tensor * embed_linear_b = nullptr;
+ ggml_tensor * embed_norm_w = nullptr; // fenc.embed.out.1
+ ggml_tensor * embed_norm_b = nullptr;
+ ggml_tensor * pre_conv1_w = nullptr; // pre_lookahead_layer
+ ggml_tensor * pre_conv1_b = nullptr;
+ ggml_tensor * pre_conv2_w = nullptr;
+ ggml_tensor * pre_conv2_b = nullptr;
+ std::vector<enc_layer> enc;
+ ggml_tensor * up_conv_w = nullptr; // fenc.up_layer.conv
+ ggml_tensor * up_conv_b = nullptr;
+ ggml_tensor * up_embed_linear_w = nullptr; // fenc.up_embed.out.0
+ ggml_tensor * up_embed_linear_b = nullptr;
+ ggml_tensor * up_embed_norm_w = nullptr; // fenc.up_embed.out.1
+ ggml_tensor * up_embed_norm_b = nullptr;
+ std::vector<enc_layer> up_enc;
+ ggml_tensor * after_norm_w = nullptr;
+ ggml_tensor * after_norm_b = nullptr;
+ ggml_tensor * encoder_proj_w = nullptr; // flow.encoder_proj
+ ggml_tensor * encoder_proj_b = nullptr;
+
+ // cfm estimator
+ ggml_tensor * time_mlp_1_w = nullptr;
+ ggml_tensor * time_mlp_1_b = nullptr;
+ ggml_tensor * time_mlp_2_w = nullptr;
+ ggml_tensor * time_mlp_2_b = nullptr;
+ ggml_tensor * time_embed_mixer_w = nullptr; // meanflow variant only
+ est_stage est_down;
+ std::vector<est_stage> est_mid;
+ est_stage est_up;
+ causal_block est_final_block;
+ ggml_tensor * est_final_proj_w = nullptr;
+ ggml_tensor * est_final_proj_b = nullptr;
+
+ // hift vocoder
+ ggml_tensor * hift_pre_w = nullptr;
+ ggml_tensor * hift_pre_b = nullptr;
+ std::vector<hift_up> hift_ups;
+ ggml_tensor * hift_post_w = nullptr;
+ ggml_tensor * hift_post_b = nullptr;
+ std::vector<f0_conv> f0_condnet;
+ ggml_tensor * f0_classifier_w = nullptr;
+ ggml_tensor * f0_classifier_b = nullptr;
+
+ // s3 tokenizer
+ ggml_tensor * s3tok_conv1_w = nullptr;
+ ggml_tensor * s3tok_conv1_b = nullptr;
+ ggml_tensor * s3tok_conv2_w = nullptr;
+ ggml_tensor * s3tok_conv2_b = nullptr;
+ std::vector<s3tok_block> s3tok_blocks;
+ ggml_tensor * s3tok_down_w = nullptr; // quantizer._codebook.project_down
+ ggml_tensor * s3tok_down_b = nullptr;
+
+ // CAMPPlus speaker encoder
+ ggml_tensor * spk_conv1_w = nullptr;
+ bn spk_bn1;
+ spk_res2d spk_layer1_0;
+ spk_res2d spk_layer1_1;
+ spk_res2d spk_layer2_0;
+ spk_res2d spk_layer2_1;
+ ggml_tensor * spk_conv2_w = nullptr;
+ bn spk_bn2;
+ ggml_tensor * spk_tdnn_w = nullptr;
+ bn spk_tdnn_bn;
+ spk_cam_block spk_block1;
+ spk_cam_block spk_block2;
+ spk_cam_block spk_block3;
+ bn spk_out_bn;
+ ggml_tensor * spk_dense_w = nullptr;
+ bn spk_dense_bn;
+ ggml_tensor * spk_affine_w = nullptr;
+ ggml_tensor * spk_affine_b = nullptr;
+};
+
struct clip_vocos {
ggml_tensor * input_w = nullptr;
ggml_tensor * input_b = nullptr;
@@ -896,8 +1126,13 @@ struct clip_model {
ggml_tensor * gen_upsample_w = nullptr; // depthwise convtr, frame rate -> encoder frame rate
std::vector<clip_layer> gen_tfm_layers; // mimi decoder_transformer
+ // chatterbox flow encoder, cfm estimator, hift vocoder, s3 tokenizer and CAMPPlus speaker encoder
+ clip_chatterbox cbx;
+
+ // chatterbox tensors read on the host, see clip_cbx_read_tensor()
+ std::map<std::string, ggml_tensor *> cbx_tensors;
// predefined voices, no rows if the voice is named in the prompt (lfm3a)
- std::map<std::string, std::vector<float>> gen_voices;
+ std::map<std::string, clip_gen_voice> gen_voices;
// soprano: Vocos decoder, hidden states -> STFT coefficients
clip_vocos vocos;
diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp
index f2c5c31b8..e6d931e69 100644
--- a/tools/mtmd/clip.cpp
+++ b/tools/mtmd/clip.cpp
@@ -1,3 +1,4 @@
+#define _USE_MATH_DEFINES // for M_PI
#include "clip.h"
#include "clip-impl.h"
#include "clip-model.h"
@@ -1200,6 +1201,24 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
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_CHATTERBOX_SPKENC:
+ {
+ builder = std::make_unique<clip_graph_chatterbox_spkenc>(ctx, img);
+ } break;
+ case PROJECTOR_TYPE_CHATTERBOX:
+ {
+ const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
+ int n_tokens = params && params->codes ? (int) params->codes->size() : 0;
+ int n_prompt_mel = 0;
+ if (n_tokens > 0) {
+ GGML_ASSERT(params->ref);
+ n_tokens += (int) params->ref->tokens.size();
+ n_prompt_mel = (int) (params->ref->feat.size() / 80);
+ }
+ const int vnm = params ? params->vocode_n_mel : 0;
+ const int vns = params ? params->vocode_n_stft : 0;
+ builder = std::make_unique<clip_graph_chatterbox>(ctx, img, gen_process, n_tokens, n_prompt_mel, vnm, vns);
+ } break;
case PROJECTOR_TYPE_LFM3A_GEN:
{
const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
@@ -2184,6 +2203,14 @@ struct clip_model_loader {
hparams.audio_window_len = 1024;
hparams.audio_hop_len = 256;
} break;
+ case PROJECTOR_TYPE_CHATTERBOX_SPKENC:
+ {
+ // kaldi fbank front-end of CAMPPlus: 25 ms / 10 ms framing, 512-point spectrum
+ hparams.audio_sample_rate = 16000;
+ hparams.audio_n_fft = 512;
+ hparams.audio_window_len = 400;
+ hparams.audio_hop_len = 160;
+ } break;
case PROJECTOR_TYPE_LFM3A_GEN:
{
// fixed by LiquidCodecConfig, head counts come from the tensors
@@ -2246,6 +2273,11 @@ struct clip_model_loader {
hparams.flow_n_step = 1;
hparams.gen_eos_threshold = -4.0f;
} break;
+ case PROJECTOR_TYPE_CHATTERBOX:
+ {
+ // s3gen output rate
+ hparams.audio_sample_rate = 24000;
+ } break;
case PROJECTOR_TYPE_SOPRANO:
{
// SopranoDecoder defaults, n_fft comes from the head tensor
@@ -2649,7 +2681,9 @@ struct clip_model_loader {
model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC &&
model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN &&
model.proj_type != PROJECTOR_TYPE_SOPRANO &&
- model.proj_type != PROJECTOR_TYPE_MUSE_SPARK_AUDIO);
+ model.proj_type != PROJECTOR_TYPE_MUSE_SPARK_AUDIO &&
+ model.proj_type != PROJECTOR_TYPE_CHATTERBOX &&
+ model.proj_type != PROJECTOR_TYPE_CHATTERBOX_SPKENC);
// layers
const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0;
@@ -3537,7 +3571,7 @@ struct clip_model_loader {
if (t->type != GGML_TYPE_F32 || ggml_n_dims(t) > 2 || t->ne[0] != model.gen_input_lin_w->ne[1]) {
throw std::runtime_error(string_format("%s: invalid voice tensor %s\n", __func__, name));
}
- auto & embd = model.gen_voices[name + 6];
+ auto & embd = model.gen_voices[name + 6].rows;
embd.resize(ggml_nelements(t));
fin.seekg(tensor_offset.at(name), std::ios::beg);
if (!fin.read(reinterpret_cast<char *>(embd.data()), ggml_nbytes(t))) {
@@ -3708,6 +3742,326 @@ struct clip_model_loader {
c2w.dac_post_conv_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_CONV, "bias"));
}
} break;
+ case PROJECTOR_TYPE_CHATTERBOX:
+ {
+ auto & c = model.cbx;
+ auto has = [&](const std::string & name) {
+ return gguf_find_tensor(ctx_gguf.get(), name.c_str()) >= 0;
+ };
+ auto load_enc_layer = [&](const std::string & p, clip_chatterbox::enc_layer & l) {
+ l.norm_mha_w = get_tensor(p + ".norm_mha.weight");
+ l.norm_mha_b = get_tensor(p + ".norm_mha.bias");
+ l.attn_q_w = get_tensor(p + ".self_attn.linear_q.weight");
+ l.attn_q_b = get_tensor(p + ".self_attn.linear_q.bias");
+ l.attn_k_w = get_tensor(p + ".self_attn.linear_k.weight");
+ l.attn_k_b = get_tensor(p + ".self_attn.linear_k.bias");
+ l.attn_v_w = get_tensor(p + ".self_attn.linear_v.weight");
+ l.attn_v_b = get_tensor(p + ".self_attn.linear_v.bias");
+ l.attn_pos_w = get_tensor(p + ".self_attn.linear_pos.weight");
+ l.attn_pos_bias_u = get_tensor(p + ".self_attn.pos_bias_u");
+ l.attn_pos_bias_v = get_tensor(p + ".self_attn.pos_bias_v");
+ l.attn_out_w = get_tensor(p + ".self_attn.linear_out.weight");
+ l.attn_out_b = get_tensor(p + ".self_attn.linear_out.bias");
+ l.norm_ff_w = get_tensor(p + ".norm_ff.weight");
+ l.norm_ff_b = get_tensor(p + ".norm_ff.bias");
+ l.ffn_1_w = get_tensor(p + ".feed_forward.w_1.weight");
+ l.ffn_1_b = get_tensor(p + ".feed_forward.w_1.bias");
+ l.ffn_2_w = get_tensor(p + ".feed_forward.w_2.weight");
+ l.ffn_2_b = get_tensor(p + ".feed_forward.w_2.bias");
+ };
+ auto load_causal = [&](const std::string & p, clip_chatterbox::causal_block & b) {
+ b.conv_w = get_tensor(p + ".block.0.weight");
+ b.conv_b = get_tensor(p + ".block.0.bias");
+ b.norm_w = get_tensor(p + ".block.2.weight");
+ b.norm_b = get_tensor(p + ".block.2.bias");
+ };
+ auto load_resnet = [&](const std::string & p, clip_chatterbox::resnet & r) {
+ load_causal(p + ".block1", r.block1);
+ load_causal(p + ".block2", r.block2);
+ r.mlp_w = get_tensor(p + ".mlp.1.weight");
+ r.mlp_b = get_tensor(p + ".mlp.1.bias");
+ r.res_conv_w = get_tensor(p + ".res_conv.weight");
+ r.res_conv_b = get_tensor(p + ".res_conv.bias");
+ };
+ auto load_tfm = [&](const std::string & p, clip_chatterbox::tfm_block & b) {
+ b.norm1_w = get_tensor(p + ".norm1.weight");
+ b.norm1_b = get_tensor(p + ".norm1.bias");
+ b.attn_q_w = get_tensor(p + ".attn1.to_q.weight");
+ b.attn_k_w = get_tensor(p + ".attn1.to_k.weight");
+ b.attn_v_w = get_tensor(p + ".attn1.to_v.weight");
+ b.attn_out_w = get_tensor(p + ".attn1.to_out.0.weight");
+ b.attn_out_b = get_tensor(p + ".attn1.to_out.0.bias");
+ b.norm3_w = get_tensor(p + ".norm3.weight");
+ b.norm3_b = get_tensor(p + ".norm3.bias");
+ b.ff_in_w = get_tensor(p + ".ff.net.0.proj.weight");
+ b.ff_in_b = get_tensor(p + ".ff.net.0.proj.bias");
+ b.ff_out_w = get_tensor(p + ".ff.net.2.weight");
+ b.ff_out_b = get_tensor(p + ".ff.net.2.bias");
+ };
+ auto load_stage = [&](const std::string & p, clip_chatterbox::est_stage & s, bool boundary) {
+ load_resnet(p + ".0", s.res);
+ for (int j = 0; has(p + ".1." + std::to_string(j) + ".norm1.weight"); j++) {
+ clip_chatterbox::tfm_block b;
+ load_tfm(p + ".1." + std::to_string(j), b);
+ s.tfm.push_back(b);
+ }
+ if (boundary) {
+ s.conv_w = get_tensor(p + ".2.weight");
+ s.conv_b = get_tensor(p + ".2.bias");
+ }
+ };
+ auto load_hres = [&](const std::string & p, clip_chatterbox::hift_res & r) {
+ for (int j = 0; has(p + ".convs1." + std::to_string(j) + ".weight"); j++) {
+ const std::string js = std::to_string(j);
+ clip_chatterbox::hift_res_unit u;
+ u.act1_alpha = get_tensor(p + ".activations1." + js + ".alpha");
+ u.act2_alpha = get_tensor(p + ".activations2." + js + ".alpha");
+ u.conv1_w = get_tensor(p + ".convs1." + js + ".weight");
+ u.conv1_b = get_tensor(p + ".convs1." + js + ".bias");
+ u.conv2_w = get_tensor(p + ".convs2." + js + ".weight");
+ u.conv2_b = get_tensor(p + ".convs2." + js + ".bias");
+ r.units.push_back(u);
+ }
+ };
+
+ // flow encoder
+ c.input_embedding_w = get_tensor("a.gen.flow.input_embedding.weight");
+ c.spk_affine_w = get_tensor("a.gen.flow.spk_embed_affine_layer.weight");
+ c.spk_affine_b = get_tensor("a.gen.flow.spk_embed_affine_layer.bias");
+ c.embed_linear_w = get_tensor("a.gen.fenc.embed.out.0.weight");
+ c.embed_linear_b = get_tensor("a.gen.fenc.embed.out.0.bias");
+ c.embed_norm_w = get_tensor("a.gen.fenc.embed.out.1.weight");
+ c.embed_norm_b = get_tensor("a.gen.fenc.embed.out.1.bias");
+ c.pre_conv1_w = get_tensor("a.gen.fenc.pre_lookahead_layer.conv1.weight");
+ c.pre_conv1_b = get_tensor("a.gen.fenc.pre_lookahead_layer.conv1.bias");
+ c.pre_conv2_w = get_tensor("a.gen.fenc.pre_lookahead_layer.conv2.weight");
+ c.pre_conv2_b = get_tensor("a.gen.fenc.pre_lookahead_layer.conv2.bias");
+ for (int i = 0; has("a.gen.fenc.encoders." + std::to_string(i) + ".norm_mha.weight"); i++) {
+ clip_chatterbox::enc_layer l;
+ load_enc_layer("a.gen.fenc.encoders." + std::to_string(i), l);
+ c.enc.push_back(l);
+ }
+ c.up_conv_w = get_tensor("a.gen.fenc.up_layer.conv.weight");
+ c.up_conv_b = get_tensor("a.gen.fenc.up_layer.conv.bias");
+ c.up_embed_linear_w = get_tensor("a.gen.fenc.up_embed.out.0.weight");
+ c.up_embed_linear_b = get_tensor("a.gen.fenc.up_embed.out.0.bias");
+ c.up_embed_norm_w = get_tensor("a.gen.fenc.up_embed.out.1.weight");
+ c.up_embed_norm_b = get_tensor("a.gen.fenc.up_embed.out.1.bias");
+ for (int i = 0; has("a.gen.fenc.up_encoders." + std::to_string(i) + ".norm_mha.weight"); i++) {
+ clip_chatterbox::enc_layer l;
+ load_enc_layer("a.gen.fenc.up_encoders." + std::to_string(i), l);
+ c.up_enc.push_back(l);
+ }
+ c.after_norm_w = get_tensor("a.gen.fenc.after_norm.weight");
+ c.after_norm_b = get_tensor("a.gen.fenc.after_norm.bias");
+ c.encoder_proj_w = get_tensor("a.gen.flow.encoder_proj.weight");
+ c.encoder_proj_b = get_tensor("a.gen.flow.encoder_proj.bias");
+
+ // cfm estimator
+ c.time_mlp_1_w = get_tensor("a.gen.est.time_mlp.linear_1.weight");
+ c.time_mlp_1_b = get_tensor("a.gen.est.time_mlp.linear_1.bias");
+ c.time_mlp_2_w = get_tensor("a.gen.est.time_mlp.linear_2.weight");
+ c.time_mlp_2_b = get_tensor("a.gen.est.time_mlp.linear_2.bias");
+ c.time_embed_mixer_w = get_tensor("a.gen.est.time_embed_mixer.weight", false);
+ load_stage("a.gen.est.down_blocks.0", c.est_down, true);
+ for (int i = 0; has("a.gen.est.mid_blocks." + std::to_string(i) + ".0.block1.block.0.weight"); i++) {
+ clip_chatterbox::est_stage s;
+ load_stage("a.gen.est.mid_blocks." + std::to_string(i), s, false);
+ c.est_mid.push_back(std::move(s));
+ }
+ load_stage("a.gen.est.up_blocks.0", c.est_up, true);
+ load_causal("a.gen.est.final_block", c.est_final_block);
+ c.est_final_proj_w = get_tensor("a.gen.est.final_proj.weight");
+ c.est_final_proj_b = get_tensor("a.gen.est.final_proj.bias");
+
+ // hift vocoder
+ c.hift_pre_w = get_tensor("a.gen.hift.conv_pre.weight");
+ c.hift_pre_b = get_tensor("a.gen.hift.conv_pre.bias");
+ c.hift_post_w = get_tensor("a.gen.hift.conv_post.weight");
+ c.hift_post_b = get_tensor("a.gen.hift.conv_post.bias");
+ for (int i = 0; has("a.gen.hift.ups." + std::to_string(i) + ".weight"); i++) {
+ const std::string is = std::to_string(i);
+ clip_chatterbox::hift_up up;
+ up.up_w = get_tensor("a.gen.hift.ups." + is + ".weight");
+ up.up_b = get_tensor("a.gen.hift.ups." + is + ".bias");
+ up.source_down_w = get_tensor("a.gen.hift.source_downs." + is + ".weight");
+ up.source_down_b = get_tensor("a.gen.hift.source_downs." + is + ".bias");
+ load_hres("a.gen.hift.source_resblocks." + is, up.source_res);
+ load_hres("a.gen.hift.resblocks." + std::to_string(3 * i), up.res_0);
+ load_hres("a.gen.hift.resblocks." + std::to_string(3 * i + 1), up.res_1);
+ load_hres("a.gen.hift.resblocks." + std::to_string(3 * i + 2), up.res_2);
+ c.hift_ups.push_back(std::move(up));
+ }
+ for (int i = 0; has("a.gen.hift.f0_predictor.condnet." + std::to_string(i) + ".weight"); i += 2) {
+ const std::string is = std::to_string(i);
+ clip_chatterbox::f0_conv fc;
+ fc.w = get_tensor("a.gen.hift.f0_predictor.condnet." + is + ".weight");
+ fc.b = get_tensor("a.gen.hift.f0_predictor.condnet." + is + ".bias");
+ c.f0_condnet.push_back(fc);
+ }
+ c.f0_classifier_w = get_tensor("a.gen.hift.f0_predictor.classifier.weight");
+ c.f0_classifier_b = get_tensor("a.gen.hift.f0_predictor.classifier.bias");
+
+ // s3 tokenizer
+ c.s3tok_conv1_w = get_tensor("a.s3tok.encoder.conv1.weight");
+ c.s3tok_conv1_b = get_tensor("a.s3tok.encoder.conv1.bias");
+ c.s3tok_conv2_w = get_tensor("a.s3tok.encoder.conv2.weight");
+ c.s3tok_conv2_b = get_tensor("a.s3tok.encoder.conv2.bias");
+ for (int i = 0; has("a.s3tok.encoder.blocks." + std::to_string(i) + ".attn_ln.weight"); i++) {
+ const std::string p = "a.s3tok.encoder.blocks." + std::to_string(i);
+ clip_chatterbox::s3tok_block b;
+ b.attn_ln_w = get_tensor(p + ".attn_ln.weight");
+ b.attn_ln_b = get_tensor(p + ".attn_ln.bias");
+ b.attn_q_w = get_tensor(p + ".attn.query.weight");
+ b.attn_q_b = get_tensor(p + ".attn.query.bias");
+ b.attn_k_w = get_tensor(p + ".attn.key.weight");
+ b.attn_v_w = get_tensor(p + ".attn.value.weight");
+ b.attn_v_b = get_tensor(p + ".attn.value.bias");
+ b.fsmn_w = get_tensor(p + ".attn.fsmn_block.weight");
+ b.attn_out_w = get_tensor(p + ".attn.out.weight");
+ b.attn_out_b = get_tensor(p + ".attn.out.bias");
+ b.mlp_ln_w = get_tensor(p + ".mlp_ln.weight");
+ b.mlp_ln_b = get_tensor(p + ".mlp_ln.bias");
+ b.mlp_in_w = get_tensor(p + ".mlp.0.weight");
+ b.mlp_in_b = get_tensor(p + ".mlp.0.bias");
+ b.mlp_out_w = get_tensor(p + ".mlp.2.weight");
+ b.mlp_out_b = get_tensor(p + ".mlp.2.bias");
+ c.s3tok_blocks.push_back(b);
+ }
+ c.s3tok_down_w = get_tensor("a.s3tok.quantizer._codebook.project_down.weight");
+ c.s3tok_down_b = get_tensor("a.s3tok.quantizer._codebook.project_down.bias");
+
+ // read on the host: embedding tables, source module
+ for (ggml_tensor * t = ggml_get_first_tensor(ctx_meta.get()); t; t = ggml_get_next_tensor(ctx_meta.get(), t)) {
+ const std::string name = t->name;
+ if (name.rfind("a.gen.code.", 0) == 0 || name.rfind("a.gen.t3.", 0) == 0 ||
+ name.rfind("a.gen.hift.m_source.", 0) == 0) {
+ model.cbx_tensors[name] = get_tensor(name);
+ }
+ }
+
+ // voice.<name> holds the talker rows, voice.<name>.ref_* the flow decoder reference
+ const int64_t n_embd_talker = model.cbx_tensors.at("a.gen.code.out_embd.weight")->ne[0];
+ for (int64_t i = 0; i < gguf_get_n_tensors(ctx_gguf.get()); ++i) {
+ const std::string name = gguf_get_tensor_name(ctx_gguf.get(), i);
+ if (name.rfind("voice.", 0) != 0) {
+ continue;
+ }
+ const ggml_tensor * t = ggml_get_tensor(ctx_meta.get(), name.c_str());
+ std::string voice = name.substr(6);
+ std::vector<float> ids; // the converter stores 1D tensors as F32, ids are exact
+ void * dst = nullptr;
+ bool ok = false;
+ auto take = [&](const char * suffix) {
+ const size_t n = strlen(suffix);
+ if (voice.size() > n && voice.compare(voice.size() - n, n, suffix) == 0) {
+ voice.resize(voice.size() - n);
+ return true;
+ }
+ return false;
+ };
+ if (take(".ref_tokens")) {
+ ids.resize(ggml_nelements(t));
+ dst = ids.data();
+ ok = t->type == GGML_TYPE_F32;
+ } else if (take(".ref_feat")) {
+ auto & v = model.gen_voices[voice].ref.feat;
+ v.resize(ggml_nelements(t));
+ dst = v.data();
+ ok = t->type == GGML_TYPE_F32 && t->ne[0] == 80;
+ } else if (take(".ref_spk")) {
+ auto & v = model.gen_voices[voice].ref.spk;
+ v.resize(ggml_nelements(t));
+ dst = v.data();
+ ok = t->type == GGML_TYPE_F32 && t->ne[0] == 192;
+ } else {
+ auto & v = model.gen_voices[voice].rows;
+ v.resize(ggml_nelements(t));
+ dst = v.data();
+ ok = t->type == GGML_TYPE_F32 && ggml_n_dims(t) <= 2 && t->ne[0] == n_embd_talker;
+ }
+ if (!ok) {
+ throw std::runtime_error(string_format("%s: invalid voice tensor %s\n", __func__, name.c_str()));
+ }
+ fin.seekg(tensor_offset.at(name), std::ios::beg);
+ if (!fin.read(reinterpret_cast<char *>(dst), ggml_nbytes(t))) {
+ throw std::runtime_error(string_format("%s: failed to read %s\n", __func__, name.c_str()));
+ }
+ if (!ids.empty()) {
+ model.gen_voices[voice].ref.tokens.assign(ids.begin(), ids.end());
+ }
+ }
+ for (const auto & [voice, v] : model.gen_voices) {
+ if (v.rows.empty() || v.ref.tokens.empty() || v.ref.feat.size() != 2 * 80 * v.ref.tokens.size() || v.ref.spk.empty()) {
+ throw std::runtime_error(string_format("%s: incomplete voice %s\n", __func__, voice.c_str()));
+ }
+ }
+ } break;
+ case PROJECTOR_TYPE_CHATTERBOX_SPKENC:
+ {
+ auto & c = model.cbx;
+ auto has = [&](const std::string & name) {
+ return gguf_find_tensor(ctx_gguf.get(), name.c_str()) >= 0;
+ };
+ auto load_bn = [&](const std::string & p, clip_chatterbox::bn & n) {
+ n.mean = get_tensor(p + ".running_mean");
+ n.var = get_tensor(p + ".running_var");
+ n.w = get_tensor(p + ".weight", false);
+ n.b = get_tensor(p + ".bias", false);
+ };
+ auto load_res2d = [&](const std::string & p, clip_chatterbox::spk_res2d & r) {
+ r.conv1_w = get_tensor(p + ".conv1.weight");
+ load_bn(p + ".bn1", r.bn1);
+ r.conv2_w = get_tensor(p + ".conv2.weight");
+ load_bn(p + ".bn2", r.bn2);
+ r.shortcut_w = get_tensor(p + ".shortcut.0.weight", false);
+ if (r.shortcut_w) {
+ load_bn(p + ".shortcut.1", r.shortcut_bn);
+ }
+ };
+ auto load_cam_block = [&](const std::string & bp, const std::string & tp, clip_chatterbox::spk_cam_block & blk) {
+ for (int li = 1; has(bp + ".tdnnd" + std::to_string(li) + ".linear1.weight"); li++) {
+ const std::string p = bp + ".tdnnd" + std::to_string(li);
+ clip_chatterbox::spk_cam_layer l;
+ load_bn(p + ".nonlinear1.batchnorm", l.nl1_bn);
+ l.linear1_w = get_tensor(p + ".linear1.weight");
+ load_bn(p + ".nonlinear2.batchnorm", l.nl2_bn);
+ l.local_w = get_tensor(p + ".cam_layer.linear_local.weight");
+ l.ctx1_w = get_tensor(p + ".cam_layer.linear1.weight");
+ l.ctx1_b = get_tensor(p + ".cam_layer.linear1.bias");
+ l.ctx2_w = get_tensor(p + ".cam_layer.linear2.weight");
+ l.ctx2_b = get_tensor(p + ".cam_layer.linear2.bias");
+ blk.layers.push_back(l);
+ }
+ load_bn(tp + ".nonlinear.batchnorm", blk.transit_bn);
+ blk.transit_w = get_tensor(tp + ".linear.weight");
+ };
+
+ c.spk_conv1_w = get_tensor("a.spk.head.conv1.weight");
+ load_bn("a.spk.head.bn1", c.spk_bn1);
+ load_res2d("a.spk.head.layer1.0", c.spk_layer1_0);
+ load_res2d("a.spk.head.layer1.1", c.spk_layer1_1);
+ load_res2d("a.spk.head.layer2.0", c.spk_layer2_0);
+ load_res2d("a.spk.head.layer2.1", c.spk_layer2_1);
+ c.spk_conv2_w = get_tensor("a.spk.head.conv2.weight");
+ load_bn("a.spk.head.bn2", c.spk_bn2);
+ c.spk_tdnn_w = get_tensor("a.spk.xvector.tdnn.linear.weight");
+ load_bn("a.spk.xvector.tdnn.nonlinear.batchnorm", c.spk_tdnn_bn);
+ load_cam_block("a.spk.xvector.block1", "a.spk.xvector.transit1", c.spk_block1);
+ load_cam_block("a.spk.xvector.block2", "a.spk.xvector.transit2", c.spk_block2);
+ load_cam_block("a.spk.xvector.block3", "a.spk.xvector.transit3", c.spk_block3);
+ load_bn("a.spk.xvector.out_nonlinear.batchnorm", c.spk_out_bn);
+ c.spk_dense_w = get_tensor("a.spk.xvector.dense.linear.weight");
+ load_bn("a.spk.xvector.dense.nonlinear.batchnorm", c.spk_dense_bn);
+
+ // read on the host: voice encoder lstm, conditioning encoder
+ for (ggml_tensor * t = ggml_get_first_tensor(ctx_meta.get()); t; t = ggml_get_next_tensor(ctx_meta.get(), t)) {
+ const std::string name = t->name;
+ if (name.rfind("a.ve.", 0) == 0 || name.rfind("a.cenc.", 0) == 0) {
+ model.cbx_tensors[name] = get_tensor(name);
+ }
+ }
+ } break;
case PROJECTOR_TYPE_SOPRANO:
{
auto & v = model.vocos;
@@ -4741,8 +5095,19 @@ struct clip_init_result clip_init(const char * fname, struct clip_context_params
ctx_gen_audio->flash_attn_type = CLIP_FLASH_ATTN_TYPE_DISABLED;
}
loader.load_tensors(*ctx_gen_audio);
+ if (ctx_gen_audio->model.proj_type == PROJECTOR_TYPE_CHATTERBOX) {
+ // the classic cfm solver unrolls up to 10 steps x 2 cfg estimator evaluations in one graph
+ ctx_gen_audio->max_nodes = 65536;
+ ctx_gen_audio->sched.reset(
+ ggml_backend_sched_new(ctx_gen_audio->backend_ptrs.data(), ctx_gen_audio->backend_buft.data(),
+ ctx_gen_audio->backend_ptrs.size(), ctx_gen_audio->max_nodes, false, true));
+ if (ctx_params.cb_eval != nullptr) {
+ ggml_backend_sched_set_eval_callback(ctx_gen_audio->sched.get(), ctx_params.cb_eval, ctx_params.cb_eval_user_data);
+ }
+ }
// TODO: fix warmup
- ctx_gen_audio->buf_compute_meta.resize(ctx_gen_audio->max_nodes * ggml_tensor_overhead() + ggml_graph_overhead());
+ ctx_gen_audio->buf_compute_meta.resize(ctx_gen_audio->max_nodes * ggml_tensor_overhead() +
+ ggml_graph_overhead_custom(ctx_gen_audio->max_nodes, false));
}
} catch (const std::exception & e) {
@@ -5148,6 +5513,11 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
// pooling gives one speaker embedding, whatever the clip length is
n_patches = 1;
} break;
+ case PROJECTOR_TYPE_CHATTERBOX_SPKENC:
+ {
+ // statistics pooling gives one speaker embedding, whatever the clip length is
+ n_patches = 1;
+ } break;
case PROJECTOR_TYPE_QWEN3TTS_GEN:
case PROJECTOR_TYPE_LFM3A_GEN:
{
@@ -5276,7 +5646,7 @@ void clip_gen_state_reset(clip_gen_state * state) {
}
}
-const std::map<std::string, std::vector<float>> & clip_gen_get_voices(const clip_ctx * ctx) {
+const std::map<std::string, clip_gen_voice> & clip_gen_get_voices(const clip_ctx * ctx) {
return ctx->model.gen_voices;
}
@@ -5354,6 +5724,32 @@ void clip_enc_state_free(clip_enc_state * state) {
delete state;
}
+size_t clip_cbx_read_tensor(struct clip_ctx * ctx, const char * name, float * out, size_t n_max) {
+ const auto & tensors = ctx->model.cbx_tensors;
+ auto it = tensors.find(name);
+ if (it == tensors.end() || !it->second) {
+ return 0;
+ }
+ ggml_tensor * t = it->second;
+ const size_t n = (size_t) ggml_nelements(t);
+ if (!out) {
+ return n;
+ }
+ if (n_max < n) {
+ return 0;
+ }
+ if (t->type == GGML_TYPE_F32) {
+ ggml_backend_tensor_get(t, out, 0, n * sizeof(float));
+ } else if (t->type == GGML_TYPE_F16) {
+ std::vector<ggml_fp16_t> tmp(n);
+ ggml_backend_tensor_get(t, tmp.data(), 0, n * sizeof(ggml_fp16_t));
+ for (size_t i = 0; i < n; i++) out[i] = ggml_fp16_to_fp32(tmp[i]);
+ } else {
+ return 0;
+ }
+ return n;
+}
+
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();
@@ -5369,7 +5765,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
clip_model_loader::warmup(*ctx, *params->imgs);
}
- if (params->gen_process != CLIP_GEN_PROCESS_GEN_UNKNOWN && params->gen_state == nullptr) {
+ const bool needs_gen_state = params->gen_process == CLIP_GEN_PROCESS_GEN_CODE ||
+ params->gen_process == CLIP_GEN_PROCESS_GEN_WAV ||
+ params->gen_process == CLIP_GEN_PROCESS_TTS;
+ if (needs_gen_state && params->gen_state == nullptr) {
LOG_ERR("%s: audio gen requires a gen state\n", __func__);
return false;
}
@@ -5543,8 +5942,11 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
}
set_input_f32("inp_raw", inp_raw);
- } else if (params->gen_process != CLIP_GEN_PROCESS_GEN_WAV && ggml_graph_get_tensor(gf, "inp_raw") != nullptr) {
- // audio input. GEN_WAV is not here: it takes codes or feats, set in the switch below.
+ } else if (params->gen_process != CLIP_GEN_PROCESS_GEN_WAV &&
+ params->gen_process != CLIP_GEN_PROCESS_TTS &&
+ params->gen_process != CLIP_GEN_PROCESS_TTS_VOCODE &&
+ ggml_graph_get_tensor(gf, "inp_raw") != nullptr) {
+ // audio input. GEN_WAV and the chatterbox TTS stages are not here: they take codes or feats, set in the switch below.
// a GEN_CODE graph may have no hidden-state input (lfm3a_gen feedback embedding)
GGML_ASSERT(imgs.entries.size() == 1);
@@ -6337,6 +6739,84 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
set_input_attn_mask(qwen2_mask);
}
} break;
+ case PROJECTOR_TYPE_CHATTERBOX:
+ {
+ if (params->gen_process == CLIP_GEN_PROCESS_TTS_VOCODE) {
+ set_input_f32("inp_mel", *params->mel_in);
+ set_input_f32("inp_sstft", *params->sstft_in);
+ break;
+ }
+ if (params->gen_process == CLIP_GEN_PROCESS_TOKENIZE) {
+ const int T1 = (imgs.entries[0].nx() - 1) / 2 + 1;
+ const int T2 = (T1 - 1) / 2 + 1;
+ std::vector<int32_t> pos(T2);
+ for (int i = 0; i < T2; i++) {
+ pos[(size_t) i] = i;
+ }
+ set_input_i32("inp_pos", pos);
+ break;
+ }
+ if (params->gen_process != CLIP_GEN_PROCESS_TTS) {
+ break;
+ }
+ const int n_gen = (int) params->codes->size();
+ GGML_ASSERT(params->ref);
+ const int n_prompt = (int) params->ref->tokens.size();
+ std::vector<int32_t> tokens = params->ref->tokens;
+ tokens.resize((size_t) n_prompt + n_gen);
+ memcpy(tokens.data() + n_prompt, params->codes->data(), (size_t) n_gen * sizeof(int32_t));
+ const int64_t n_codes = model.cbx.input_embedding_w->ne[1];
+ for (int32_t t : tokens) {
+ if (t < 0 || t >= n_codes) {
+ LOG_ERR("%s: speech token out of range (%d, codebook size %lld)\n", __func__, t, (long long) n_codes);
+ return false;
+ }
+ }
+ const int T1 = n_prompt + n_gen;
+ const int T2 = 2 * T1;
+ set_input_i32("inp_tokens", tokens);
+
+ // mel-rate reference conditioning of the voice
+ set_input_f32("inp_prompt_feat", params->ref->feat);
+ set_input_f32("inp_xvec", params->ref->spk);
+
+ // espnet relative positional encoding, entry k holds the sinusoid of relative position (T-1) - k
+ auto fill_pos = [&](const char * name, int T) {
+ const int d = 512;
+ std::vector<float> pos((size_t) (2 * T - 1) * d);
+ for (int k = 0; k < 2 * T - 1; k++) {
+ const double rel = (double) (T - 1 - k);
+ for (int i = 0; i < d / 2; i++) {
+ const double div = exp(-(double) (2 * i) * log(10000.0) / d);
+ pos[(size_t) k * d + 2 * i ] = (float) sin(rel * div);
+ pos[(size_t) k * d + 2 * i + 1] = (float) cos(rel * div);
+ }
+ }
+ set_input_f32(name, pos);
+ };
+ fill_pos("inp_pos1", T1);
+ fill_pos("inp_pos2", T2);
+
+ // noise and the time embeddings of the solver span points, matcha layout (sines then cosines, scale 1000)
+ std::vector<float> noise((size_t) 80 * T2);
+ std::normal_distribution<float> nd(0.0f, 1.0f);
+ for (auto & f : noise) f = nd(params->gen_state->rng);
+ set_input_f32("inp_noise", noise);
+
+ const bool meanflow = model.cbx.time_embed_mixer_w != nullptr;
+ const int n_steps = meanflow ? 2 : 10;
+ std::vector<float> temb((size_t) 320 * (n_steps + 1));
+ for (int s = 0; s <= n_steps; s++) {
+ const double u = (double) s / n_steps;
+ const double t = meanflow ? u : 1.0 - cos(u * M_PI / 2.0);
+ for (int i = 0; i < 160; i++) {
+ const double div = exp(-(double) i * log(10000.0) / 159.0);
+ temb[(size_t) s * 320 + i ] = (float) sin(1000.0 * t * div);
+ temb[(size_t) s * 320 + i + 160] = (float) cos(1000.0 * t * div);
+ }
+ }
+ set_input_f32("inp_temb", temb);
+ } break;
case PROJECTOR_TYPE_GEMMA3:
case PROJECTOR_TYPE_GEMMA3NV:
case PROJECTOR_TYPE_IDEFICS3:
@@ -6363,6 +6843,16 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
{
// do nothing
} break;
+ case PROJECTOR_TYPE_CHATTERBOX_SPKENC:
+ {
+ // the ceil-mode seg pooling divides by the full segment length, rescale the last partial segment
+ const int T1 = (imgs.entries[0].nx() - 1) / 2 + 1;
+ const int S = (T1 + 99) / 100;
+ std::vector<float> segfix((size_t) S, 1.0f);
+ const int last = T1 - (S - 1) * 100;
+ segfix[(size_t) S - 1] = 100.0f / (float) last;
+ set_input_f32("inp_segfix", segfix);
+ } break;
case PROJECTOR_TYPE_LFM3A_GEN:
{
const int64_t n_cb = model.gen_code_embd_w->ne[2];
@@ -6966,6 +7456,210 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
// for audio gen models
//
+ if (ctx->proj_type() == PROJECTOR_TYPE_CHATTERBOX && params->gen_process == CLIP_GEN_PROCESS_TOKENIZE) {
+ ggml_tensor * fsq = ggml_graph_get_tensor(gf, "out_fsq");
+ GGML_ASSERT(fsq != nullptr && params->out_codes);
+ const int n_tok = (int) fsq->ne[1];
+ std::vector<float> h((size_t) 8 * n_tok);
+ ggml_backend_tensor_get(fsq, h.data(), 0, ggml_nbytes(fsq));
+
+ // fsq round to base 3: h in (-1, 1) maps to digits {0, 1, 2}
+ params->out_codes->resize(n_tok);
+ for (int t = 0; t < n_tok; t++) {
+ int32_t code = 0;
+ for (int i = 7; i >= 0; i--) {
+ const int32_t d = (int32_t) roundf(h[(size_t) t * 8 + i] * 0.9990000128746033f) + 1;
+ code = code * 3 + d;
+ }
+ (*params->out_codes)[(size_t) t] = code;
+ }
+
+ // speech table rows of the codes, the multilingual variant adds its learned positions as its perceiver expects
+ if (params->out_code_embd) {
+ const auto & tensors = ctx->model.cbx_tensors;
+ auto tab_it = tensors.find("a.gen.code.out_embd.weight");
+ GGML_ASSERT(tab_it != tensors.end());
+ ggml_tensor * tab = tab_it->second;
+ const int n_e = (int) tab->ne[0];
+ auto pos_it = tensors.find("a.gen.t3.speech_pos_emb");
+ ggml_tensor * pos = pos_it == tensors.end() ? nullptr : pos_it->second;
+ if ((tab->type != GGML_TYPE_F16 && tab->type != GGML_TYPE_F32) ||
+ (pos && (pos->type != GGML_TYPE_F32 || pos->ne[0] != n_e || pos->ne[1] < n_tok))) {
+ LOG_ERR("%s: unexpected speech table or positions\n", __func__);
+ return false;
+ }
+
+ params->out_code_embd->resize((size_t) n_tok * n_e);
+ std::vector<ggml_fp16_t> h16(n_e);
+ std::vector<float> pr(n_e);
+ for (int t = 0; t < n_tok; t++) {
+ float * dst = params->out_code_embd->data() + (size_t) t * n_e;
+ if ((*params->out_codes)[(size_t) t] >= tab->ne[1]) {
+ LOG_ERR("%s: speech code beyond the speech table\n", __func__);
+ return false;
+ }
+ const size_t r = (size_t) (*params->out_codes)[(size_t) t] * n_e;
+ if (tab->type == GGML_TYPE_F16) {
+ ggml_backend_tensor_get(tab, h16.data(), r * sizeof(ggml_fp16_t), (size_t) n_e * sizeof(ggml_fp16_t));
+ for (int j = 0; j < n_e; j++) {
+ dst[j] = ggml_fp16_to_fp32(h16[j]);
+ }
+ } else {
+ ggml_backend_tensor_get(tab, dst, r * sizeof(float), (size_t) n_e * sizeof(float));
+ }
+ if (pos) {
+ ggml_backend_tensor_get(pos, pr.data(), (size_t) t * n_e * sizeof(float), (size_t) n_e * sizeof(float));
+ for (int j = 0; j < n_e; j++) {
+ dst[j] += pr[j];
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ if (ctx->proj_type() == PROJECTOR_TYPE_CHATTERBOX && params->gen_process == CLIP_GEN_PROCESS_TTS_VOCODE) {
+ ggml_tensor * sp = ggml_graph_get_tensor(gf, "out_spec");
+ GGML_ASSERT(sp != nullptr && params->out_spec);
+ params->out_spec->resize(ggml_nelements(sp));
+ ggml_backend_tensor_get(sp, params->out_spec->data(), 0, ggml_nbytes(sp));
+ return true;
+ }
+
+ if (ctx->proj_type() == PROJECTOR_TYPE_CHATTERBOX && params->gen_process == CLIP_GEN_PROCESS_TTS) {
+ ggml_tensor * mel = ggml_graph_get_tensor(gf, "out_mel");
+ GGML_ASSERT(mel != nullptr);
+ std::vector<float> mel_data(ggml_nelements(mel));
+ ggml_backend_tensor_get(mel, mel_data.data(), 0, ggml_nbytes(mel));
+
+ // hift: f0 -> harmonic source -> source stft on the host, then the vocoder graph and the istft
+ ggml_tensor * f0_t = ggml_graph_get_tensor(gf, "out_f0");
+ GGML_ASSERT(f0_t != nullptr);
+ const int n_mel_out = (int) ggml_nelements(f0_t);
+ std::vector<float> f0(n_mel_out);
+ ggml_backend_tensor_get(f0_t, f0.data(), 0, ggml_nbytes(f0_t));
+
+ // SineGen: f0 upsampled x480, 9 harmonics with cumulative phase, uv gating and noise, merged by l_linear + tanh
+ const int ups_total = 480;
+ const double sr = 24000.0;
+ const int64_t n_wav = (int64_t) n_mel_out * ups_total;
+ std::vector<float> lw(9);
+ float lb = 0.0f;
+ if (clip_cbx_read_tensor(ctx, "a.gen.hift.m_source.l_linear.weight", lw.data(), lw.size()) != 9 ||
+ clip_cbx_read_tensor(ctx, "a.gen.hift.m_source.l_linear.bias", &lb, 1) != 1) {
+ LOG_ERR("%s: unexpected hift source module\n", __func__);
+ return false;
+ }
+
+ std::mt19937 & srng = params->gen_state->rng;
+ std::uniform_real_distribution<double> ud(-M_PI, M_PI);
+ std::normal_distribution<float> snd(0.0f, 1.0f);
+ double phase[9];
+ for (int h = 0; h < 9; h++) {
+ phase[h] = h == 0 ? 0.0 : ud(srng);
+ }
+ std::vector<float> src((size_t) n_wav);
+ double cum[9] = {0.0};
+ for (int64_t t = 0; t < n_wav; t++) {
+ const float f = f0[(size_t) (t / ups_total)];
+ const float uv = f > 10.0f ? 1.0f : 0.0f; // nsf_voiced_threshold
+ const float namp = uv * 0.003f + (1.0f - uv) * 0.1f / 3.0f;
+ float merged = lb;
+ for (int h = 0; h < 9; h++) {
+ cum[h] += (double) f * (h + 1) / sr;
+ cum[h] -= floor(cum[h]);
+ float sine = 0.1f * (float) sin(2.0 * M_PI * cum[h] + phase[h]);
+ sine = sine * uv + namp * snd(srng);
+ merged += lw[(size_t) h] * sine;
+ }
+ src[(size_t) t] = tanhf(merged);
+ }
+
+ // stft of the source: n_fft 16, hop 4, hann window, centered
+ const int n_fft = 16, hop = 4, n_bins = 9;
+ std::vector<float> win(n_fft);
+ for (int i = 0; i < n_fft; i++) win[(size_t) i] = 0.5f - 0.5f * cosf(2.0f * (float) M_PI * i / n_fft);
+ const int n_stft = (int) (n_wav / hop) + 1;
+ std::vector<float> sstft((size_t) n_stft * 18);
+ for (int fr = 0; fr < n_stft; fr++) {
+ const int64_t c0 = (int64_t) fr * hop - n_fft / 2; // centered, reflect padded
+ for (int k = 0; k < n_bins; k++) {
+ double re = 0.0, im = 0.0;
+ for (int i = 0; i < n_fft; i++) {
+ int64_t idx = c0 + i;
+ if (idx < 0) idx = -idx;
+ if (idx >= n_wav) idx = 2 * (n_wav - 1) - idx;
+ const double v = (double) src[(size_t) idx] * win[(size_t) i];
+ const double a = 2.0 * M_PI * k * i / n_fft;
+ re += v * cos(a);
+ im -= v * sin(a);
+ }
+ sstft[(size_t) fr * 18 + k ] = (float) re;
+ sstft[(size_t) fr * 18 + 9 + k] = (float) im;
+ }
+ }
+
+ // vocoder graph on mel + source stft
+ std::vector<float> spec;
+ {
+ clip_encode_params vp = *params;
+ vp.gen_process = CLIP_GEN_PROCESS_TTS_VOCODE;
+ vp.mel_in = &mel_data;
+ vp.sstft_in = &sstft;
+ vp.vocode_n_mel = n_mel_out;
+ vp.vocode_n_stft = n_stft;
+ vp.out_audio = nullptr;
+ vp.out_spec = &spec;
+ if (!clip_encode(ctx, &vp)) {
+ LOG_ERR("%s: vocoder stage failed\n", __func__);
+ return false;
+ }
+ }
+
+ // istft: mag = clipped exp, phase = sin, hann overlap-add
+ const int n_frames_out = (int) (spec.size() / 18);
+ const int64_t n_out = (int64_t) (n_frames_out - 1) * hop;
+ std::vector<double> acc((size_t) n_out + n_fft, 0.0);
+ std::vector<double> wsum((size_t) n_out + n_fft, 0.0);
+ for (int fr = 0; fr < n_frames_out; fr++) {
+ double frame[16];
+ for (int i = 0; i < n_fft; i++) {
+ double v = 0.0;
+ for (int k = 0; k < n_bins; k++) {
+ const double mag = fmin(exp((double) spec[(size_t) fr * 18 + k]), 1e2);
+ const double ph = sin((double) spec[(size_t) fr * 18 + 9 + k]);
+ const double re = mag * cos(ph), im = mag * sin(ph);
+ const double a = 2.0 * M_PI * k * i / n_fft;
+ const double w = (k == 0 || k == n_fft / 2) ? 1.0 : 2.0;
+ v += w * (re * cos(a) - im * sin(a));
+ }
+ frame[i] = v / n_fft;
+ }
+ const int64_t o = (int64_t) fr * hop;
+ for (int i = 0; i < n_fft; i++) {
+ acc [(size_t) (o + i)] += frame[i] * win[(size_t) i];
+ wsum[(size_t) (o + i)] += (double) win[(size_t) i] * win[(size_t) i];
+ }
+ }
+ GGML_ASSERT(params->out_audio);
+ auto & out_audio = *params->out_audio;
+ out_audio.resize((size_t) std::min<int64_t>(n_out - n_fft / 2, n_wav));
+ for (size_t i = 0; i < out_audio.size(); i++) {
+ const size_t j = i + n_fft / 2; // drop the centering pad
+ const double v = wsum[j] > 1e-11 ? acc[j] / wsum[j] : 0.0;
+ out_audio[(size_t) i] = (float) fmax(-0.99, fmin(0.99, v));
+ }
+
+ // trim_fade of the reference: silence 20 ms then fade in 20 ms, hides the artifact of the prompt boundary
+ const size_t n_trim = 24000 / 50;
+ for (size_t i = 0; i < 2 * n_trim && i < out_audio.size(); i++) {
+ const double g = i < n_trim ? 0.0
+ : (cos(M_PI * (1.0 - (double) (i - n_trim) / n_trim)) + 1.0) / 2.0;
+ out_audio[(size_t) i] *= (float) g;
+ }
+ return true;
+ }
+
// 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");
@@ -7235,6 +7929,12 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
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_CHATTERBOX:
+ // talker width, carried by the speech table
+ return ctx->model.cbx_tensors.at("a.gen.code.out_embd.weight")->ne[0];
+ case PROJECTOR_TYPE_CHATTERBOX_SPKENC:
+ // talker conditioning rows, the speaker projection carries their width
+ return ctx->model.cbx_tensors.at("a.cenc.spkr_enc.weight")->ne[1];
case PROJECTOR_TYPE_SOPRANO:
return ctx->model.hparams.projection_dim;
case PROJECTOR_TYPE_PARAKEET:
diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h
index b2c9e0759..6c8dcf31c 100644
--- a/tools/mtmd/clip.h
+++ b/tools/mtmd/clip.h
@@ -96,7 +96,18 @@ enum clip_gen_process_type {
CLIP_GEN_PROCESS_GEN_UNKNOWN,
CLIP_GEN_PROCESS_GEN_CODE, // h_state to codes
CLIP_GEN_PROCESS_GEN_WAV, // codes to raw PCM audio
+ CLIP_GEN_PROCESS_TTS, // full utterance of codes to raw PCM audio
+ CLIP_GEN_PROCESS_TTS_VOCODE, // internal: mel + source stft to istft input
+ CLIP_GEN_PROCESS_TOKENIZE, // raw PCM audio to semantic speech tokens
};
+
+// gen-side voice reference (chatterbox): speech tokens, mel-rate features and the speaker vector
+struct clip_gen_ref {
+ std::vector<int32_t> tokens;
+ std::vector<float> feat;
+ std::vector<float> spk;
+};
+
struct clip_encode_params {
int n_threads = 1;
const clip_image_f32_batch * imgs = nullptr;
@@ -114,10 +125,20 @@ struct clip_encode_params {
float temp = 0.0f; // sampling temperature, noise scale for flow-matching decoders
bool * out_is_eos = nullptr;
+ // TOKENIZE: speech embedding rows of the produced codes
+ std::vector<float> * out_code_embd = 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 clip_gen_ref * ref = nullptr; // TTS reference conditioning of the voice
+ // TTS_VOCODE internal stage inputs
+ const std::vector<float> * mel_in = nullptr;
+ const std::vector<float> * sstft_in = nullptr;
+ int vocode_n_mel = 0;
+ int vocode_n_stft = 0;
+ std::vector<float> * out_spec = nullptr;
// audio gen: state of one sequence, clip_encode() updates it in place
struct clip_gen_state * gen_state = nullptr;
@@ -139,7 +160,16 @@ struct clip_gen_state * clip_gen_state_init(struct clip_ctx * ctx, uint32_t seed
void clip_gen_state_free(struct clip_gen_state * state);
void clip_gen_state_reset(struct clip_gen_state * state);
-const std::map<std::string, std::vector<float>> & clip_gen_get_voices(const struct clip_ctx * ctx);
+// read a chatterbox tensor by name as F32, returns the element count (0 if not found), out may be null to query the size
+size_t clip_cbx_read_tensor(struct clip_ctx * ctx, const char * name, float * out, size_t n_max);
+
+// a predefined voice: backbone rows, plus the gen-side reference when the gen stack needs one
+struct clip_gen_voice {
+ std::vector<float> rows;
+ clip_gen_ref ref;
+};
+
+const std::map<std::string, clip_gen_voice> & clip_gen_get_voices(const struct clip_ctx * ctx);
bool clip_is_llava(const struct clip_ctx * ctx);
// note for contributor: this clip_is_(model) pattern is deprecated
diff --git a/tools/mtmd/models/chatterbox-gen.cpp b/tools/mtmd/models/chatterbox-gen.cpp
new file mode 100644
index 000000000..d4c4f77bf
--- /dev/null
+++ b/tools/mtmd/models/chatterbox-gen.cpp
@@ -0,0 +1,542 @@
+#define _USE_MATH_DEFINES // for M_PI
+#include "models.h"
+
+// Chatterbox generation graphs: flow encoder, cfm estimator, s3 tokenizer and hift vocoder
+
+// x [C, T]: y = W x + b with torch Linear weights stored as [in, out]
+ggml_tensor * clip_graph_chatterbox_base::cbx_linear(ggml_tensor * w, ggml_tensor * b, ggml_tensor * x) const {
+ ggml_tensor * y = ggml_mul_mat(ctx0, w, x);
+ if (b) {
+ y = ggml_add(ctx0, y, b);
+ }
+ return y;
+}
+
+// conv1d over time with asymmetric zero padding: x [C, T] -> [OC, T_out], kernel [K, IC, OC]
+ggml_tensor * clip_graph_chatterbox_base::cbx_conv1d(ggml_tensor * k, ggml_tensor * b, ggml_tensor * x,
+ int stride, int pad_l, int pad_r) const {
+ ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [T, C]
+ if (pad_l > 0) {
+ ggml_tensor * z = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, pad_l, xt->ne[1]);
+ z = ggml_scale(ctx0, z, 0.0f);
+ xt = ggml_concat(ctx0, z, xt, 0);
+ }
+ if (pad_r > 0) {
+ ggml_tensor * z = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, pad_r, xt->ne[1]);
+ z = ggml_scale(ctx0, z, 0.0f);
+ xt = ggml_concat(ctx0, xt, z, 0);
+ }
+ ggml_tensor * y = ggml_conv_1d(ctx0, k, xt, stride, 0, 1); // [T_out, OC]
+ y = ggml_cont(ctx0, ggml_transpose(ctx0, y)); // [OC, T_out]
+ if (b) {
+ y = ggml_add(ctx0, y, b);
+ }
+ return y;
+}
+
+// x [C, T] -> symmetric-padded dilated conv -> [OC, T]
+ggml_tensor * clip_graph_chatterbox_base::cbx_conv1d_dil(ggml_tensor * k, ggml_tensor * b, ggml_tensor * x,
+ int pad, int dil) const {
+ ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x));
+ ggml_tensor * y = ggml_conv_1d(ctx0, k, xt, 1, pad, dil);
+ y = ggml_cont(ctx0, ggml_transpose(ctx0, y));
+ if (b) {
+ y = ggml_add(ctx0, y, b);
+ }
+ return y;
+}
+
+// Transformer-XL relative shift: bd [2T-1, T, H] -> [T, T, H] with out[j, i, h] = bd[(T-1) - i + j, i, h]
+// same buffer walk as the espnet rel_shift: pad one column, drop the first row, keep the first T columns
+static ggml_tensor * cbx_rel_shift(ggml_context * ctx0, ggml_tensor * bd, int T) {
+ const int H = (int) bd->ne[2];
+ ggml_tensor * z = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, T, H);
+ z = ggml_scale(ctx0, z, 0.0f);
+ ggml_tensor * p = ggml_cont(ctx0, ggml_concat(ctx0, z, bd, 0)); // [2T, T, H]
+ p = ggml_reshape_3d(ctx0, p, T, 2 * T, H); // [T, 2T, H]
+ p = ggml_view_3d(ctx0, p, T, 2 * T - 1, H, p->nb[1], p->nb[2], p->nb[1]); // drop first row
+ p = ggml_cont(ctx0, p);
+ p = ggml_reshape_3d(ctx0, p, 2 * T - 1, T, H); // [2T-1, T, H]
+ p = ggml_view_3d(ctx0, p, T, T, H, p->nb[1], p->nb[2], 0); // first T columns
+ return ggml_cont(ctx0, p);
+}
+
+// espnet rel-pos self attention block, pre-norm, x [512, T], pos [512, 2T-1]
+ggml_tensor * clip_graph_chatterbox::enc_layer(const clip_chatterbox::enc_layer & l, ggml_tensor * x,
+ ggml_tensor * pos, int T) {
+ const int n_head = 8;
+ const int d_head = 64;
+ const float scale = 1.0f / sqrtf((float) d_head);
+
+ ggml_tensor * res = x;
+ ggml_tensor * cur = build_norm(x, l.norm_mha_w, l.norm_mha_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+
+ ggml_tensor * q = cbx_linear(l.attn_q_w, l.attn_q_b, cur);
+ ggml_tensor * k = cbx_linear(l.attn_k_w, l.attn_k_b, cur);
+ ggml_tensor * v = cbx_linear(l.attn_v_w, l.attn_v_b, cur);
+ ggml_tensor * pe = cbx_linear(l.attn_pos_w, nullptr, pos); // [512, 2T-1]
+
+ q = ggml_reshape_3d(ctx0, q, d_head, n_head, T);
+ k = ggml_reshape_3d(ctx0, k, d_head, n_head, T);
+ v = ggml_reshape_3d(ctx0, v, d_head, n_head, T);
+ pe = ggml_reshape_3d(ctx0, pe, d_head, n_head, 2 * T - 1);
+
+ ggml_tensor * u = l.attn_pos_bias_u; // [64, 8]
+ ggml_tensor * w = l.attn_pos_bias_v;
+
+ ggml_tensor * qu = ggml_add(ctx0, q, ggml_reshape_3d(ctx0, u, d_head, n_head, 1));
+ ggml_tensor * qv = ggml_add(ctx0, q, ggml_reshape_3d(ctx0, w, d_head, n_head, 1));
+
+ // per head: [64, T] tensors, scores [T(k), T(q)]
+ qu = ggml_cont(ctx0, ggml_permute(ctx0, qu, 0, 2, 1, 3)); // [64, T, 8]
+ qv = ggml_cont(ctx0, ggml_permute(ctx0, qv, 0, 2, 1, 3));
+ k = ggml_cont(ctx0, ggml_permute(ctx0, k, 0, 2, 1, 3));
+ v = ggml_cont(ctx0, ggml_permute(ctx0, v, 0, 2, 1, 3));
+ pe = ggml_cont(ctx0, ggml_permute(ctx0, pe, 0, 2, 1, 3)); // [64, 2T-1, 8]
+
+ ggml_tensor * ac = ggml_mul_mat(ctx0, k, qu); // [T(k), T(q), 8]
+ ggml_tensor * bd = ggml_mul_mat(ctx0, pe, qv); // [2T-1, T(q), 8]
+ bd = cbx_rel_shift(ctx0, bd, T); // [T(k), T(q), 8]
+
+ ggml_tensor * scores = ggml_scale(ctx0, ggml_add(ctx0, ac, bd), scale);
+ ggml_tensor * probs = ggml_soft_max(ctx0, scores);
+
+ ggml_tensor * o = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, v)), probs); // [64, T, 8]
+ o = ggml_cont(ctx0, ggml_permute(ctx0, o, 0, 2, 1, 3)); // [64, 8, T]
+ o = ggml_reshape_2d(ctx0, o, n_head * d_head, T);
+ o = cbx_linear(l.attn_out_w, l.attn_out_b, o);
+ x = ggml_add(ctx0, res, o);
+
+ res = x;
+ cur = build_norm(x, l.norm_ff_w, l.norm_ff_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+ cur = cbx_linear(l.ffn_1_w, l.ffn_1_b, cur);
+ cur = ggml_silu(ctx0, cur); // swish
+ cur = cbx_linear(l.ffn_2_w, l.ffn_2_b, cur);
+ x = ggml_add(ctx0, res, cur);
+ return x;
+}
+
+// mish = x * tanh(softplus(x))
+static ggml_tensor * cbx_mish(ggml_context * ctx0, ggml_tensor * x) {
+ return ggml_mul(ctx0, x, ggml_tanh(ctx0, ggml_softplus(ctx0, x)));
+}
+
+// causal block: conv k3 left-padded, layer norm over channels, mish; x [C, T]
+ggml_tensor * clip_graph_chatterbox::causal_block(const clip_chatterbox::causal_block & b, ggml_tensor * x) {
+ x = cbx_conv1d(b.conv_w, b.conv_b, x, 1, (int) b.conv_w->ne[0] - 1, 0);
+ x = build_norm(x, b.norm_w, b.norm_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+ return cbx_mish(ctx0, x);
+}
+
+// resnet block with time conditioning; x [C, T], temb [1024]
+ggml_tensor * clip_graph_chatterbox::resnet(const clip_chatterbox::resnet & r, ggml_tensor * x, ggml_tensor * temb) {
+ ggml_tensor * h = causal_block(r.block1, x);
+ ggml_tensor * tproj = cbx_linear(r.mlp_w, r.mlp_b, cbx_mish(ctx0, temb));
+ h = ggml_add(ctx0, h, tproj); // broadcast [256, 1] over T
+ h = causal_block(r.block2, h);
+ ggml_tensor * res = cbx_conv1d(r.res_conv_w, r.res_conv_b, x, 1, 0, 0);
+ return ggml_add(ctx0, h, res);
+}
+
+// diffusers-style transformer block, full attention; x [256, T]
+ggml_tensor * clip_graph_chatterbox::tfm_block(const clip_chatterbox::tfm_block & b, ggml_tensor * x) {
+ const int n_head = 8;
+ const int d_head = 64;
+ const int T = (int) x->ne[1];
+
+ ggml_tensor * cur = build_norm(x, b.norm1_w, b.norm1_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+ ggml_tensor * q = ggml_reshape_3d(ctx0, ggml_mul_mat(ctx0, b.attn_q_w, cur), d_head, n_head, T);
+ ggml_tensor * k = ggml_reshape_3d(ctx0, ggml_mul_mat(ctx0, b.attn_k_w, cur), d_head, n_head, T);
+ ggml_tensor * v = ggml_reshape_3d(ctx0, ggml_mul_mat(ctx0, b.attn_v_w, cur), d_head, n_head, T);
+ x = ggml_add(ctx0, x, build_attn(b.attn_out_w, b.attn_out_b, q, k, v, nullptr, 1.0f / sqrtf((float) d_head), -1));
+
+ cur = build_norm(x, b.norm3_w, b.norm3_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+ cur = cbx_linear(b.ff_in_w, b.ff_in_b, cur);
+ cur = ggml_gelu_erf(ctx0, cur);
+ cur = cbx_linear(b.ff_out_w, b.ff_out_b, cur);
+ return ggml_add(ctx0, x, cur);
+}
+
+// one estimator evaluation; x_noise, mu, cond [80, T], spks [80], temb [1024]
+ggml_tensor * clip_graph_chatterbox::estimator(ggml_tensor * x_noise, ggml_tensor * mu, ggml_tensor * spks,
+ ggml_tensor * cond, ggml_tensor * temb, int T) {
+ const auto & c = model.cbx;
+
+ // channels live on ne0, time on ne1: pack along ne0
+ ggml_tensor * x = ggml_concat(ctx0, x_noise, mu, 0); // [160, T]
+ ggml_tensor * spks_b = ggml_repeat(ctx0, ggml_reshape_2d(ctx0, spks, 80, 1), ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 80, T));
+ x = ggml_concat(ctx0, x, spks_b, 0); // [240, T]
+ x = ggml_concat(ctx0, x, cond, 0); // [320, T]
+ x = ggml_cont(ctx0, x);
+
+ // down
+ ggml_tensor * skip;
+ x = resnet(c.est_down.res, x, temb);
+ for (const auto & b : c.est_down.tfm) {
+ x = tfm_block(b, x);
+ }
+ skip = x;
+ x = cbx_conv1d(c.est_down.conv_w, c.est_down.conv_b, x, 1, (int) c.est_down.conv_w->ne[0] - 1, 0);
+
+ // mid
+ for (const auto & m : c.est_mid) {
+ x = resnet(m.res, x, temb);
+ for (const auto & b : m.tfm) {
+ x = tfm_block(b, x);
+ }
+ }
+
+ // up with skip
+ x = ggml_concat(ctx0, x, skip, 0); // [512, T]
+ x = resnet(c.est_up.res, x, temb);
+ for (const auto & b : c.est_up.tfm) {
+ x = tfm_block(b, x);
+ }
+ x = cbx_conv1d(c.est_up.conv_w, c.est_up.conv_b, x, 1, (int) c.est_up.conv_w->ne[0] - 1, 0);
+
+ x = causal_block(c.est_final_block, x);
+ x = cbx_conv1d(c.est_final_proj_w, c.est_final_proj_b, x, 1, 0, 0); // [80, T]
+ return x;
+}
+
+// s3 tokenizer encoder: log-mel [T, 128] -> fsq code [8, T / 4], rounded to base 3 tokens on the host
+ggml_cgraph * clip_graph_chatterbox::build_s3tok(int T) {
+ const auto & c = model.cbx;
+ const int n_head = 20;
+ const int d_head = 64;
+ const int T1 = (T - 1) / 2 + 1;
+ const int T2 = (T1 - 1) / 2 + 1;
+
+ ggml_tensor * inp = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, T, 128);
+ ggml_set_name(inp, "inp_raw");
+ ggml_set_input(inp);
+
+ ggml_tensor * pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, T2);
+ ggml_set_name(pos, "inp_pos");
+ ggml_set_input(pos);
+
+ ggml_tensor * x = ggml_cont(ctx0, ggml_transpose(ctx0, inp)); // [128, T]
+ x = cbx_conv1d(c.s3tok_conv1_w, c.s3tok_conv1_b, x, 2, 1, 1);
+ x = ggml_gelu_erf(ctx0, x);
+ x = cbx_conv1d(c.s3tok_conv2_w, c.s3tok_conv2_b, x, 2, 1, 1);
+ x = ggml_gelu_erf(ctx0, x); // [1280, T2]
+
+ for (const auto & blk : c.s3tok_blocks) {
+ ggml_tensor * res = x;
+ ggml_tensor * cur = build_norm(x, blk.attn_ln_w, blk.attn_ln_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+ ggml_tensor * q = cbx_linear(blk.attn_q_w, blk.attn_q_b, cur);
+ ggml_tensor * k = ggml_mul_mat(ctx0, blk.attn_k_w, cur);
+ ggml_tensor * v = cbx_linear(blk.attn_v_w, blk.attn_v_b, cur);
+
+ // fsmn memory: depthwise conv over time on v plus residual, added to the attention output
+ ggml_tensor * fsm = ggml_cont(ctx0, ggml_transpose(ctx0, v)); // [T2, 1280]
+ {
+ ggml_tensor * w = blk.fsmn_w;
+ ggml_tensor * m = ggml_conv_1d_dw(ctx0, w, fsm, 1, ((int) w->ne[0] - 1) / 2, 1);
+ fsm = ggml_add(ctx0, ggml_reshape_2d(ctx0, m, fsm->ne[0], fsm->ne[1]), fsm);
+ }
+ fsm = ggml_cont(ctx0, ggml_transpose(ctx0, fsm)); // [1280, T2]
+
+ q = ggml_reshape_3d(ctx0, q, d_head, n_head, T2);
+ k = ggml_reshape_3d(ctx0, k, d_head, n_head, T2);
+ v = ggml_reshape_3d(ctx0, v, d_head, n_head, T2);
+ q = ggml_rope_ext(ctx0, q, pos, nullptr, d_head, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+ k = ggml_rope_ext(ctx0, k, pos, nullptr, d_head, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+ ggml_tensor * o = build_attn(blk.attn_out_w, blk.attn_out_b, q, k, v, nullptr, 1.0f / sqrtf((float) d_head), -1);
+ x = ggml_add(ctx0, res, ggml_add(ctx0, o, fsm));
+
+ res = x;
+ cur = build_norm(x, blk.mlp_ln_w, blk.mlp_ln_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+ cur = cbx_linear(blk.mlp_in_w, blk.mlp_in_b, cur);
+ cur = ggml_gelu_erf(ctx0, cur);
+ cur = cbx_linear(blk.mlp_out_w, blk.mlp_out_b, cur);
+ x = ggml_add(ctx0, res, cur);
+ }
+
+ x = cbx_linear(c.s3tok_down_w, c.s3tok_down_b, x); // [8, T2]
+ x = ggml_tanh(ctx0, x);
+
+ ggml_set_name(x, "out_fsq");
+ ggml_set_output(x);
+ ggml_build_forward_expand(gf, x);
+ return gf;
+}
+
+ggml_cgraph * clip_graph_chatterbox::build() {
+ if (gen_process == CLIP_GEN_PROCESS_TTS_VOCODE) {
+ return build_vocoder(vocode_n_mel, vocode_n_stft);
+ }
+ if (gen_process == CLIP_GEN_PROCESS_TOKENIZE) {
+ return build_s3tok(img.nx());
+ }
+ if (gen_process != CLIP_GEN_PROCESS_TTS) {
+ // load-time buffer sizing path
+ ggml_tensor * inp = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);
+ ggml_set_name(inp, "inp_stub");
+ ggml_set_input(inp);
+ ggml_tensor * cur = ggml_dup(ctx0, inp);
+ ggml_set_name(cur, "out_stub");
+ ggml_set_output(cur);
+ ggml_build_forward_expand(gf, cur);
+ return gf;
+ }
+
+ const auto & c = model.cbx;
+
+ const int T1 = n_tokens; // token-rate length (prompt + generated)
+ const int T2 = 2 * n_tokens; // mel-rate length after the x2 upsample
+
+ ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, T1);
+ ggml_set_name(inp_tokens, "inp_tokens");
+ ggml_set_input(inp_tokens);
+
+ ggml_tensor * pos1 = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 512, 2 * T1 - 1);
+ ggml_set_name(pos1, "inp_pos1");
+ ggml_set_input(pos1);
+
+ ggml_tensor * pos2 = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 512, 2 * T2 - 1);
+ ggml_set_name(pos2, "inp_pos2");
+ ggml_set_input(pos2);
+
+ // token embedding
+ ggml_tensor * x = ggml_get_rows(ctx0, c.input_embedding_w, inp_tokens); // [512, T1]
+
+ // embed: linear + layer norm, then the espnet xscale
+ x = cbx_linear(c.embed_linear_w, c.embed_linear_b, x);
+ x = build_norm(x, c.embed_norm_w, c.embed_norm_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+ x = ggml_scale(ctx0, x, sqrtf(512.0f));
+ cb(x, "fenc_embd", -1);
+
+ // pre-lookahead: conv k=4 right-padded 3, leaky 0.01, conv k=3 left-padded 2, residual
+ {
+ ggml_tensor * res = x;
+ ggml_tensor * cur = cbx_conv1d(c.pre_conv1_w, c.pre_conv1_b, x, 1, 0, 3);
+ cur = ggml_leaky_relu(ctx0, cur, 0.01f, false);
+ cur = cbx_conv1d(c.pre_conv2_w, c.pre_conv2_b, cur, 1, 2, 0);
+ x = ggml_add(ctx0, res, cur);
+ cb(x, "fenc_pre_lookahead", -1);
+ }
+
+ for (size_t i = 0; i < c.enc.size(); i++) {
+ x = enc_layer(c.enc[i], x, pos1, T1);
+ cb(x, "fenc_enc", (int) i);
+ }
+
+ // upsample x2: nearest repeat, left pad 4, conv k=5
+ {
+ ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [T1, 512]
+ xt = ggml_interpolate(ctx0, xt, 2 * T1, 512, 1, 1, GGML_SCALE_MODE_NEAREST);
+ ggml_tensor * z = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 4, 512);
+ z = ggml_scale(ctx0, z, 0.0f);
+ xt = ggml_concat(ctx0, z, xt, 0);
+ ggml_tensor * y = ggml_conv_1d(ctx0, c.up_conv_w, xt, 1, 0, 1);
+ x = ggml_cont(ctx0, ggml_transpose(ctx0, y)); // [512, T2]
+ x = ggml_add(ctx0, x, c.up_conv_b);
+ cb(x, "fenc_upsample", -1);
+ }
+
+ // up embed: linear + layer norm + xscale
+ x = cbx_linear(c.up_embed_linear_w, c.up_embed_linear_b, x);
+ x = build_norm(x, c.up_embed_norm_w, c.up_embed_norm_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+ x = ggml_scale(ctx0, x, sqrtf(512.0f));
+
+ for (size_t i = 0; i < c.up_enc.size(); i++) {
+ x = enc_layer(c.up_enc[i], x, pos2, T2);
+ cb(x, "fenc_up_enc", (int) i);
+ }
+
+ x = build_norm(x, c.after_norm_w, c.after_norm_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+
+ // encoder projection to the mel channel count
+ ggml_tensor * mu = cbx_linear(c.encoder_proj_w, c.encoder_proj_b, x); // [80, T2]
+ cb(mu, "flow_mu", -1);
+
+ // cfm solver unrolled in the graph: meanflow takes 2 euler steps without cfg, its time embeds mix t and r
+ // the classic solver takes 10 euler steps on the cosine schedule with cfg 0.7
+ const bool meanflow = c.time_embed_mixer_w != nullptr;
+ const int n_steps = meanflow ? 2 : 10;
+
+ // span points, same schedule as the host side sinusoid fill in clip.cpp
+ float span[11];
+ for (int i = 0; i <= n_steps; i++) {
+ const float u = (float) i / n_steps;
+ span[i] = meanflow ? u : 1.0f - cosf(u * (float) M_PI / 2.0f);
+ }
+
+ ggml_tensor * noise = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 80, T2);
+ ggml_set_name(noise, "inp_noise");
+ ggml_set_input(noise);
+ // sinusoidal time embeddings, one row per span point
+ ggml_tensor * temb_sin = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 320, n_steps + 1);
+ ggml_set_name(temb_sin, "inp_temb");
+ ggml_set_input(temb_sin);
+
+ auto time_mlp = [&](ggml_tensor * e) {
+ e = cbx_linear(c.time_mlp_1_w, c.time_mlp_1_b, e);
+ e = ggml_silu(ctx0, e);
+ e = cbx_linear(c.time_mlp_2_w, c.time_mlp_2_b, e);
+ return e;
+ };
+ auto span_emb = [&](int i) {
+ return ggml_view_2d(ctx0, temb_sin, 320, 1, temb_sin->nb[1], (size_t) i * temb_sin->nb[1]);
+ };
+ auto step_temb = [&](int i) {
+ if (!meanflow) {
+ return time_mlp(span_emb(i));
+ }
+ ggml_tensor * e = ggml_concat(ctx0, time_mlp(span_emb(i)), time_mlp(span_emb(i + 1)), 0); // [2048, 1]
+ return ggml_mul_mat(ctx0, c.time_embed_mixer_w, e); // [1024, 1]
+ };
+
+ // mel-rate conditions of the voice: prompt features then zeros, and the speaker vector
+ ggml_tensor * pf = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 80, n_prompt_mel);
+ ggml_set_name(pf, "inp_prompt_feat");
+ ggml_set_input(pf);
+ ggml_tensor * zc = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 80, T2 - n_prompt_mel);
+ zc = ggml_scale(ctx0, zc, 0.0f);
+ ggml_tensor * cond = ggml_concat(ctx0, pf, zc, 1); // [80, T2]
+ // raw 192-dim campplus x-vector, normalized then through the flow speaker affine
+ ggml_tensor * xvec = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 192);
+ ggml_set_name(xvec, "inp_xvec");
+ ggml_set_input(xvec);
+ ggml_tensor * n2 = ggml_sqrt(ctx0, ggml_sum(ctx0, ggml_mul(ctx0, xvec, xvec)));
+ ggml_tensor * unit = ggml_div(ctx0, xvec, n2);
+ ggml_tensor * spks = cbx_linear(c.spk_affine_w, c.spk_affine_b,
+ ggml_reshape_2d(ctx0, unit, 192, 1));
+ spks = ggml_cont(ctx0, ggml_reshape_1d(ctx0, spks, 80));
+
+ const float cfg = 0.7f;
+ ggml_tensor * mu_zero = meanflow ? nullptr : ggml_scale(ctx0, mu, 0.0f);
+ ggml_tensor * cond_zero = meanflow ? nullptr : ggml_scale(ctx0, cond, 0.0f);
+ ggml_tensor * spks_zero = meanflow ? nullptr : ggml_scale(ctx0, spks, 0.0f);
+
+ ggml_tensor * mx = noise;
+ for (int i = 0; i < n_steps; i++) {
+ ggml_tensor * temb = step_temb(i);
+ ggml_tensor * d = estimator(mx, mu, spks, cond, temb, T2);
+ if (!meanflow) {
+ ggml_tensor * du = estimator(mx, mu_zero, spks_zero, cond_zero, temb, T2);
+ d = ggml_add(ctx0, ggml_scale(ctx0, d, 1.0f + cfg), ggml_scale(ctx0, du, -cfg));
+ }
+ mx = ggml_add(ctx0, mx, ggml_scale(ctx0, d, span[i + 1] - span[i]));
+ cb(mx, "cfm_step", i);
+ }
+ ggml_tensor * mel = mx;
+
+ // trim the prompt frames at mel rate
+ mel = ggml_view_2d(ctx0, mel, 80, T2 - n_prompt_mel, mel->nb[1], (size_t) n_prompt_mel * mel->nb[1]);
+ mel = ggml_cont(ctx0, mel);
+ ggml_set_name(mel, "out_mel");
+ ggml_set_output(mel);
+ ggml_build_forward_expand(gf, mel);
+
+ // f0 predictor on the trimmed mel: 5x (conv k3 same-pad + elu), abs(linear)
+ {
+ ggml_tensor * fx = mel;
+ for (const auto & fc : c.f0_condnet) {
+ fx = cbx_conv1d(fc.w, fc.b, fx, 1, 1, 1);
+ fx = ggml_elu(ctx0, fx);
+ }
+ fx = cbx_linear(c.f0_classifier_w, c.f0_classifier_b, fx);
+ fx = ggml_abs(ctx0, fx); // [1, T]
+ ggml_set_name(fx, "out_f0");
+ ggml_set_output(fx);
+ ggml_build_forward_expand(gf, fx);
+ }
+ return gf;
+}
+
+// snake activation with per-channel alpha: x + sin^2(alpha x) / alpha
+static ggml_tensor * cbx_snake(ggml_context * ctx0, ggml_tensor * x, ggml_tensor * alpha) {
+ ggml_tensor * sx = ggml_sin(ctx0, ggml_mul(ctx0, x, alpha));
+ sx = ggml_mul(ctx0, sx, sx);
+ sx = ggml_div(ctx0, sx, alpha);
+ return ggml_add(ctx0, x, sx);
+}
+
+// hifigan-snake resblock; kernels with dilations 1/3/5 on convs1, 1 on convs2
+ggml_tensor * clip_graph_chatterbox::hift_resblock(const clip_chatterbox::hift_res & r, ggml_tensor * x) {
+ static const int dil[3] = {1, 3, 5};
+ for (size_t j = 0; j < r.units.size(); j++) {
+ const auto & u = r.units[j];
+ const int d = dil[j % 3];
+ const int p1 = (int) (u.conv1_w->ne[0] - 1) / 2 * d;
+ const int p2 = (int) (u.conv2_w->ne[0] - 1) / 2;
+ ggml_tensor * xt = cbx_snake(ctx0, x, u.act1_alpha);
+ xt = cbx_conv1d_dil(u.conv1_w, u.conv1_b, xt, p1, d);
+ xt = cbx_snake(ctx0, xt, u.act2_alpha);
+ xt = cbx_conv1d_dil(u.conv2_w, u.conv2_b, xt, p2, 1);
+ x = ggml_add(ctx0, x, xt);
+ }
+ return x;
+}
+
+// mel [80, T] + source stft [18, T_stft] -> conv_post output [18, T_stft2]
+ggml_cgraph * clip_graph_chatterbox::build_vocoder(int n_mel, int n_stft) {
+ const auto & c = model.cbx;
+
+ ggml_tensor * mel = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 80, n_mel);
+ ggml_set_name(mel, "inp_mel");
+ ggml_set_input(mel);
+ ggml_tensor * sstft = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 18, n_stft);
+ ggml_set_name(sstft, "inp_sstft");
+ ggml_set_input(sstft);
+
+ ggml_tensor * x = cbx_conv1d_dil(c.hift_pre_w, c.hift_pre_b, mel, 3, 1);
+
+ for (size_t i = 0; i < c.hift_ups.size(); i++) {
+ const auto & up = c.hift_ups[i];
+ ggml_tensor * uk = up.up_w;
+ const int K = (int) uk->ne[0];
+ const int S = K / 2;
+ const int P = (K - S) / 2;
+
+ x = ggml_leaky_relu(ctx0, x, 0.1f, false);
+ // conv transpose then trim the torch padding P on both sides
+ ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x));
+ xt = ggml_conv_transpose_1d(ctx0, uk, xt, S, 0, 1);
+ xt = ggml_cont(ctx0, ggml_view_2d(ctx0, xt, xt->ne[0] - 2 * P, xt->ne[1], xt->nb[1], (size_t) P * ggml_element_size(xt)));
+ x = ggml_cont(ctx0, ggml_transpose(ctx0, xt));
+ x = ggml_add(ctx0, x, up.up_b);
+
+ const bool is_last = i + 1 == c.hift_ups.size();
+ if (is_last) {
+ ggml_tensor * xr = ggml_cont(ctx0, ggml_transpose(ctx0, x));
+ xr = ggml_pad_reflect_1d(ctx0, xr, 1, 0);
+ x = ggml_cont(ctx0, ggml_transpose(ctx0, xr));
+ }
+
+ // source injection: strided conv on the source stft, one resblock
+ ggml_tensor * sk = up.source_down_w;
+ const int SK = (int) sk->ne[0];
+ const int SS = SK > 1 ? SK / 2 : 1;
+ const int SP = SK > 1 ? SS / 2 : 0;
+ ggml_tensor * si;
+ {
+ ggml_tensor * st = ggml_cont(ctx0, ggml_transpose(ctx0, sstft));
+ st = ggml_conv_1d(ctx0, sk, st, SS, SP, 1);
+ si = ggml_cont(ctx0, ggml_transpose(ctx0, st));
+ si = ggml_add(ctx0, si, up.source_down_b);
+ }
+ si = hift_resblock(up.source_res, si);
+ // align lengths: the reflection pad on the last stage adds one step
+ if ((int) si->ne[1] != (int) x->ne[1]) {
+ const int n = (int) std::min(si->ne[1], x->ne[1]);
+ si = ggml_cont(ctx0, ggml_view_2d(ctx0, si, si->ne[0], n, si->nb[1], 0));
+ x = ggml_cont(ctx0, ggml_view_2d(ctx0, x, x->ne[0], n, x->nb[1], 0));
+ }
+ x = ggml_add(ctx0, x, si);
+
+ ggml_tensor * acc = hift_resblock(up.res_0, x);
+ acc = ggml_add(ctx0, acc, hift_resblock(up.res_1, x));
+ acc = ggml_add(ctx0, acc, hift_resblock(up.res_2, x));
+ x = ggml_scale(ctx0, acc, 1.0f / 3.0f);
+ }
+
+ x = ggml_leaky_relu(ctx0, x, 0.01f, false);
+ x = cbx_conv1d_dil(c.hift_post_w, c.hift_post_b, x, 3, 1);
+ ggml_set_name(x, "out_spec");
+ ggml_set_output(x);
+ ggml_build_forward_expand(gf, x);
+ return gf;
+}
diff --git a/tools/mtmd/models/chatterbox-spkenc.cpp b/tools/mtmd/models/chatterbox-spkenc.cpp
new file mode 100644
index 000000000..afd9f0c29
--- /dev/null
+++ b/tools/mtmd/models/chatterbox-spkenc.cpp
@@ -0,0 +1,156 @@
+#include "models.h"
+
+// Chatterbox speaker encoder: CAMPPlus x-vector on kaldi fbank features, see s3gen/xvector.py
+
+static constexpr float bn_eps = 1e-5f;
+
+// per-channel batchnorm on x [C, T], w/b are null on the affine=False variant
+ggml_tensor * clip_graph_chatterbox_spkenc::bn1d(const clip_chatterbox::bn & n, ggml_tensor * x) {
+ ggml_tensor * sd = ggml_sqrt(ctx0, ggml_scale_bias(ctx0, n.var, 1.0f, bn_eps));
+ if (!n.w) {
+ return ggml_div(ctx0, ggml_sub(ctx0, x, n.mean), sd);
+ }
+ ggml_tensor * a = ggml_div(ctx0, n.w, sd);
+ ggml_tensor * shift = ggml_sub(ctx0, n.b, ggml_mul(ctx0, n.mean, a));
+ return ggml_add(ctx0, ggml_mul(ctx0, x, a), shift);
+}
+
+// batchnorm on a conv2d activation [W=T, H=F, C, 1], stats on ne2
+static ggml_tensor * cbx_bn2d(ggml_context * ctx0, const clip_chatterbox::bn & n, ggml_tensor * x) {
+ const int C = (int) x->ne[2];
+ ggml_tensor * mean = ggml_reshape_4d(ctx0, n.mean, 1, 1, C, 1);
+ ggml_tensor * var = ggml_reshape_4d(ctx0, n.var, 1, 1, C, 1);
+ ggml_tensor * w = ggml_reshape_4d(ctx0, n.w, 1, 1, C, 1);
+ ggml_tensor * b = ggml_reshape_4d(ctx0, n.b, 1, 1, C, 1);
+ ggml_tensor * a = ggml_div(ctx0, w, ggml_sqrt(ctx0, ggml_scale_bias(ctx0, var, 1.0f, bn_eps)));
+ ggml_tensor * shift = ggml_sub(ctx0, b, ggml_mul(ctx0, mean, a));
+ return ggml_add(ctx0, ggml_mul(ctx0, x, a), shift);
+}
+
+ggml_tensor * clip_graph_chatterbox_spkenc::bn2d_relu(const clip_chatterbox::bn & n, ggml_tensor * x) {
+ return ggml_relu(ctx0, cbx_bn2d(ctx0, n, x));
+}
+
+// fcm residual 2d block, stride on the frequency axis only
+ggml_tensor * clip_graph_chatterbox_spkenc::res2d(const clip_chatterbox::spk_res2d & r, ggml_tensor * x,
+ int stride) {
+ ggml_tensor * cur = ggml_conv_2d(ctx0, r.conv1_w, x, 1, stride, 1, 1, 1, 1);
+ cur = bn2d_relu(r.bn1, cur);
+ cur = ggml_conv_2d(ctx0, r.conv2_w, cur, 1, 1, 1, 1, 1, 1);
+ // bn2 without the relu, applied before the residual add
+ cur = cbx_bn2d(ctx0, r.bn2, cur);
+ ggml_tensor * res = x;
+ if (r.shortcut_w) {
+ res = ggml_conv_2d(ctx0, r.shortcut_w, x, 1, stride, 0, 0, 1, 1);
+ res = cbx_bn2d(ctx0, r.shortcut_bn, res);
+ }
+ return ggml_relu(ctx0, ggml_add(ctx0, cur, res));
+}
+
+// cam dense tdnn layer: bottleneck then context-gated conv; x [C_in, T] -> [growth, T]
+ggml_tensor * clip_graph_chatterbox_spkenc::cam_layer(const clip_chatterbox::spk_cam_layer & l, ggml_tensor * x,
+ int dil, ggml_tensor * segfix) {
+ ggml_tensor * h = ggml_relu(ctx0, bn1d(l.nl1_bn, x));
+ h = cbx_conv1d(l.linear1_w, nullptr, h, 1, 0, 0);
+ h = ggml_relu(ctx0, bn1d(l.nl2_bn, h));
+
+ ggml_tensor * k = l.local_w;
+ const int pad = ((int) k->ne[0] - 1) / 2 * dil;
+ ggml_tensor * y = cbx_conv1d_dil(k, nullptr, h, pad, dil);
+
+ // context: global mean plus ceil-mode segment means of length 100
+ const int T = (int) h->ne[1];
+ const int C = (int) h->ne[0];
+ const int S = (T + 99) / 100;
+ ggml_tensor * ht = ggml_cont(ctx0, ggml_transpose(ctx0, h)); // [T, C]
+ ggml_tensor * gmean = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_mean(ctx0, ht))); // [C, 1]
+ ggml_tensor * seg;
+ {
+ ggml_tensor * padded = ht;
+ if (S * 100 != T) {
+ ggml_tensor * z = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, S * 100 - T, C);
+ z = ggml_scale(ctx0, z, 0.0f);
+ padded = ggml_concat(ctx0, ht, z, 0);
+ }
+ ggml_tensor * pooled = ggml_pool_1d(ctx0, padded, GGML_OP_POOL_AVG, 100, 100, 0); // [S, C]
+ pooled = ggml_mul(ctx0, pooled, segfix);
+ ggml_tensor * exp = ggml_interpolate(ctx0, pooled, S * 100, C, 1, 1, GGML_SCALE_MODE_NEAREST);
+ exp = ggml_cont(ctx0, ggml_view_2d(ctx0, exp, T, C, exp->nb[1], 0));
+ seg = ggml_cont(ctx0, ggml_transpose(ctx0, exp)); // [C, T]
+ }
+ ggml_tensor * context = ggml_add(ctx0, seg, gmean);
+ context = cbx_conv1d(l.ctx1_w, l.ctx1_b, context, 1, 0, 0);
+ context = ggml_relu(ctx0, context);
+ context = cbx_conv1d(l.ctx2_w, l.ctx2_b, context, 1, 0, 0);
+ ggml_tensor * m = ggml_sigmoid(ctx0, context);
+ return ggml_mul(ctx0, y, m);
+}
+
+ggml_cgraph * clip_graph_chatterbox_spkenc::build() {
+ const auto & c = model.cbx;
+
+ const int T = img.nx();
+ const int T1 = (T - 1) / 2 + 1;
+ const int S = (T1 + 99) / 100;
+ ggml_tensor * segfix = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, S);
+ ggml_set_name(segfix, "inp_segfix");
+ ggml_set_input(segfix);
+
+ // fbank features [T, 80] from the preprocessor
+ ggml_tensor * inp = build_inp_raw(1);
+
+ // fcm 2d front: [W=T, H=F=80, C=1] -> [T, 10, 32] -> [320, T]
+ ggml_tensor * x = ggml_reshape_4d(ctx0, inp, T, 80, 1, 1);
+ x = ggml_conv_2d(ctx0, c.spk_conv1_w, x, 1, 1, 1, 1, 1, 1);
+ x = bn2d_relu(c.spk_bn1, x);
+ x = res2d(c.spk_layer1_0, x, 2);
+ x = res2d(c.spk_layer1_1, x, 1);
+ x = res2d(c.spk_layer2_0, x, 2);
+ x = res2d(c.spk_layer2_1, x, 1);
+ x = ggml_conv_2d(ctx0, c.spk_conv2_w, x, 1, 2, 1, 1, 1, 1);
+ x = bn2d_relu(c.spk_bn2, x);
+ x = ggml_reshape_2d(ctx0, x, T, 320);
+ x = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [320, T]
+ cb(x, "spk_fcm", -1);
+
+ // tdnn k5 stride 2 over time, then the three cam dense blocks
+ x = cbx_conv1d(c.spk_tdnn_w, nullptr, x, 2, 2, 2); // [128, T1]
+ x = ggml_relu(ctx0, bn1d(c.spk_tdnn_bn, x));
+ cb(x, "spk_tdnn", -1);
+
+ static const int block_dil[3] = {1, 2, 2};
+ const clip_chatterbox::spk_cam_block * blocks[3] = { &c.spk_block1, &c.spk_block2, &c.spk_block3 };
+ for (int bi = 1; bi <= 3; bi++) {
+ const auto & blk = *blocks[bi - 1];
+ for (const auto & l : blk.layers) {
+ ggml_tensor * out = cam_layer(l, x, block_dil[bi - 1], segfix);
+ x = ggml_concat(ctx0, x, out, 0);
+ }
+ x = ggml_relu(ctx0, bn1d(blk.transit_bn, x));
+ x = cbx_conv1d(blk.transit_w, nullptr, x, 1, 0, 0);
+ cb(x, "spk_block", bi);
+ }
+ x = ggml_relu(ctx0, bn1d(c.spk_out_bn, x)); // [512, T1]
+
+ // statistics pooling: mean and unbiased std over time -> [1024, 1]
+ ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [T1, 512]
+ ggml_tensor * mean = ggml_mean(ctx0, xt); // [1, 512]
+ ggml_tensor * m2 = ggml_mean(ctx0, ggml_mul(ctx0, xt, xt));
+ ggml_tensor * var = ggml_sub(ctx0, m2, ggml_mul(ctx0, mean, mean));
+ var = ggml_scale(ctx0, var, (float) T1 / (float) (T1 - 1));
+ ggml_tensor * sd = ggml_sqrt(ctx0, ggml_relu(ctx0, var));
+ ggml_tensor * stats = ggml_concat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, mean)),
+ ggml_cont(ctx0, ggml_transpose(ctx0, sd)), 0); // [1024, 1]
+ cb(stats, "spk_stats_pool", -1);
+
+ // the x-vector stays raw, the gen graph normalizes it
+ ggml_tensor * dw = ggml_reshape_2d(ctx0, c.spk_dense_w, 1024, 192);
+ ggml_tensor * emb = ggml_mul_mat(ctx0, dw, stats); // [192, 1]
+ emb = bn1d(c.spk_dense_bn, emb);
+ emb = ggml_cont(ctx0, ggml_reshape_1d(ctx0, emb, 192));
+ cb(emb, "spk_embd", -1);
+ ggml_set_name(emb, "out_xvec");
+ ggml_set_output(emb);
+ ggml_build_forward_expand(gf, emb);
+ return gf;
+}
diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h
index dc816fa87..c7520468f 100644
--- a/tools/mtmd/models/models.h
+++ b/tools/mtmd/models/models.h
@@ -206,6 +206,53 @@ struct clip_graph_conformer : clip_graph {
ggml_cgraph * build() override;
};
+// builders shared by the chatterbox gen and spkenc graphs, defined in chatterbox-gen.cpp
+struct clip_graph_chatterbox_base : clip_graph {
+ using clip_graph::clip_graph;
+
+ protected:
+ ggml_tensor * cbx_linear(ggml_tensor * w, ggml_tensor * b, ggml_tensor * x) const;
+ ggml_tensor * cbx_conv1d(ggml_tensor * k, ggml_tensor * b, ggml_tensor * x,
+ int stride, int pad_l, int pad_r) const;
+ ggml_tensor * cbx_conv1d_dil(ggml_tensor * k, ggml_tensor * b, ggml_tensor * x,
+ int pad, int dil) const;
+};
+
+struct clip_graph_chatterbox_spkenc : clip_graph_chatterbox_base {
+ clip_graph_chatterbox_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph_chatterbox_base(ctx, img) {}
+ ggml_cgraph * build() override;
+
+ private:
+ ggml_tensor * bn1d(const clip_chatterbox::bn & n, ggml_tensor * x);
+ ggml_tensor * bn2d_relu(const clip_chatterbox::bn & n, ggml_tensor * x);
+ ggml_tensor * res2d(const clip_chatterbox::spk_res2d & r, ggml_tensor * x, int stride);
+ ggml_tensor * cam_layer(const clip_chatterbox::spk_cam_layer & l, ggml_tensor * x, int dil, ggml_tensor * segfix);
+};
+
+struct clip_graph_chatterbox : clip_graph_chatterbox_base {
+ clip_gen_process_type gen_process = CLIP_GEN_PROCESS_GEN_CODE;
+ int n_tokens = 0;
+ int n_prompt_mel = 0;
+ int vocode_n_mel = 0;
+ int vocode_n_stft = 0;
+ clip_graph_chatterbox(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int n_tokens,
+ int n_prompt_mel, int vocode_n_mel, int vocode_n_stft)
+ : clip_graph_chatterbox_base(ctx, img), gen_process(gen_process), n_tokens(n_tokens), n_prompt_mel(n_prompt_mel),
+ vocode_n_mel(vocode_n_mel), vocode_n_stft(vocode_n_stft) {}
+ ggml_cgraph * build() override;
+
+ private:
+ ggml_tensor * enc_layer(const clip_chatterbox::enc_layer & l, ggml_tensor * x, ggml_tensor * pos, int T);
+ ggml_tensor * causal_block(const clip_chatterbox::causal_block & b, ggml_tensor * x);
+ ggml_tensor * resnet(const clip_chatterbox::resnet & r, ggml_tensor * x, ggml_tensor * temb);
+ ggml_tensor * tfm_block(const clip_chatterbox::tfm_block & b, ggml_tensor * x);
+ ggml_tensor * estimator(ggml_tensor * x_noise, ggml_tensor * mu, ggml_tensor * spks,
+ ggml_tensor * cond, ggml_tensor * temb, int T);
+ ggml_tensor * hift_resblock(const clip_chatterbox::hift_res & r, ggml_tensor * x);
+ ggml_cgraph * build_s3tok(int T);
+ ggml_cgraph * build_vocoder(int n_mel, int n_stft);
+};
+
struct clip_graph_lfm3a : clip_graph {
clip_graph_lfm3a(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
ggml_cgraph * build() override;
diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp
index 6e1ee2324..74d265526 100644
--- a/tools/mtmd/mtmd-audio.cpp
+++ b/tools/mtmd/mtmd-audio.cpp
@@ -1,6 +1,7 @@
#include "mtmd-audio.h"
#define _USE_MATH_DEFINES // for M_PI
+#include <cfloat>
#include <cmath>
#include <cstdint>
#include <cstring>
@@ -970,6 +971,480 @@ bool mtmd_audio_preprocessor_qwen3tts_spk::preprocess(const float *
return true;
}
+// windowed frame of a centered, reflect padded signal: dst[i] = x[start + i] * win[i], zeros past the window
+static void mtmd_audio_reflect_frame(const float * x, int n, int start, const std::vector<float> & win, std::vector<float> & dst) {
+ std::fill(dst.begin(), dst.end(), 0.0f);
+ for (int i = 0; i < (int) win.size(); i++) {
+ int idx = start + i;
+ if (idx < 0) {
+ idx = -idx;
+ }
+ if (idx >= n) {
+ idx = 2 * n - 2 - idx;
+ }
+ dst[(size_t) i] = x[idx] * win[(size_t) i];
+ }
+}
+
+// mel energies of the frame in fft_in: power (or magnitude) spectrum through cache.filters
+static void mtmd_audio_mel_energies(const mtmd_audio_cache & cache, std::vector<float> & fft_in, std::vector<float> & fft_out,
+ bool magnitude, std::vector<double> & mel) {
+ const int n_fft = (int) cache.filters.n_fft;
+ const int n_bins = n_fft / 2 + 1;
+ fft(cache, fft_in.data(), n_fft, fft_out.data());
+ for (int k = 0; k < n_bins; k++) {
+ const float p = fft_out[2 * k + 0] * fft_out[2 * k + 0] + fft_out[2 * k + 1] * fft_out[2 * k + 1];
+ fft_out[k] = magnitude ? sqrtf(p + 1e-9f) : p;
+ }
+ mel.assign(cache.filters.n_mel, 0.0);
+ for (int64_t m = 0; m < cache.filters.n_mel; m++) {
+ const float * w = cache.filters.data.data() + (size_t) m * n_bins;
+ for (int k = 0; k < n_bins; k++) {
+ mel[(size_t) m] += w[k] * fft_out[k];
+ }
+ }
+}
+
+// whisper style log-mel of the chatterbox s3 tokenizer (s3tokenizer.py), mel major
+static bool mtmd_audio_s3tok_log_mel(const mtmd_audio_cache & cache, const float * samples, size_t n_samples,
+ std::vector<float> & out, int & n_frames) {
+ const int n_fft = 400;
+ const int hop = 160;
+ const int n_mel = (int) cache.filters.n_mel;
+ const int n = (int) n_samples;
+
+ n_frames = n / hop;
+ if (n_frames <= 0) {
+ return false;
+ }
+
+ std::vector<float> fft_in((size_t) n_fft * 2), fft_out((size_t) n_fft * 8);
+ std::vector<double> e;
+ std::vector<double> mel((size_t) n_mel * n_frames);
+ for (int fr = 0; fr < n_frames; fr++) {
+ mtmd_audio_reflect_frame(samples, n, fr * hop - n_fft / 2, cache.hann_window, fft_in);
+ mtmd_audio_mel_energies(cache, fft_in, fft_out, false, e);
+ for (int m = 0; m < n_mel; m++) {
+ mel[(size_t) m * n_frames + fr] = log10(std::max(e[(size_t) m], 1e-10));
+ }
+ }
+
+ const double mx = *std::max_element(mel.begin(), mel.end());
+ out.resize(mel.size());
+ for (size_t i = 0; i < mel.size(); i++) {
+ out[i] = (float) ((std::max(mel[i], mx - 8.0) + 4.0) / 4.0);
+ }
+ return true;
+}
+
+// 16 kHz -> 24 kHz upsampler: hann windowed sinc cut just under the source nyquist
+static void mtmd_audio_upsample_3_2(const float * samples, size_t n_samples, std::vector<float> & out) {
+ const int W = 16; // sinc half width in source samples
+ const double fc = 0.495; // cutoff, normalized to the source rate
+
+ const size_t n_out = n_samples * 3 / 2;
+ out.assign(n_out, 0.0f);
+ for (size_t n = 0; n < n_out; n++) {
+ const double t = (double) (2 * n) / 3.0;
+ const int k0 = (int) floor(t) - W + 1;
+ double acc = 0.0;
+ for (int k = k0; k < k0 + 2 * W; k++) {
+ if (k < 0 || k >= (int) n_samples) {
+ continue;
+ }
+ const double x = t - k;
+ const double s = x == 0.0 ? 1.0 : sin(2.0 * M_PI * fc * x) / (M_PI * x);
+ acc += samples[k] * s * 0.5 * (1.0 + cos(M_PI * x / (W + 1)));
+ }
+ out[n] = (float) acc;
+ }
+}
+
+// matcha style log-mel of the s3gen prompt features (s3gen/utils/mel.py), frame major
+static bool mtmd_audio_matcha_log_mel(const mtmd_audio_cache & cache, const float * samples, size_t n_samples,
+ std::vector<float> & out, int & n_frames) {
+ const int n_fft = 1920;
+ const int hop = 480;
+ const int pad = (n_fft - hop) / 2;
+ const int n_mel = 80;
+ const int n = (int) n_samples;
+
+ n_frames = 1 + (n + 2 * pad - n_fft) / hop;
+ if (n <= pad || n_frames <= 0) {
+ return false;
+ }
+
+ std::vector<float> fft_in((size_t) n_fft * 2), fft_out((size_t) n_fft * 8);
+ std::vector<double> e;
+ out.resize((size_t) n_mel * n_frames);
+ for (int fr = 0; fr < n_frames; fr++) {
+ mtmd_audio_reflect_frame(samples, n, fr * hop - pad, cache.hann_window, fft_in);
+ mtmd_audio_mel_energies(cache, fft_in, fft_out, true, e);
+ for (int m = 0; m < n_mel; m++) {
+ out[(size_t) fr * n_mel + m] = logf(std::max((float) e[(size_t) m], 1e-5f));
+ }
+ }
+ return true;
+}
+
+// librosa.effects.trim: non-silent where the rms window is within top_db of the peak
+static void mtmd_audio_trim_silence(const float * samples, size_t n_samples, float top_db,
+ size_t & start, size_t & end) {
+ const int win = 2048;
+ const int hop = 512;
+ const int n = (int) n_samples;
+
+ const int n_fr = 1 + n / hop;
+ std::vector<double> rms((size_t) n_fr);
+ double mx = 0.0;
+ for (int fr = 0; fr < n_fr; fr++) {
+ double acc = 0.0;
+ for (int i = 0; i < win; i++) {
+ const int idx = fr * hop - win / 2 + i;
+ if (idx >= 0 && idx < n) {
+ acc += (double) samples[idx] * samples[idx];
+ }
+ }
+ rms[(size_t) fr] = sqrt(acc / win);
+ mx = std::max(mx, rms[(size_t) fr]);
+ }
+
+ const double thr = mx * pow(10.0, -top_db / 20.0);
+ int first = -1, last = -1;
+ for (int fr = 0; fr < n_fr; fr++) {
+ if (rms[(size_t) fr] > thr) {
+ if (first < 0) {
+ first = fr;
+ }
+ last = fr;
+ }
+ }
+ if (first < 0) {
+ start = end = 0;
+ return;
+ }
+ start = (size_t) first * hop;
+ end = std::min((size_t) (last + 1) * hop, n_samples);
+}
+
+// power mel of the voice encoder front-end, frame major, no log
+static bool mtmd_audio_ve_mel(const mtmd_audio_cache & cache, const float * samples, size_t n_samples,
+ std::vector<float> & out, int & n_frames) {
+ const int n_fft = 400;
+ const int hop = 160;
+ const int n_mel = 40;
+ const int n = (int) n_samples;
+
+ n_frames = 1 + n / hop;
+ if (n < 2) {
+ return false;
+ }
+
+ std::vector<float> fft_in((size_t) n_fft * 2), fft_out((size_t) n_fft * 8);
+ std::vector<double> e;
+ out.resize((size_t) n_mel * n_frames);
+ for (int fr = 0; fr < n_frames; fr++) {
+ mtmd_audio_reflect_frame(samples, n, fr * hop - n_fft / 2, cache.hann_window, fft_in);
+ mtmd_audio_mel_energies(cache, fft_in, fft_out, false, e);
+ for (int m = 0; m < n_mel; m++) {
+ out[(size_t) fr * n_mel + m] = (float) e[(size_t) m];
+ }
+ }
+ return true;
+}
+
+// ITU-R BS.1770 integrated loudness of a mono signal, matching pyloudnorm
+static float mtmd_audio_lufs(const float * samples, size_t n_samples, int sample_rate) {
+ std::vector<double> y(samples, samples + n_samples);
+
+ auto biquad = [&](double b0, double b1, double b2, double a1, double a2) {
+ double x1 = 0.0, x2 = 0.0, y1 = 0.0, y2 = 0.0;
+ for (double & v : y) {
+ const double x0 = v;
+ v = b0 * x0 + b1 * x1 + b2 * x2 - a1 * y1 - a2 * y2;
+ x2 = x1; x1 = x0;
+ y2 = y1; y1 = v;
+ }
+ };
+
+ // stage 1: high shelf, f0 1681.9744509555319 Hz, +3.99984385397 dB, Q 0.7071752369554196
+ {
+ const double A = pow(10.0, 3.99984385397 / 40.0);
+ const double w0 = 2.0 * M_PI * 1681.9744509555319 / sample_rate;
+ const double alpha = sin(w0) / (2.0 * 0.7071752369554196);
+ const double c = cos(w0);
+ const double sq = 2.0 * sqrt(A) * alpha;
+ const double b0 = A * ((A + 1.0) + (A - 1.0) * c + sq);
+ const double b1 = -2.0 * A * ((A - 1.0) + (A + 1.0) * c);
+ const double b2 = A * ((A + 1.0) + (A - 1.0) * c - sq);
+ const double a0 = (A + 1.0) - (A - 1.0) * c + sq;
+ const double a1 = 2.0 * ((A - 1.0) - (A + 1.0) * c);
+ const double a2 = (A + 1.0) - (A - 1.0) * c - sq;
+ biquad(b0 / a0, b1 / a0, b2 / a0, a1 / a0, a2 / a0);
+ }
+ // stage 2: high pass, f0 38.13547087602444 Hz, Q 0.5003270373238773
+ {
+ const double w0 = 2.0 * M_PI * 38.13547087602444 / sample_rate;
+ const double alpha = sin(w0) / (2.0 * 0.5003270373238773);
+ const double c = cos(w0);
+ const double b0 = (1.0 + c) / 2.0;
+ const double b1 = -(1.0 + c);
+ const double b2 = (1.0 + c) / 2.0;
+ const double a0 = 1.0 + alpha;
+ const double a1 = -2.0 * c;
+ const double a2 = 1.0 - alpha;
+ biquad(b0 / a0, b1 / a0, b2 / a0, a1 / a0, a2 / a0);
+ }
+
+ const int win = (int) (0.4 * sample_rate);
+ const int hop = win / 4;
+ if ((int) n_samples < win) {
+ return -HUGE_VALF;
+ }
+ const int n_blocks = 1 + ((int) n_samples - win) / hop;
+ std::vector<double> z((size_t) n_blocks);
+ for (int b = 0; b < n_blocks; b++) {
+ double acc = 0.0;
+ for (int i = 0; i < win; i++) {
+ acc += y[(size_t) b * hop + i] * y[(size_t) b * hop + i];
+ }
+ z[(size_t) b] = acc / win;
+ }
+
+ auto gated_mean = [&](double thr_lufs) {
+ double acc = 0.0;
+ int n = 0;
+ for (const double v : z) {
+ if (-0.691 + 10.0 * log10(std::max(v, 1e-30)) > thr_lufs) {
+ acc += v;
+ n++;
+ }
+ }
+ return n > 0 ? acc / n : 0.0;
+ };
+
+ const double z_abs = gated_mean(-70.0);
+ if (z_abs <= 0.0) {
+ return -HUGE_VALF;
+ }
+ const double thr_rel = -0.691 + 10.0 * log10(z_abs) - 10.0;
+ const double z_rel = gated_mean(thr_rel);
+ if (z_rel <= 0.0) {
+ return -HUGE_VALF;
+ }
+ return (float) (-0.691 + 10.0 * log10(z_rel));
+}
+
+//
+// mtmd_audio_preprocessor_chatterbox_ref
+//
+
+void mtmd_audio_preprocessor_chatterbox_ref::initialize() {
+ // kaldi fbank: povey window, kaldi mel banks over the 512-point spectrum, the nyquist bin left out
+ {
+ const int frame_len = 400;
+ const int n_fft = 512;
+ const int n_bins = n_fft / 2 + 1;
+ const int n_mel = hparams.n_mel_bins;
+ const double sr = (double) hparams.audio_sample_rate;
+
+ fbank_cache.fill_sin_cos_table(n_fft);
+ fbank_cache.hann_window.resize(frame_len);
+ for (int i = 0; i < frame_len; i++) {
+ fbank_cache.hann_window[(size_t) i] = (float) pow(0.5 - 0.5 * cos(2.0 * M_PI * i / (frame_len - 1)), 0.85);
+ }
+
+ auto mel = [](double f) { return 1127.0 * log(1.0 + f / 700.0); };
+ const double mel_lo = mel(20.0);
+ const double mel_hi = mel(sr / 2.0);
+ const double delta = (mel_hi - mel_lo) / (n_mel + 1);
+ const double bin_hz = sr / n_fft;
+
+ auto & filters = fbank_cache.filters;
+ filters.n_mel = n_mel;
+ filters.n_fft = n_fft;
+ filters.data.assign((size_t) n_mel * n_bins, 0.0f);
+ for (int m = 0; m < n_mel; m++) {
+ const double left = mel_lo + m * delta;
+ const double center = left + delta;
+ const double right = center + delta;
+ for (int i = 0; i < n_bins - 1; i++) {
+ const double f = mel(bin_hz * i);
+ if (f > left && f < right) {
+ const double w = f <= center ? (f - left) / (center - left)
+ : (right - f) / (right - center);
+ filters.data[(size_t) m * n_bins + i] = (float) w;
+ }
+ }
+ }
+ }
+
+ // s3 tokenizer: librosa slaney mel, 128 bands
+ s3tok_cache.fill_sin_cos_table(400);
+ s3tok_cache.fill_hann_window(400, true);
+ s3tok_cache.fill_mel_filterbank_matrix(128, 400, 16000);
+
+ // voice encoder: slaney mel, 40 bands up to 8 kHz
+ ve_cache.fill_sin_cos_table(400);
+ ve_cache.fill_hann_window(400, true);
+ ve_cache.fill_mel_filterbank_matrix(40, 400, 16000, 0.0f, 8000.0f);
+
+ // s3gen prompt features at the 24 kHz decoder rate
+ matcha_cache.fill_sin_cos_table(1920);
+ matcha_cache.fill_hann_window(1920, true);
+ matcha_cache.fill_mel_filterbank_matrix(80, 1920, 24000, 0.0f, 8000.0f);
+}
+
+bool mtmd_audio_preprocessor_chatterbox_ref::preprocess(const float * samples,
+ size_t n_samples,
+ std::vector<mtmd_audio_mel> & output) const {
+ output.clear();
+ const int sr = (int) hparams.audio_sample_rate;
+
+ // the turbo variant loudness-normalizes the whole clip before any feature
+ std::vector<float> pcm(samples, samples + n_samples);
+ if (!is_mtl) {
+ const float lufs = mtmd_audio_lufs(pcm.data(), pcm.size(), sr);
+ if (lufs != -HUGE_VALF) {
+ const float gain = powf(10.0f, (-27.0f - lufs) / 20.0f);
+ if (std::isfinite(gain) && gain > 0.0f) {
+ for (float & v : pcm) {
+ v *= gain;
+ }
+ }
+ }
+ }
+
+ // entry 0: torchaudio kaldi.fbank of CAMPPlus (snip edges, dc removal, preemphasis 0.97), mean subtracted over time
+ {
+ const int frame_len = 400;
+ const int hop = 160;
+ const int n_fft = 512;
+ const int n_mel = hparams.n_mel_bins;
+
+ if ((int) pcm.size() < frame_len) {
+ return false;
+ }
+ const int n_frames = 1 + ((int) pcm.size() - frame_len) / hop;
+
+ mtmd_audio_mel out;
+ out.n_len = n_frames;
+ out.n_len_org = n_frames;
+ out.n_mel = n_mel;
+ out.data.assign((size_t) n_mel * n_frames, 0.0f);
+
+ const std::vector<float> & win = fbank_cache.hann_window;
+ std::vector<float> fft_in((size_t) n_fft * 2), fft_out((size_t) n_fft * 8);
+ std::vector<double> e;
+ for (int fr = 0; fr < n_frames; fr++) {
+ const float * x = pcm.data() + (size_t) fr * hop;
+
+ double mean = 0.0;
+ for (int i = 0; i < frame_len; i++) {
+ mean += x[i];
+ }
+ mean /= frame_len;
+
+ std::fill(fft_in.begin(), fft_in.end(), 0.0f);
+ fft_in[0] = (float) ((x[0] - mean) * (1.0 - 0.97) * win[0]);
+ for (int i = 1; i < frame_len; i++) {
+ fft_in[(size_t) i] = (float) (((x[i] - mean) - 0.97 * (x[i - 1] - mean)) * win[(size_t) i]);
+ }
+ mtmd_audio_mel_energies(fbank_cache, fft_in, fft_out, false, e);
+ for (int m = 0; m < n_mel; m++) {
+ out.data[(size_t) m * n_frames + fr] = (float) log(std::max(e[(size_t) m], (double) FLT_EPSILON));
+ }
+ }
+
+ for (int m = 0; m < n_mel; m++) {
+ float * row = out.data.data() + (size_t) m * n_frames;
+ double mean = 0.0;
+ for (int fr = 0; fr < n_frames; fr++) {
+ mean += row[fr];
+ }
+ mean /= n_frames;
+ for (int fr = 0; fr < n_frames; fr++) {
+ row[fr] -= (float) mean;
+ }
+ }
+
+ output.push_back(std::move(out));
+ }
+
+ // entries 1 and 2: s3 tokenizer log-mels at the flow and t3 caps, padded to whole 40 ms tokens
+ auto s3tok_entry = [&](size_t cap) -> bool {
+ std::vector<float> clip(pcm.begin(), pcm.begin() + std::min(pcm.size(), cap));
+ clip.resize((clip.size() + 639) / 640 * 640, 0.0f);
+ std::vector<float> mel;
+ int n_frames = 0;
+ if (!mtmd_audio_s3tok_log_mel(s3tok_cache, clip.data(), clip.size(), mel, n_frames)) {
+ return false;
+ }
+ mtmd_audio_mel out;
+ out.n_len = n_frames;
+ out.n_len_org = n_frames;
+ out.n_mel = s3tok_cache.filters.n_mel;
+ out.data = std::move(mel);
+ output.push_back(std::move(out));
+ return true;
+ };
+ const size_t gen_cap = (size_t) 10 * sr;
+ const size_t t3_cap = (size_t) (is_mtl ? 6 : 15) * sr;
+ if (!s3tok_entry(gen_cap) || !s3tok_entry(t3_cap)) {
+ return false;
+ }
+
+ // entry 3: s3gen prompt features of the flow-capped clip at 24 kHz
+ {
+ std::vector<float> clip(pcm.begin(), pcm.begin() + std::min(pcm.size(), gen_cap));
+ clip.resize((clip.size() + 639) / 640 * 640, 0.0f);
+ std::vector<float> pcm24;
+ mtmd_audio_upsample_3_2(clip.data(), clip.size(), pcm24);
+ std::vector<float> feat;
+ int n_frames = 0;
+ if (!mtmd_audio_matcha_log_mel(matcha_cache, pcm24.data(), pcm24.size(), feat, n_frames)) {
+ return false;
+ }
+ mtmd_audio_mel out;
+ out.n_len = n_frames;
+ out.n_len_org = n_frames;
+ out.n_mel = 80;
+ out.data = std::move(feat);
+ output.push_back(std::move(out));
+ }
+
+ // entry 4: voice encoder mel of the silence-trimmed clip, padded to the 160-frame partial grid (rate 1.3)
+ {
+ size_t t0 = 0, t1 = 0;
+ mtmd_audio_trim_silence(pcm.data(), pcm.size(), 20.0f, t0, t1);
+ if (t1 <= t0) {
+ return false;
+ }
+ std::vector<float> mel;
+ int n_frames = 0;
+ if (!mtmd_audio_ve_mel(ve_cache, pcm.data() + t0, t1 - t0, mel, n_frames)) {
+ return false;
+ }
+ const int n_mel = 40;
+ const int n_partial = 160;
+ const int step = (int) lround((16000.0 / 1.3) / n_partial);
+ int n_wins = std::max(n_frames - n_partial + step, 0) / step;
+ const int rem = std::max(n_frames - n_partial + step, 0) % step;
+ if (n_wins == 0 || (double) (rem + n_partial - step) / n_partial >= 0.8) {
+ n_wins++;
+ }
+ const int target = n_partial + step * (n_wins - 1);
+ mel.resize((size_t) target * n_mel, 0.0f);
+ mtmd_audio_mel out;
+ out.n_len = target;
+ out.n_len_org = target;
+ out.n_mel = n_mel;
+ out.data = std::move(mel);
+ output.push_back(std::move(out));
+ }
+ return true;
+}
+
//
// mtmd_audio_preprocessor_conformer
//
diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h
index 087e4b8c0..ca95b217f 100644
--- a/tools/mtmd/mtmd-audio.h
+++ b/tools/mtmd/mtmd-audio.h
@@ -171,6 +171,26 @@ struct mtmd_audio_preprocessor_pockettts : mtmd_audio_preprocessor {
bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) const override;
};
+// chatterbox reference front-end: one clip gives five entries, used by position in mtmd.cpp
+// 0: CAMPPlus kaldi fbank [80 x n_frames] (mel major)
+// 1: s3 tokenizer log-mel, flow cap [n_mel x n_frames] (mel major)
+// 2: s3 tokenizer log-mel, t3 cap [n_mel x n_frames] (mel major)
+// 3: s3gen prompt features [n_frames x 80] (frame major, 24 kHz mel rate)
+// 4: voice encoder power mel [n_frames x 40] (frame major, trimmed, partial grid)
+// turbo loudness-normalizes the clip to -27 LUFS and caps the t3 clip at 15 s instead of 6 s
+struct mtmd_audio_preprocessor_chatterbox_ref : mtmd_audio_preprocessor {
+ mtmd_audio_preprocessor_chatterbox_ref(const clip_ctx * ctx, bool is_mtl) : mtmd_audio_preprocessor(ctx), is_mtl(is_mtl) {}
+ void initialize() override;
+ bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) const override;
+
+ private:
+ bool is_mtl;
+ mtmd_audio_cache fbank_cache; // entry 0, the window slot holds the povey window
+ mtmd_audio_cache s3tok_cache; // entries 1 and 2
+ mtmd_audio_cache matcha_cache; // entry 3
+ mtmd_audio_cache ve_cache; // entry 4
+};
+
struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor {
mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { }
void initialize() override;
diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp
index bed8f3d8c..0d7901575 100644
--- a/tools/mtmd/mtmd-helper-gen.cpp
+++ b/tools/mtmd/mtmd-helper-gen.cpp
@@ -124,6 +124,39 @@ protected:
return true;
}
+ // encodes a speaker clip or a voice, out receives the encoder embeddings
+ bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
+ if (!mtmd_bitmap_is_voice(bitmap) && !mtmd_support_audio(mctx)) {
+ LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n");
+ return false;
+ }
+ const std::string marker = mtmd_get_marker(mctx);
+ 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: speaker 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;
+ }
+
llama_context * lctx;
mtmd_context * mctx;
const llama_model * model;
@@ -428,39 +461,6 @@ private:
return true;
}
- // runs the reference wav through the speaker encoder, returns one x-vector embedding row
- bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
- if (!mtmd_bitmap_is_voice(bitmap) && !mtmd_support_audio(mctx)) {
- LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n");
- return false;
- }
- const std::string marker = mtmd_get_marker(mctx);
- 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: speaker 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;
- }
-
// one GEN_WAV process() call over the buffered codes, state is carried across batches
bool flush_gen_wav() {
if (codes_buf.empty()) {
@@ -996,39 +996,6 @@ private:
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_bitmap_is_voice(bitmap) && !mtmd_support_audio(mctx)) {
- LOG_ERR("mtmd_helper_gen_audio: mmproj has no voice encoder\n");
- return false;
- }
- const std::string marker = mtmd_get_marker(mctx);
- 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()) {
@@ -1718,6 +1685,398 @@ private:
bool wav_header_sent = false;
};
+// Chatterbox: the backbone emits s3 speech tokens, then one flow-matching mel decode and the NSF-iSTFT vocoder make the PCM
+// the multilingual variant runs cfg with an uncond sequence (text rows zeroed), paired as seq_id and n_seq_max / 2 + seq_id
+class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
+public:
+ using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline;
+
+ void reset() override {
+ seq_id = 0;
+ seq_uncond = 0;
+ pos = 0;
+ ar_idx = 0;
+ codes_buf.clear();
+ audio_pcm.clear();
+ h_state_buf.clear();
+ out_buf.clear();
+ prompt_embd_buf.clear();
+ uncond_embd_buf.clear();
+ cond_logits.clear();
+ n_prompt = 0;
+ prompt_pos = 0;
+ decoded = false;
+ stream = false;
+ pcm_sent = 0;
+ wav_header_sent = false;
+ gen_state.reset();
+ }
+
+ int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
+ reset();
+ seq_id = inp->seq_id;
+
+ if (!ensure_cache()) {
+ return 1;
+ }
+
+ out_type = inp->out_type;
+ stream = inp->stream;
+ if (!init_gen_state(inp->seed)) {
+ return 1;
+ }
+
+ if (mtl) {
+ // the uncond sequence lives half a seq space away from the cond one
+ const int32_t n_seq_max = (int32_t) llama_n_seq_max(lctx);
+ seq_uncond = (llama_seq_id) (n_seq_max / 2 + seq_id);
+ if (n_seq_max < 2 || seq_id >= n_seq_max / 2) {
+ LOG_ERR("mtmd_helper_gen_audio: multilingual chatterbox needs a paired cfg sequence "
+ "(n_seq_max = %d, seq_id = %d)\n", n_seq_max, seq_id);
+ return 1;
+ }
+ }
+ llama_memory_t mem = llama_get_memory(lctx);
+ if (!llama_memory_seq_rm(mem, seq_id, 0, -1) || (mtl && !llama_memory_seq_rm(mem, seq_uncond, 0, -1))) {
+ LOG_ERR("mtmd_helper_gen_audio: cannot reset chatterbox sequences\n");
+ return 1;
+ }
+
+ // the encode gives the talker conditioning rows, set_ref keeps the flow decoder reference
+ if (!inp->speaker_ref) {
+ LOG_ERR("mtmd_helper_gen_audio: chatterbox requires a speaker reference or a voice\n");
+ return 1;
+ }
+ if (!encode_speaker(inp->speaker_ref, prompt_embd_buf) ||
+ mtmd_gen_state_set_ref(mctx, gen_state.get()) != 0) {
+ LOG_ERR("mtmd_helper_gen_audio: speaker reference encoding failed\n");
+ return 1;
+ }
+
+ const int n_e = n_embd;
+ auto push_row = [&](const float * src) {
+ prompt_embd_buf.insert(prompt_embd_buf.end(), src, src + n_e);
+ };
+ auto add_pos = [&](float * r, const std::vector<float> & tab, int idx) {
+ const float * p = tab.data() + (size_t) idx * n_e;
+ for (int j = 0; j < n_e; j++) {
+ r[j] += p[j];
+ }
+ };
+
+ // prompt rows: [conditioning, text, speech bos]
+ // no text preprocessing here, the [SPACE] substitution and the case folding are baked in at conversion
+ std::string txt(inp->prompt, inp->prompt_len);
+ if (mtl && inp->lang && inp->lang[0]) {
+ // the reference prefixes the language tag, e.g. "[fr]"
+ std::string lang = inp->lang;
+ std::transform(lang.begin(), lang.end(), lang.begin(), [](unsigned char c) { return (char) std::tolower(c); });
+ txt = "[" + lang + "]" + txt;
+ }
+ std::vector<llama_token> ids(txt.size() + 16);
+ int n_ids = llama_tokenize(vocab, txt.c_str(), (int32_t) txt.size(), ids.data(), (int32_t) ids.size(),
+ false, true);
+ if (n_ids < 1) {
+ LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n");
+ return 1;
+ }
+ ids.resize((size_t) n_ids);
+ if (mtl) {
+ // only the multilingual prompt wraps the text in start/stop tokens
+ ids.insert(ids.begin(), text_start);
+ ids.push_back(text_stop);
+ if (ids.size() > text_pos.size() / (size_t) n_e) {
+ LOG_ERR("mtmd_helper_gen_audio: prompt too long (%zu tokens, max %zu)\n", ids.size(), text_pos.size() / (size_t) n_e);
+ return 1;
+ }
+ }
+ const size_t text0 = prompt_embd_buf.size() / (size_t) n_e;
+ for (size_t i = 0; i < ids.size(); i++) {
+ push_row(tok_embd.data() + (size_t) ids[i] * n_e);
+ if (mtl) {
+ add_pos(prompt_embd_buf.data() + (text0 + i) * n_e, text_pos, (int) i);
+ }
+ }
+
+ // speech bos opens the AR stream
+ push_row(tok_embd.data() + (size_t) llama_vocab_bos(vocab) * n_e);
+ if (mtl) {
+ add_pos(prompt_embd_buf.data() + prompt_embd_buf.size() - n_e, speech_pos, 0);
+ }
+ ar_idx = 1;
+ n_prompt = (int) (prompt_embd_buf.size() / (size_t) n_e);
+
+ if (mtl) {
+ // cfg uncond prompt: text rows are zeroed but keep their learned positions, as in the reference
+ uncond_embd_buf = prompt_embd_buf;
+ for (size_t i = 0; i < ids.size(); i++) {
+ float * r = uncond_embd_buf.data() + (text0 + i) * n_e;
+ std::fill(r, r + n_e, 0.0f);
+ add_pos(r, text_pos, (int) i);
+ }
+ }
+
+ prompt_pos = 0;
+ return 0;
+ }
+
+ int32_t step_prompt(int32_t n_batch) override {
+ GGML_ASSERT(n_batch > 0);
+ // the uncond prompt follows the cond one on mtl
+ const int n_total = mtl ? 2 * n_prompt : n_prompt;
+ if (prompt_pos >= n_total) {
+ return 0;
+ }
+ const bool uncond = prompt_pos >= n_prompt;
+ const int p0 = uncond ? prompt_pos - n_prompt : prompt_pos;
+ const int32_t n_tokens_batch = std::min(n_batch, n_prompt - p0);
+ const bool last = p0 + n_tokens_batch == n_prompt;
+
+ decode_embd_batch batch((uncond ? uncond_embd_buf : prompt_embd_buf).data(), n_prompt, 1, n_embd);
+ batch.set_position_normal(0, uncond ? seq_uncond : seq_id);
+ if (last) {
+ batch.logits[n_prompt - 1] = 1;
+ }
+ if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.render(lctx, p0, n_tokens_batch)) != 0) {
+ LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n");
+ return -1;
+ }
+ prompt_pos += n_tokens_batch;
+
+ if (last && mtl) {
+ if (!uncond) {
+ cfg_read();
+ } else {
+ cfg_apply();
+ }
+ }
+
+ if (prompt_pos >= n_total) {
+ pos = n_prompt;
+ prompt_embd_buf.clear();
+ uncond_embd_buf.clear();
+ return 0;
+ }
+ return n_total - prompt_pos;
+ }
+
+ int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override {
+ GGML_UNUSED(h_state_in);
+ if (sampled < 0 || sampled >= llama_vocab_n_tokens(vocab)) {
+ LOG_ERR("mtmd_helper_gen_audio: chatterbox requires a token sampled from the backbone\n");
+ return 1;
+ }
+
+ // the speech stop token is the backbone eos, the multilingual speech positions end the stream too
+ if (llama_vocab_is_eog(vocab, sampled) || (mtl && (size_t) ar_idx >= speech_pos.size() / (size_t) n_embd)) {
+ *out_stop = true;
+ *h_state_out = nullptr;
+ return 0;
+ }
+
+ // keep only the 6561 s3gen codes, dropping start/stop and oov ids
+ if (sampled >= speech_base && sampled - speech_base < 6561) {
+ codes_buf.push_back(sampled - speech_base);
+ }
+
+ std::vector<float> e(tok_embd.begin() + (size_t) sampled * n_embd,
+ tok_embd.begin() + (size_t) (sampled + 1) * n_embd);
+ if (mtl) {
+ // the multilingual backbone takes the token row plus the learned speech position
+ const float * p = speech_pos.data() + (size_t) ar_idx * n_embd;
+ for (int j = 0; j < n_embd; j++) {
+ e[(size_t) j] += p[j];
+ }
+ ar_idx++;
+ }
+
+ if (!decode_one(e.data(), seq_id)) {
+ return 1;
+ }
+ if (mtl) {
+ // cfg: the sampled token feeds both branches
+ cfg_read();
+ if (!decode_one(e.data(), seq_uncond)) {
+ return 1;
+ }
+ cfg_apply();
+ }
+ pos++;
+
+ const float * h = llama_get_embeddings_ith(lctx, -1);
+ h_state_buf.assign(h, h + 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 {
+ *out_sample_rate = info.sample_rate;
+
+ if (!stream) {
+ if (!flush_gen_wav()) {
+ return 1;
+ }
+ 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;
+ }
+
+ // streaming: the flow decode is one-shot, so all the audio arrives after flush()
+ const size_t n_new = audio_pcm.size() - pcm_sent;
+ if (out_n_samples) {
+ *out_n_samples = (int64_t) n_new;
+ }
+ if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
+ *out_data = (const char *) (audio_pcm.data() + pcm_sent);
+ *out_data_len = n_new * sizeof(float);
+ pcm_sent = audio_pcm.size();
+ return 0;
+ }
+ out_buf.clear();
+ if (!wav_header_sent) {
+ write_wav16_header(out_buf, UINT32_MAX, info.sample_rate);
+ wav_header_sent = true;
+ }
+ append_wav16_pcm(out_buf, audio_pcm.data() + pcm_sent, n_new);
+ pcm_sent = audio_pcm.size();
+ *out_data = out_buf.data();
+ *out_data_len = out_buf.size();
+ return 0;
+ }
+
+ int32_t flush() override {
+ return flush_gen_wav() ? 0 : 1;
+ }
+
+private:
+ bool decode_one(const float * e, llama_seq_id seq) {
+ decode_embd_batch batch(e, 1, 1, n_embd);
+ batch.set_position_normal(pos, seq);
+ batch.logits[0] = 1;
+ if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.render(lctx, 0, 1)) != 0) {
+ LOG_ERR("mtmd_helper_gen_audio: step decode failed\n");
+ return false;
+ }
+ return true;
+ }
+
+ // cfg of the reference: logits = cond + 0.5 * (cond - uncond)
+ // the mix goes into the last output row (the uncond one), the caller samples from it
+ void cfg_read() {
+ const float * c = llama_get_logits_ith(lctx, -1);
+ cond_logits.assign(c, c + llama_vocab_n_tokens(vocab));
+ }
+ void cfg_apply() {
+ float * u = llama_get_logits_ith(lctx, -1);
+ for (size_t i = 0; i < cond_logits.size(); i++) {
+ u[i] = cond_logits[i] + 0.5f * (cond_logits[i] - u[i]);
+ }
+ }
+
+ // one GEN_WAV call over the whole utterance, the flow decode is not streamed
+ bool flush_gen_wav() {
+ if (decoded || codes_buf.empty()) {
+ return true;
+ }
+ if (!mtl) {
+ // turbo appends a short silence tail before vocoding
+ codes_buf.insert(codes_buf.end(), 3, 4299);
+ }
+ 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();
+ mtmd_gen_out out{};
+ if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) {
+ LOG_ERR("mtmd_helper_gen_audio: gen_wav decode failed\n");
+ return false;
+ }
+ audio_pcm.assign(out.audio, out.audio + out.n_samples);
+ decoded = true;
+ return true;
+ }
+
+ bool ensure_cache() {
+ if (!tok_embd.empty()) {
+ return true;
+ }
+ // fused vocab layout: [text 0..speech_base) then the speech tokens
+ speech_base = find_special_token(vocab, "<|speech_0|>");
+ if (speech_base == LLAMA_TOKEN_NULL) {
+ LOG_ERR("mtmd_helper_gen_audio: fused speech tokens not found in vocab\n");
+ return false;
+ }
+
+ const uint32_t n_tok_embd = llama_model_get_tok_embd(model, nullptr);
+ if (n_tok_embd != (uint32_t) llama_vocab_n_tokens(vocab) * (uint32_t) n_embd) {
+ LOG_ERR("mtmd_helper_gen_audio: unexpected token embedding size\n");
+ return false;
+ }
+ tok_embd.resize(n_tok_embd);
+ llama_model_get_tok_embd(model, tok_embd.data());
+
+ auto read_table = [&](const char * name, std::vector<float> & dst) {
+ const size_t n = mtmd_gen_audio_read_tensor(mctx, name, nullptr, 0);
+ dst.resize(n);
+ if (n == 0 || mtmd_gen_audio_read_tensor(mctx, name, dst.data(), n) != n) {
+ LOG_ERR("mtmd_helper_gen_audio: %s read failed\n", name);
+ return false;
+ }
+ return true;
+ };
+
+ // multilingual: the learned position tables
+ mtl = mtmd_gen_audio_read_tensor(mctx, "a.gen.t3.text_pos_emb", nullptr, 0) > 0;
+ if (mtl && (!read_table("a.gen.t3.text_pos_emb", text_pos) || !read_table("a.gen.t3.speech_pos_emb", speech_pos))) {
+ return false;
+ }
+ return true;
+ }
+
+ // vocab and mmproj data fixed across the whole session, read once
+ std::vector<float> tok_embd;
+ bool mtl = false;
+ std::vector<float> text_pos;
+ std::vector<float> speech_pos;
+ llama_token speech_base = LLAMA_TOKEN_NULL;
+ // start_text_token and stop_text_token of the reference, only the multilingual prompt uses them
+ static constexpr llama_token text_start = 255;
+ static constexpr llama_token text_stop = 0;
+
+ // per-generation state, cleared by reset()
+ llama_seq_id seq_id = 0;
+ llama_seq_id seq_uncond = 0;
+ llama_pos pos = 0;
+ int ar_idx = 0;
+ std::vector<float> prompt_embd_buf;
+ std::vector<float> uncond_embd_buf;
+ int n_prompt = 0;
+ int prompt_pos = 0;
+ std::vector<float> cond_logits;
+ std::vector<int32_t> codes_buf;
+ bool decoded = false;
+ 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;
+ bool stream = false;
+ size_t pcm_sent = 0; // samples already returned by get_output()
+ bool wav_header_sent = false;
+};
+
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:
@@ -1728,6 +2087,8 @@ static std::unique_ptr<mtmd_gen_audio_pipeline> make_pipeline(llama_context * lc
return std::unique_ptr<mtmd_gen_audio_pipeline>(new soprano_gen_audio_pipeline(lctx, mctx));
case MTMD_GEN_AUDIO_TYPE_LFM3A:
return std::unique_ptr<mtmd_gen_audio_pipeline>(new lfm3a_gen_audio_pipeline(lctx, mctx));
+ case MTMD_GEN_AUDIO_TYPE_CHATTERBOX:
+ return std::unique_ptr<mtmd_gen_audio_pipeline>(new chatterbox_gen_audio_pipeline(lctx, mctx));
default:
return nullptr;
}
diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp
index c1b024a88..1c89eabfa 100644
--- a/tools/mtmd/mtmd.cpp
+++ b/tools/mtmd/mtmd.cpp
@@ -176,6 +176,7 @@ struct mtmd_bitmap {
std::string id; // optional user-defined id, for ex: can be set to image hash, useful for KV cache tracking
bool is_audio = false; // true if the bitmap is audio
bool is_voice = false;
+ std::shared_ptr<const clip_gen_ref> gen_ref; // gen-side reference of a predefined voice, if any
bool mergeable = false; // [QWEN_VIDEO] set only on frames of the same video
bool is_video = false; // see mtmd_bitmap_set_is_video()
@@ -201,9 +202,14 @@ struct mtmd_bitmap {
}
}
- static mtmd_bitmap * init_voice(const std::string & name, const float * rows, uint32_t n_embd, uint32_t n_rows) {
+ static mtmd_bitmap * init_voice(const std::string & name, const clip_gen_voice & voice, uint32_t n_embd) {
+ const float * rows = voice.rows.data();
+ const uint32_t n_rows = voice.rows.size() / n_embd;
auto * bitmap = new mtmd_bitmap(nullptr, n_embd, n_rows);
bitmap->is_voice = true;
+ if (!voice.ref.tokens.empty()) {
+ bitmap->gen_ref = std::make_shared<const clip_gen_ref>(voice.ref);
+ }
bitmap->id = name;
bitmap->data.assign((const unsigned char *) rows, (const unsigned char *) (rows + (size_t) n_embd * n_rows));
return bitmap;
@@ -348,6 +354,7 @@ struct mtmd_audio_track {
clip_image_f32_batch batch_f32; // the mel of the whole track
uint32_t n_tokens = 0;
std::vector<float> embd; // filled by the first encode
+ std::shared_ptr<const clip_gen_ref> gen_ref; // gen-side reference of a predefined voice, if any
};
struct mtmd_audio_tokens {
@@ -564,6 +571,11 @@ mtmd_context_params mtmd_context_params_default() {
return params;
}
+// chatterbox multilingual has a perceiver in its conditioning encoder (speaker encoding context)
+static bool cbx_is_mtl(clip_ctx * ctx_a) {
+ return clip_cbx_read_tensor(ctx_a, "a.cenc.perceiver.pre_attention_query", nullptr, 0) > 0;
+}
+
struct mtmd_context {
struct clip_ctx * ctx_v; // vision
struct clip_ctx * ctx_a; // audio
@@ -572,6 +584,9 @@ struct mtmd_context {
// generation context
struct clip_ctx * ctx_gen_a; // audio
+ // gen-side reference data of the last encoded chunk (chatterbox), see mtmd_gen_state_set_ref()
+ clip_gen_ref gen_ref;
+
bool print_timings;
int n_threads;
std::string media_marker;
@@ -1184,6 +1199,10 @@ struct mtmd_context {
{
audio_preproc = std::make_unique<mtmd_audio_preprocessor_pockettts>(ctx_a);
} break;
+ case PROJECTOR_TYPE_CHATTERBOX_SPKENC:
+ {
+ audio_preproc = std::make_unique<mtmd_audio_preprocessor_chatterbox_ref>(ctx_a, cbx_is_mtl(ctx_a));
+ } break;
default:
throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj));
}
@@ -1845,6 +1864,38 @@ struct mtmd_tokenizer {
}
}
+ // one chunk for the five entries of a chatterbox clip, one token per talker conditioning row
+ if (clip_get_projector_type(ctx->ctx_a) == PROJECTOR_TYPE_CHATTERBOX_SPKENC) {
+ GGML_ASSERT(mel_spec_chunks.size() == 5);
+ // mtl: spkr, perceiver query rows, emotion; turbo: spkr, t3 speech tokens
+ const size_t n_query = clip_cbx_read_tensor(ctx->ctx_a, "a.cenc.perceiver.pre_attention_query", nullptr, 0) / ctx->n_embd_out();
+ const size_t n_tokens = cbx_is_mtl(ctx->ctx_a) ? 2 + n_query : 1 + (size_t) mel_spec_chunks[2].n_len / 4;
+
+ clip_image_f32_batch batch_f32;
+ batch_f32.is_audio = true;
+ for (auto & mel_spec : mel_spec_chunks) {
+ GGML_ASSERT(mel_spec.n_len <= INT32_MAX && mel_spec.n_len >= 0);
+ GGML_ASSERT(mel_spec.n_mel <= INT32_MAX && mel_spec.n_mel >= 0);
+ clip_image_f32 mel_f32;
+ mel_f32.set_size({(int) mel_spec.n_len, (int) mel_spec.n_mel},
+ mel_spec.data.empty(), /* is_audio */ true);
+ mel_f32.cpy_buf(mel_spec.data);
+ batch_f32.entries.push_back(std::move(mel_f32));
+ }
+
+ mtmd_audio_tokens_ptr audio_tokens(new mtmd_audio_tokens);
+ audio_tokens->n_tokens = (uint32_t) n_tokens;
+ audio_tokens->batch_f32 = std::move(batch_f32);
+ audio_tokens->id = bitmap->id;
+
+ mtmd_input_chunk chunk{
+ MTMD_INPUT_CHUNK_TYPE_AUDIO,
+ {}, // text tokens
+ nullptr, // image tokens
+ std::move(audio_tokens),
+ };
+ cur.entries.emplace_back(std::move(chunk));
+ } else
// consider each mel_spec as a separate audio chunk
// TODO: maybe support batching, but this may come with memory cost
for (auto & mel_spec : mel_spec_chunks) {
@@ -1898,6 +1949,7 @@ struct mtmd_tokenizer {
track->n_tokens = bitmap.ny;
const auto & buf = bitmap.get_ro_buf();
track->embd.assign((const float *) buf.data(), (const float *) (buf.data() + buf.size()));
+ track->gen_ref = bitmap.gen_ref;
mtmd_audio_tokens_ptr audio_tokens(new mtmd_audio_tokens);
audio_tokens->n_tokens = bitmap.ny;
@@ -2189,7 +2241,10 @@ static int32_t mtmd_encode_impl(mtmd_context * ctx, const mtmd_image_tokens * im
return ok ? 0 : 1;
}
+static int32_t cbx_ref_encode(mtmd_context * ctx, const mtmd_audio_tokens * audio_tokens, std::vector<float> & out_embd);
+
static int32_t mtmd_encode_chunk_impl(mtmd_context * ctx, const mtmd_input_chunk * chunk, std::vector<float> & out_embd) {
+ ctx->gen_ref = {}; // see mtmd_gen_state_set_ref()
if (chunk->type == MTMD_INPUT_CHUNK_TYPE_TEXT) {
LOG_WRN("mtmd_encode_chunk has no effect for text chunks\n");
return 0;
@@ -2234,8 +2289,14 @@ static int32_t mtmd_encode_chunk_impl(mtmd_context * ctx, const mtmd_input_chunk
const size_t off = (size_t) chunk->tokens_audio->row_begin * n_mmproj_embd;
out_embd.assign(track->embd.begin() + off,
track->embd.begin() + off + (size_t) chunk->tokens_audio->n_tokens * n_mmproj_embd);
+ if (track->gen_ref) {
+ ctx->gen_ref = *track->gen_ref;
+ }
return 0;
}
+ if (clip_get_projector_type(ctx->ctx_a) == PROJECTOR_TYPE_CHATTERBOX_SPKENC) {
+ return cbx_ref_encode(ctx, chunk->tokens_audio.get(), out_embd);
+ }
out_embd.resize((size_t)chunk->tokens_audio->n_tokens * ctx->n_embd_out());
bool ok = clip_image_batch_encode(
ctx->ctx_a,
@@ -2302,6 +2363,10 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
info.type = MTMD_GEN_AUDIO_TYPE_LFM3A;
info.sample_rate = 24000;
break;
+ case PROJECTOR_TYPE_CHATTERBOX:
+ info.type = MTMD_GEN_AUDIO_TYPE_CHATTERBOX;
+ info.sample_rate = 24000;
+ break;
default:
info.type = MTMD_GEN_AUDIO_TYPE_NONE;
break;
@@ -2309,6 +2374,355 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
return info;
}
+size_t mtmd_gen_audio_read_tensor(mtmd_context * ctx, const char * name, float * out, size_t n_max) {
+ if (!ctx->ctx_gen_a) {
+ return 0;
+ }
+ return clip_cbx_read_tensor(ctx->ctx_gen_a, name, out, n_max);
+}
+
+// raw 192-dim CAMPPlus x-vector of the fbank entry
+static bool cbx_ref_xvec(mtmd_context * ctx, const clip_image_f32 & fbank, std::vector<float> & xvec) {
+ clip_image_f32 mel_img;
+ mel_img.set_size(fbank.get_size(), false, true);
+ mel_img.cpy_buf(fbank.get_ro_buf());
+ clip_image_f32_batch batch;
+ batch.is_audio = true;
+ batch.entries.push_back(std::move(mel_img));
+ xvec.resize(192);
+ if (!clip_image_batch_encode(ctx->ctx_a, ctx->n_threads, &batch, xvec)) {
+ LOG_ERR("%s: speaker encoder failed\n", __func__);
+ return false;
+ }
+ return true;
+}
+
+// s3 speech tokens of a log-mel entry, rows optionally receives their speech embedding rows
+static int32_t cbx_ref_tokenize(mtmd_context * ctx, const clip_image_f32 & mel,
+ std::vector<int32_t> & codes, std::vector<float> * rows) {
+ clip_image_f32 mel_img;
+ mel_img.set_size(mel.get_size(), false, true);
+ mel_img.cpy_buf(mel.get_ro_buf());
+
+ clip_image_f32_batch batch;
+ batch.is_audio = true;
+ batch.entries.push_back(std::move(mel_img));
+
+ std::vector<int32_t> out_codes;
+ std::vector<float> out_embd;
+
+ clip_encode_params params;
+ params.imgs = &batch;
+ params.n_threads = ctx->n_threads;
+ params.gen_process = CLIP_GEN_PROCESS_TOKENIZE;
+ params.out_codes = &out_codes;
+ params.out_code_embd = rows ? &out_embd : nullptr;
+
+ if (!clip_encode(ctx->ctx_gen_a, ¶ms)) {
+ LOG_ERR("%s: clip_encode failed (tokenize)\n", __func__);
+ return 1;
+ }
+
+ codes = std::move(out_codes);
+ if (rows) {
+ *rows = std::move(out_embd);
+ }
+ return 0;
+}
+
+// talker conditioning rows: the speaker row, then on mtl the perceiver block and the emotion row
+static int32_t cbx_ref_cond(mtmd_context * ctx, const clip_image_f32 & ve_mel,
+ const std::vector<float> & pse, std::vector<float> & out_rows) {
+ clip_ctx * ctx_clip = ctx->ctx_a;
+ // n_expect 0 accepts any non-empty size
+ auto read_t = [&](const char * name, std::vector<float> & v, size_t n_expect) -> bool {
+ const size_t n = clip_cbx_read_tensor(ctx_clip, name, nullptr, 0);
+ if (n == 0 || (n_expect > 0 && n != n_expect)) {
+ return false;
+ }
+ v.resize(n);
+ return clip_cbx_read_tensor(ctx_clip, name, v.data(), n) == n;
+ };
+
+ // voice encoder (embeds_from_wavs): 3-layer lstm per 160-frame partial, the normalized partial embeddings are averaged
+ const std::vector<float> & mel = ve_mel.get_ro_buf();
+ const int n_mel = 40;
+ const int n_partial = 160;
+ const int n_h = 256;
+ const int step = (int) lround((16000.0 / 1.3) / n_partial); // reference rate 1.3
+ if (ve_mel.nx() < n_partial || mel.size() < (size_t) ve_mel.nx() * n_mel) {
+ LOG_ERR("%s: voice encoder mel too short\n", __func__);
+ return 1;
+ }
+ const int n_wins = (ve_mel.nx() - n_partial) / step + 1;
+
+ std::vector<float> w_ih[3], w_hh[3], b_ih[3], b_hh[3];
+ std::vector<float> w_proj, b_proj;
+ for (int l = 0; l < 3; l++) {
+ const std::string s = std::to_string(l);
+ const size_t n_in = l == 0 ? n_mel : n_h;
+ if (!read_t(("a.ve.lstm.weight_ih_l" + s).c_str(), w_ih[l], 4 * n_h * n_in) ||
+ !read_t(("a.ve.lstm.weight_hh_l" + s).c_str(), w_hh[l], 4 * n_h * n_h) ||
+ !read_t(("a.ve.lstm.bias_ih_l" + s).c_str(), b_ih[l], 4 * n_h) ||
+ !read_t(("a.ve.lstm.bias_hh_l" + s).c_str(), b_hh[l], 4 * n_h)) {
+ LOG_ERR("%s: model has no voice encoder\n", __func__);
+ return 1;
+ }
+ }
+ if (!read_t("a.ve.proj.weight", w_proj, n_h * n_h) || !read_t("a.ve.proj.bias", b_proj, n_h)) {
+ LOG_ERR("%s: model has no voice encoder projection\n", __func__);
+ return 1;
+ }
+
+ std::vector<float> ve(n_h, 0.0f);
+ std::vector<float> h((size_t) 3 * n_h), c((size_t) 3 * n_h), x(n_h), g((size_t) 4 * n_h);
+ for (int p = 0; p < n_wins; p++) {
+ std::fill(h.begin(), h.end(), 0.0f);
+ std::fill(c.begin(), c.end(), 0.0f);
+ for (int t = 0; t < n_partial; t++) {
+ const float * in = mel.data() + (size_t) (p * step + t) * n_mel;
+ int n_in = n_mel;
+ for (int l = 0; l < 3; l++) {
+ float * hl = h.data() + (size_t) l * n_h;
+ float * cl = c.data() + (size_t) l * n_h;
+ for (int j = 0; j < 4 * n_h; j++) {
+ double acc = b_ih[l][(size_t) j] + b_hh[l][(size_t) j];
+ const float * wi = w_ih[l].data() + (size_t) j * n_in;
+ for (int i = 0; i < n_in; i++) {
+ acc += (double) wi[i] * in[i];
+ }
+ const float * wh = w_hh[l].data() + (size_t) j * n_h;
+ for (int i = 0; i < n_h; i++) {
+ acc += (double) wh[i] * hl[i];
+ }
+ g[(size_t) j] = (float) acc;
+ }
+ // torch gate order: input, forget, cell, output
+ for (int i = 0; i < n_h; i++) {
+ const float gi = 1.0f / (1.0f + expf(-g[(size_t) i]));
+ const float gf = 1.0f / (1.0f + expf(-g[(size_t) i + n_h]));
+ const float gc = tanhf(g[(size_t) i + 2 * n_h]);
+ const float go = 1.0f / (1.0f + expf(-g[(size_t) i + 3 * n_h]));
+ cl[i] = gf * cl[i] + gi * gc;
+ x[(size_t) i] = go * tanhf(cl[i]);
+ }
+ memcpy(hl, x.data(), (size_t) n_h * sizeof(float));
+ in = hl;
+ n_in = n_h;
+ }
+ }
+ // projected, relu'd, normalized partial embedding
+ std::vector<float> e(n_h);
+ double norm = 0.0;
+ for (int o = 0; o < n_h; o++) {
+ double acc = b_proj[(size_t) o];
+ const float * w = w_proj.data() + (size_t) o * n_h;
+ const float * hl = h.data() + (size_t) 2 * n_h;
+ for (int i = 0; i < n_h; i++) {
+ acc += (double) w[i] * hl[i];
+ }
+ e[(size_t) o] = (float) std::max(acc, 0.0);
+ norm += (double) e[(size_t) o] * e[(size_t) o];
+ }
+ norm = sqrt(norm);
+ for (int o = 0; o < n_h; o++) {
+ ve[(size_t) o] += (float) (e[(size_t) o] / norm);
+ }
+ }
+ double norm = 0.0;
+ for (float v : ve) {
+ norm += (double) v * v;
+ }
+ norm = sqrt(norm);
+ for (float & v : ve) {
+ v = (float) (v / norm);
+ }
+
+ // speaker projection row
+ const int n_e = (int) ctx->n_embd_out();
+ std::vector<float> w_spkr, b_spkr;
+ if (!read_t("a.cenc.spkr_enc.weight", w_spkr, (size_t) n_e * n_h) || !read_t("a.cenc.spkr_enc.bias", b_spkr, n_e)) {
+ LOG_ERR("%s: model has no speaker conditioning projection\n", __func__);
+ return 1;
+ }
+ std::vector<float> rows((size_t) n_e);
+ for (int o = 0; o < n_e; o++) {
+ double acc = b_spkr[(size_t) o];
+ const float * w = w_spkr.data() + (size_t) o * n_h;
+ for (int i = 0; i < n_h; i++) {
+ acc += (double) w[i] * ve[(size_t) i];
+ }
+ rows[(size_t) o] = (float) acc;
+ }
+
+ // the perceiver runs its shared attention block as cross then self attention
+ std::vector<float> query;
+ if (read_t("a.cenc.perceiver.pre_attention_query", query, 0)) {
+ const int n_head = 4;
+ const size_t n_ee = (size_t) n_e * n_e;
+ if (pse.empty() || pse.size() % n_e != 0 || query.size() % n_e != 0 || n_e % n_head != 0) {
+ LOG_ERR("%s: reference speech embeddings required for the perceiver\n", __func__);
+ return 1;
+ }
+ std::vector<float> ln_w, ln_b, wq, bq, wk, bk, wv, bv, wo, bo, emo;
+ if (!read_t("a.cenc.perceiver.attn.norm.weight", ln_w, n_e) || !read_t("a.cenc.perceiver.attn.norm.bias", ln_b, n_e) ||
+ !read_t("a.cenc.perceiver.attn.to_q.weight", wq, n_ee) || !read_t("a.cenc.perceiver.attn.to_q.bias", bq, n_e) ||
+ !read_t("a.cenc.perceiver.attn.to_k.weight", wk, n_ee) || !read_t("a.cenc.perceiver.attn.to_k.bias", bk, n_e) ||
+ !read_t("a.cenc.perceiver.attn.to_v.weight", wv, n_ee) || !read_t("a.cenc.perceiver.attn.to_v.bias", bv, n_e) ||
+ !read_t("a.cenc.perceiver.attn.proj_out.weight", wo, n_ee) || !read_t("a.cenc.perceiver.attn.proj_out.bias", bo, n_e) ||
+ !read_t("a.cenc.emotion_adv_fc.weight", emo, n_e)) {
+ LOG_ERR("%s: model has an incomplete perceiver\n", __func__);
+ return 1;
+ }
+ const int d_head = n_e / n_head;
+
+ auto layer_norm = [&](const float * in, float * out) {
+ double mean = 0.0, var = 0.0;
+ for (int i = 0; i < n_e; i++) {
+ mean += in[i];
+ }
+ mean /= n_e;
+ for (int i = 0; i < n_e; i++) {
+ var += (in[i] - mean) * (in[i] - mean);
+ }
+ const double sd = sqrt(var / n_e + 1e-5);
+ for (int i = 0; i < n_e; i++) {
+ out[i] = (float) ((in[i] - mean) / sd * ln_w[(size_t) i] + ln_b[(size_t) i]);
+ }
+ };
+ auto linear = [&](const std::vector<float> & w, const std::vector<float> & b,
+ const std::vector<float> & in, int n_rows, std::vector<float> & out) {
+ out.resize((size_t) n_rows * n_e);
+ for (int r = 0; r < n_rows; r++) {
+ for (int o = 0; o < n_e; o++) {
+ double acc = b[(size_t) o];
+ const float * wr = w.data() + (size_t) o * n_e;
+ const float * ir = in.data() + (size_t) r * n_e;
+ for (int i = 0; i < n_e; i++) {
+ acc += (double) wr[i] * ir[i];
+ }
+ out[(size_t) r * n_e + o] = (float) acc;
+ }
+ }
+ };
+ auto attn_block = [&](const std::vector<float> & x1, int n1,
+ const std::vector<float> & x2, int n2, std::vector<float> & out) {
+ std::vector<float> nx1((size_t) n1 * n_e), nx2((size_t) n2 * n_e);
+ for (int r = 0; r < n1; r++) {
+ layer_norm(x1.data() + (size_t) r * n_e, nx1.data() + (size_t) r * n_e);
+ }
+ for (int r = 0; r < n2; r++) {
+ layer_norm(x2.data() + (size_t) r * n_e, nx2.data() + (size_t) r * n_e);
+ }
+ std::vector<float> q, k, v;
+ linear(wq, bq, nx1, n1, q);
+ linear(wk, bk, nx2, n2, k);
+ linear(wv, bv, nx2, n2, v);
+
+ std::vector<float> ctxt((size_t) n1 * n_e);
+ std::vector<double> sc((size_t) n2);
+ for (int hd = 0; hd < n_head; hd++) {
+ const int off = hd * d_head;
+ for (int t = 0; t < n1; t++) {
+ double mx = -1e30;
+ for (int s = 0; s < n2; s++) {
+ double acc = 0.0;
+ for (int i = 0; i < d_head; i++) {
+ acc += (double) q[(size_t) t * n_e + off + i] * k[(size_t) s * n_e + off + i];
+ }
+ sc[(size_t) s] = acc / sqrt((double) d_head);
+ mx = std::max(mx, sc[(size_t) s]);
+ }
+ double sum = 0.0;
+ for (int s = 0; s < n2; s++) {
+ sc[(size_t) s] = exp(sc[(size_t) s] - mx);
+ sum += sc[(size_t) s];
+ }
+ for (int i = 0; i < d_head; i++) {
+ double acc = 0.0;
+ for (int s = 0; s < n2; s++) {
+ acc += sc[(size_t) s] * v[(size_t) s * n_e + off + i];
+ }
+ ctxt[(size_t) t * n_e + off + i] = (float) (acc / sum);
+ }
+ }
+ }
+ linear(wo, bo, ctxt, n1, out);
+ for (size_t i = 0; i < out.size(); i++) {
+ out[i] += x1[i];
+ }
+ };
+
+ const int n_q = (int) (query.size() / n_e);
+ std::vector<float> pre, p32;
+ attn_block(query, n_q, pse, (int) (pse.size() / (size_t) n_e), pre);
+ attn_block(pre, n_q, pre, n_q, p32);
+
+ rows.insert(rows.end(), p32.begin(), p32.end());
+ const float exaggeration = 0.5f; // reference default
+ for (int i = 0; i < n_e; i++) {
+ rows.push_back(emo[(size_t) i] * exaggeration);
+ }
+ }
+
+ out_rows = std::move(rows);
+ return 0;
+}
+
+// talker conditioning rows go to out_embd, the flow decoder reference to ctx->gen_ref
+static int32_t cbx_ref_encode(mtmd_context * ctx, const mtmd_audio_tokens * audio_tokens, std::vector<float> & out_embd) {
+ if (!ctx->ctx_gen_a) {
+ LOG_ERR("%s: reference encoding needs the audio generation context\n", __func__);
+ return 1;
+ }
+ const auto & entries = audio_tokens->batch_f32.entries;
+ GGML_ASSERT(entries.size() == 5);
+
+ std::vector<float> xvec;
+ if (!cbx_ref_xvec(ctx, entries[0], xvec)) {
+ return 1;
+ }
+
+ std::vector<int32_t> flow_codes, t3_codes;
+ std::vector<float> t3_rows;
+ if (cbx_ref_tokenize(ctx, entries[1], flow_codes, nullptr) != 0 ||
+ cbx_ref_tokenize(ctx, entries[2], t3_codes, &t3_rows) != 0) {
+ return 1;
+ }
+
+ const size_t n_feat = (size_t) entries[3].nx() * 80;
+ if ((size_t) entries[3].nx() != 2 * flow_codes.size()) {
+ LOG_ERR("%s: reference mel length %d does not match %zu codes\n",
+ __func__, entries[3].nx(), flow_codes.size());
+ return 1;
+ }
+
+ // [spkr, perceiver block] on mtl, [spkr, reference speech rows] on turbo
+ const bool mtl = cbx_is_mtl(ctx->ctx_a);
+ std::vector<float> pse;
+ if (mtl) {
+ pse = std::move(t3_rows);
+ }
+ std::vector<float> rows;
+ if (cbx_ref_cond(ctx, entries[4], pse, rows) != 0) {
+ return 1;
+ }
+ if (!mtl) {
+ rows.insert(rows.end(), t3_rows.begin(), t3_rows.end());
+ }
+ if (rows.size() != (size_t) audio_tokens->n_tokens * (size_t) ctx->n_embd_out()) {
+ LOG_ERR("%s: unexpected number of conditioning rows\n", __func__);
+ return 1;
+ }
+
+ ctx->gen_ref.tokens = std::move(flow_codes);
+ ctx->gen_ref.feat.assign(entries[3].get_ro_buf().begin(), entries[3].get_ro_buf().begin() + n_feat);
+ ctx->gen_ref.spk = std::move(xvec);
+
+ out_embd = std::move(rows);
+ return 0;
+}
+
const char * mtmd_gen_audio_voice_name(const mtmd_context * ctx, int32_t idx) {
if (!ctx->ctx_gen_a || idx < 0) {
return nullptr;
@@ -2330,8 +2744,7 @@ mtmd_bitmap * mtmd_gen_audio_get_voice(const mtmd_context * ctx, const char * na
if (it == voices.end()) {
return nullptr;
}
- const uint32_t n_embd = clip_n_mmproj_embd(ctx->ctx_gen_a);
- return mtmd_bitmap::init_voice(it->first, it->second.data(), n_embd, it->second.size() / n_embd);
+ return mtmd_bitmap::init_voice(it->first, it->second, clip_n_mmproj_embd(ctx->ctx_gen_a));
}
mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) {
@@ -2382,6 +2795,9 @@ struct mtmd_gen_state {
std::vector<float> out_embd; // next-step hidden state fed back to backbone (GEN_CODE)
std::vector<float> out_audio; // decoded PCM samples for the current frame (GEN_WAV)
+ // voice reference for GEN_WAV, from a speaker clip or a predefined voice (chatterbox)
+ clip_gen_ref ref;
+
// lfm3a: the codec graph yields STFT frames, the ISTFT runs on the host
std::unique_ptr<mtmd_audio_streaming_istft> istft;
@@ -2488,6 +2904,19 @@ static int32_t mtmd_gen_wav_lfm3a(mtmd_context * ctx, mtmd_gen_state * state, co
return 0;
}
+int32_t mtmd_gen_state_set_ref(mtmd_context * ctx, mtmd_gen_state * state) {
+ if (!state || state->ctx_clip != ctx->ctx_gen_a) {
+ LOG_ERR("%s: invalid gen state\n", __func__);
+ return 1;
+ }
+ if (ctx->gen_ref.tokens.empty()) {
+ LOG_ERR("%s: the last encode produced no reference data\n", __func__);
+ return 1;
+ }
+ state->ref = ctx->gen_ref;
+ return 0;
+}
+
static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, mtmd_gen_state * state, const mtmd_gen_inp * inp, mtmd_gen_out * out) {
clip_ctx * ctx_clip = ctx->ctx_gen_a;
if (!ctx_clip) {
@@ -2569,6 +2998,44 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, mtmd_gen_state *
}
// MTMD_GEN_PROCESS_TYPE_GEN_WAV
+ if (clip_get_projector_type(ctx_clip) == PROJECTOR_TYPE_CHATTERBOX) {
+ if (!inp->codes || inp->n_codes == 0) {
+ LOG_ERR("%s: codes required for gen_wav\n", __func__);
+ return 1;
+ }
+ if (state->ref.tokens.empty()) {
+ LOG_ERR("%s: gen_wav needs a voice reference, see mtmd_gen_state_set_ref()\n", __func__);
+ return 1;
+ }
+ std::vector<int32_t> in_codes(inp->codes, inp->codes + inp->n_codes);
+
+ // the batch entry is unused, present to satisfy the encode interface
+ clip_image_f32 dummy;
+ dummy.set_size({1, 1}, false, true);
+ dummy.cpy_buf(std::vector<float>(1, 0.0f));
+ clip_image_f32_batch batch;
+ batch.is_audio = true;
+ batch.entries.push_back(std::move(dummy));
+
+ clip_encode_params params;
+ params.imgs = &batch;
+ params.n_threads = ctx->n_threads;
+ params.gen_process = CLIP_GEN_PROCESS_TTS;
+ params.codes = &in_codes;
+ params.out_audio = &state->out_audio;
+ params.ref = &state->ref;
+ params.gen_state = state->clip_state;
+
+ if (!clip_encode(ctx_clip, ¶ms)) {
+ LOG_ERR("%s: clip_encode failed (gen_wav)\n", __func__);
+ return 1;
+ }
+
+ out->audio = state->out_audio.data();
+ out->n_samples = state->out_audio.size();
+ return 0;
+ }
+
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) {
diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h
index 4f69a644d..d490e9668 100644
--- a/tools/mtmd/mtmd.h
+++ b/tools/mtmd/mtmd.h
@@ -465,6 +465,7 @@ enum mtmd_gen_audio_type {
MTMD_GEN_AUDIO_TYPE_POCKETTTS,
MTMD_GEN_AUDIO_TYPE_SOPRANO,
MTMD_GEN_AUDIO_TYPE_LFM3A,
+ MTMD_GEN_AUDIO_TYPE_CHATTERBOX,
};
struct mtmd_gen_audio_info {
@@ -476,6 +477,8 @@ struct mtmd_gen_audio_info {
MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx);
+// read a named tensor of the gen audio model as F32, returns the element count (0 if not found), out may be null to query the size
+MTMD_API size_t mtmd_gen_audio_read_tensor(mtmd_context * ctx, const char * name, float * out, size_t n_max);
// name of the idx-th predefined voice, nullptr when idx is out of range
MTMD_API const char * mtmd_gen_audio_voice_name(const mtmd_context * ctx, int32_t idx);
@@ -543,6 +546,10 @@ MTMD_API void mtmd_gen_state_free(mtmd_gen_state * state);
// clear the decoder state (cold start), keep the RNG stream
MTMD_API void mtmd_gen_state_reset(mtmd_gen_state * state);
+// copy the gen-side reference data of the last mtmd_encode_chunk() call into the state (speaker clip or predefined voice)
+// returns non-zero if the last encode produced no such data
+MTMD_API int32_t mtmd_gen_state_set_ref(mtmd_context * ctx, mtmd_gen_state * state);
+
// 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);
diff --git a/tools/tts/README.md b/tools/tts/README.md
index 9bec17624..82c23473c 100644
--- a/tools/tts/README.md
+++ b/tools/tts/README.md
@@ -121,3 +121,36 @@ The mmproj holds the audio encoder (for ASR), the depthformer and the codec deco
python convert_hf_to_gguf.py path/to/lfm3-audio --outfile lfm3-audio.gguf
python convert_hf_to_gguf.py path/to/lfm3-audio --mmproj --outfile mmproj-lfm3-audio.gguf
```
+
+## Chatterbox
+
+Two variants are supported: [Chatterbox Turbo](https://huggingface.co/ResembleAI/chatterbox-turbo) (English) and [Chatterbox Multilingual V3](https://huggingface.co/ResembleAI/chatterbox) (23 languages).
+
+Available params:
+- `--tts-speaker-file` should point to a speaker reference audio file (wav, mp3). Optional, the predefined voice is used without it
+- `--tts-voice` can be `default`, the voice shipped with the model (`conds.pt`)
+- `--tts-lang` sets the language of the multilingual variant, for example `fr` (default: none). It is not used by Turbo
+
+Example usage:
+
+```sh
+llama-tts -m chatterbox-mtl.gguf \
+ -mm mmproj-chatterbox-mtl.gguf \
+ -p "Bonjour tout le monde" \
+ --tts-lang fr \
+ --tts-speaker-file speaker.mp3 \
+ --output out.wav
+```
+
+Note: the multilingual variant runs classifier-free guidance with a second sequence. `llama-tts` sets this up by itself, but `llama-server` does not support it yet.
+
+**Note for GGUF conversion:**
+
+The upstream repositories ship raw checkpoints without `config.json`. Download the files below into a directory, then convert it as usual (the variant is detected from the talker file):
+- Turbo: `t3_turbo_v1.safetensors`, `s3gen_meanflow.safetensors`, `ve.safetensors`, `conds.pt` and the GPT-2 tokenizer files (`vocab.json`, `merges.txt`, `added_tokens.json`, `special_tokens_map.json`, `tokenizer_config.json`)
+- Multilingual V3: `t3_mtl23ls_v3.safetensors`, `s3gen_v3.safetensors`, `ve.safetensors`, `conds.pt`, `mtl_tokenizer.json`
+
+```sh
+python convert_hf_to_gguf.py path/to/chatterbox --outfile chatterbox.gguf
+python convert_hf_to_gguf.py path/to/chatterbox --mmproj --outfile mmproj-chatterbox.gguf
+```
diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp
index 8578f8869..6354be6d6 100644
--- a/tools/tts/tts.cpp
+++ b/tools/tts/tts.cpp
@@ -69,6 +69,11 @@ int main(int argc, char ** argv) {
// always enable embd, so that we can pass hidden states to the audio generation helper
params.embedding = true;
+ // for chatterbox multilingual: its cfg pairs slot i with sequence n_parallel + i in a unified kv cache
+ // TODO: this applies to every model, find a better way for a pipeline to request extra sequences
+ params.n_parallel *= 2;
+ params.kv_unified = true;
+
llama_backend_init();
llama_numa_init(params.numa);