diff --git a/conversion/__init__.py b/conversion/__init__.py index 31da4963cf86..1a5546cc5821 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -37,6 +37,7 @@ "BloomModel": "bloom", "CamembertModel": "bert", "ChameleonForCausalLM": "chameleon", + "ChatterboxModel": "chatterbox", "ChameleonForConditionalGeneration": "chameleon", "ChatGLMForConditionalGeneration": "chatglm", "ChatGLMModel": "chatglm", @@ -260,6 +261,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = { + "ChatterboxModel": "chatterbox", "AudioFlamingo3ForConditionalGeneration": "ultravox", "CogVLMForCausalLM": "cogvlm", "DeepseekOCR2ForCausalLM": "deepseek", diff --git a/conversion/chatterbox.py b/conversion/chatterbox.py new file mode 100644 index 000000000000..54a2bf6008db --- /dev/null +++ b/conversion/chatterbox.py @@ -0,0 +1,581 @@ +# Chatterbox (ResembleAI) conversion: turbo and multilingual v3 variants. +# +# Checkpoint layout is the official repo layout (raw safetensors, no HF weights): +# - turbo (ResembleAI/chatterbox-turbo): t3_turbo_v1.safetensors (GPT-2 medium talker), +# s3gen_meanflow.safetensors, ve.safetensors, conds.pt, GPT-2 BPE tokenizer files +# - multilingual v3 (ResembleAI/chatterbox): t3_mtl23ls_v3.safetensors (Llama 520M talker), +# s3gen_v3.safetensors, ve.safetensors, conds.pt, mtl_tokenizer.json +# The variant is detected by which talker file is present. A minimal config.json with +# architectures ["ChatterboxModel"] routes the directory to these classes. +# +# Talker: the transformer input embeddings (wte / embed_tokens) are dead in the +# reference (inputs_embeds everywhere); the live tables are text_emb and speech_emb, +# fused here into one [text | speech] vocab. Text tokens keep their ids, speech token i +# becomes <|speech_i|> at text_vocab + i. The speech start/stop tokens map to bos/eos. +# +# Mmproj: the whole s3gen sidecar (flow encoder, CFM estimator, HiFT vocoder, CAMPPlus +# speaker encoder, S3 tokenizer), the voice encoder, the talker conditioning encoder, +# the learned position tables, precomputed default-voice conditioning from conds.pt, +# and the talker speech embedding table so that reference speech tokens can be turned +# into talker-space embeddings without a lookup on the text model side. + +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 + +TURBO_TALKER = "t3_turbo_v1.safetensors" +MTL_TALKER = "t3_mtl23ls_v3.safetensors" +TURBO_S3GEN = "s3gen_meanflow.safetensors" +MTL_S3GEN = "s3gen_v3.safetensors" + +# relative speech ids shared by both variants (start/stop_speech_token in the reference) +SPEECH_BOS = 6561 +SPEECH_EOS = 6562 + +# reference sampling defaults (tts_turbo.py generate / mtl tts.py); samplers +# absent from the reference carry their explicit disable value, so that the +# common defaults never leak in when tools apply these as model defaults +TURBO_SAMPLING = {"top_k": 1000, "min_p": 0.0, "top_p": 0.95, "temp": 0.8, "penalty_repeat": 1.2} +MTL_SAMPLING = {"top_k": 0, "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 _s3tok_mel_filters() -> Tensor: + # slaney-normalized mel filterbank of the s3 tokenizer front end (librosa + # defaults: sr 16000, n_fft 400, 128 bands); not every checkpoint ships + # it, so it is synthesized here for both variants + sr, n_fft, n_mels = 16000, 400, 128 + + def hz_to_mel(f: Tensor) -> Tensor: + lin = f / (200.0 / 3.0) + logstep = math.log(6.4) / 27.0 + return torch.where(f >= 1000.0, 15.0 + torch.log(f.clamp(min=1000.0) / 1000.0) / logstep, lin) + + def mel_to_hz(m: Tensor) -> Tensor: + logstep = math.log(6.4) / 27.0 + return torch.where(m >= 15.0, 1000.0 * torch.exp(logstep * (m - 15.0)), m * (200.0 / 3.0)) + + fftfreqs = torch.arange(n_fft // 2 + 1, dtype=torch.float64) * (sr / n_fft) + bounds = hz_to_mel(torch.tensor([0.0, sr / 2.0], dtype=torch.float64)) + mel_f = mel_to_hz(torch.linspace(bounds[0], bounds[1], n_mels + 2, dtype=torch.float64)) + fdiff = mel_f.diff() + ramps = mel_f[:, None] - fftfreqs[None, :] + lower = -ramps[:n_mels] / fdiff[:n_mels, None] + upper = ramps[2:] / fdiff[1:, None] + weights = torch.minimum(lower, upper).clamp(min=0.0) + weights *= (2.0 / (mel_f[2:] - mel_f[:n_mels]))[:, None] + return weights.float() + + +def _is_turbo(dir_model: Path) -> bool: + if (dir_model / TURBO_TALKER).is_file(): + return True + if (dir_model / MTL_TALKER).is_file(): + return False + raise FileNotFoundError(f"no chatterbox talker checkpoint in {dir_model}") + + +def _index_safetensors(path: Path, lazy: bool, rename: Callable[[str], str | None]) -> dict[str, Callable[[], Tensor]]: + tensors: dict[str, Callable[[], Tensor]] = {} + with gguf.utility.SafetensorsLocal(path) as model_part: + for name in model_part.keys(): + new_name = rename(name) + if new_name is None: + continue + data: gguf.utility.LocalTensor = model_part[name] + if lazy: + data_gen = lambda data=data: LazyTorchTensor.from_local_tensor(data) # noqa: E731 + else: + dtype = LazyTorchTensor._dtype_str_map[data.dtype] + data_gen = lambda data=data, dtype=dtype: torch.from_numpy(data.mmap_bytes()).view(dtype).reshape(data.shape) # noqa: E731 + tensors[new_name] = data_gen + return tensors + + +@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]]: + talker = TURBO_TALKER if self.is_turbo else MTL_TALKER + + def rename(name: str) -> str | None: + # transformer input embeddings are dead in the reference (inputs_embeds + # everywhere); the conditioning encoder and position tables go to the mmproj + 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 name + + return _index_safetensors(self.dir_model / talker, self.lazy, rename) + + 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.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: suppress the text zone + # so the sampling chain can never pick a text token + self.gguf_writer.add_suppress_tokens(list(range(n_text))) + + def _set_vocab_mtl(self): + # custom multilingual BPE (mtl_tokenizer.json), extended with the speech + # tokens. the reference tokenizer is char-level (raw unicode chars in + # the vocab), while the gpt2 tokenizer of llama.cpp is byte-level: the + # vocab and merges are re-encoded through the gpt2 byte-to-unicode map, + # and synthetic merges rebuild each multi-byte char from its bytes so + # that the byte-level closure reproduces the char-level tokenization. + # the reference also lowercases and NFKD-normalizes its input; this + # port keeps composed characters (they live in the vocab with merges) + # and folds case in the embedding table instead, so the raw prompt + # needs no runtime text preprocessing at all + 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 raw space char is a dead token (the reference substitutes + # [SPACE] before encoding), its slot stays unused + 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) + + # the reference replaces ' ' with the [SPACE] token before encoding: + # that substitution is baked into the vocab by giving the [SPACE] slot + # the byte-level space as content, so raw spaces tokenize to it 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 + + # char-building merges rank first: chars are atomic in the reference, + # they must form before any of its merges apply + 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.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre("default") + 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: suppress the text zone + # so the sampling chain can never pick a text token + self.gguf_writer.add_suppress_tokens(list(range(n_text))) + + def _mtl_case_fold_pairs(self) -> list[tuple[int, int]]: + # single-char uppercase vocab entries whose lowercase form is also a + # single-char vocab entry, as (upper_id, lower_id) pairs. covers every + # script in the vocab (latin, latin-1 accented, cyrillic, greek, ...) + 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): + if self.is_turbo: + self.gguf_writer.add_block_count(self.hparams["n_layer"]) + self.gguf_writer.add_context_length(self.hparams["n_ctx"]) + self.gguf_writer.add_embedding_length(self.hparams["n_embd"]) + self.gguf_writer.add_feed_forward_length(4 * self.hparams["n_embd"]) + self.gguf_writer.add_head_count(self.hparams["n_head"]) + self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"]) + else: + self.gguf_writer.add_block_count(self.hparams["num_hidden_layers"]) + self.gguf_writer.add_context_length(self.hparams["max_position_embeddings"]) + self.gguf_writer.add_embedding_length(self.hparams["hidden_size"]) + self.gguf_writer.add_feed_forward_length(self.hparams["intermediate_size"]) + self.gguf_writer.add_head_count(self.hparams["num_attention_heads"]) + self.gguf_writer.add_head_count_kv(self.hparams["num_key_value_heads"]) + self.gguf_writer.add_rope_freq_base(self.hparams["rope_theta"]) + self.gguf_writer.add_rope_dimension_count(self.hparams["head_dim"]) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams["rms_norm_eps"]) + self.gguf_writer.add_file_type(self.ftype) + + sampling = TURBO_SAMPLING if self.is_turbo else MTL_SAMPLING + self.gguf_writer.add_sampling_top_k(sampling["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)) + + @staticmethod + def permute(weights: Tensor, n_head: int) -> Tensor: + # HF half-split rope layout to the interleaved layout of the llama arch + return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:]) + .swapaxes(1, 2) + .reshape(weights.shape)) + + 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: + # case folding baked into the embedding table: uppercase rows + # are replaced by their lowercase rows (the reference lowercases + # before encoding, so the uppercase rows are never trained) + 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 gpt2 arch has no output bias tensor; the constant speech logit + # bias is dropped, matching the validated behavior of this port + 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 = self.permute(data_torch, self.hparams["num_attention_heads"]) + if name.endswith("k_proj.weight"): + data_torch = self.permute(data_torch, 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]]: + talker = TURBO_TALKER if self.is_turbo else MTL_TALKER + s3gen = TURBO_S3GEN if self.is_turbo else MTL_S3GEN + + def rename_s3gen(name: str) -> str | None: + # batchnorm bookkeeping, unused in inference (and over the gguf name length cap) + if name.endswith("num_batches_tracked"): + return None + # dsp buffers; the mel filterbank is synthesized in + # generate_extra_tensors, the window is rebuilt at runtime + if 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."), + ): + if name.startswith(src): + return dst + name[len(src):] + if name.startswith("flow."): + # input_embedding, encoder_proj, spk_embed_affine_layer + return "a.gen." + name + return None + + def rename_ve(name: str) -> str | None: + if name.startswith("similarity_"): + return None + return "a.ve." + name + + def rename_talker(name: str) -> str | None: + # conditioning encoder, learned position tables and the speech + # embedding table live on the mmproj side + if name.startswith("cond_enc."): + return "a.cenc." + name[len("cond_enc."):] + if name == "text_pos_emb.emb.weight": + return "a.gen.t3.text_pos_emb" + if name == "speech_pos_emb.emb.weight": + return "a.gen.t3.speech_pos_emb" + if name == "speech_emb.weight": + return self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_OUT_EMBD) + return None + + tensors = _index_safetensors(self.dir_model / s3gen, self.lazy, rename_s3gen) + tensors.update(_index_safetensors(self.dir_model / "ve.safetensors", self.lazy, rename_ve)) + tensors.update(_index_safetensors(self.dir_model / talker, self.lazy, rename_talker)) + return tensors + + def set_gguf_parameters(self): + self.gguf_writer.add_file_type(self.ftype) + + # speaker encoder (CAMPPlus, chatterbox_spkenc projector); the DSP front-end + # hparams are fixed by the projector type on the C++ side + 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) + + self.gguf_writer.add_uint32("chatterbox.n_mels", 80) + self.gguf_writer.add_uint32("chatterbox.sample_rate", 24000) + self.gguf_writer.add_uint32("chatterbox.speech_vocab", self.global_config["text_config"]["speech_vocab_size"]) + self.gguf_writer.add_uint32("chatterbox.meanflow", 1 if self.is_turbo else 0) + + 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 raw on the host (voice encoder, conditioning encoder, + # precomputed conditioning, position tables, mel filterbank, source + # module) must stay F32: the reader handles F32/F16/I32 only + if new_name.startswith(("a.ve.", "a.cenc.", "a.gen.cond.", "a.gen.t3.", "a.gen.hift.m_source.")) or new_name == "a.s3tok.mel_filters": + return gguf.GGMLQuantizationType.F32 + # conv kernels of the graphs (ggml_conv_1d/_2d/_dw and the transposed + # convs of the vocoder) have no BF16 kernels; F16 is the graph-side + # storage type for everything large, F32 for the rest + 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]]: + talker = TURBO_TALKER if self.is_turbo else MTL_TALKER + conds = torch.load(self.dir_model / "conds.pt", map_location="cpu", weights_only=False) + 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: + with gguf.utility.SafetensorsLocal(self.dir_model / talker) as parts: + data = parts[name] + dtype = LazyTorchTensor._dtype_str_map[data.dtype] + return torch.from_numpy(data.mmap_bytes()).view(dtype).reshape(data.shape).clone() + + yield ("a.s3tok.mel_filters", _s3tok_mel_filters()) + + # default voice: precomputed s3gen conditioning from conds.pt + yield ("a.gen.cond.gen_prompt_token", genc["prompt_token"][0].to(torch.int32)) + yield ("a.gen.cond.gen_prompt_feat", genc["prompt_feat"][0].float()) + # raw 192-dim campplus x-vector of the default voice: normalization + # and the flow speaker affine run inside the code2wav graph + yield ("a.gen.cond.gen_embedding", genc["embedding"][0].float()) + + spkr_w = talker_tensor("cond_enc.spkr_enc.weight").float() + spkr_b = talker_tensor("cond_enc.spkr_enc.bias").float() + spkr_row = spkr_w @ t3c["speaker_emb"][0].float() + spkr_b + + if self.is_turbo: + # default talker conditioning: projected speaker row + speech token ids, + # resolved through the speech embedding table at inference time + yield ("a.gen.cond.spkr_default", spkr_row) + yield ("a.gen.cond.prompt_speech_tokens", t3c["cond_prompt_speech_tokens"][0].to(torch.int32)) + return + + # multilingual default talker conditioning: [spkr, perceiver x32, emotion] + # block precomputed by running the reference perceiver over the embedded + # default cond speech tokens (flash attention path of AttentionBlock2) + with torch.no_grad(): + speech_emb = talker_tensor("speech_emb.weight").float() + speech_pos = talker_tensor("speech_pos_emb.emb.weight").float() + cond_tokens = t3c["cond_prompt_speech_tokens"][0] + pse = speech_emb[cond_tokens] + speech_pos[: cond_tokens.shape[0]] + + ln_w = talker_tensor("cond_enc.perceiver.attn.norm.weight").float() + ln_b = talker_tensor("cond_enc.perceiver.attn.norm.bias").float() + wq = talker_tensor("cond_enc.perceiver.attn.to_q.weight").float() + bq = talker_tensor("cond_enc.perceiver.attn.to_q.bias").float() + wk = talker_tensor("cond_enc.perceiver.attn.to_k.weight").float() + bk = talker_tensor("cond_enc.perceiver.attn.to_k.bias").float() + wv = talker_tensor("cond_enc.perceiver.attn.to_v.weight").float() + bv = talker_tensor("cond_enc.perceiver.attn.to_v.bias").float() + wo = talker_tensor("cond_enc.perceiver.attn.proj_out.weight").float() + bo = talker_tensor("cond_enc.perceiver.attn.proj_out.bias").float() + query = talker_tensor("cond_enc.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("cond_enc.emotion_adv_fc.weight").float() + emo_row = emo[:, 0] * t3c["emotion_adv"].reshape(-1)[0] + yield ("a.gen.cond.t3_cond", 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 e72b6564ff02..7ef4844b688d 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -5084,6 +5084,8 @@ class VisionProjectorType: NEMOTRON_V2_VL = "nemotron_v2_vl" QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor + CHATTERBOX_SPKENC = "chatterbox_spkenc" # audio: CAMPPlus speaker encoder + CHATTERBOX = "chatterbox" # audio generation: s3gen flow matching + HiFT vocoder HUNYUANVL = "hunyuanvl" PARAKEET = "parakeet" # audio MINIMAXM3 = "minimax_m3" diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 4675fb9a97b6..516a9b47898b 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -27,6 +27,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/dotsocr.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 7222660c7797..f7d03ad81b43 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -455,6 +455,8 @@ enum projector_type { PROJECTOR_TYPE_MIMO_AUDIO, PROJECTOR_TYPE_QWEN3TTS_SPKENC, PROJECTOR_TYPE_QWEN3TTS_GEN, + PROJECTOR_TYPE_CHATTERBOX, + PROJECTOR_TYPE_CHATTERBOX_SPKENC, PROJECTOR_TYPE_UNKNOWN, }; @@ -514,6 +516,8 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_PARAKEET, "parakeet"}, { PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"}, { PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"}, + { PROJECTOR_TYPE_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 30014f1f506c..27ff5eef3ec7 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -456,6 +457,237 @@ 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 + // (the boundary conv stays null on mid stages) + struct est_stage { + resnet res; + std::vector 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 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 + // (shortcut stays null on the 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 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; + 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 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_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_ups; + ggml_tensor * hift_post_w = nullptr; + ggml_tensor * hift_post_b = nullptr; + std::vector 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_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_model { clip_modality modality = CLIP_MODALITY_VISION; projector_type proj_type = PROJECTOR_TYPE_MLP; @@ -687,6 +919,15 @@ struct clip_model { // qwen3tts code2wav: RVQ codes -> raw PCM clip_code2wav c2w; + // chatterbox flow encoder, cfm estimator, hift vocoder, s3 tokenizer + // and CAMPPlus speaker encoder + clip_chatterbox cbx; + + // chatterbox host-read side data (conditioning defaults, embedding + // tables, filterbanks, voice/conditioning encoders run on the host), + // accessed by name through clip_cbx_read_tensor + std::map cbx_tensors; + // cogvlm ggml_tensor * mm_post_fc_norm_w = nullptr; ggml_tensor * mm_post_fc_norm_b = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 50644bf1e526..b575e60278a4 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1054,6 +1054,35 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + { + builder = std::make_unique(ctx, img); + } break; + case PROJECTOR_TYPE_CHATTERBOX: + { + const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_CODE_GEN; + int n_tokens = params && params->codes ? (int) params->codes->size() : 0; + int n_prompt_mel = 0; + if (n_tokens > 0) { + if (params->ref_tokens) { + n_tokens += (int) params->ref_tokens->size(); + } else { + auto it = ctx->model.cbx_tensors.find("a.gen.cond.gen_prompt_token"); + GGML_ASSERT(it != ctx->model.cbx_tensors.end()); + n_tokens += (int) it->second->ne[0]; + } + if (params->ref_feat) { + n_prompt_mel = (int) (params->ref_feat->size() / 80); + } else { + auto it = ctx->model.cbx_tensors.find("a.gen.cond.gen_prompt_feat"); + GGML_ASSERT(it != ctx->model.cbx_tensors.end()); + n_prompt_mel = (int) it->second->ne[1]; + } + } + const int vnm = params ? params->vocode_n_mel : 0; + const int vns = params ? params->vocode_n_stft : 0; + builder = std::make_unique(ctx, img, gen_process, n_tokens, n_prompt_mel, vnm, vns); + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_CODE_GEN; @@ -1708,6 +1737,15 @@ struct clip_model_loader { hparams.audio_window_len = 1024; hparams.audio_hop_len = 256; } break; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + { + // CAMPPlus x-vector; kaldi fbank front-end (povey + // window, 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_QWEN3TTS_GEN: { // discrete-token autoregressive predictor, no mel-frontend needed @@ -1728,6 +1766,11 @@ struct clip_model_loader { // matches the reference decoder's sliding_window (speech_tokenizer/config.json) hparams.wav_tfm_swa = 72; } break; + case PROJECTOR_TYPE_CHATTERBOX: + { + // s3gen output rate; the s3 tokenizer mel front-end runs at 16 kHz + get_u32("chatterbox.sample_rate", hparams.audio_sample_rate); + } break; case PROJECTOR_TYPE_PADDLEOCR: { hparams.n_merge = 2; @@ -2057,7 +2100,9 @@ struct clip_model_loader { const bool has_standard_layers = ( model.proj_type != PROJECTOR_TYPE_GEMMA3NV && - model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC); + model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC && + model.proj_type != PROJECTOR_TYPE_CHATTERBOX && + model.proj_type != PROJECTOR_TYPE_CHATTERBOX_SPKENC); // layers const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0; @@ -2823,6 +2868,274 @@ 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"); + + // host-read side data, accessed by name through + // clip_cbx_read_tensor: conditioning defaults, embedding + // tables, source module, filterbank + 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.cond.", 0) == 0 || name.rfind("a.gen.code.", 0) == 0 || + name.rfind("a.gen.t3.", 0) == 0 || name.rfind("a.gen.hift.m_source.", 0) == 0 || + name == "a.s3tok.mel_filters") { + model.cbx_tensors[name] = get_tensor(name); + } + } + } 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); + + // host-read side data, accessed by name through + // clip_cbx_read_tensor: 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_VOXTRAL: { model.conv1d_1_w = get_tensor(string_format(TN_CONV1D, 1, "weight")); @@ -3688,6 +4001,17 @@ struct clip_init_result clip_init(const char * fname, struct clip_context_params ctx_gen_audio = new clip_ctx(ctx_params); loader.load_hparams(ctx_gen_audio->model, CLIP_MODALITY_GEN_AUDIO); 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()); } @@ -4024,6 +4348,12 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { // a single speaker embedding vector, regardless of its length n_patches = 1; } break; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + { + // statistics pooling collapses the whole clip into a single + // speaker embedding vector, regardless of its length + n_patches = 1; + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { // one hidden-state vector fed back to the talker per call @@ -4070,7 +4400,38 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 return clip_encode(ctx, ¶ms); } +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 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 if (t->type == GGML_TYPE_I32) { + std::vector tmp(n); + ggml_backend_tensor_get(t, tmp.data(), 0, n * sizeof(int32_t)); + for (size_t i = 0; i < n; i++) out[i] = (float) 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(); @@ -4172,7 +4533,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } set_input_f32("inp_raw", inp_raw); - } else if (!(ctx->proj_type() == PROJECTOR_TYPE_QWEN3TTS_GEN && params->gen_process == CLIP_GEN_PROCESS_CODE2WAV)) { + } else if (!(ctx->proj_type() == PROJECTOR_TYPE_QWEN3TTS_GEN && params->gen_process == CLIP_GEN_PROCESS_CODE2WAV) && + !(ctx->proj_type() == PROJECTOR_TYPE_CHATTERBOX && (params->gen_process == CLIP_GEN_PROCESS_TTS || params->gen_process == CLIP_GEN_PROCESS_TTS_VOCODE))) { // audio input (code2wav has no hidden-state/raw input at all, its only input is the "inp_codes" tensor handled in the switch below) GGML_ASSERT(imgs.entries.size() == 1); @@ -4711,6 +5073,109 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { set_input_f32("qwen2_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 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(); + int n_prompt = 0; + std::vector tokens; + if (params->ref_tokens) { + n_prompt = (int) params->ref_tokens->size(); + tokens = *params->ref_tokens; + tokens.resize((size_t) n_prompt + n_gen); + } else { + // precomputed prompt ids, stored as floats and converted + // through the typed accessor like the other sidecar data + n_prompt = (int) clip_cbx_read_tensor(ctx, "a.gen.cond.gen_prompt_token", nullptr, 0); + tokens.resize((size_t) n_prompt + n_gen); + std::vector ids((size_t) n_prompt); + clip_cbx_read_tensor(ctx, "a.gen.cond.gen_prompt_token", ids.data(), ids.size()); + for (int i = 0; i < n_prompt; i++) { + tokens[(size_t) i] = (int32_t) ids[(size_t) i]; + } + } + memcpy(tokens.data() + n_prompt, params->codes->data(), (size_t) n_gen * sizeof(int32_t)); + const int T1 = n_prompt + n_gen; + const int T2 = 2 * T1; + set_input_i32("inp_tokens", tokens); + + // mel-rate reference conditioning, from the clip or the + // precomputed defaults shipped in the mmproj + if (params->ref_feat) { + set_input_f32("inp_prompt_feat", *params->ref_feat); + } else { + ggml_tensor * pf = model.cbx_tensors.at("a.gen.cond.gen_prompt_feat"); + std::vector feat(ggml_nelements(pf)); + ggml_backend_tensor_get(pf, feat.data(), 0, ggml_nbytes(pf)); + set_input_f32("inp_prompt_feat", feat); + } + if (params->ref_spk) { + set_input_f32("inp_xvec", *params->ref_spk); + } else { + ggml_tensor * sp = model.cbx_tensors.at("a.gen.cond.gen_embedding"); + std::vector spk(ggml_nelements(sp)); + ggml_backend_tensor_get(sp, spk.data(), 0, ggml_nbytes(sp)); + set_input_f32("inp_xvec", 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 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); + + // meanflow inputs: gaussian noise and the sinusoidal time + // embeddings of the solver span points, same schedule as the + // graph builder (matcha layout: sines then cosines, scale 1000) + std::vector noise((size_t) 80 * T2); + std::mt19937 rng(42); + std::normal_distribution nd(0.0f, 1.0f); + for (auto & f : noise) f = nd(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 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: @@ -4733,6 +5198,19 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { { // do nothing } break; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + { + // batchnorm epsilon and the ceil-mode correction of the cam + // seg pooling: pooled sums are divided by the full segment + // length, the last partial segment gets rescaled + set_input_f32("inp_eps", {1e-5f}); + const int T1 = (imgs.entries[0].nx() - 1) / 2 + 1; + const int S = (T1 + 99) / 100; + std::vector 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_QWEN3TTS_GEN: { if (params->gen_process == CLIP_GEN_PROCESS_CODE2WAV) { @@ -5252,6 +5730,205 @@ 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 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; + } + + // embedding rows of the produced codes, gathered from the talker + // speech table shipped in the mmproj; the multilingual variant adds + // its learned speech positions so the rows feed the conditioning + // perceiver directly + 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; + GGML_ASSERT(tab->type == GGML_TYPE_F16 || tab->type == GGML_TYPE_F32); + 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; + + params->out_code_embd->resize((size_t) n_tok * n_e); + std::vector h16(n_e); + std::vector pr(n_e); + for (int t = 0; t < n_tok; t++) { + float * dst = params->out_code_embd->data() + (size_t) t * n_e; + 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 mel_data(ggml_nelements(mel)); + ggml_backend_tensor_get(mel, mel_data.data(), 0, ggml_nbytes(mel)); + + // hift bridge: f0 -> harmonic source -> source stft on the host, + // then the vocoder graph, then 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 f0(n_mel_out); + ggml_backend_tensor_get(f0_t, f0.data(), 0, ggml_nbytes(f0_t)); + + // source: f0 upsampled x480 nearest, 9 harmonics, cumulative phase, + // uv gating and noise as in SineGen, 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 lw(9); + float lb = 0.0f; + GGML_ASSERT(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); + + std::mt19937 srng(1234); + std::uniform_real_distribution ud(-M_PI, M_PI); + std::normal_distribution 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 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 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 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 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 acc((size_t) n_out + n_fft, 0.0); + std::vector 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(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)); + } + + // the reference silences the first 20 ms and fades the next 20 ms in + // to hide the onset artifact of the flow prompt boundary (trim_fade) + 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; + } + if (params->out_codes != nullptr) { ggml_tensor * codes = ggml_graph_get_tensor(gf, "out_codes"); if (codes == nullptr) { @@ -5434,6 +6111,18 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_fc_w->ne[2]; case PROJECTOR_TYPE_QWEN3TTS_GEN: return ctx->model.gen_code_out_embd_w->ne[0]; + case PROJECTOR_TYPE_CHATTERBOX: + // gen-only stack, no input projection into the backbone; the mel + // channel count stands in for the interface dimension + return 80; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + { + // the encoder emits talker conditioning rows in the backbone + // embedding space, whose dim the speaker projection carries + auto it = ctx->model.cbx_tensors.find("a.cenc.spkr_enc.weight"); + GGML_ASSERT(it != ctx->model.cbx_tensors.end()); + return it->second->ne[1]; + } case PROJECTOR_TYPE_PARAKEET: return ctx->model.mm_1_w->ne[1]; default: diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index e969b2b9b192..ace45fd021e2 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -89,6 +89,9 @@ bool clip_image_batch_encode(struct clip_ctx * ctx, int n_threads, const struct enum clip_gen_process_type { CLIP_GEN_PROCESS_CODE_GEN, // h_state to codes CLIP_GEN_PROCESS_CODE2WAV, // 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 }; struct clip_encode_params { int n_threads = 1; @@ -106,18 +109,41 @@ struct clip_encode_params { int32_t top_k = 50; float top_p = 1.0f; std::vector * out_codes = nullptr; + // TOKENIZE: out_code_embd receives the speech embedding rows of the + // produced codes (kept apart from out_embd, which is reserved for graphs + // whose last node is the embedding tensor) + std::vector * out_code_embd = nullptr; // CODE2WAV: codes holds this frame's 16 RVQ codes, out_audio receives the // decoded PCM samples (F32). state_in is the state from the previous // call (null or wrong size means cold start, state is zero-filled). // state_out receives the state to pass into the next call. const std::vector * codes = nullptr; + // TOKENIZE input + const float * pcm_in = nullptr; + size_t n_pcm = 0; + // TTS reference conditioning overriding the precomputed cond.gen_* + // defaults: speech tokens, mel-rate features and the 80-dim speaker + // vector of the reference clip (null means default) + const std::vector * ref_tokens = nullptr; + const std::vector * ref_feat = nullptr; + const std::vector * ref_spk = nullptr; + // TTS_VOCODE internal stage inputs + const std::vector * mel_in = nullptr; + const std::vector * sstft_in = nullptr; + int vocode_n_mel = 0; + int vocode_n_stft = 0; + std::vector * out_spec = nullptr; std::vector * out_audio = nullptr; const std::vector * state_in = nullptr; std::vector * state_out = nullptr; }; bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params); +// read a chatterbox tensor by source name, converted to F32. returns the +// element count, 0 if not found. out may be null to only query the size. +size_t clip_cbx_read_tensor(struct clip_ctx * ctx, const char * name, float * out, size_t n_max); + bool clip_is_llava(const struct clip_ctx * ctx); // note for contributor: this clip_is_(model) pattern is deprecated // do NOT add new functions like this diff --git a/tools/mtmd/models/chatterbox-gen.cpp b/tools/mtmd/models/chatterbox-gen.cpp new file mode 100644 index 000000000000..a552d9835917 --- /dev/null +++ b/tools/mtmd/models/chatterbox-gen.cpp @@ -0,0 +1,584 @@ +#include "models.h" + +// Chatterbox generation graphs: flow encoder, cfm estimator, s3 tokenizer +// and hift vocoder. Weights come from the nested model.cbx structs; the +// speaker encoder lives in chatterbox-spkenc.cpp. + +// 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; +} + +static ggml_tensor * cbx_layer_norm(ggml_context * ctx0, ggml_tensor * w, ggml_tensor * b, ggml_tensor * x, float eps) { + x = ggml_norm(ctx0, x, eps); + x = ggml_mul(ctx0, x, w); + x = ggml_add(ctx0, x, b); + return x; +} + +// x [C, T] -> conv1d over time -> [OC, T_out]; kernel [K, IC, OC], explicit +// host-side asymmetric padding is applied by the caller through pad_l/pad_r +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] where +// out[j, i, h] = bd[(T-1) - i + j, i, h] (ggml ne0 is the fastest dim). +// Same buffer walk as the espnet rel_shift: left-pad one column, reinterpret +// rows/cols, drop the first row, reinterpret back, 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 = cbx_layer_norm(ctx0, l.norm_mha_w, l.norm_mha_b, x, 1e-5f); + + 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 = cbx_layer_norm(ctx0, l.norm_ff_w, l.norm_ff_b, x, 1e-5f); + 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 = cbx_layer_norm(ctx0, b.norm_w, b.norm_b, x, 1e-5f); + 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 * res = x; + ggml_tensor * cur = cbx_layer_norm(ctx0, b.norm1_w, b.norm1_b, x, 1e-5f); + ggml_tensor * q = ggml_mul_mat(ctx0, b.attn_q_w, cur); + ggml_tensor * k = ggml_mul_mat(ctx0, b.attn_k_w, cur); + ggml_tensor * v = ggml_mul_mat(ctx0, b.attn_v_w, cur); + q = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, q, d_head, n_head, T), 0, 2, 1, 3)); // [64, T, 8] + k = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, k, d_head, n_head, T), 0, 2, 1, 3)); + v = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, v, d_head, n_head, T), 0, 2, 1, 3)); + ggml_tensor * scores = ggml_scale(ctx0, ggml_mul_mat(ctx0, k, q), 1.0f / sqrtf((float) d_head)); + 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)); + o = ggml_reshape_2d(ctx0, o, n_head * d_head, T); + o = cbx_linear(b.attn_out_w, b.attn_out_b, o); + x = ggml_add(ctx0, res, o); + + res = x; + cur = cbx_layer_norm(ctx0, b.norm3_w, b.norm3_b, x, 1e-5f); + 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, res, cur); +} + +// one estimator evaluation; x_noise [80, T], mu [80, T], spks [80], +// cond [80, T], 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] + // conv layout is [C, T] with channels contiguous per step; conv helpers + // transpose internally, the pack above must land on the channel dim + 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: whisper style log-mel [T, 128] in, two stride 2 +// convs to token rate, 6 pre-norm attention blocks with neox rope on q/k and +// an fsmn memory over the value projection, then the fsq down projection. +// output is the post-tanh 8-dim 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 = cbx_layer_norm(ctx0, blk.attn_ln_w, blk.attn_ln_b, x, 1e-5f); + 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 k31 over time on the value projection, + // residual, added to the projected attention context + 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); + 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); + q = ggml_cont(ctx0, ggml_permute(ctx0, q, 0, 2, 1, 3)); // [64, T2, 20] + k = ggml_cont(ctx0, ggml_permute(ctx0, k, 0, 2, 1, 3)); + v = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, v, d_head, n_head, T2), 0, 2, 1, 3)); + ggml_tensor * scores = ggml_scale(ctx0, ggml_mul_mat(ctx0, k, q), 1.0f / sqrtf((float) d_head)); + ggml_tensor * probs = ggml_soft_max(ctx0, scores); + ggml_tensor * o = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, v)), probs); // [64, T2, 20] + o = ggml_cont(ctx0, ggml_permute(ctx0, o, 0, 2, 1, 3)); + o = ggml_reshape_2d(ctx0, o, n_head * d_head, T2); + o = cbx_linear(blk.attn_out_w, blk.attn_out_b, o); + x = ggml_add(ctx0, res, ggml_add(ctx0, o, fsm)); + + res = x; + cur = cbx_layer_norm(ctx0, blk.mlp_ln_w, blk.mlp_ln_b, x, 1e-5f); + 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 = cbx_layer_norm(ctx0, c.embed_norm_w, c.embed_norm_b, x, 1e-5f); + 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 = cbx_layer_norm(ctx0, c.up_embed_norm_w, c.up_embed_norm_b, x, 1e-5f); + 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 = cbx_layer_norm(ctx0, c.after_norm_w, c.after_norm_b, x, 1e-5f); + + // 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 (distilled): 2 euler steps + // over t = 0 -> 0.5 -> 1, no cfg, time embeds mix t and r. classic: 10 + // euler steps on the cosine schedule with cfg 0.7, time embeds on t only. + 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: prompt features then zeros, and the 80-dim + // speaker vector; both are fed by the host from either the reference + // clip or the precomputed defaults + 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 from the speaker encoder or the + // precomputed default, 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 000000000000..7220319b08e3 --- /dev/null +++ b/tools/mtmd/models/chatterbox-spkenc.cpp @@ -0,0 +1,161 @@ +#include "models.h" + +// Chatterbox speaker encoder: CAMPPlus x-vector on kaldi fbank features, +// output raw. Mirrors s3gen/xvector.py. + +// per-channel batchnorm on x [C, T]; scale = w / sqrt(var + eps), shift folds +// the running mean. w/b stay null on the affine=False variant +ggml_tensor * clip_graph_chatterbox_spkenc::bn1d(const clip_chatterbox::bn & n, ggml_tensor * x, ggml_tensor * eps) { + ggml_tensor * sd = ggml_sqrt(ctx0, ggml_add(ctx0, n.var, 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, ggml_tensor * eps) { + 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_add(ctx0, var, 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, ggml_tensor * eps) { + return ggml_relu(ctx0, cbx_bn2d(ctx0, n, x, eps)); +} + +// 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 * eps) { + ggml_tensor * cur = ggml_conv_2d(ctx0, r.conv1_w, x, 1, stride, 1, 1, 1, 1); + cur = bn2d_relu(r.bn1, cur, eps); + 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, eps); + 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, eps); + } + 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 * eps, ggml_tensor * segfix) { + ggml_tensor * h = ggml_relu(ctx0, bn1d(l.nl1_bn, x, eps)); + h = cbx_conv1d(l.linear1_w, nullptr, h, 1, 0, 0); + h = ggml_relu(ctx0, bn1d(l.nl2_bn, h, eps)); + + 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; + + ggml_tensor * eps = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1); + ggml_set_name(eps, "inp_eps"); + ggml_set_input(eps); + + 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, eps); + x = res2d(c.spk_layer1_0, x, 2, eps); + x = res2d(c.spk_layer1_1, x, 1, eps); + x = res2d(c.spk_layer2_0, x, 2, eps); + x = res2d(c.spk_layer2_1, x, 1, eps); + x = ggml_conv_2d(ctx0, c.spk_conv2_w, x, 1, 2, 1, 1, 1, 1); + x = bn2d_relu(c.spk_bn2, x, eps); + 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, eps)); + 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], eps, segfix); + x = ggml_concat(ctx0, x, out, 0); + } + x = ggml_relu(ctx0, bn1d(blk.transit_bn, x, eps)); + 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, eps)); // [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); + + // dense 1024 -> 192, batchnorm without affine, into the raw x-vector: + // normalization and the s3gen speaker affine belong to the code2wav graph + 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, eps); + 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 67b10b2a4610..b89e34dd213e 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -151,6 +151,55 @@ struct clip_graph_conformer : clip_graph { ggml_cgraph * build() override; }; +// linear/conv builders shared between 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 * eps); + ggml_tensor * bn2d_relu(const clip_chatterbox::bn & n, ggml_tensor * x, ggml_tensor * eps); + ggml_tensor * res2d(const clip_chatterbox::spk_res2d & r, ggml_tensor * x, int stride, ggml_tensor * eps); + ggml_tensor * cam_layer(const clip_chatterbox::spk_cam_layer & l, ggml_tensor * x, + int dil, ggml_tensor * eps, ggml_tensor * segfix); +}; + +struct clip_graph_chatterbox : clip_graph_chatterbox_base { + clip_gen_process_type gen_process = CLIP_GEN_PROCESS_CODE_GEN; + 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_granite_speech : clip_graph { clip_graph_granite_speech(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 2811d24df764..10e36fc1c62b 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 #include #include #include @@ -855,6 +856,550 @@ bool mtmd_audio_preprocessor_qwen3tts_spk::preprocess(const float * return true; } +// whisper style log-mel of the chatterbox s3 tokenizer (s3tokenizer.py) +static bool mtmd_audio_s3tok_log_mel(const float * samples, size_t n_samples, + const float * filters, int n_mel, + std::vector & out, int & n_frames) { + const int n_fft = 400; + const int hop = 160; + const int n_bins = n_fft / 2 + 1; + const int half = n_fft / 2; + const int n = (int) n_samples; + + n_frames = n / hop; + if (n_frames <= 0) { + return false; + } + + std::vector window(n_fft); + for (int i = 0; i < n_fft; i++) { + window[(size_t) i] = 0.5 * (1.0 - cos(2.0 * M_PI * i / n_fft)); + } + + std::vector mel((size_t) n_mel * n_frames, 0.0); + std::vector frame(n_fft); + std::vector power(n_bins); + for (int fr = 0; fr < n_frames; fr++) { + for (int i = 0; i < n_fft; i++) { + int idx = fr * hop - half + i; + if (idx < 0) { + idx = -idx; + } + if (idx >= n) { + idx = 2 * n - 2 - idx; + } + frame[(size_t) i] = samples[idx] * window[(size_t) i]; + } + + for (int k = 0; k < n_bins; k++) { + double re = 0.0, im = 0.0; + for (int i = 0; i < n_fft; i++) { + const double a = 2.0 * M_PI * k * i / n_fft; + re += frame[(size_t) i] * cos(a); + im -= frame[(size_t) i] * sin(a); + } + power[(size_t) k] = re * re + im * im; + } + + for (int m = 0; m < n_mel; m++) { + double e = 0.0; + const float * w = filters + (size_t) m * n_bins; + for (int k = 0; k < n_bins; k++) { + e += w[k] * power[(size_t) k]; + } + mel[(size_t) m * n_frames + fr] = log10(std::max(e, 1e-10)); + } + } + + double mx = mel[0]; + for (const double v : mel) { + mx = std::max(mx, v); + } + 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; +} + +// rational 3/2 upsampler: every output sample sits at source position +// 2 n / 3, interpolated by a hann windowed sinc cut just under the source +// nyquist. edges are zero extended. +static void mtmd_audio_upsample_3_2(const float * samples, size_t n_samples, std::vector & 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) +static bool mtmd_audio_matcha_log_mel(const float * samples, size_t n_samples, + std::vector & out, int & n_frames) { + const int n_fft = 1920; + const int hop = 480; + const int n_bins = n_fft / 2 + 1; + 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; + } + + mtmd_audio_cache cache; + cache.fill_sin_cos_table(n_fft); + cache.fill_hann_window(n_fft, true); + cache.fill_mel_filterbank_matrix(n_mel, n_fft, 24000, 0.0f, 8000.0f); + + std::vector fft_in((size_t) n_fft * 2, 0.0f); + std::vector fft_out((size_t) n_fft * 8); + std::vector mag(n_bins); + out.resize((size_t) n_mel * n_frames); + for (int fr = 0; fr < n_frames; fr++) { + for (int i = 0; i < n_fft; i++) { + int idx = fr * hop - pad + i; + if (idx < 0) { + idx = -idx; + } + if (idx >= n) { + idx = 2 * n - 2 - idx; + } + fft_in[(size_t) i] = samples[idx] * cache.hann_window[(size_t) i]; + } + fft(cache, fft_in.data(), n_fft, fft_out.data()); + + for (int k = 0; k < n_bins; k++) { + const float re = fft_out[2 * k + 0]; + const float im = fft_out[2 * k + 1]; + mag[(size_t) k] = sqrtf(re * re + im * im + 1e-9f); + } + for (int m = 0; m < n_mel; m++) { + float e = 0.0f; + const float * w = cache.filters.data.data() + (size_t) m * n_bins; + for (int k = 0; k < n_bins; k++) { + e += w[k] * mag[(size_t) k]; + } + out[(size_t) fr * n_mel + m] = logf(std::max(e, 1e-5f)); + } + } + return true; +} + +// librosa.effects.trim replica, rms windows 2048/512 centered with zero +// padding, non-silent where the window sits less than top_db under 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 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: centered reflect padded frames, +// hann 400 periodic, hop 160, squared magnitude, slaney mel 40 bins, no log +static bool mtmd_audio_ve_mel(const float * samples, size_t n_samples, + std::vector & out, int & n_frames) { + const int n_fft = 400; + const int hop = 160; + const int n_bins = n_fft / 2 + 1; + const int half = n_fft / 2; + const int n_mel = 40; + const int n = (int) n_samples; + + n_frames = 1 + n / hop; + if (n < 2) { + return false; + } + + mtmd_audio_cache cache; + cache.fill_mel_filterbank_matrix(n_mel, n_fft, 16000, 0.0f, 8000.0f); + + std::vector window(n_fft); + for (int i = 0; i < n_fft; i++) { + window[(size_t) i] = 0.5 * (1.0 - cos(2.0 * M_PI * i / n_fft)); + } + + std::vector frame(n_fft); + std::vector power(n_bins); + out.resize((size_t) n_mel * n_frames); + for (int fr = 0; fr < n_frames; fr++) { + for (int i = 0; i < n_fft; i++) { + int idx = fr * hop - half + i; + if (idx < 0) { + idx = -idx; + } + if (idx >= n) { + idx = 2 * n - 2 - idx; + } + frame[(size_t) i] = samples[idx] * window[(size_t) i]; + } + for (int k = 0; k < n_bins; k++) { + double re = 0.0, im = 0.0; + for (int i = 0; i < n_fft; i++) { + const double a = 2.0 * M_PI * k * i / n_fft; + re += frame[(size_t) i] * cos(a); + im -= frame[(size_t) i] * sin(a); + } + power[(size_t) k] = re * re + im * im; + } + for (int m = 0; m < n_mel; m++) { + double e = 0.0; + const float * w = cache.filters.data.data() + (size_t) m * n_bins; + for (int k = 0; k < n_bins; k++) { + e += w[k] * power[(size_t) k]; + } + out[(size_t) fr * n_mel + m] = (float) e; + } + } + 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 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 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 +// +// Mirrors torchaudio.compliance.kaldi.fbank(wav, num_mel_bins=80) at 16 kHz as +// used by the CAMPPlus x-vector front-end (s3gen/xvector.py extract_feature): +// snip_edges framing 400/160, per-frame dc removal, preemphasis 0.97, povey +// window, 512-point power spectrum, kaldi mel banks (low 20 Hz, high +// nyquist, nyquist fft bin excluded), log with float-eps floor, then the +// reference's own cepstral mean subtraction over time. +// + +void mtmd_audio_preprocessor_chatterbox_ref::initialize() { + const int frame_len = 400; + const int n_fft = 512; + const int n_bins = n_fft / 2; + const int n_mel = hparams.n_mel_bins; + const double sr = (double) hparams.audio_sample_rate; + + window.resize(frame_len); + for (int i = 0; i < frame_len; i++) { + 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; + + filters.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; 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[(size_t) m * n_bins + i] = (float) w; + } + } + } +} + +bool mtmd_audio_preprocessor_chatterbox_ref::preprocess(const float * samples, + size_t n_samples, + std::vector & output) { + output.clear(); + const int sr = (int) hparams.audio_sample_rate; + + // the turbo variant loudness-normalizes the whole clip before any feature + std::vector 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: CAMPPlus kaldi fbank, per-channel mean subtracted over time + { + const int frame_len = 400; + const int hop = 160; + const int n_fft = 512; + const int n_bins = n_fft / 2; + 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; + + GGML_ASSERT(!window.empty()); + GGML_ASSERT(!filters.empty()); + + 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); + + std::vector frame(n_fft); + std::vector power(n_bins); + 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; + + frame[0] = (x[0] - mean) * (1.0 - 0.97) * window[0]; + for (int i = 1; i < frame_len; i++) { + frame[(size_t) i] = ((x[i] - mean) - 0.97 * (x[i - 1] - mean)) * window[(size_t) i]; + } + std::fill(frame.begin() + frame_len, frame.end(), 0.0); + + for (int k = 0; k < n_bins; k++) { + double re = 0.0, im = 0.0; + for (int i = 0; i < frame_len; i++) { + const double a = 2.0 * M_PI * k * i / n_fft; + re += frame[(size_t) i] * cos(a); + im -= frame[(size_t) i] * sin(a); + } + power[(size_t) k] = re * re + im * im; + } + + for (int m = 0; m < n_mel; m++) { + double e = 0.0; + const float * w = filters.data() + (size_t) m * n_bins; + for (int k = 0; k < n_bins; k++) { + e += w[k] * power[(size_t) k]; + } + out.data[(size_t) m * n_frames + fr] = (float) log(std::max(e, (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, each + // clip padded to whole 40 ms tokens so the mel stays twice the token grid + const int n_mel_s3 = (int) (s3tok_filters.size() / (400 / 2 + 1)); + auto s3tok_entry = [&](size_t cap) -> bool { + std::vector clip(pcm.begin(), pcm.begin() + std::min(pcm.size(), cap)); + clip.resize((clip.size() + 639) / 640 * 640, 0.0f); + std::vector mel; + int n_frames = 0; + if (!mtmd_audio_s3tok_log_mel(clip.data(), clip.size(), s3tok_filters.data(), n_mel_s3, mel, n_frames)) { + return false; + } + mtmd_audio_mel out; + out.n_len = n_frames; + out.n_len_org = n_frames; + out.n_mel = n_mel_s3; + 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 (n_mel_s3 == 0 || !s3tok_entry(gen_cap) || !s3tok_entry(t3_cap)) { + return false; + } + + // entry 3: s3gen prompt features, the flow-capped clip upsampled to the + // 24 kHz decoder rate then through the matcha log-mel + { + std::vector clip(pcm.begin(), pcm.begin() + std::min(pcm.size(), gen_cap)); + clip.resize((clip.size() + 639) / 640 * 640, 0.0f); + std::vector pcm24; + mtmd_audio_upsample_3_2(clip.data(), clip.size(), pcm24); + std::vector feat; + int n_frames = 0; + if (!mtmd_audio_matcha_log_mel(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 power mel of the silence-trimmed clip, padded + // (or trimmed) to the 160-frame partial grid at the reference 1.3 rate + { + 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 mel; + int n_frames = 0; + if (!mtmd_audio_ve_mel(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 b4d6f7259808..acf7c0605425 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -50,6 +50,7 @@ struct mtmd_audio_cache { ); }; + struct mtmd_audio_preprocessor { const clip_hparams & hparams; @@ -129,6 +130,29 @@ struct mtmd_audio_preprocessor_qwen3tts_spk : mtmd_audio_preprocessor { mtmd_audio_cache cache; }; +// full reference chain front-end of the chatterbox speaker encoder: one +// preprocess call emits the host DSP products of a reference clip as five +// entries, consumed positionally by the encoder orchestration 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) +// the turbo variant loudness-normalizes the clip to -27 LUFS first 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, std::vector s3tok_filters) + : mtmd_audio_preprocessor(ctx), is_mtl(is_mtl), s3tok_filters(std::move(s3tok_filters)) {} + void initialize() override; + bool preprocess(const float * samples, size_t n_samples, std::vector & output) override; + + private: + bool is_mtl; + std::vector s3tok_filters; // s3 tokenizer filterbank [n_mel x (400 / 2 + 1)], from the mmproj + std::vector window; // povey window of the fbank front-end, frame_length points + std::vector filters; // kaldi mel filterbank of the fbank front-end, n_mel x (n_fft / 2) dense +}; + 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 3b7762c2145f..2a792ee98a7b 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -91,6 +91,40 @@ class mtmd_gen_audio_pipeline { virtual int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) = 0; protected: + // encodes a speaker reference clip through the standard audio path: + // bitmap to chunks to chunk encode, out receives the encoder embeddings + bool encode_speaker(mtmd_bitmap * bitmap, std::vector & out) { + if (!mtmd_support_audio(mctx)) { + LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n"); + return false; + } + const std::string marker = mtmd_default_marker(); + mtmd_input_text text{ marker.c_str(), marker.size(), false, true }; + mtmd_input_chunks * chunks = mtmd_input_chunks_init(); + const mtmd_bitmap * bptr = bitmap; + bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0; + if (ok) { + ok = false; + for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) { + const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i); + if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) { + continue; + } + if (mtmd_encode_chunk(mctx, chunk) != 0) { + LOG_ERR("mtmd_helper_gen_audio: 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; @@ -325,38 +359,6 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { // encodes a reference wav (already loaded as a bitmap) through the mmproj's // speaker encoder, returning the single x-vector embedding row it produces - bool encode_speaker(mtmd_bitmap * bitmap, std::vector & out) { - if (!mtmd_support_audio(mctx)) { - LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n"); - return false; - } - const std::string marker = mtmd_default_marker(); - mtmd_input_text text{ marker.c_str(), marker.size(), false, true }; - mtmd_input_chunks * chunks = mtmd_input_chunks_init(); - const mtmd_bitmap * bptr = bitmap; - bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0; - if (ok) { - ok = false; - for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) { - const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i); - if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) { - continue; - } - if (mtmd_encode_chunk(mctx, chunk) != 0) { - LOG_ERR("mtmd_helper_gen_audio: 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; - } - // runs one CODE2WAV process() call on whatever is currently buffered, carrying // the persisted state (KV cache + conv left-context) across batches bool flush_c2w() { @@ -413,10 +415,387 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::vector out_buf; }; + +// Chatterbox: single-track discrete AR (the backbone emits the s3 speech tokens +// directly) into a one-shot flow-matching mel decode and NSF-iSTFT vocoder. +// The prompt is pure embedding concat [spkr, cond speech, text, speech bos], +// positions are handled by the backbone (gpt2 wpe / llama learned pos). +class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { +public: + using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline; + + void reset() override { + pos = 0; + ar_idx = 0; + codes_buf.clear(); + audio_pcm.clear(); + h_state_buf.clear(); + out_buf.clear(); + ref_cond.clear(); + ref_codes.clear(); + ref_feat.clear(); + ref_spk.clear(); + } + + int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override { + reset(); + + if (!ensure_cache()) { + return 1; + } + + if (inp->speaker_ref) { + // the reference clip goes through the standard audio encode path: + // the chunk encodes to the talker conditioning rows, the flow + // decoder reference comes back through the typed side outputs + if (!encode_speaker(inp->speaker_ref, ref_cond)) { + LOG_ERR("mtmd_helper_gen_audio: speaker reference encoding failed\n"); + return 1; + } + size_t n = 0; + const float * p = mtmd_get_output_typed_embd(mctx, MTMD_EMBD_OUT_TYPE_REF_CODES, &n); + ref_codes.assign(p, p + n); + p = mtmd_get_output_typed_embd(mctx, MTMD_EMBD_OUT_TYPE_REF_FEAT, &n); + ref_feat.assign(p, p + n); + p = mtmd_get_output_typed_embd(mctx, MTMD_EMBD_OUT_TYPE_REF_SPK, &n); + ref_spk.assign(p, p + n); + } + + const int n_e = n_embd; + auto row = [&](llama_token t) { + return std::vector(tok_embd.begin() + (size_t) t * n_e, + tok_embd.begin() + (size_t) (t + 1) * n_e); + }; + auto add_pos = [&](std::vector & r, const std::vector & tab, int idx) { + const float * p = tab.data() + (size_t) idx * n_e; + for (int j = 0; j < n_e; j++) { + r[(size_t) j] += p[j]; + } + }; + const bool mtl = !t3_cond.empty(); + + std::vector> prompt; + + if (mtl) { + // conditioning: [spkr, perceiver, emotion] block, cloned from the + // reference clip when present, precomputed default otherwise + const auto & cond = ref_cond.empty() ? t3_cond : ref_cond; + for (size_t i = 0; i < cond.size() / (size_t) n_e; i++) { + prompt.emplace_back(cond.begin() + i * (size_t) n_e, cond.begin() + (i + 1) * (size_t) n_e); + } + } else { + // conditioning rows: the reference block from the spk-ref stage, + // or the precomputed default speaker row followed by the + // table-resolved prompt ids + if (ref_cond.empty()) { + prompt.push_back(cond_spkr); + for (size_t i = 0; i < cond_speech_rows.size() / (size_t) n_e; i++) { + prompt.emplace_back(cond_speech_rows.begin() + i * (size_t) n_e, + cond_speech_rows.begin() + (i + 1) * (size_t) n_e); + } + } else { + for (size_t i = 0; i < ref_cond.size() / (size_t) n_e; i++) { + prompt.emplace_back(ref_cond.begin() + i * (size_t) n_e, + ref_cond.begin() + (i + 1) * (size_t) n_e); + } + } + } + + // the raw prompt goes straight to the tokenizer: the multilingual + // vocab carries the [SPACE] substitution and the case folding lives + // in the embedding table, both baked in at conversion + const std::string txt(inp->prompt, inp->prompt_len); + std::vector 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); + } + const size_t text0 = prompt.size(); + for (size_t i = 0; i < ids.size(); i++) { + prompt.push_back(row(ids[i])); + if (mtl) { + add_pos(prompt.back(), text_pos, (int) i); + } + } + + // speech bos opens the AR stream + prompt.push_back(row(llama_vocab_bos(vocab))); + if (mtl) { + add_pos(prompt.back(), speech_pos, 0); + } + ar_idx = 1; + + const int n_prompt = (int) prompt.size(); + std::vector embd_buf((size_t) n_prompt * (size_t) n_e); + for (int i = 0; i < n_prompt; i++) { + memcpy(embd_buf.data() + (size_t) i * n_e, prompt[(size_t) i].data(), (size_t) n_e * sizeof(float)); + } + + decode_embd_batch batch_embd(embd_buf.data(), n_prompt, 1, n_e); + batch_embd.set_position_normal(0, 0); + batch_embd.batch.logits[n_prompt - 1] = 1; + + if (llama_decode(lctx, batch_embd.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: prefill decode failed\n"); + return 1; + } + + if (mtl) { + // cfg second sequence: the same prompt with the text embeddings + // zeroed, keeping their learned positions (reference prepares the + // uncond branch before the position add) + std::vector cond_logits; + cfg_read(cond_logits); + for (size_t i = 0; i < ids.size(); i++) { + std::vector u((size_t) n_e, 0.0f); + add_pos(u, text_pos, (int) i); + memcpy(embd_buf.data() + (text0 + i) * (size_t) n_e, u.data(), (size_t) n_e * sizeof(float)); + } + decode_embd_batch batch_uncond(embd_buf.data(), n_prompt, 1, n_e); + // the uncond branch runs in the paired sequence half a seq space + // away from the cond one (slot i pairs with n_seq_max / 2 + i) + const llama_seq_id seq_uncond = (llama_seq_id) (llama_n_seq_max(lctx) / 2); + batch_uncond.set_position_normal(0, seq_uncond); + batch_uncond.batch.logits[n_prompt - 1] = 1; + if (llama_decode(lctx, batch_uncond.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: cfg prefill decode failed\n"); + return 1; + } + cfg_apply(cond_logits); + } + + pos = n_prompt; + out_type = inp->out_type; + return 0; + } + + int32_t step(llama_token sampled, const float * h_state_in, const float ** h_state_out) override { + GGML_UNUSED(h_state_in); + + // 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); + } + + if (!t3_cond.empty()) { + // multilingual backbone reads embeddings with the learned speech + // position added on top of the token row + std::vector e(tok_embd.begin() + (size_t) sampled * n_embd, + tok_embd.begin() + (size_t) (sampled + 1) * n_embd); + 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++; + decode_embd_batch batch_embd(e.data(), 1, 1, n_embd); + batch_embd.set_position_normal(pos, 0); + batch_embd.batch.logits[0] = 1; + if (llama_decode(lctx, batch_embd.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: step decode failed\n"); + return 1; + } + // cfg second sequence: the sampled token feeds both branches + std::vector cond_logits; + cfg_read(cond_logits); + decode_embd_batch batch_uncond(e.data(), 1, 1, n_embd); + const llama_seq_id seq_uncond = (llama_seq_id) (llama_n_seq_max(lctx) / 2); + batch_uncond.set_position_normal(pos, seq_uncond); + batch_uncond.batch.logits[0] = 1; + if (llama_decode(lctx, batch_uncond.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: cfg step decode failed\n"); + return 1; + } + cfg_apply(cond_logits); + } else { + llama_batch batch = llama_batch_get_one(&sampled, 1); + if (llama_decode(lctx, batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: step decode failed\n"); + return 1; + } + } + 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 { + if (codes_buf.empty()) { + LOG_ERR("mtmd_helper_gen_audio: no speech tokens generated\n"); + return 1; + } + + if (t3_cond.empty()) { + // turbo appends a short silence tail before vocoding + codes_buf.insert(codes_buf.end(), 3, 4299); + } + + mtmd_gen_inp gen_inp{}; + gen_inp.type = MTMD_GEN_PROCESS_TYPE_CODE2WAV; + gen_inp.codes = codes_buf.data(); + gen_inp.n_codes = codes_buf.size(); + if (!ref_codes.empty()) { + gen_inp.ref_codes = ref_codes.data(); gen_inp.n_ref_codes = ref_codes.size(); + gen_inp.ref_feat = ref_feat.data(); gen_inp.n_ref_feat = ref_feat.size(); + gen_inp.ref_spk = ref_spk.data(); gen_inp.n_ref_spk = ref_spk.size(); + } + mtmd_gen_out gen_out{}; + if (mtmd_gen_audio_process(mctx, &gen_inp, &gen_out) != 0) { + LOG_ERR("mtmd_helper_gen_audio: code2wav decode failed\n"); + return 1; + } + audio_pcm.assign(gen_out.audio, gen_out.audio + gen_out.n_samples); + + *out_sample_rate = info.sample_rate; + *out_n_samples = (int64_t) audio_pcm.size(); + out_buf.clear(); + if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV) { + write_wav16(out_buf, audio_pcm, info.sample_rate); + } else { + out_buf.resize(audio_pcm.size() * sizeof(float)); + memcpy(out_buf.data(), audio_pcm.data(), out_buf.size()); + } + *out_data = out_buf.data(); + *out_data_len = out_buf.size(); + return 0; + } + +private: + // multilingual cfg, reference t3 combine: logits = cond + w * (cond - uncond) + // with the reference default weight 0.5. the cond row is saved after the + // first decode, the mix lands in the row the tool samples from (the last + // one with logits enabled, which the second decode produced) + void cfg_read(std::vector & cond_logits) { + const float * c = llama_get_logits_ith(lctx, -1); + cond_logits.assign(c, c + llama_vocab_n_tokens(vocab)); + } + void cfg_apply(const std::vector & cond_logits) { + 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]); + } + } + + 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; + } + n_speech = llama_vocab_n_tokens(vocab) - speech_base; + + // reference config: start_text_token = 255, stop_text_token = 0 + // (only the multilingual prompt wraps the text with them) + text_start = 255; + text_stop = 0; + + 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()); + + // multilingual variant: the mmproj ships a precomputed t3 conditioning + // block [spkr, perceiver, emotion] and the learned positional tables + // that the backbone needs added to its input embeddings + size_t n_t3 = mtmd_gen_audio_read_tensor(mctx, "a.gen.cond.t3_cond", nullptr, 0); + if (n_t3 > 0) { + t3_cond.resize(n_t3); + if (mtmd_gen_audio_read_tensor(mctx, "a.gen.cond.t3_cond", t3_cond.data(), n_t3) != n_t3 || + n_t3 % (size_t) n_embd != 0) { + LOG_ERR("mtmd_helper_gen_audio: a.gen.cond.t3_cond read failed\n"); + return false; + } + auto read_table = [&](const char * name, std::vector & dst) { + 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 || + n % (size_t) n_embd != 0) { + LOG_ERR("mtmd_helper_gen_audio: %s read failed\n", name); + return false; + } + return true; + }; + if (!read_table("a.gen.t3.text_pos_emb", text_pos) || !read_table("a.gen.t3.speech_pos_emb", speech_pos)) { + return false; + } + return true; + } + + cond_spkr.resize((size_t) n_embd); + if (mtmd_gen_audio_read_tensor(mctx, "a.gen.cond.spkr_default", cond_spkr.data(), cond_spkr.size()) != (size_t) n_embd) { + LOG_ERR("mtmd_helper_gen_audio: a.gen.cond.spkr_default missing\n"); + return false; + } + // resolve the precomputed conditioning ids through the speech + // embedding table shipped in the mmproj + size_t n_ct = mtmd_gen_audio_read_tensor(mctx, "a.gen.cond.prompt_speech_tokens", nullptr, 0); + std::vector cond_ids(n_ct); + if (n_ct == 0 || mtmd_gen_audio_read_tensor(mctx, "a.gen.cond.prompt_speech_tokens", cond_ids.data(), n_ct) != n_ct) { + LOG_ERR("mtmd_helper_gen_audio: a.gen.cond.prompt_speech_tokens missing\n"); + return false; + } + const size_t n_tab = mtmd_gen_audio_read_tensor(mctx, "a.gen.code.out_embd.weight", nullptr, 0); + std::vector table(n_tab); + if (n_tab == 0 || n_tab % (size_t) n_embd != 0 || + mtmd_gen_audio_read_tensor(mctx, "a.gen.code.out_embd.weight", table.data(), n_tab) != n_tab) { + LOG_ERR("mtmd_helper_gen_audio: a.gen.code.out_embd.weight read failed\n"); + return false; + } + cond_speech_rows.resize(n_ct * (size_t) n_embd); + for (size_t i = 0; i < n_ct; i++) { + const size_t r = (size_t) cond_ids[i] * (size_t) n_embd; + memcpy(cond_speech_rows.data() + i * (size_t) n_embd, table.data() + r, (size_t) n_embd * sizeof(float)); + } + return true; + } + + std::vector tok_embd; + std::vector cond_spkr; + std::vector cond_speech_rows; // default conditioning ids resolved through the mmproj speech table + std::vector t3_cond; + std::vector text_pos; + std::vector speech_pos; + std::vector ref_cond; + std::vector ref_codes; // flow reference codes carried as exact float values + std::vector ref_feat; + std::vector ref_spk; + llama_token speech_base = LLAMA_TOKEN_NULL; + int n_speech = 0; + llama_token text_start = LLAMA_TOKEN_NULL; + llama_token text_stop = LLAMA_TOKEN_NULL; + int ar_idx = 0; + + llama_pos pos = 0; + std::vector codes_buf; + std::vector audio_pcm; + std::vector h_state_buf; + std::vector out_buf; + mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; +}; + static std::unique_ptr make_pipeline(llama_context * lctx, mtmd_context * mctx) { switch (mtmd_gen_audio_get_info(mctx).type) { case MTMD_GEN_AUDIO_TYPE_QWEN3TTS: return std::unique_ptr(new qwen3tts_gen_audio_pipeline(lctx, mctx)); + case MTMD_GEN_AUDIO_TYPE_CHATTERBOX: + return std::unique_ptr(new chatterbox_gen_audio_pipeline(lctx, mctx)); default: return nullptr; } diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 361af6dfb51c..af329c0c413f 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -61,6 +61,10 @@ struct mtmd_bitmap { return data; } + std::vector & get_rw_buf() { + return data; + } + bool is_placeholder() const { return data.empty(); } @@ -269,6 +273,13 @@ struct mtmd_context { std::vector gen_out_audio; // decoded PCM samples for the current frame (CODE2WAV) std::vector gen_out_state; // state to feed into the next CODE2WAV call + // typed side outputs of the last encoded reference chunk (chatterbox): + // flow codes carried as exact float values, mel-rate features, speaker + // vector; read back through mtmd_get_output_typed_embd + std::vector gen_out_ref_codes; + std::vector gen_out_ref_feat; + std::vector gen_out_ref_spk; + bool print_timings; int n_threads; std::string media_marker; @@ -362,7 +373,7 @@ struct mtmd_context { ctx_v = res.ctx_v; ctx_a = res.ctx_a; ctx_gen_a = res.ctx_gen_a; - if (!ctx_v && !ctx_a) { + if (!ctx_v && !ctx_a && !ctx_gen_a) { throw std::runtime_error(string_format("Failed to load CLIP model from %s\n", mmproj_fname)); } @@ -379,14 +390,17 @@ struct mtmd_context { // since we already validate n_embd of vision and audio mmproj, // we can safely assume that they are the same - int n_embd_clip = clip_n_mmproj_embd(ctx_v ? ctx_v : ctx_a); - if (n_embd_text > 0 && n_embd_text != n_embd_clip) { - throw std::runtime_error(string_format( - "mismatch between text model (n_embd = %d) and mmproj (n_embd = %d)\n" - "hint: you may be using wrong mmproj\n", - n_embd_text, n_embd_clip)); - } - if (ctx_gen_a) { + if (ctx_v || ctx_a) { + int n_embd_clip = clip_n_mmproj_embd(ctx_v ? ctx_v : ctx_a); + if (n_embd_text > 0 && n_embd_text != n_embd_clip) { + throw std::runtime_error(string_format( + "mismatch between text model (n_embd = %d) and mmproj (n_embd = %d)\n" + "hint: you may be using wrong mmproj\n", + n_embd_text, n_embd_clip)); + } + } + if (ctx_gen_a && clip_get_projector_type(ctx_gen_a) != PROJECTOR_TYPE_CHATTERBOX) { + // the chatterbox gen mmproj emits mel channels, not text embeddings int n_embd_gen = clip_n_mmproj_embd(ctx_gen_a); if (n_embd_text > 0 && n_embd_text != n_embd_gen) { throw std::runtime_error(string_format( @@ -761,6 +775,25 @@ struct mtmd_context { { audio_preproc = std::make_unique(ctx_a); } break; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + { + { + // the ref-chain front-end needs the s3 tokenizer + // filterbank of the generation context and the + // variant to pick its caps and loudness handling + std::vector filt; + bool is_mtl = false; + if (ctx_gen_a) { + const size_t n_filt = clip_cbx_read_tensor(ctx_gen_a, "a.s3tok.mel_filters", nullptr, 0); + if (n_filt > 0) { + filt.resize(n_filt); + clip_cbx_read_tensor(ctx_gen_a, "a.s3tok.mel_filters", filt.data(), n_filt); + } + is_mtl = clip_cbx_read_tensor(ctx_gen_a, "a.gen.t3.speech_pos_emb", nullptr, 0) > 0; + } + audio_preproc = std::make_unique(ctx_a, is_mtl, std::move(filt)); + } + } break; default: throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj)); } @@ -1337,6 +1370,41 @@ struct mtmd_tokenizer { } } + // the chatterbox reference encoder consumes the five DSP entries + // of one clip as a single chunk; its token count is the number of + // talker conditioning rows the clip encodes to (fixed on the + // multilingual variant, spkr row plus the t3 token grid on turbo) + if (clip_get_projector_type(ctx->ctx_a) == PROJECTOR_TYPE_CHATTERBOX_SPKENC) { + GGML_ASSERT(mel_spec_chunks.size() == 5); + const bool mtl = ctx->ctx_gen_a && + clip_cbx_read_tensor(ctx->ctx_gen_a, "a.gen.t3.speech_pos_emb", nullptr, 0) > 0; + const size_t n_tokens = mtl ? 34 : 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) { @@ -1507,6 +1575,10 @@ static int32_t mtmd_encode_impl(mtmd_context * ctx, const mtmd_image_tokens * im return ok ? 0 : 1; } +// defined with the reference chain below; encodes a chatterbox reference +// chunk into talker conditioning rows and the typed decoder reference outputs +static int32_t cbx_ref_encode(mtmd_context * ctx, const mtmd_audio_tokens * audio_tokens, std::vector & out_embd); + static int32_t mtmd_encode_chunk_impl(mtmd_context * ctx, const mtmd_input_chunk * chunk, std::vector & out_embd) { if (chunk->type == MTMD_INPUT_CHUNK_TYPE_TEXT) { LOG_WRN("mtmd_encode_chunk has no effect for text chunks\n"); @@ -1538,6 +1610,9 @@ static int32_t mtmd_encode_chunk_impl(mtmd_context * ctx, const mtmd_input_chunk LOG_ERR("%s: audio tokens batch is placeholder\n", __func__); return 1; } + if (clip_get_projector_type(ctx->ctx_a) == PROJECTOR_TYPE_CHATTERBOX_SPKENC) { + return cbx_ref_encode(ctx, chunk->tokens_audio.get(), out_embd); + } int n_mmproj_embd = ctx->n_embd_out(); out_embd.resize((size_t)chunk->tokens_audio->n_tokens * n_mmproj_embd); bool ok = clip_image_batch_encode( @@ -1575,6 +1650,23 @@ float * mtmd_get_output_embd(mtmd_context * ctx) { return ctx->out_embd.data(); } +const float * mtmd_get_output_typed_embd(mtmd_context * ctx, + enum mtmd_embd_out_type type, + size_t * n_elements) { + const std::vector * buf = nullptr; + switch (type) { + case MTMD_EMBD_OUT_TYPE_REF_CODES: buf = &ctx->gen_out_ref_codes; break; + case MTMD_EMBD_OUT_TYPE_REF_FEAT: buf = &ctx->gen_out_ref_feat; break; + case MTMD_EMBD_OUT_TYPE_REF_SPK: buf = &ctx->gen_out_ref_spk; break; + } + if (!buf || buf->empty()) { + *n_elements = 0; + return nullptr; + } + *n_elements = buf->size(); + return buf->data(); +} + // // audio generation // @@ -1590,6 +1682,10 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) { info.type = MTMD_GEN_AUDIO_TYPE_QWEN3TTS; info.sample_rate = 24000; break; + case PROJECTOR_TYPE_CHATTERBOX: + info.type = MTMD_GEN_AUDIO_TYPE_CHATTERBOX; + info.sample_rate = 24000; + break; default: info.type = MTMD_GEN_AUDIO_TYPE_NONE; break; @@ -1597,6 +1693,368 @@ 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 reference clip: the precomputed fbank +// entry through the CAMPPlus graph of the speaker encoding context +static bool cbx_ref_xvec(mtmd_context * ctx, const clip_image_f32 & fbank, std::vector & xvec) { + if (!ctx->ctx_a || clip_get_projector_type(ctx->ctx_a) != PROJECTOR_TYPE_CHATTERBOX_SPKENC) { + LOG_ERR("%s: mmproj has no speaker encoder\n", __func__); + return false; + } + 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 precomputed reference log-mel entry; rows, when +// requested, receives the speech embedding rows of the codes (with the +// learned speech positions added on the multilingual variant) +static int32_t cbx_ref_tokenize(mtmd_context * ctx, const clip_image_f32 & mel, + std::vector & codes, std::vector * 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 out_codes; + std::vector 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 of a reference clip: voice encoder chain and +// speaker projection row; the multilingual variant appends its perceiver +// output over the reference speech embedding rows and the emotion row +static int32_t cbx_ref_cond(mtmd_context * ctx, const clip_image_f32 & ve_mel, + const std::vector & pse, std::vector & out_rows) { + clip_ctx * ctx_clip = ctx->ctx_a; + auto read_t = [&](const char * name, std::vector & v) -> bool { + const size_t n = clip_cbx_read_tensor(ctx_clip, name, nullptr, 0); + if (n == 0) { + return false; + } + v.resize(n); + return clip_cbx_read_tensor(ctx_clip, name, v.data(), n) == n; + }; + + // voice encoder reference chain (embeds_from_wavs) over the precomputed + // power mel entry, already on the 160-frame partial grid: 3-layer lstm + // per partial, projected/relu/normalized embeddings averaged into the + // utterance embedding + const std::vector & mel = ve_mel.get_ro_buf(); + const int n_mel = 40; + const int n_partial = 160; + const int step = (int) lround((16000.0 / 1.3) / n_partial); // reference rate 1.3 + const int n_wins = (ve_mel.nx() - n_partial) / step + 1; + + std::vector w_ih[3], w_hh[3], b_ih[3], b_hh[3]; + std::vector w_proj, b_proj; + for (int l = 0; l < 3; l++) { + const std::string s = std::to_string(l); + if (!read_t(("a.ve.lstm.weight_ih_l" + s).c_str(), w_ih[l]) || + !read_t(("a.ve.lstm.weight_hh_l" + s).c_str(), w_hh[l]) || + !read_t(("a.ve.lstm.bias_ih_l" + s).c_str(), b_ih[l]) || + !read_t(("a.ve.lstm.bias_hh_l" + s).c_str(), b_hh[l])) { + LOG_ERR("%s: model has no voice encoder\n", __func__); + return 1; + } + } + if (!read_t("a.ve.proj.weight", w_proj) || !read_t("a.ve.proj.bias", b_proj)) { + LOG_ERR("%s: model has no voice encoder projection\n", __func__); + return 1; + } + + const int n_h = 256; + std::vector ve(n_h, 0.0f); + std::vector 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 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 + std::vector w_spkr, b_spkr; + if (!read_t("a.cenc.spkr_enc.weight", w_spkr) || !read_t("a.cenc.spkr_enc.bias", b_spkr)) { + LOG_ERR("%s: model has no speaker conditioning projection\n", __func__); + return 1; + } + const int n_e = (int) b_spkr.size(); + std::vector 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; + } + + // multilingual variant: [spkr, perceiver x32, emotion] block, the + // perceiver runs its shared attention block as cross then self + // attention over the reference speech embedding rows + std::vector query; + if (read_t("a.cenc.perceiver.pre_attention_query", query)) { + if (pse.empty()) { + LOG_ERR("%s: reference speech embeddings required for the perceiver\n", __func__); + return 1; + } + std::vector ln_w, ln_b, wq, bq, wk, bk, wv, bv, wo, bo, emo; + if (!read_t("a.cenc.perceiver.attn.norm.weight", ln_w) || !read_t("a.cenc.perceiver.attn.norm.bias", ln_b) || + !read_t("a.cenc.perceiver.attn.to_q.weight", wq) || !read_t("a.cenc.perceiver.attn.to_q.bias", bq) || + !read_t("a.cenc.perceiver.attn.to_k.weight", wk) || !read_t("a.cenc.perceiver.attn.to_k.bias", bk) || + !read_t("a.cenc.perceiver.attn.to_v.weight", wv) || !read_t("a.cenc.perceiver.attn.to_v.bias", bv) || + !read_t("a.cenc.perceiver.attn.proj_out.weight", wo) || !read_t("a.cenc.perceiver.attn.proj_out.bias", bo) || + !read_t("a.cenc.emotion_adv_fc.weight", emo)) { + LOG_ERR("%s: model has an incomplete perceiver\n", __func__); + return 1; + } + const int n_head = 4; + 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 & w, const std::vector & b, + const std::vector & in, int n_rows, std::vector & 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 & x1, int n1, + const std::vector & x2, int n2, std::vector & out) { + std::vector 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 q, k, v; + linear(wq, bq, nx1, n1, q); + linear(wk, bk, nx2, n2, k); + linear(wv, bv, nx2, n2, v); + + std::vector ctxt((size_t) n1 * n_e); + std::vector 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); + const std::vector & x2 = pse; + std::vector pre, p32; + attn_block(query, n_q, x2, (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; +} + +// encodes the five DSP entries of a chatterbox reference chunk (fbank, s3 +// log-mels at the flow and t3 caps, s3gen prompt features, voice encoder +// mel): the primary output is the talker conditioning rows in the backbone +// embedding space, the flow decoder reference (codes, mel-rate features, +// speaker vector) lands in the typed side outputs +static int32_t cbx_ref_encode(mtmd_context * ctx, const mtmd_audio_tokens * audio_tokens, std::vector & 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 xvec; + if (!cbx_ref_xvec(ctx, entries[0], xvec)) { + return 1; + } + + std::vector flow_codes, t3_codes; + std::vector 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; + } + + // conditioning rows: [spkr] then, multilingual, the perceiver block over + // the reference rows, or, turbo, the raw reference rows + const bool mtl = clip_cbx_read_tensor(ctx->ctx_a, "a.cenc.perceiver.pre_attention_query", nullptr, 0) > 0; + std::vector pse; + if (mtl) { + pse = std::move(t3_rows); + } + std::vector 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()); + } + GGML_ASSERT(rows.size() == (size_t) audio_tokens->n_tokens * (size_t) ctx->n_embd_out()); + + ctx->gen_out_ref_codes.assign(flow_codes.begin(), flow_codes.end()); + ctx->gen_out_ref_feat.assign(entries[3].get_ro_buf().begin(), + entries[3].get_ro_buf().begin() + n_feat); + ctx->gen_out_ref_spk = std::move(xvec); + + out_embd = std::move(rows); + return 0; +} + + static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) { clip_ctx * ctx_clip = ctx->ctx_gen_a; if (!ctx_clip) { @@ -1643,6 +2101,63 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in } // MTMD_GEN_PROCESS_TYPE_CODE2WAV + if (clip_get_projector_type(ctx_clip) == PROJECTOR_TYPE_CHATTERBOX) { + if (!inp->codes || inp->n_codes == 0) { + LOG_ERR("%s: codes required for code2wav\n", __func__); + return 1; + } + std::vector in_codes(inp->codes, inp->codes + inp->n_codes); + std::vector out_audio; + + // optional voice cloning reference from the encoded reference chunk: + // flow prompt codes, mel-rate prompt features, raw x-vector; + // all null selects the model's precomputed default voice + std::vector ref_tokens; + std::vector ref_feat; + std::vector ref_spk; + const bool has_ref = inp->ref_codes && inp->n_ref_codes > 0; + if (has_ref) { + if (!inp->ref_feat || inp->n_ref_feat == 0 || !inp->ref_spk || inp->n_ref_spk == 0) { + LOG_ERR("%s: incomplete cloning reference\n", __func__); + return 1; + } + ref_tokens.resize(inp->n_ref_codes); + for (size_t i = 0; i < inp->n_ref_codes; i++) { + ref_tokens[i] = (int32_t) lroundf(inp->ref_codes[i]); + } + ref_feat.assign(inp->ref_feat, inp->ref_feat + inp->n_ref_feat); + ref_spk.assign(inp->ref_spk, inp->ref_spk + inp->n_ref_spk); + } + + // 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(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 = &out_audio; + params.ref_tokens = has_ref ? &ref_tokens : nullptr; + params.ref_feat = has_ref ? &ref_feat : nullptr; + params.ref_spk = has_ref ? &ref_spk : nullptr; + + if (!clip_encode(ctx_clip, ¶ms)) { + LOG_ERR("%s: clip_encode failed (code2wav)\n", __func__); + return 1; + } + + ctx->gen_out_audio = std::move(out_audio); + out->audio = ctx->gen_out_audio.data(); + out->n_samples = ctx->gen_out_audio.size(); + return 0; + } + if (!inp->codes || inp->n_codes == 0) { LOG_ERR("%s: codes required for code2wav\n", __func__); return 1; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index f7d1fc65b79c..d228b65c3013 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -295,6 +295,20 @@ MTMD_API int32_t mtmd_encode_chunk(mtmd_context * ctx, // llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk) * sizeof(float) MTMD_API float * mtmd_get_output_embd(mtmd_context * ctx); +// typed side outputs of the last encoded chunk, for encoders that produce +// more than the backbone-space embeddings; entries are model-defined and +// valid until the next encode call. n_elements receives the element count, +// the return is null when the encoder has no output of the requested type. +// integer outputs (audio codes) are carried as exact float values. +enum mtmd_embd_out_type { + MTMD_EMBD_OUT_TYPE_REF_CODES, // speech codes of the reference clip + MTMD_EMBD_OUT_TYPE_REF_FEAT, // mel-rate features of the reference clip + MTMD_EMBD_OUT_TYPE_REF_SPK, // raw x-vector of the reference clip +}; +MTMD_API const float * mtmd_get_output_typed_embd(mtmd_context * ctx, + enum mtmd_embd_out_type type, + size_t * n_elements); + // batch encoding API // chunks are not owned by the batch, they will not be freed by mtmd_batch_free() @@ -334,6 +348,7 @@ MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname); enum mtmd_gen_audio_type { MTMD_GEN_AUDIO_TYPE_NONE, // not supported MTMD_GEN_AUDIO_TYPE_QWEN3TTS, + MTMD_GEN_AUDIO_TYPE_CHATTERBOX, }; struct mtmd_gen_audio_info { enum mtmd_gen_audio_type type; @@ -341,6 +356,12 @@ struct mtmd_gen_audio_info { }; MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx); +// read a named conditioning tensor from the gen audio model, converted to 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); + + + enum mtmd_gen_process_type { MTMD_GEN_PROCESS_TYPE_GEN_CODE, // h_state to codes MTMD_GEN_PROCESS_TYPE_CODE2WAV, // codes to raw PCM audio @@ -357,8 +378,15 @@ struct mtmd_gen_inp { // for MTMD_GEN_PROCESS_TYPE_CODE2WAV int32_t * codes; size_t n_codes; + // opaque state: the decoder carry-over between calls const char * state_data; size_t state_size; + // voice cloning reference from mtmd_get_output_typed_embd after encoding + // the reference clip; all null selects the model's precomputed default + // voice + const float * ref_codes; size_t n_ref_codes; + const float * ref_feat; size_t n_ref_feat; + const float * ref_spk; size_t n_ref_spk; }; struct mtmd_gen_out { // note: output memory is allocated by the context, valid until next process() call @@ -368,10 +396,12 @@ struct mtmd_gen_out { size_t n_codes; const float * embd; // the generated hidden state, to be fed back to backbone // it must have n_text_embd elements + size_t n_embd; // for MTMD_GEN_PROCESS_TYPE_CODE2WAV const float * audio; size_t n_samples; + // opaque state: the decoder carry-over to pass into the next CODE2WAV call const char * state_data; size_t state_size; }; diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index 23e797dc1255..7392c8c4be35 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -1,3 +1,4 @@ +#include #include "arg.h" #include "common.h" #include "sampling.h" @@ -65,6 +66,12 @@ int main(int argc, char ** argv) { // always enable embd, so that we can pass hidden states to the audio generation helper params.embedding = true; + // provision the cfg pair sequences for pipelines that decode one: each + // slot i pairs with the uncond sequence n_parallel + i, and the unified + // kv cache shares the window across sequences instead of splitting it + params.n_parallel *= 2; + params.kv_unified = true; + llama_backend_init(); llama_numa_init(params.numa); @@ -114,17 +121,10 @@ int main(int argc, char ** argv) { return 1; } - // codec_0 (backbone) EOS token: ordinary LLM sampling concern, kept out of the - // model-agnostic audio-generation helper + // codec_0 (backbone) generation ends on the model's eog tokens; sampling is + // restricted to the audio zone by the tokenizer.ggml.suppress_tokens GGUF + // metadata, merged into the sampling chain by common_sampler_init const llama_vocab * vocab = llama_model_get_vocab(model); - llama_token codec_eos_tok = LLAMA_TOKEN_NULL; - for (llama_token t = 0; t < llama_vocab_n_tokens(vocab); t++) { - if (!strcmp(llama_vocab_get_text(vocab, t), "<|codec_eos_token|>")) { codec_eos_tok = t; break; } - } - if (codec_eos_tok == LLAMA_TOKEN_NULL) { - LOG_ERR("missing codec eos token in vocab\n"); - return 1; - } auto sample_codec0 = [&]() -> llama_token { llama_token t = common_sampler_sample(smpl, lctx, -1); @@ -140,7 +140,7 @@ int main(int argc, char ** argv) { tts_timings timings; const int64_t t_gen_start_us = ggml_time_us(); - for (; n_frames < max_new && sampled != codec_eos_tok; n_frames++) { + for (; n_frames < max_new && !llama_vocab_is_eog(vocab, sampled); n_frames++) { const float * h_next = nullptr; if (gen.step(sampled, h_state, &h_next) != 0) { LOG_ERR("step failed at frame %d\n", n_frames);