mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-10-28 17:56:13 +08:00
add chatterbox TTS support
Co-authored-by: Pascal <admin@serveurperso.com>
This commit is contained in:
Xuan Son NguyenandPascal committed
1 parent
0a1ef309cd
commit
0353975120
18 files changed
+3971
-79
No files matched your search
@@ -40,6 +40,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
|
||||
"BloomModel": "bloom",
|
||||
"CamembertModel": "bert",
|
||||
"ChameleonForCausalLM": "chameleon",
|
||||
"ChatterboxModel": "chatterbox",
|
||||
"ChameleonForConditionalGeneration": "chameleon",
|
||||
"ChatGLMForConditionalGeneration": "chatglm",
|
||||
"ChatGLMModel": "chatglm",
|
||||
@@ -314,6 +315,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
|
||||
|
||||
|
||||
MMPROJ_MODEL_MAP: dict[str, str] = {
|
||||
"ChatterboxModel": "chatterbox",
|
||||
"AudioFlamingo3ForConditionalGeneration": "ultravox",
|
||||
"ClefModel": "clef",
|
||||
"CogVLMForCausalLM": "cogvlm",
|
||||
|
||||
@@ -0,0 +1,612 @@
|
||||
# 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. The checkpoints ship no
|
||||
# config.json, the hparams below come from the reference code (llama_configs.py, t3_config.py).
|
||||
#
|
||||
# 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 logging
|
||||
import math
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Callable, Iterable
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch import Tensor
|
||||
|
||||
from .base import ModelBase, TextModel, MmprojModel, LazyTorchTensor, gguf
|
||||
|
||||
logger = logging.getLogger("hf-to-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}")
|
||||
|
||||
|
||||
# talker hparams of the reference (GPT2_MEDIUM_CONFIG / LLAMA_520M_CONFIG_DICT, T3Config)
|
||||
TURBO_TEXT_CONFIG = {
|
||||
"n_layer": 24, "n_embd": 1024, "n_head": 16, "n_ctx": 8196, "layer_norm_epsilon": 1e-05,
|
||||
"vocab_size": 50276, "speech_vocab_size": 6563,
|
||||
}
|
||||
MTL_TEXT_CONFIG = {
|
||||
"num_hidden_layers": 30, "hidden_size": 1024, "intermediate_size": 4096, "num_attention_heads": 16,
|
||||
"num_key_value_heads": 16, "head_dim": 64, "max_position_embeddings": 131072, "rms_norm_eps": 1e-05,
|
||||
"rope_theta": 500000.0,
|
||||
"rope_scaling": {"factor": 8.0, "high_freq_factor": 4.0, "low_freq_factor": 1.0, "original_max_position_embeddings": 8192},
|
||||
"vocab_size": 2454, "speech_vocab_size": 8194,
|
||||
}
|
||||
|
||||
|
||||
def _is_chatterbox_checkpoint(dir_model: Path) -> bool:
|
||||
return (dir_model / "ve.safetensors").is_file() and \
|
||||
((dir_model / TURBO_TALKER).is_file() or (dir_model / MTL_TALKER).is_file())
|
||||
|
||||
|
||||
@ModelBase.register_hparams_loader(_is_chatterbox_checkpoint)
|
||||
def _load_chatterbox_hparams(dir_model: Path) -> dict[str, Any]:
|
||||
turbo = _is_turbo(dir_model)
|
||||
logger.info(f"gguf: detected Chatterbox {'turbo' if turbo else 'multilingual v3'} checkpoint")
|
||||
text_config = dict(TURBO_TEXT_CONFIG if turbo else MTL_TEXT_CONFIG)
|
||||
# the mmproj tensors are named explicitly, they do not go through the block tensor map
|
||||
return {"architectures": ["ChatterboxModel"], "text_config": text_config, "audio_config": {"num_hidden_layers": 0}}
|
||||
|
||||
|
||||
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=True)
|
||||
t3c = conds["t3"] if isinstance(conds, dict) else conds.t3
|
||||
genc = conds["gen"] if isinstance(conds, dict) else conds.gen
|
||||
t3c = vars(t3c) if not isinstance(t3c, dict) else t3c
|
||||
genc = vars(genc) if not isinstance(genc, dict) else genc
|
||||
|
||||
def talker_tensor(name: str) -> Tensor:
|
||||
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))
|
||||
@@ -6394,6 +6394,8 @@ class VisionProjectorType:
|
||||
QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor
|
||||
POCKETTTS_SPKENC = "pockettts_spkenc" # audio: mimi encoder as voice-prompt encoder
|
||||
POCKETTTS_GEN = "pockettts_gen" # audio generation: flow-matching decoder + mimi decoder
|
||||
CHATTERBOX_SPKENC = "chatterbox_spkenc" # audio: CAMPPlus speaker encoder
|
||||
CHATTERBOX = "chatterbox" # audio generation: s3gen flow matching + HiFT vocoder
|
||||
HUNYUANVL = "hunyuanvl"
|
||||
PARAKEET = "parakeet" # audio
|
||||
MINIMAXM3 = "minimax_m3"
|
||||
|
||||
@@ -28,6 +28,8 @@ add_library(mtmd
|
||||
clip-model.h
|
||||
clip-graph.h
|
||||
models/models.h
|
||||
models/chatterbox-gen.cpp
|
||||
models/chatterbox-spkenc.cpp
|
||||
models/cogvlm.cpp
|
||||
models/conformer.cpp
|
||||
models/deepseek4v.cpp
|
||||
|
||||
@@ -549,6 +549,8 @@ enum projector_type {
|
||||
PROJECTOR_TYPE_MUSE_GLIMMER,
|
||||
PROJECTOR_TYPE_COHERE2V,
|
||||
PROJECTOR_TYPE_MUSE_SPARK_AUDIO,
|
||||
PROJECTOR_TYPE_CHATTERBOX,
|
||||
PROJECTOR_TYPE_CHATTERBOX_SPKENC,
|
||||
PROJECTOR_TYPE_UNKNOWN,
|
||||
};
|
||||
|
||||
@@ -623,6 +625,8 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
|
||||
{ PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"},
|
||||
{ PROJECTOR_TYPE_COHERE2V, "cohere2v"},
|
||||
{ PROJECTOR_TYPE_MUSE_SPARK_AUDIO, "muse-spark-audio"},
|
||||
{ PROJECTOR_TYPE_CHATTERBOX, "chatterbox"},
|
||||
{ PROJECTOR_TYPE_CHATTERBOX_SPKENC, "chatterbox_spkenc"},
|
||||
};
|
||||
|
||||
static projector_type clip_projector_type_from_string(const std::string & str) {
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
|
||||
#include <algorithm>
|
||||
#include <array>
|
||||
#include <map>
|
||||
#include <vector>
|
||||
#include <unordered_set>
|
||||
#include <cstdint>
|
||||
@@ -579,6 +580,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_block> tfm;
|
||||
ggml_tensor * conv_w = nullptr;
|
||||
ggml_tensor * conv_b = nullptr;
|
||||
};
|
||||
|
||||
// snake resblock unit of the hift vocoder
|
||||
struct hift_res_unit {
|
||||
ggml_tensor * act1_alpha = nullptr;
|
||||
ggml_tensor * act2_alpha = nullptr;
|
||||
ggml_tensor * conv1_w = nullptr;
|
||||
ggml_tensor * conv1_b = nullptr;
|
||||
ggml_tensor * conv2_w = nullptr;
|
||||
ggml_tensor * conv2_b = nullptr;
|
||||
};
|
||||
struct hift_res {
|
||||
std::vector<hift_res_unit> units; // dilations 1/3/5 on conv1
|
||||
};
|
||||
|
||||
// one hift upsample stage: conv transpose, source injection, 3 resblocks
|
||||
struct hift_up {
|
||||
ggml_tensor * up_w = nullptr;
|
||||
ggml_tensor * up_b = nullptr;
|
||||
ggml_tensor * source_down_w = nullptr;
|
||||
ggml_tensor * source_down_b = nullptr;
|
||||
hift_res source_res;
|
||||
hift_res res_0;
|
||||
hift_res res_1;
|
||||
hift_res res_2;
|
||||
};
|
||||
|
||||
// f0 predictor conv of the hift vocoder
|
||||
struct f0_conv {
|
||||
ggml_tensor * w = nullptr;
|
||||
ggml_tensor * b = nullptr;
|
||||
};
|
||||
|
||||
// s3 tokenizer attention block (neox rope on q/k, fsmn memory on v)
|
||||
struct s3tok_block {
|
||||
ggml_tensor * attn_ln_w = nullptr;
|
||||
ggml_tensor * attn_ln_b = nullptr;
|
||||
ggml_tensor * attn_q_w = nullptr;
|
||||
ggml_tensor * attn_q_b = nullptr;
|
||||
ggml_tensor * attn_k_w = nullptr; // no bias
|
||||
ggml_tensor * attn_v_w = nullptr;
|
||||
ggml_tensor * attn_v_b = nullptr;
|
||||
ggml_tensor * fsmn_w = nullptr;
|
||||
ggml_tensor * attn_out_w = nullptr;
|
||||
ggml_tensor * attn_out_b = nullptr;
|
||||
ggml_tensor * mlp_ln_w = nullptr;
|
||||
ggml_tensor * mlp_ln_b = nullptr;
|
||||
ggml_tensor * mlp_in_w = nullptr;
|
||||
ggml_tensor * mlp_in_b = nullptr;
|
||||
ggml_tensor * mlp_out_w = nullptr;
|
||||
ggml_tensor * mlp_out_b = nullptr;
|
||||
};
|
||||
|
||||
// folded batchnorm (w/b stay null on the affine free variant)
|
||||
struct bn {
|
||||
ggml_tensor * mean = nullptr;
|
||||
ggml_tensor * var = nullptr;
|
||||
ggml_tensor * w = nullptr;
|
||||
ggml_tensor * b = nullptr;
|
||||
};
|
||||
|
||||
// fcm residual 2d block of the speaker encoder
|
||||
// (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<spk_cam_layer> layers;
|
||||
bn transit_bn;
|
||||
ggml_tensor * transit_w = nullptr;
|
||||
};
|
||||
|
||||
// flow encoder
|
||||
ggml_tensor * input_embedding_w = nullptr; // flow.input_embedding
|
||||
ggml_tensor * embed_linear_w = nullptr; // fenc.embed.out.0
|
||||
ggml_tensor * embed_linear_b = nullptr;
|
||||
ggml_tensor * embed_norm_w = nullptr; // fenc.embed.out.1
|
||||
ggml_tensor * embed_norm_b = nullptr;
|
||||
ggml_tensor * pre_conv1_w = nullptr; // pre_lookahead_layer
|
||||
ggml_tensor * pre_conv1_b = nullptr;
|
||||
ggml_tensor * pre_conv2_w = nullptr;
|
||||
ggml_tensor * pre_conv2_b = nullptr;
|
||||
std::vector<enc_layer> enc;
|
||||
ggml_tensor * up_conv_w = nullptr; // fenc.up_layer.conv
|
||||
ggml_tensor * up_conv_b = nullptr;
|
||||
ggml_tensor * up_embed_linear_w = nullptr; // fenc.up_embed.out.0
|
||||
ggml_tensor * up_embed_linear_b = nullptr;
|
||||
ggml_tensor * up_embed_norm_w = nullptr; // fenc.up_embed.out.1
|
||||
ggml_tensor * up_embed_norm_b = nullptr;
|
||||
std::vector<enc_layer> up_enc;
|
||||
ggml_tensor * after_norm_w = nullptr;
|
||||
ggml_tensor * after_norm_b = nullptr;
|
||||
ggml_tensor * encoder_proj_w = nullptr; // flow.encoder_proj
|
||||
ggml_tensor * encoder_proj_b = nullptr;
|
||||
|
||||
// cfm estimator
|
||||
ggml_tensor * time_mlp_1_w = nullptr;
|
||||
ggml_tensor * time_mlp_1_b = nullptr;
|
||||
ggml_tensor * time_mlp_2_w = nullptr;
|
||||
ggml_tensor * time_mlp_2_b = nullptr;
|
||||
ggml_tensor * time_embed_mixer_w = nullptr; // meanflow variant only
|
||||
est_stage est_down;
|
||||
std::vector<est_stage> est_mid;
|
||||
est_stage est_up;
|
||||
causal_block est_final_block;
|
||||
ggml_tensor * est_final_proj_w = nullptr;
|
||||
ggml_tensor * est_final_proj_b = nullptr;
|
||||
|
||||
// hift vocoder
|
||||
ggml_tensor * hift_pre_w = nullptr;
|
||||
ggml_tensor * hift_pre_b = nullptr;
|
||||
std::vector<hift_up> hift_ups;
|
||||
ggml_tensor * hift_post_w = nullptr;
|
||||
ggml_tensor * hift_post_b = nullptr;
|
||||
std::vector<f0_conv> f0_condnet;
|
||||
ggml_tensor * f0_classifier_w = nullptr;
|
||||
ggml_tensor * f0_classifier_b = nullptr;
|
||||
|
||||
// s3 tokenizer
|
||||
ggml_tensor * s3tok_conv1_w = nullptr;
|
||||
ggml_tensor * s3tok_conv1_b = nullptr;
|
||||
ggml_tensor * s3tok_conv2_w = nullptr;
|
||||
ggml_tensor * s3tok_conv2_b = nullptr;
|
||||
std::vector<s3tok_block> s3tok_blocks;
|
||||
ggml_tensor * s3tok_down_w = nullptr; // quantizer._codebook.project_down
|
||||
ggml_tensor * s3tok_down_b = nullptr;
|
||||
|
||||
// CAMPPlus speaker encoder
|
||||
ggml_tensor * spk_conv1_w = nullptr;
|
||||
bn spk_bn1;
|
||||
spk_res2d spk_layer1_0;
|
||||
spk_res2d spk_layer1_1;
|
||||
spk_res2d spk_layer2_0;
|
||||
spk_res2d spk_layer2_1;
|
||||
ggml_tensor * spk_conv2_w = nullptr;
|
||||
bn spk_bn2;
|
||||
ggml_tensor * spk_tdnn_w = nullptr;
|
||||
bn spk_tdnn_bn;
|
||||
spk_cam_block spk_block1;
|
||||
spk_cam_block spk_block2;
|
||||
spk_cam_block spk_block3;
|
||||
bn spk_out_bn;
|
||||
ggml_tensor * spk_dense_w = nullptr;
|
||||
bn spk_dense_bn;
|
||||
ggml_tensor * spk_affine_w = nullptr;
|
||||
ggml_tensor * spk_affine_b = nullptr;
|
||||
};
|
||||
|
||||
struct clip_model {
|
||||
clip_modality modality = CLIP_MODALITY_VISION;
|
||||
projector_type proj_type = PROJECTOR_TYPE_MLP;
|
||||
@@ -845,6 +1077,15 @@ struct clip_model {
|
||||
ggml_tensor * gen_upsample_w = nullptr; // depthwise convtr, frame rate -> encoder frame rate
|
||||
std::vector<clip_layer> gen_tfm_layers; // mimi decoder_transformer
|
||||
|
||||
// chatterbox flow encoder, cfm estimator, hift vocoder, s3 tokenizer
|
||||
// and CAMPPlus speaker encoder
|
||||
clip_chatterbox cbx;
|
||||
|
||||
// chatterbox 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<std::string, ggml_tensor *> cbx_tensors;
|
||||
|
||||
// cogvlm
|
||||
ggml_tensor * mm_post_fc_norm_w = nullptr;
|
||||
ggml_tensor * mm_post_fc_norm_b = nullptr;
|
||||
|
||||
+698
-5
@@ -1156,6 +1156,35 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
|
||||
const int n_frames = params && params->feats ? (int) (params->feats->size() / n_latent) : 1;
|
||||
builder = std::make_unique<clip_graph_pockettts_gen>(ctx, img, gen_process, n_step, n_frames);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_CHATTERBOX_SPKENC:
|
||||
{
|
||||
builder = std::make_unique<clip_graph_chatterbox_spkenc>(ctx, img);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_CHATTERBOX:
|
||||
{
|
||||
const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
|
||||
int n_tokens = params && params->codes ? (int) params->codes->size() : 0;
|
||||
int n_prompt_mel = 0;
|
||||
if (n_tokens > 0) {
|
||||
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<clip_graph_chatterbox>(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_GEN_CODE;
|
||||
@@ -2018,6 +2047,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:
|
||||
{
|
||||
// TODO: hardcoded for now, read from code_predictor_config instead
|
||||
@@ -2053,6 +2091,11 @@ struct clip_model_loader {
|
||||
hparams.flow_n_step = 1;
|
||||
hparams.gen_eos_threshold = -4.0f;
|
||||
} 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;
|
||||
@@ -2432,7 +2475,9 @@ struct clip_model_loader {
|
||||
model.proj_type != PROJECTOR_TYPE_GEMMA3NV &&
|
||||
model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC &&
|
||||
model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN &&
|
||||
model.proj_type != PROJECTOR_TYPE_MUSE_SPARK_AUDIO);
|
||||
model.proj_type != PROJECTOR_TYPE_MUSE_SPARK_AUDIO &&
|
||||
model.proj_type != PROJECTOR_TYPE_CHATTERBOX &&
|
||||
model.proj_type != PROJECTOR_TYPE_CHATTERBOX_SPKENC);
|
||||
|
||||
// layers
|
||||
const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0;
|
||||
@@ -3369,6 +3414,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:
|
||||
case PROJECTOR_TYPE_VOXTRAL_RT:
|
||||
{
|
||||
@@ -4254,8 +4567,20 @@ 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());
|
||||
ctx_gen_audio->buf_compute_meta.resize(ctx_gen_audio->max_nodes * ggml_tensor_overhead() +
|
||||
ggml_graph_overhead_custom(ctx_gen_audio->max_nodes, false));
|
||||
}
|
||||
|
||||
} catch (const std::exception & e) {
|
||||
@@ -4643,6 +4968,12 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
|
||||
// pooling gives one speaker embedding, whatever the clip length is
|
||||
n_patches = 1;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_CHATTERBOX_SPKENC:
|
||||
{
|
||||
// statistics pooling 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
|
||||
@@ -4816,7 +5147,38 @@ void clip_enc_state_free(clip_enc_state * state) {
|
||||
delete state;
|
||||
}
|
||||
|
||||
size_t clip_cbx_read_tensor(struct clip_ctx * ctx, const char * name, float * out, size_t n_max) {
|
||||
const auto & tensors = ctx->model.cbx_tensors;
|
||||
auto it = tensors.find(name);
|
||||
if (it == tensors.end() || !it->second) {
|
||||
return 0;
|
||||
}
|
||||
ggml_tensor * t = it->second;
|
||||
const size_t n = (size_t) ggml_nelements(t);
|
||||
if (!out) {
|
||||
return n;
|
||||
}
|
||||
if (n_max < n) {
|
||||
return 0;
|
||||
}
|
||||
if (t->type == GGML_TYPE_F32) {
|
||||
ggml_backend_tensor_get(t, out, 0, n * sizeof(float));
|
||||
} else if (t->type == GGML_TYPE_F16) {
|
||||
std::vector<ggml_fp16_t> tmp(n);
|
||||
ggml_backend_tensor_get(t, tmp.data(), 0, n * sizeof(ggml_fp16_t));
|
||||
for (size_t i = 0; i < n; i++) out[i] = ggml_fp16_to_fp32(tmp[i]);
|
||||
} else if (t->type == GGML_TYPE_I32) {
|
||||
std::vector<int32_t> 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();
|
||||
|
||||
@@ -4831,7 +5193,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
clip_model_loader::warmup(*ctx, *params->imgs);
|
||||
}
|
||||
|
||||
if (params->gen_process != CLIP_GEN_PROCESS_GEN_UNKNOWN && params->gen_state == nullptr) {
|
||||
const bool needs_gen_state = params->gen_process == CLIP_GEN_PROCESS_GEN_CODE ||
|
||||
params->gen_process == CLIP_GEN_PROCESS_GEN_WAV ||
|
||||
params->gen_process == CLIP_GEN_PROCESS_TTS;
|
||||
if (needs_gen_state && params->gen_state == nullptr) {
|
||||
LOG_ERR("%s: audio gen requires a gen state\n", __func__);
|
||||
return false;
|
||||
}
|
||||
@@ -4978,8 +5343,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
}
|
||||
set_input_f32("inp_raw", inp_raw);
|
||||
|
||||
} else if (params->gen_process != CLIP_GEN_PROCESS_GEN_WAV) {
|
||||
// audio input. GEN_WAV is not here: it takes codes or feats, set in the switch below
|
||||
} else if (params->gen_process != CLIP_GEN_PROCESS_GEN_WAV &&
|
||||
params->gen_process != CLIP_GEN_PROCESS_TTS &&
|
||||
params->gen_process != CLIP_GEN_PROCESS_TTS_VOCODE) {
|
||||
// audio input. GEN_WAV and the chatterbox TTS stages are not here: they take codes or feats, set in the switch below
|
||||
GGML_ASSERT(imgs.entries.size() == 1);
|
||||
|
||||
const auto & mel_inp = imgs.entries[0];
|
||||
@@ -5701,6 +6068,108 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
set_input_attn_mask(qwen2_mask);
|
||||
}
|
||||
} break;
|
||||
case PROJECTOR_TYPE_CHATTERBOX:
|
||||
{
|
||||
if (params->gen_process == CLIP_GEN_PROCESS_TTS_VOCODE) {
|
||||
set_input_f32("inp_mel", *params->mel_in);
|
||||
set_input_f32("inp_sstft", *params->sstft_in);
|
||||
break;
|
||||
}
|
||||
if (params->gen_process == CLIP_GEN_PROCESS_TOKENIZE) {
|
||||
const int T1 = (imgs.entries[0].nx() - 1) / 2 + 1;
|
||||
const int T2 = (T1 - 1) / 2 + 1;
|
||||
std::vector<int32_t> pos(T2);
|
||||
for (int i = 0; i < T2; i++) {
|
||||
pos[(size_t) i] = i;
|
||||
}
|
||||
set_input_i32("inp_pos", pos);
|
||||
break;
|
||||
}
|
||||
if (params->gen_process != CLIP_GEN_PROCESS_TTS) {
|
||||
break;
|
||||
}
|
||||
const int n_gen = (int) params->codes->size();
|
||||
int n_prompt = 0;
|
||||
std::vector<int32_t> 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<float> 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<float> 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<float> 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<float> pos((size_t) (2 * T - 1) * d);
|
||||
for (int k = 0; k < 2 * T - 1; k++) {
|
||||
const double rel = (double) (T - 1 - k);
|
||||
for (int i = 0; i < d / 2; i++) {
|
||||
const double div = exp(-(double) (2 * i) * log(10000.0) / d);
|
||||
pos[(size_t) k * d + 2 * i ] = (float) sin(rel * div);
|
||||
pos[(size_t) k * d + 2 * i + 1] = (float) cos(rel * div);
|
||||
}
|
||||
}
|
||||
set_input_f32(name, pos);
|
||||
};
|
||||
fill_pos("inp_pos1", T1);
|
||||
fill_pos("inp_pos2", T2);
|
||||
|
||||
// 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<float> noise((size_t) 80 * T2);
|
||||
std::normal_distribution<float> nd(0.0f, 1.0f);
|
||||
for (auto & f : noise) f = nd(params->gen_state->rng);
|
||||
set_input_f32("inp_noise", noise);
|
||||
|
||||
const bool meanflow = model.cbx.time_embed_mixer_w != nullptr;
|
||||
const int n_steps = meanflow ? 2 : 10;
|
||||
std::vector<float> temb((size_t) 320 * (n_steps + 1));
|
||||
for (int s = 0; s <= n_steps; s++) {
|
||||
const double u = (double) s / n_steps;
|
||||
const double t = meanflow ? u : 1.0 - cos(u * M_PI / 2.0);
|
||||
for (int i = 0; i < 160; i++) {
|
||||
const double div = exp(-(double) i * log(10000.0) / 159.0);
|
||||
temb[(size_t) s * 320 + i ] = (float) sin(1000.0 * t * div);
|
||||
temb[(size_t) s * 320 + i + 160] = (float) cos(1000.0 * t * div);
|
||||
}
|
||||
}
|
||||
set_input_f32("inp_temb", temb);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_GEMMA3:
|
||||
case PROJECTOR_TYPE_GEMMA3NV:
|
||||
case PROJECTOR_TYPE_IDEFICS3:
|
||||
@@ -5727,6 +6196,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<float> segfix((size_t) S, 1.0f);
|
||||
const int last = T1 - (S - 1) * 100;
|
||||
segfix[(size_t) S - 1] = 100.0f / (float) last;
|
||||
set_input_f32("inp_segfix", segfix);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
{
|
||||
if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) {
|
||||
@@ -6247,6 +6729,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<float> h((size_t) 8 * n_tok);
|
||||
ggml_backend_tensor_get(fsq, h.data(), 0, ggml_nbytes(fsq));
|
||||
|
||||
// fsq round to base 3: h in (-1, 1) maps to digits {0, 1, 2}
|
||||
params->out_codes->resize(n_tok);
|
||||
for (int t = 0; t < n_tok; t++) {
|
||||
int32_t code = 0;
|
||||
for (int i = 7; i >= 0; i--) {
|
||||
const int32_t d = (int32_t) roundf(h[(size_t) t * 8 + i] * 0.9990000128746033f) + 1;
|
||||
code = code * 3 + d;
|
||||
}
|
||||
(*params->out_codes)[(size_t) t] = code;
|
||||
}
|
||||
|
||||
// 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<ggml_fp16_t> h16(n_e);
|
||||
std::vector<float> pr(n_e);
|
||||
for (int t = 0; t < n_tok; t++) {
|
||||
float * dst = params->out_code_embd->data() + (size_t) t * n_e;
|
||||
const size_t r = (size_t) (*params->out_codes)[(size_t) t] * n_e;
|
||||
if (tab->type == GGML_TYPE_F16) {
|
||||
ggml_backend_tensor_get(tab, h16.data(), r * sizeof(ggml_fp16_t), (size_t) n_e * sizeof(ggml_fp16_t));
|
||||
for (int j = 0; j < n_e; j++) {
|
||||
dst[j] = ggml_fp16_to_fp32(h16[j]);
|
||||
}
|
||||
} else {
|
||||
ggml_backend_tensor_get(tab, dst, r * sizeof(float), (size_t) n_e * sizeof(float));
|
||||
}
|
||||
if (pos) {
|
||||
ggml_backend_tensor_get(pos, pr.data(), (size_t) t * n_e * sizeof(float), (size_t) n_e * sizeof(float));
|
||||
for (int j = 0; j < n_e; j++) {
|
||||
dst[j] += pr[j];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
if (ctx->proj_type() == PROJECTOR_TYPE_CHATTERBOX && params->gen_process == CLIP_GEN_PROCESS_TTS_VOCODE) {
|
||||
ggml_tensor * sp = ggml_graph_get_tensor(gf, "out_spec");
|
||||
GGML_ASSERT(sp != nullptr && params->out_spec);
|
||||
params->out_spec->resize(ggml_nelements(sp));
|
||||
ggml_backend_tensor_get(sp, params->out_spec->data(), 0, ggml_nbytes(sp));
|
||||
return true;
|
||||
}
|
||||
|
||||
if (ctx->proj_type() == PROJECTOR_TYPE_CHATTERBOX && params->gen_process == CLIP_GEN_PROCESS_TTS) {
|
||||
ggml_tensor * mel = ggml_graph_get_tensor(gf, "out_mel");
|
||||
GGML_ASSERT(mel != nullptr);
|
||||
std::vector<float> mel_data(ggml_nelements(mel));
|
||||
ggml_backend_tensor_get(mel, mel_data.data(), 0, ggml_nbytes(mel));
|
||||
|
||||
// hift 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<float> 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<float> 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 = params->gen_state->rng;
|
||||
std::uniform_real_distribution<double> ud(-M_PI, M_PI);
|
||||
std::normal_distribution<float> snd(0.0f, 1.0f);
|
||||
double phase[9];
|
||||
for (int h = 0; h < 9; h++) {
|
||||
phase[h] = h == 0 ? 0.0 : ud(srng);
|
||||
}
|
||||
std::vector<float> src((size_t) n_wav);
|
||||
double cum[9] = {0.0};
|
||||
for (int64_t t = 0; t < n_wav; t++) {
|
||||
const float f = f0[(size_t) (t / ups_total)];
|
||||
const float uv = f > 10.0f ? 1.0f : 0.0f; // nsf_voiced_threshold
|
||||
const float namp = uv * 0.003f + (1.0f - uv) * 0.1f / 3.0f;
|
||||
float merged = lb;
|
||||
for (int h = 0; h < 9; h++) {
|
||||
cum[h] += (double) f * (h + 1) / sr;
|
||||
cum[h] -= floor(cum[h]);
|
||||
float sine = 0.1f * (float) sin(2.0 * M_PI * cum[h] + phase[h]);
|
||||
sine = sine * uv + namp * snd(srng);
|
||||
merged += lw[(size_t) h] * sine;
|
||||
}
|
||||
src[(size_t) t] = tanhf(merged);
|
||||
}
|
||||
|
||||
// stft of the source: n_fft 16, hop 4, hann window, centered
|
||||
const int n_fft = 16, hop = 4, n_bins = 9;
|
||||
std::vector<float> win(n_fft);
|
||||
for (int i = 0; i < n_fft; i++) win[(size_t) i] = 0.5f - 0.5f * cosf(2.0f * (float) M_PI * i / n_fft);
|
||||
const int n_stft = (int) (n_wav / hop) + 1;
|
||||
std::vector<float> sstft((size_t) n_stft * 18);
|
||||
for (int fr = 0; fr < n_stft; fr++) {
|
||||
const int64_t c0 = (int64_t) fr * hop - n_fft / 2; // centered, reflect padded
|
||||
for (int k = 0; k < n_bins; k++) {
|
||||
double re = 0.0, im = 0.0;
|
||||
for (int i = 0; i < n_fft; i++) {
|
||||
int64_t idx = c0 + i;
|
||||
if (idx < 0) idx = -idx;
|
||||
if (idx >= n_wav) idx = 2 * (n_wav - 1) - idx;
|
||||
const double v = (double) src[(size_t) idx] * win[(size_t) i];
|
||||
const double a = 2.0 * M_PI * k * i / n_fft;
|
||||
re += v * cos(a);
|
||||
im -= v * sin(a);
|
||||
}
|
||||
sstft[(size_t) fr * 18 + k ] = (float) re;
|
||||
sstft[(size_t) fr * 18 + 9 + k] = (float) im;
|
||||
}
|
||||
}
|
||||
|
||||
// vocoder graph on mel + source stft
|
||||
std::vector<float> spec;
|
||||
{
|
||||
clip_encode_params vp = *params;
|
||||
vp.gen_process = CLIP_GEN_PROCESS_TTS_VOCODE;
|
||||
vp.mel_in = &mel_data;
|
||||
vp.sstft_in = &sstft;
|
||||
vp.vocode_n_mel = n_mel_out;
|
||||
vp.vocode_n_stft = n_stft;
|
||||
vp.out_audio = nullptr;
|
||||
vp.out_spec = &spec;
|
||||
if (!clip_encode(ctx, &vp)) {
|
||||
LOG_ERR("%s: vocoder stage failed\n", __func__);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// istft: mag = clipped exp, phase = sin, hann overlap-add
|
||||
const int n_frames_out = (int) (spec.size() / 18);
|
||||
const int64_t n_out = (int64_t) (n_frames_out - 1) * hop;
|
||||
std::vector<double> acc((size_t) n_out + n_fft, 0.0);
|
||||
std::vector<double> wsum((size_t) n_out + n_fft, 0.0);
|
||||
for (int fr = 0; fr < n_frames_out; fr++) {
|
||||
double frame[16];
|
||||
for (int i = 0; i < n_fft; i++) {
|
||||
double v = 0.0;
|
||||
for (int k = 0; k < n_bins; k++) {
|
||||
const double mag = fmin(exp((double) spec[(size_t) fr * 18 + k]), 1e2);
|
||||
const double ph = sin((double) spec[(size_t) fr * 18 + 9 + k]);
|
||||
const double re = mag * cos(ph), im = mag * sin(ph);
|
||||
const double a = 2.0 * M_PI * k * i / n_fft;
|
||||
const double w = (k == 0 || k == n_fft / 2) ? 1.0 : 2.0;
|
||||
v += w * (re * cos(a) - im * sin(a));
|
||||
}
|
||||
frame[i] = v / n_fft;
|
||||
}
|
||||
const int64_t o = (int64_t) fr * hop;
|
||||
for (int i = 0; i < n_fft; i++) {
|
||||
acc [(size_t) (o + i)] += frame[i] * win[(size_t) i];
|
||||
wsum[(size_t) (o + i)] += (double) win[(size_t) i] * win[(size_t) i];
|
||||
}
|
||||
}
|
||||
GGML_ASSERT(params->out_audio);
|
||||
auto & out_audio = *params->out_audio;
|
||||
out_audio.resize((size_t) std::min<int64_t>(n_out - n_fft / 2, n_wav));
|
||||
for (size_t i = 0; i < out_audio.size(); i++) {
|
||||
const size_t j = i + n_fft / 2; // drop the centering pad
|
||||
const double v = wsum[j] > 1e-11 ? acc[j] / wsum[j] : 0.0;
|
||||
out_audio[(size_t) i] = (float) fmax(-0.99, fmin(0.99, v));
|
||||
}
|
||||
|
||||
// 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;
|
||||
}
|
||||
|
||||
// optional outputs: a pipeline yields codes or feats, and not all have an eos head
|
||||
if (params->out_codes != nullptr) {
|
||||
ggml_tensor * codes = ggml_graph_get_tensor(gf, "out_codes");
|
||||
@@ -6486,6 +7167,18 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
|
||||
return ctx->model.spk_proj_w->ne[1];
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN:
|
||||
return ctx->model.gen_input_lin_w->ne[1];
|
||||
case PROJECTOR_TYPE_CHATTERBOX:
|
||||
// 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:
|
||||
|
||||
@@ -96,6 +96,9 @@ enum clip_gen_process_type {
|
||||
CLIP_GEN_PROCESS_GEN_UNKNOWN,
|
||||
CLIP_GEN_PROCESS_GEN_CODE, // h_state to codes
|
||||
CLIP_GEN_PROCESS_GEN_WAV, // codes to raw PCM audio
|
||||
CLIP_GEN_PROCESS_TTS, // full utterance of codes to raw PCM audio
|
||||
CLIP_GEN_PROCESS_TTS_VOCODE, // internal: mel + source stft to istft input
|
||||
CLIP_GEN_PROCESS_TOKENIZE, // raw PCM audio to semantic speech tokens
|
||||
};
|
||||
struct clip_encode_params {
|
||||
int n_threads = 1;
|
||||
@@ -114,10 +117,30 @@ struct clip_encode_params {
|
||||
float temp = 0.0f; // sampling temperature, noise scale for flow-matching decoders
|
||||
bool * out_is_eos = 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<float> * out_code_embd = nullptr;
|
||||
|
||||
// GEN_WAV
|
||||
const std::vector<int32_t> * codes = nullptr; // this frame's 16 RVQ codes
|
||||
const std::vector<float> * feats = nullptr; // continuous counterpart of codes
|
||||
std::vector<float> * out_audio = nullptr; // decoded PCM samples, F32
|
||||
// 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<int32_t> * ref_tokens = nullptr;
|
||||
const std::vector<float> * ref_feat = nullptr;
|
||||
const std::vector<float> * ref_spk = nullptr;
|
||||
// TTS_VOCODE internal stage inputs
|
||||
const std::vector<float> * mel_in = nullptr;
|
||||
const std::vector<float> * sstft_in = nullptr;
|
||||
int vocode_n_mel = 0;
|
||||
int vocode_n_stft = 0;
|
||||
std::vector<float> * out_spec = nullptr;
|
||||
|
||||
// audio gen: state of one sequence, clip_encode() updates it in place
|
||||
struct clip_gen_state * gen_state = nullptr;
|
||||
@@ -139,6 +162,10 @@ struct clip_gen_state * clip_gen_state_init(struct clip_ctx * ctx, uint32_t seed
|
||||
void clip_gen_state_free(struct clip_gen_state * state);
|
||||
void clip_gen_state_reset(struct clip_gen_state * state);
|
||||
|
||||
// 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
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
@@ -180,6 +180,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_GEN_CODE;
|
||||
int n_tokens = 0;
|
||||
int n_prompt_mel = 0;
|
||||
int vocode_n_mel = 0;
|
||||
int vocode_n_stft = 0;
|
||||
clip_graph_chatterbox(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int n_tokens,
|
||||
int n_prompt_mel, int vocode_n_mel, int vocode_n_stft)
|
||||
: clip_graph_chatterbox_base(ctx, img), gen_process(gen_process), n_tokens(n_tokens), n_prompt_mel(n_prompt_mel),
|
||||
vocode_n_mel(vocode_n_mel), vocode_n_stft(vocode_n_stft) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
private:
|
||||
ggml_tensor * enc_layer(const clip_chatterbox::enc_layer & l, ggml_tensor * x, ggml_tensor * pos, int T);
|
||||
ggml_tensor * causal_block(const clip_chatterbox::causal_block & b, ggml_tensor * x);
|
||||
ggml_tensor * resnet(const clip_chatterbox::resnet & r, ggml_tensor * x, ggml_tensor * temb);
|
||||
ggml_tensor * tfm_block(const clip_chatterbox::tfm_block & b, ggml_tensor * x);
|
||||
ggml_tensor * estimator(ggml_tensor * x_noise, ggml_tensor * mu, ggml_tensor * spks,
|
||||
ggml_tensor * cond, ggml_tensor * temb, int T);
|
||||
ggml_tensor * hift_resblock(const clip_chatterbox::hift_res & r, ggml_tensor * x);
|
||||
ggml_cgraph * build_s3tok(int T);
|
||||
ggml_cgraph * build_vocoder(int n_mel, int n_stft);
|
||||
};
|
||||
|
||||
struct clip_graph_granite_speech : clip_graph {
|
||||
clip_graph_granite_speech(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
#include "mtmd-audio.h"
|
||||
|
||||
#define _USE_MATH_DEFINES // for M_PI
|
||||
#include <cfloat>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
@@ -970,6 +971,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<float> & 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<double> 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<double> mel((size_t) n_mel * n_frames, 0.0);
|
||||
std::vector<double> frame(n_fft);
|
||||
std::vector<double> 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<float> & out) {
|
||||
const int W = 16; // sinc half width in source samples
|
||||
const double fc = 0.495; // cutoff, normalized to the source rate
|
||||
|
||||
const size_t n_out = n_samples * 3 / 2;
|
||||
out.assign(n_out, 0.0f);
|
||||
for (size_t n = 0; n < n_out; n++) {
|
||||
const double t = (double) (2 * n) / 3.0;
|
||||
const int k0 = (int) floor(t) - W + 1;
|
||||
double acc = 0.0;
|
||||
for (int k = k0; k < k0 + 2 * W; k++) {
|
||||
if (k < 0 || k >= (int) n_samples) {
|
||||
continue;
|
||||
}
|
||||
const double x = t - k;
|
||||
const double s = x == 0.0 ? 1.0 : sin(2.0 * M_PI * fc * x) / (M_PI * x);
|
||||
acc += samples[k] * s * 0.5 * (1.0 + cos(M_PI * x / (W + 1)));
|
||||
}
|
||||
out[n] = (float) acc;
|
||||
}
|
||||
}
|
||||
|
||||
// matcha style log-mel of the s3gen prompt features (s3gen/utils/mel.py)
|
||||
static bool mtmd_audio_matcha_log_mel(const float * samples, size_t n_samples,
|
||||
std::vector<float> & out, int & n_frames) {
|
||||
const int n_fft = 1920;
|
||||
const int hop = 480;
|
||||
const int 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<float> fft_in((size_t) n_fft * 2, 0.0f);
|
||||
std::vector<float> fft_out((size_t) n_fft * 8);
|
||||
std::vector<float> 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<double> rms((size_t) n_fr);
|
||||
double mx = 0.0;
|
||||
for (int fr = 0; fr < n_fr; fr++) {
|
||||
double acc = 0.0;
|
||||
for (int i = 0; i < win; i++) {
|
||||
const int idx = fr * hop - win / 2 + i;
|
||||
if (idx >= 0 && idx < n) {
|
||||
acc += (double) samples[idx] * samples[idx];
|
||||
}
|
||||
}
|
||||
rms[(size_t) fr] = sqrt(acc / win);
|
||||
mx = std::max(mx, rms[(size_t) fr]);
|
||||
}
|
||||
|
||||
const double thr = mx * pow(10.0, -top_db / 20.0);
|
||||
int first = -1, last = -1;
|
||||
for (int fr = 0; fr < n_fr; fr++) {
|
||||
if (rms[(size_t) fr] > thr) {
|
||||
if (first < 0) {
|
||||
first = fr;
|
||||
}
|
||||
last = fr;
|
||||
}
|
||||
}
|
||||
if (first < 0) {
|
||||
start = end = 0;
|
||||
return;
|
||||
}
|
||||
start = (size_t) first * hop;
|
||||
end = std::min((size_t) (last + 1) * hop, n_samples);
|
||||
}
|
||||
|
||||
// power mel of the voice encoder front-end: 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<float> & 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<double> 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<double> frame(n_fft);
|
||||
std::vector<double> 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<double> y(samples, samples + n_samples);
|
||||
|
||||
auto biquad = [&](double b0, double b1, double b2, double a1, double a2) {
|
||||
double x1 = 0.0, x2 = 0.0, y1 = 0.0, y2 = 0.0;
|
||||
for (double & v : y) {
|
||||
const double x0 = v;
|
||||
v = b0 * x0 + b1 * x1 + b2 * x2 - a1 * y1 - a2 * y2;
|
||||
x2 = x1; x1 = x0;
|
||||
y2 = y1; y1 = v;
|
||||
}
|
||||
};
|
||||
|
||||
// stage 1: high shelf, f0 1681.9744509555319 Hz, +3.99984385397 dB, Q 0.7071752369554196
|
||||
{
|
||||
const double A = pow(10.0, 3.99984385397 / 40.0);
|
||||
const double w0 = 2.0 * M_PI * 1681.9744509555319 / sample_rate;
|
||||
const double alpha = sin(w0) / (2.0 * 0.7071752369554196);
|
||||
const double c = cos(w0);
|
||||
const double sq = 2.0 * sqrt(A) * alpha;
|
||||
const double b0 = A * ((A + 1.0) + (A - 1.0) * c + sq);
|
||||
const double b1 = -2.0 * A * ((A - 1.0) + (A + 1.0) * c);
|
||||
const double b2 = A * ((A + 1.0) + (A - 1.0) * c - sq);
|
||||
const double a0 = (A + 1.0) - (A - 1.0) * c + sq;
|
||||
const double a1 = 2.0 * ((A - 1.0) - (A + 1.0) * c);
|
||||
const double a2 = (A + 1.0) - (A - 1.0) * c - sq;
|
||||
biquad(b0 / a0, b1 / a0, b2 / a0, a1 / a0, a2 / a0);
|
||||
}
|
||||
// stage 2: high pass, f0 38.13547087602444 Hz, Q 0.5003270373238773
|
||||
{
|
||||
const double w0 = 2.0 * M_PI * 38.13547087602444 / sample_rate;
|
||||
const double alpha = sin(w0) / (2.0 * 0.5003270373238773);
|
||||
const double c = cos(w0);
|
||||
const double b0 = (1.0 + c) / 2.0;
|
||||
const double b1 = -(1.0 + c);
|
||||
const double b2 = (1.0 + c) / 2.0;
|
||||
const double a0 = 1.0 + alpha;
|
||||
const double a1 = -2.0 * c;
|
||||
const double a2 = 1.0 - alpha;
|
||||
biquad(b0 / a0, b1 / a0, b2 / a0, a1 / a0, a2 / a0);
|
||||
}
|
||||
|
||||
const int win = (int) (0.4 * sample_rate);
|
||||
const int hop = win / 4;
|
||||
if ((int) n_samples < win) {
|
||||
return -HUGE_VALF;
|
||||
}
|
||||
const int n_blocks = 1 + ((int) n_samples - win) / hop;
|
||||
std::vector<double> z((size_t) n_blocks);
|
||||
for (int b = 0; b < n_blocks; b++) {
|
||||
double acc = 0.0;
|
||||
for (int i = 0; i < win; i++) {
|
||||
acc += y[(size_t) b * hop + i] * y[(size_t) b * hop + i];
|
||||
}
|
||||
z[(size_t) b] = acc / win;
|
||||
}
|
||||
|
||||
auto gated_mean = [&](double thr_lufs) {
|
||||
double acc = 0.0;
|
||||
int n = 0;
|
||||
for (const double v : z) {
|
||||
if (-0.691 + 10.0 * log10(std::max(v, 1e-30)) > thr_lufs) {
|
||||
acc += v;
|
||||
n++;
|
||||
}
|
||||
}
|
||||
return n > 0 ? acc / n : 0.0;
|
||||
};
|
||||
|
||||
const double z_abs = gated_mean(-70.0);
|
||||
if (z_abs <= 0.0) {
|
||||
return -HUGE_VALF;
|
||||
}
|
||||
const double thr_rel = -0.691 + 10.0 * log10(z_abs) - 10.0;
|
||||
const double z_rel = gated_mean(thr_rel);
|
||||
if (z_rel <= 0.0) {
|
||||
return -HUGE_VALF;
|
||||
}
|
||||
return (float) (-0.691 + 10.0 * log10(z_rel));
|
||||
}
|
||||
|
||||
//
|
||||
// mtmd_audio_preprocessor_chatterbox_ref
|
||||
//
|
||||
// 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<mtmd_audio_mel> & output) const {
|
||||
output.clear();
|
||||
const int sr = (int) hparams.audio_sample_rate;
|
||||
|
||||
// the turbo variant loudness-normalizes the whole clip before any feature
|
||||
std::vector<float> pcm(samples, samples + n_samples);
|
||||
if (!is_mtl) {
|
||||
const float lufs = mtmd_audio_lufs(pcm.data(), pcm.size(), sr);
|
||||
if (lufs != -HUGE_VALF) {
|
||||
const float gain = powf(10.0f, (-27.0f - lufs) / 20.0f);
|
||||
if (std::isfinite(gain) && gain > 0.0f) {
|
||||
for (float & v : pcm) {
|
||||
v *= gain;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// entry 0: 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<double> frame(n_fft);
|
||||
std::vector<double> 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<float> clip(pcm.begin(), pcm.begin() + std::min(pcm.size(), cap));
|
||||
clip.resize((clip.size() + 639) / 640 * 640, 0.0f);
|
||||
std::vector<float> mel;
|
||||
int n_frames = 0;
|
||||
if (!mtmd_audio_s3tok_log_mel(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<float> clip(pcm.begin(), pcm.begin() + std::min(pcm.size(), gen_cap));
|
||||
clip.resize((clip.size() + 639) / 640 * 640, 0.0f);
|
||||
std::vector<float> pcm24;
|
||||
mtmd_audio_upsample_3_2(clip.data(), clip.size(), pcm24);
|
||||
std::vector<float> feat;
|
||||
int n_frames = 0;
|
||||
if (!mtmd_audio_matcha_log_mel(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<float> 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
|
||||
//
|
||||
|
||||
@@ -162,6 +162,29 @@ struct mtmd_audio_preprocessor_pockettts : mtmd_audio_preprocessor {
|
||||
bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) const override;
|
||||
};
|
||||
|
||||
// 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<float> 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<mtmd_audio_mel> & output) const override;
|
||||
|
||||
private:
|
||||
bool is_mtl;
|
||||
std::vector<float> s3tok_filters; // s3 tokenizer filterbank [n_mel x (400 / 2 + 1)], from the mmproj
|
||||
std::vector<float> window; // povey window of the fbank front-end, frame_length points
|
||||
std::vector<float> 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;
|
||||
|
||||
+464
-66
@@ -117,6 +117,40 @@ protected:
|
||||
return true;
|
||||
}
|
||||
|
||||
// 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<float> & 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_get_marker(mctx);
|
||||
mtmd_input_text text{ marker.c_str(), marker.size(), false, true };
|
||||
mtmd_input_chunks * chunks = mtmd_input_chunks_init();
|
||||
const mtmd_bitmap * bptr = bitmap;
|
||||
bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0;
|
||||
if (ok) {
|
||||
ok = false;
|
||||
for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {
|
||||
const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
|
||||
if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) {
|
||||
continue;
|
||||
}
|
||||
if (mtmd_encode_chunk(mctx, chunk) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: speaker encode failed\n");
|
||||
break;
|
||||
}
|
||||
const float * embd = mtmd_get_output_embd(mctx);
|
||||
const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk);
|
||||
out.assign(embd, embd + n);
|
||||
ok = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
mtmd_input_chunks_free(chunks);
|
||||
return ok;
|
||||
}
|
||||
|
||||
llama_context * lctx;
|
||||
mtmd_context * mctx;
|
||||
const llama_model * model;
|
||||
@@ -421,39 +455,6 @@ private:
|
||||
return true;
|
||||
}
|
||||
|
||||
// runs the reference wav through the speaker encoder, returns one x-vector embedding row
|
||||
bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
|
||||
if (!mtmd_support_audio(mctx)) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n");
|
||||
return false;
|
||||
}
|
||||
const std::string marker = mtmd_get_marker(mctx);
|
||||
mtmd_input_text text{ marker.c_str(), marker.size(), false, true };
|
||||
mtmd_input_chunks * chunks = mtmd_input_chunks_init();
|
||||
const mtmd_bitmap * bptr = bitmap;
|
||||
bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0;
|
||||
if (ok) {
|
||||
ok = false;
|
||||
for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {
|
||||
const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
|
||||
if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) {
|
||||
continue;
|
||||
}
|
||||
if (mtmd_encode_chunk(mctx, chunk) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: speaker encode failed\n");
|
||||
break;
|
||||
}
|
||||
const float * embd = mtmd_get_output_embd(mctx);
|
||||
const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk);
|
||||
out.assign(embd, embd + n);
|
||||
ok = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
mtmd_input_chunks_free(chunks);
|
||||
return ok;
|
||||
}
|
||||
|
||||
// one GEN_WAV process() call over the buffered codes, state is carried across batches
|
||||
bool flush_gen_wav() {
|
||||
if (codes_buf.empty()) {
|
||||
@@ -989,39 +990,6 @@ private:
|
||||
return n;
|
||||
}
|
||||
|
||||
// runs the reference wav through the mimi encoder, returns one row per 12.5Hz frame
|
||||
bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
|
||||
if (!mtmd_support_audio(mctx)) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: mmproj has no voice encoder\n");
|
||||
return false;
|
||||
}
|
||||
const std::string marker = mtmd_get_marker(mctx);
|
||||
mtmd_input_text text{ marker.c_str(), marker.size(), false, true };
|
||||
mtmd_input_chunks * chunks = mtmd_input_chunks_init();
|
||||
const mtmd_bitmap * bptr = bitmap;
|
||||
bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0;
|
||||
if (ok) {
|
||||
ok = false;
|
||||
for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {
|
||||
const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
|
||||
if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) {
|
||||
continue;
|
||||
}
|
||||
if (mtmd_encode_chunk(mctx, chunk) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: voice encode failed\n");
|
||||
break;
|
||||
}
|
||||
const float * embd = mtmd_get_output_embd(mctx);
|
||||
const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk);
|
||||
out.assign(embd, embd + n);
|
||||
ok = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
mtmd_input_chunks_free(chunks);
|
||||
return ok;
|
||||
}
|
||||
|
||||
// decodes the buffered latents, the mimi decoder state carries over between calls
|
||||
bool flush_gen_wav() {
|
||||
if (feats_buf.empty()) {
|
||||
@@ -1076,12 +1044,442 @@ private:
|
||||
bool wav_header_sent = false;
|
||||
};
|
||||
|
||||
|
||||
// Chatterbox: the backbone emits s3 speech tokens, then one flow-matching mel decode and the NSF-iSTFT vocoder make the PCM
|
||||
// the multilingual variant runs cfg with an uncond sequence (text rows zeroed), paired as seq_id and n_seq_max / 2 + seq_id
|
||||
class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
|
||||
public:
|
||||
using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline;
|
||||
|
||||
void reset() override {
|
||||
seq_id = 0;
|
||||
seq_uncond = 0;
|
||||
pos = 0;
|
||||
ar_idx = 0;
|
||||
codes_buf.clear();
|
||||
audio_pcm.clear();
|
||||
h_state_buf.clear();
|
||||
out_buf.clear();
|
||||
prompt_embd_buf.clear();
|
||||
uncond_embd_buf.clear();
|
||||
cond_logits.clear();
|
||||
n_prompt = 0;
|
||||
prompt_pos = 0;
|
||||
decoded = false;
|
||||
stream = false;
|
||||
pcm_sent = 0;
|
||||
wav_header_sent = false;
|
||||
gen_state.reset();
|
||||
}
|
||||
|
||||
int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
|
||||
reset();
|
||||
seq_id = inp->seq_id;
|
||||
|
||||
if (!ensure_cache()) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
out_type = inp->out_type;
|
||||
stream = inp->stream;
|
||||
if (!init_gen_state(inp->seed)) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
const bool mtl = !t3_cond.empty();
|
||||
if (mtl) {
|
||||
// the uncond sequence lives half a seq space away from the cond one
|
||||
const int32_t n_seq_max = (int32_t) llama_n_seq_max(lctx);
|
||||
seq_uncond = (llama_seq_id) (n_seq_max / 2 + seq_id);
|
||||
if (n_seq_max < 2 || seq_id >= n_seq_max / 2) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: multilingual chatterbox needs a paired cfg sequence "
|
||||
"(n_seq_max = %d, seq_id = %d)\n", n_seq_max, seq_id);
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
// the encode gives the talker conditioning rows, set_ref keeps the flow decoder reference in the gen state
|
||||
std::vector<float> ref_cond;
|
||||
if (inp->speaker_ref) {
|
||||
if (!encode_speaker(inp->speaker_ref, ref_cond) ||
|
||||
mtmd_gen_state_set_ref(mctx, gen_state.get()) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: speaker reference encoding failed\n");
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
const int n_e = n_embd;
|
||||
auto push_row = [&](const float * src) {
|
||||
prompt_embd_buf.insert(prompt_embd_buf.end(), src, src + n_e);
|
||||
};
|
||||
auto add_pos = [&](float * r, const std::vector<float> & tab, int idx) {
|
||||
const float * p = tab.data() + (size_t) idx * n_e;
|
||||
for (int j = 0; j < n_e; j++) {
|
||||
r[j] += p[j];
|
||||
}
|
||||
};
|
||||
|
||||
// prompt rows: [conditioning, text, speech bos], the conditioning is the precomputed default voice when not cloning
|
||||
if (!ref_cond.empty()) {
|
||||
prompt_embd_buf = ref_cond;
|
||||
} else if (mtl) {
|
||||
prompt_embd_buf = t3_cond;
|
||||
} else {
|
||||
push_row(cond_spkr.data());
|
||||
prompt_embd_buf.insert(prompt_embd_buf.end(), cond_speech_rows.begin(), cond_speech_rows.end());
|
||||
}
|
||||
|
||||
// no text preprocessing here, the [SPACE] substitution and the case folding are baked in at conversion
|
||||
std::string txt(inp->prompt, inp->prompt_len);
|
||||
if (mtl && inp->lang && inp->lang[0]) {
|
||||
// the reference prefixes the language tag, e.g. "[fr]"
|
||||
std::string lang = inp->lang;
|
||||
std::transform(lang.begin(), lang.end(), lang.begin(), [](unsigned char c) { return (char) std::tolower(c); });
|
||||
txt = "[" + lang + "]" + txt;
|
||||
}
|
||||
std::vector<llama_token> ids(txt.size() + 16);
|
||||
int n_ids = llama_tokenize(vocab, txt.c_str(), (int32_t) txt.size(), ids.data(), (int32_t) ids.size(),
|
||||
false, true);
|
||||
if (n_ids < 1) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n");
|
||||
return 1;
|
||||
}
|
||||
ids.resize((size_t) n_ids);
|
||||
if (mtl) {
|
||||
// only the multilingual prompt wraps the text in start/stop tokens
|
||||
ids.insert(ids.begin(), text_start);
|
||||
ids.push_back(text_stop);
|
||||
}
|
||||
const size_t text0 = prompt_embd_buf.size() / (size_t) n_e;
|
||||
for (size_t i = 0; i < ids.size(); i++) {
|
||||
push_row(tok_embd.data() + (size_t) ids[i] * n_e);
|
||||
if (mtl) {
|
||||
add_pos(prompt_embd_buf.data() + (text0 + i) * n_e, text_pos, (int) i);
|
||||
}
|
||||
}
|
||||
|
||||
// speech bos opens the AR stream
|
||||
push_row(tok_embd.data() + (size_t) llama_vocab_bos(vocab) * n_e);
|
||||
if (mtl) {
|
||||
add_pos(prompt_embd_buf.data() + prompt_embd_buf.size() - n_e, speech_pos, 0);
|
||||
}
|
||||
ar_idx = 1;
|
||||
n_prompt = (int) (prompt_embd_buf.size() / (size_t) n_e);
|
||||
|
||||
if (mtl) {
|
||||
// cfg uncond prompt: text rows are zeroed but keep their learned positions, as in the reference
|
||||
uncond_embd_buf = prompt_embd_buf;
|
||||
for (size_t i = 0; i < ids.size(); i++) {
|
||||
float * r = uncond_embd_buf.data() + (text0 + i) * n_e;
|
||||
std::fill(r, r + n_e, 0.0f);
|
||||
add_pos(r, text_pos, (int) i);
|
||||
}
|
||||
}
|
||||
|
||||
prompt_pos = 0;
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t step_prompt(int32_t n_batch) override {
|
||||
GGML_ASSERT(n_batch > 0);
|
||||
const bool mtl = !t3_cond.empty();
|
||||
// the uncond prompt follows the cond one on mtl
|
||||
const int n_total = mtl ? 2 * n_prompt : n_prompt;
|
||||
if (prompt_pos >= n_total) {
|
||||
return 0;
|
||||
}
|
||||
const bool uncond = prompt_pos >= n_prompt;
|
||||
const int p0 = uncond ? prompt_pos - n_prompt : prompt_pos;
|
||||
const int32_t n_tokens_batch = std::min(n_batch, n_prompt - p0);
|
||||
const bool last = p0 + n_tokens_batch == n_prompt;
|
||||
|
||||
decode_embd_batch batch((uncond ? uncond_embd_buf : prompt_embd_buf).data(), n_prompt, 1, n_embd);
|
||||
batch.set_position_normal(0, uncond ? seq_uncond : seq_id);
|
||||
if (last) {
|
||||
batch.logits[n_prompt - 1] = 1;
|
||||
}
|
||||
if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.render(lctx, p0, n_tokens_batch)) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n");
|
||||
return -1;
|
||||
}
|
||||
prompt_pos += n_tokens_batch;
|
||||
|
||||
if (last && mtl) {
|
||||
if (!uncond) {
|
||||
cfg_read();
|
||||
} else {
|
||||
cfg_apply();
|
||||
}
|
||||
}
|
||||
|
||||
if (prompt_pos >= n_total) {
|
||||
pos = n_prompt;
|
||||
prompt_embd_buf.clear();
|
||||
uncond_embd_buf.clear();
|
||||
return 0;
|
||||
}
|
||||
return n_total - prompt_pos;
|
||||
}
|
||||
|
||||
int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override {
|
||||
GGML_UNUSED(h_state_in);
|
||||
if (sampled == LLAMA_TOKEN_NULL) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: chatterbox requires a token sampled from the backbone\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
// the speech stop token is the backbone eos
|
||||
if (llama_vocab_is_eog(vocab, sampled)) {
|
||||
*out_stop = true;
|
||||
*h_state_out = nullptr;
|
||||
return 0;
|
||||
}
|
||||
|
||||
// keep only the 6561 s3gen codes, dropping start/stop and oov ids
|
||||
if (sampled >= speech_base && sampled - speech_base < 6561) {
|
||||
codes_buf.push_back(sampled - speech_base);
|
||||
}
|
||||
|
||||
std::vector<float> e(tok_embd.begin() + (size_t) sampled * n_embd,
|
||||
tok_embd.begin() + (size_t) (sampled + 1) * n_embd);
|
||||
const bool mtl = !t3_cond.empty();
|
||||
if (mtl) {
|
||||
// the multilingual backbone takes the token row plus the learned speech position
|
||||
const float * p = speech_pos.data() + (size_t) ar_idx * n_embd;
|
||||
for (int j = 0; j < n_embd; j++) {
|
||||
e[(size_t) j] += p[j];
|
||||
}
|
||||
ar_idx++;
|
||||
}
|
||||
|
||||
if (!decode_one(e.data(), seq_id)) {
|
||||
return 1;
|
||||
}
|
||||
if (mtl) {
|
||||
// cfg: the sampled token feeds both branches
|
||||
cfg_read();
|
||||
if (!decode_one(e.data(), seq_uncond)) {
|
||||
return 1;
|
||||
}
|
||||
cfg_apply();
|
||||
}
|
||||
pos++;
|
||||
|
||||
const float * h = llama_get_embeddings_ith(lctx, -1);
|
||||
h_state_buf.assign(h, h + n_embd);
|
||||
*h_state_out = h_state_buf.data();
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override {
|
||||
*out_sample_rate = info.sample_rate;
|
||||
|
||||
if (!stream) {
|
||||
if (!flush_gen_wav()) {
|
||||
return 1;
|
||||
}
|
||||
if (out_n_samples) {
|
||||
*out_n_samples = (int64_t) audio_pcm.size();
|
||||
}
|
||||
if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
|
||||
*out_data = (const char *) audio_pcm.data();
|
||||
*out_data_len = audio_pcm.size() * sizeof(float);
|
||||
return 0;
|
||||
}
|
||||
out_buf.clear();
|
||||
if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
|
||||
return 1;
|
||||
}
|
||||
*out_data = out_buf.data();
|
||||
*out_data_len = out_buf.size();
|
||||
return 0;
|
||||
}
|
||||
|
||||
// streaming: the flow decode is one-shot, so all the audio arrives after flush()
|
||||
const size_t n_new = audio_pcm.size() - pcm_sent;
|
||||
if (out_n_samples) {
|
||||
*out_n_samples = (int64_t) n_new;
|
||||
}
|
||||
if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
|
||||
*out_data = (const char *) (audio_pcm.data() + pcm_sent);
|
||||
*out_data_len = n_new * sizeof(float);
|
||||
pcm_sent = audio_pcm.size();
|
||||
return 0;
|
||||
}
|
||||
out_buf.clear();
|
||||
if (!wav_header_sent) {
|
||||
write_wav16_header(out_buf, UINT32_MAX, info.sample_rate);
|
||||
wav_header_sent = true;
|
||||
}
|
||||
append_wav16_pcm(out_buf, audio_pcm.data() + pcm_sent, n_new);
|
||||
pcm_sent = audio_pcm.size();
|
||||
*out_data = out_buf.data();
|
||||
*out_data_len = out_buf.size();
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t flush() override {
|
||||
return flush_gen_wav() ? 0 : 1;
|
||||
}
|
||||
|
||||
private:
|
||||
bool decode_one(const float * e, llama_seq_id seq) {
|
||||
decode_embd_batch batch(e, 1, 1, n_embd);
|
||||
batch.set_position_normal(pos, seq);
|
||||
batch.logits[0] = 1;
|
||||
if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.render(lctx, 0, 1)) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: step decode failed\n");
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// cfg of the reference: logits = cond + 0.5 * (cond - uncond)
|
||||
// the mix goes into the last output row (the uncond one), the caller samples from it
|
||||
void cfg_read() {
|
||||
const float * c = llama_get_logits_ith(lctx, -1);
|
||||
cond_logits.assign(c, c + llama_vocab_n_tokens(vocab));
|
||||
}
|
||||
void cfg_apply() {
|
||||
float * u = llama_get_logits_ith(lctx, -1);
|
||||
for (size_t i = 0; i < cond_logits.size(); i++) {
|
||||
u[i] = cond_logits[i] + 0.5f * (cond_logits[i] - u[i]);
|
||||
}
|
||||
}
|
||||
|
||||
// one GEN_WAV call over the whole utterance, the flow decode is not streamed
|
||||
bool flush_gen_wav() {
|
||||
if (decoded || codes_buf.empty()) {
|
||||
return true;
|
||||
}
|
||||
if (t3_cond.empty()) {
|
||||
// turbo appends a short silence tail before vocoding
|
||||
codes_buf.insert(codes_buf.end(), 3, 4299);
|
||||
}
|
||||
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
|
||||
inp.codes = codes_buf.data();
|
||||
inp.n_codes = codes_buf.size();
|
||||
mtmd_gen_out out{};
|
||||
if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: gen_wav decode failed\n");
|
||||
return false;
|
||||
}
|
||||
audio_pcm.assign(out.audio, out.audio + out.n_samples);
|
||||
decoded = true;
|
||||
return true;
|
||||
}
|
||||
|
||||
bool ensure_cache() {
|
||||
if (!tok_embd.empty()) {
|
||||
return true;
|
||||
}
|
||||
// fused vocab layout: [text 0..speech_base) then the speech tokens
|
||||
speech_base = find_special_token(vocab, "<|speech_0|>");
|
||||
if (speech_base == LLAMA_TOKEN_NULL) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: fused speech tokens not found in vocab\n");
|
||||
return false;
|
||||
}
|
||||
|
||||
// start_text_token and stop_text_token of the reference, only the multilingual prompt uses 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());
|
||||
|
||||
auto read_table = [&](const char * name, std::vector<float> & dst) {
|
||||
const size_t n = mtmd_gen_audio_read_tensor(mctx, name, nullptr, 0);
|
||||
dst.resize(n);
|
||||
if (n == 0 || mtmd_gen_audio_read_tensor(mctx, name, dst.data(), n) != n) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: %s read failed\n", name);
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
};
|
||||
|
||||
// multilingual: precomputed [spkr, perceiver, emotion] conditioning block and the learned position tables
|
||||
if (mtmd_gen_audio_read_tensor(mctx, "a.gen.cond.t3_cond", nullptr, 0) > 0) {
|
||||
if (!read_table("a.gen.cond.t3_cond", t3_cond) ||
|
||||
!read_table("a.gen.t3.text_pos_emb", text_pos) ||
|
||||
!read_table("a.gen.t3.speech_pos_emb", speech_pos)) {
|
||||
return false;
|
||||
}
|
||||
if (t3_cond.size() % (size_t) n_embd != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: unexpected a.gen.cond.t3_cond size\n");
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// turbo: precomputed speaker row, plus the default speech ids looked up in the mmproj speech table
|
||||
std::vector<float> cond_ids, table;
|
||||
if (!read_table("a.gen.cond.spkr_default", cond_spkr) ||
|
||||
!read_table("a.gen.cond.prompt_speech_tokens", cond_ids) ||
|
||||
!read_table("a.gen.code.out_embd.weight", table)) {
|
||||
return false;
|
||||
}
|
||||
if (cond_spkr.size() != (size_t) n_embd || table.size() % (size_t) n_embd != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: unexpected default conditioning size\n");
|
||||
return false;
|
||||
}
|
||||
const size_t n_rows = table.size() / (size_t) n_embd;
|
||||
cond_speech_rows.resize(cond_ids.size() * (size_t) n_embd);
|
||||
for (size_t i = 0; i < cond_ids.size(); i++) {
|
||||
const size_t r = (size_t) cond_ids[i];
|
||||
GGML_ASSERT(r < n_rows);
|
||||
memcpy(cond_speech_rows.data() + i * (size_t) n_embd, table.data() + r * (size_t) n_embd,
|
||||
(size_t) n_embd * sizeof(float));
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// vocab and mmproj data fixed across the whole session, read once
|
||||
std::vector<float> tok_embd;
|
||||
std::vector<float> cond_spkr;
|
||||
std::vector<float> cond_speech_rows; // default conditioning ids resolved through the mmproj speech table
|
||||
std::vector<float> t3_cond;
|
||||
std::vector<float> text_pos;
|
||||
std::vector<float> speech_pos;
|
||||
llama_token speech_base = LLAMA_TOKEN_NULL;
|
||||
llama_token text_start = LLAMA_TOKEN_NULL;
|
||||
llama_token text_stop = LLAMA_TOKEN_NULL;
|
||||
|
||||
// per-generation state, cleared by reset()
|
||||
llama_seq_id seq_id = 0;
|
||||
llama_seq_id seq_uncond = 0;
|
||||
llama_pos pos = 0;
|
||||
int ar_idx = 0;
|
||||
std::vector<float> prompt_embd_buf;
|
||||
std::vector<float> uncond_embd_buf;
|
||||
int n_prompt = 0;
|
||||
int prompt_pos = 0;
|
||||
std::vector<float> cond_logits;
|
||||
std::vector<int32_t> codes_buf;
|
||||
bool decoded = false;
|
||||
std::vector<float> audio_pcm;
|
||||
std::vector<float> h_state_buf;
|
||||
mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
|
||||
std::vector<char> out_buf;
|
||||
bool stream = false;
|
||||
size_t pcm_sent = 0; // samples already returned by get_output()
|
||||
bool wav_header_sent = false;
|
||||
};
|
||||
|
||||
static std::unique_ptr<mtmd_gen_audio_pipeline> make_pipeline(llama_context * lctx, mtmd_context * mctx) {
|
||||
switch (mtmd_gen_audio_get_info(mctx).type) {
|
||||
case MTMD_GEN_AUDIO_TYPE_QWEN3TTS:
|
||||
return std::unique_ptr<mtmd_gen_audio_pipeline>(new qwen3tts_gen_audio_pipeline(lctx, mctx));
|
||||
case MTMD_GEN_AUDIO_TYPE_POCKETTTS:
|
||||
return std::unique_ptr<mtmd_gen_audio_pipeline>(new pockettts_gen_audio_pipeline(lctx, mctx));
|
||||
case MTMD_GEN_AUDIO_TYPE_CHATTERBOX:
|
||||
return std::unique_ptr<mtmd_gen_audio_pipeline>(new chatterbox_gen_audio_pipeline(lctx, mctx));
|
||||
default:
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
+509
-8
@@ -201,6 +201,10 @@ struct mtmd_bitmap {
|
||||
return data;
|
||||
}
|
||||
|
||||
std::vector<unsigned char> & get_rw_buf() {
|
||||
return data;
|
||||
}
|
||||
|
||||
bool is_placeholder() const {
|
||||
return data.empty();
|
||||
}
|
||||
@@ -524,6 +528,11 @@ struct mtmd_context {
|
||||
// generation context
|
||||
struct clip_ctx * ctx_gen_a; // audio
|
||||
|
||||
// gen-side reference data of the last encoded chunk (chatterbox), see mtmd_gen_state_set_ref()
|
||||
std::vector<int32_t> gen_ref_tokens;
|
||||
std::vector<float> gen_ref_feat;
|
||||
std::vector<float> gen_ref_spk;
|
||||
|
||||
bool print_timings;
|
||||
int n_threads;
|
||||
std::string media_marker;
|
||||
@@ -623,7 +632,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));
|
||||
}
|
||||
|
||||
@@ -640,14 +649,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_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) {
|
||||
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(
|
||||
@@ -1094,6 +1106,25 @@ struct mtmd_context {
|
||||
{
|
||||
audio_preproc = std::make_unique<mtmd_audio_preprocessor_pockettts>(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<float> 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<mtmd_audio_preprocessor_chatterbox_ref>(ctx_a, is_mtl, std::move(filt));
|
||||
}
|
||||
} break;
|
||||
default:
|
||||
throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj));
|
||||
}
|
||||
@@ -1708,6 +1739,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) {
|
||||
@@ -1976,6 +2042,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<float> & out_embd);
|
||||
|
||||
static int32_t mtmd_encode_chunk_impl(mtmd_context * ctx, const mtmd_input_chunk * chunk, std::vector<float> & out_embd) {
|
||||
if (chunk->type == MTMD_INPUT_CHUNK_TYPE_TEXT) {
|
||||
LOG_WRN("mtmd_encode_chunk has no effect for text chunks\n");
|
||||
@@ -2007,6 +2077,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(
|
||||
@@ -2023,6 +2096,9 @@ static int32_t mtmd_encode_chunk_impl(mtmd_context * ctx, const mtmd_input_chunk
|
||||
|
||||
int32_t mtmd_encode_chunk(mtmd_context * ctx, const mtmd_input_chunk * chunk) {
|
||||
// this is the non-batching version
|
||||
ctx->gen_ref_tokens.clear();
|
||||
ctx->gen_ref_feat.clear();
|
||||
ctx->gen_ref_spk.clear();
|
||||
try {
|
||||
return mtmd_encode_chunk_impl(ctx, chunk, ctx->out_embd);
|
||||
} catch (const std::exception & e) {
|
||||
@@ -2065,6 +2141,10 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS;
|
||||
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;
|
||||
@@ -2072,6 +2152,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<float> & 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<int32_t> & codes, std::vector<float> * rows) {
|
||||
clip_image_f32 mel_img;
|
||||
mel_img.set_size(mel.get_size(), false, true);
|
||||
mel_img.cpy_buf(mel.get_ro_buf());
|
||||
|
||||
clip_image_f32_batch batch;
|
||||
batch.is_audio = true;
|
||||
batch.entries.push_back(std::move(mel_img));
|
||||
|
||||
std::vector<int32_t> out_codes;
|
||||
std::vector<float> out_embd;
|
||||
|
||||
clip_encode_params params;
|
||||
params.imgs = &batch;
|
||||
params.n_threads = ctx->n_threads;
|
||||
params.gen_process = CLIP_GEN_PROCESS_TOKENIZE;
|
||||
params.out_codes = &out_codes;
|
||||
params.out_code_embd = rows ? &out_embd : nullptr;
|
||||
|
||||
if (!clip_encode(ctx->ctx_gen_a, ¶ms)) {
|
||||
LOG_ERR("%s: clip_encode failed (tokenize)\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
codes = std::move(out_codes);
|
||||
if (rows) {
|
||||
*rows = std::move(out_embd);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
||||
// talker conditioning rows 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<float> & pse, std::vector<float> & out_rows) {
|
||||
clip_ctx * ctx_clip = ctx->ctx_a;
|
||||
auto read_t = [&](const char * name, std::vector<float> & 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<float> & 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<float> w_ih[3], w_hh[3], b_ih[3], b_hh[3];
|
||||
std::vector<float> w_proj, b_proj;
|
||||
for (int l = 0; l < 3; l++) {
|
||||
const std::string s = std::to_string(l);
|
||||
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<float> ve(n_h, 0.0f);
|
||||
std::vector<float> h((size_t) 3 * n_h), c((size_t) 3 * n_h), x(n_h), g((size_t) 4 * n_h);
|
||||
for (int p = 0; p < n_wins; p++) {
|
||||
std::fill(h.begin(), h.end(), 0.0f);
|
||||
std::fill(c.begin(), c.end(), 0.0f);
|
||||
for (int t = 0; t < n_partial; t++) {
|
||||
const float * in = mel.data() + (size_t) (p * step + t) * n_mel;
|
||||
int n_in = n_mel;
|
||||
for (int l = 0; l < 3; l++) {
|
||||
float * hl = h.data() + (size_t) l * n_h;
|
||||
float * cl = c.data() + (size_t) l * n_h;
|
||||
for (int j = 0; j < 4 * n_h; j++) {
|
||||
double acc = b_ih[l][(size_t) j] + b_hh[l][(size_t) j];
|
||||
const float * wi = w_ih[l].data() + (size_t) j * n_in;
|
||||
for (int i = 0; i < n_in; i++) {
|
||||
acc += (double) wi[i] * in[i];
|
||||
}
|
||||
const float * wh = w_hh[l].data() + (size_t) j * n_h;
|
||||
for (int i = 0; i < n_h; i++) {
|
||||
acc += (double) wh[i] * hl[i];
|
||||
}
|
||||
g[(size_t) j] = (float) acc;
|
||||
}
|
||||
// torch gate order: input, forget, cell, output
|
||||
for (int i = 0; i < n_h; i++) {
|
||||
const float gi = 1.0f / (1.0f + expf(-g[(size_t) i]));
|
||||
const float gf = 1.0f / (1.0f + expf(-g[(size_t) i + n_h]));
|
||||
const float gc = tanhf(g[(size_t) i + 2 * n_h]);
|
||||
const float go = 1.0f / (1.0f + expf(-g[(size_t) i + 3 * n_h]));
|
||||
cl[i] = gf * cl[i] + gi * gc;
|
||||
x[(size_t) i] = go * tanhf(cl[i]);
|
||||
}
|
||||
memcpy(hl, x.data(), (size_t) n_h * sizeof(float));
|
||||
in = hl;
|
||||
n_in = n_h;
|
||||
}
|
||||
}
|
||||
// projected, relu'd, normalized partial embedding
|
||||
std::vector<float> e(n_h);
|
||||
double norm = 0.0;
|
||||
for (int o = 0; o < n_h; o++) {
|
||||
double acc = b_proj[(size_t) o];
|
||||
const float * w = w_proj.data() + (size_t) o * n_h;
|
||||
const float * hl = h.data() + (size_t) 2 * n_h;
|
||||
for (int i = 0; i < n_h; i++) {
|
||||
acc += (double) w[i] * hl[i];
|
||||
}
|
||||
e[(size_t) o] = (float) std::max(acc, 0.0);
|
||||
norm += (double) e[(size_t) o] * e[(size_t) o];
|
||||
}
|
||||
norm = sqrt(norm);
|
||||
for (int o = 0; o < n_h; o++) {
|
||||
ve[(size_t) o] += (float) (e[(size_t) o] / norm);
|
||||
}
|
||||
}
|
||||
double norm = 0.0;
|
||||
for (float v : ve) {
|
||||
norm += (double) v * v;
|
||||
}
|
||||
norm = sqrt(norm);
|
||||
for (float & v : ve) {
|
||||
v = (float) (v / norm);
|
||||
}
|
||||
|
||||
// speaker projection row
|
||||
std::vector<float> 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<float> rows((size_t) n_e);
|
||||
for (int o = 0; o < n_e; o++) {
|
||||
double acc = b_spkr[(size_t) o];
|
||||
const float * w = w_spkr.data() + (size_t) o * n_h;
|
||||
for (int i = 0; i < n_h; i++) {
|
||||
acc += (double) w[i] * ve[(size_t) i];
|
||||
}
|
||||
rows[(size_t) o] = (float) acc;
|
||||
}
|
||||
|
||||
// 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<float> 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<float> ln_w, ln_b, wq, bq, wk, bk, wv, bv, wo, bo, emo;
|
||||
if (!read_t("a.cenc.perceiver.attn.norm.weight", ln_w) || !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<float> & w, const std::vector<float> & b,
|
||||
const std::vector<float> & in, int n_rows, std::vector<float> & out) {
|
||||
out.resize((size_t) n_rows * n_e);
|
||||
for (int r = 0; r < n_rows; r++) {
|
||||
for (int o = 0; o < n_e; o++) {
|
||||
double acc = b[(size_t) o];
|
||||
const float * wr = w.data() + (size_t) o * n_e;
|
||||
const float * ir = in.data() + (size_t) r * n_e;
|
||||
for (int i = 0; i < n_e; i++) {
|
||||
acc += (double) wr[i] * ir[i];
|
||||
}
|
||||
out[(size_t) r * n_e + o] = (float) acc;
|
||||
}
|
||||
}
|
||||
};
|
||||
auto attn_block = [&](const std::vector<float> & x1, int n1,
|
||||
const std::vector<float> & x2, int n2, std::vector<float> & out) {
|
||||
std::vector<float> nx1((size_t) n1 * n_e), nx2((size_t) n2 * n_e);
|
||||
for (int r = 0; r < n1; r++) {
|
||||
layer_norm(x1.data() + (size_t) r * n_e, nx1.data() + (size_t) r * n_e);
|
||||
}
|
||||
for (int r = 0; r < n2; r++) {
|
||||
layer_norm(x2.data() + (size_t) r * n_e, nx2.data() + (size_t) r * n_e);
|
||||
}
|
||||
std::vector<float> q, k, v;
|
||||
linear(wq, bq, nx1, n1, q);
|
||||
linear(wk, bk, nx2, n2, k);
|
||||
linear(wv, bv, nx2, n2, v);
|
||||
|
||||
std::vector<float> ctxt((size_t) n1 * n_e);
|
||||
std::vector<double> sc((size_t) n2);
|
||||
for (int hd = 0; hd < n_head; hd++) {
|
||||
const int off = hd * d_head;
|
||||
for (int t = 0; t < n1; t++) {
|
||||
double mx = -1e30;
|
||||
for (int s = 0; s < n2; s++) {
|
||||
double acc = 0.0;
|
||||
for (int i = 0; i < d_head; i++) {
|
||||
acc += (double) q[(size_t) t * n_e + off + i] * k[(size_t) s * n_e + off + i];
|
||||
}
|
||||
sc[(size_t) s] = acc / sqrt((double) d_head);
|
||||
mx = std::max(mx, sc[(size_t) s]);
|
||||
}
|
||||
double sum = 0.0;
|
||||
for (int s = 0; s < n2; s++) {
|
||||
sc[(size_t) s] = exp(sc[(size_t) s] - mx);
|
||||
sum += sc[(size_t) s];
|
||||
}
|
||||
for (int i = 0; i < d_head; i++) {
|
||||
double acc = 0.0;
|
||||
for (int s = 0; s < n2; s++) {
|
||||
acc += sc[(size_t) s] * v[(size_t) s * n_e + off + i];
|
||||
}
|
||||
ctxt[(size_t) t * n_e + off + i] = (float) (acc / sum);
|
||||
}
|
||||
}
|
||||
}
|
||||
linear(wo, bo, ctxt, n1, out);
|
||||
for (size_t i = 0; i < out.size(); i++) {
|
||||
out[i] += x1[i];
|
||||
}
|
||||
};
|
||||
|
||||
const int n_q = (int) (query.size() / n_e);
|
||||
const std::vector<float> & x2 = pse;
|
||||
std::vector<float> 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<float> & out_embd) {
|
||||
if (!ctx->ctx_gen_a) {
|
||||
LOG_ERR("%s: reference encoding needs the audio generation context\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
const auto & entries = audio_tokens->batch_f32.entries;
|
||||
GGML_ASSERT(entries.size() == 5);
|
||||
|
||||
std::vector<float> xvec;
|
||||
if (!cbx_ref_xvec(ctx, entries[0], xvec)) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
std::vector<int32_t> flow_codes, t3_codes;
|
||||
std::vector<float> t3_rows;
|
||||
if (cbx_ref_tokenize(ctx, entries[1], flow_codes, nullptr) != 0 ||
|
||||
cbx_ref_tokenize(ctx, entries[2], t3_codes, &t3_rows) != 0) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
const size_t n_feat = (size_t) entries[3].nx() * 80;
|
||||
if ((size_t) entries[3].nx() != 2 * flow_codes.size()) {
|
||||
LOG_ERR("%s: reference mel length %d does not match %zu codes\n",
|
||||
__func__, entries[3].nx(), flow_codes.size());
|
||||
return 1;
|
||||
}
|
||||
|
||||
// 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<float> pse;
|
||||
if (mtl) {
|
||||
pse = std::move(t3_rows);
|
||||
}
|
||||
std::vector<float> rows;
|
||||
if (cbx_ref_cond(ctx, entries[4], pse, rows) != 0) {
|
||||
return 1;
|
||||
}
|
||||
if (!mtl) {
|
||||
rows.insert(rows.end(), t3_rows.begin(), t3_rows.end());
|
||||
}
|
||||
GGML_ASSERT(rows.size() == (size_t) audio_tokens->n_tokens * (size_t) ctx->n_embd_out());
|
||||
|
||||
ctx->gen_ref_tokens = std::move(flow_codes);
|
||||
ctx->gen_ref_feat.assign(entries[3].get_ro_buf().begin(),
|
||||
entries[3].get_ro_buf().begin() + n_feat);
|
||||
ctx->gen_ref_spk = std::move(xvec);
|
||||
|
||||
out_embd = std::move(rows);
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
||||
mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) {
|
||||
mtmd_gen_inp inp{};
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
|
||||
@@ -2107,6 +2549,11 @@ struct mtmd_gen_state {
|
||||
std::vector<float> out_embd; // next-step hidden state fed back to backbone (GEN_CODE)
|
||||
std::vector<float> out_audio; // decoded PCM samples for the current frame (GEN_WAV)
|
||||
|
||||
// voice cloning reference for GEN_WAV, empty selects the default voice of the model (chatterbox)
|
||||
std::vector<int32_t> ref_tokens;
|
||||
std::vector<float> ref_feat;
|
||||
std::vector<float> ref_spk;
|
||||
|
||||
~mtmd_gen_state() {
|
||||
clip_gen_state_free(clip_state);
|
||||
}
|
||||
@@ -2140,6 +2587,21 @@ void mtmd_gen_state_reset(mtmd_gen_state * state) {
|
||||
clip_gen_state_reset(state->clip_state);
|
||||
}
|
||||
|
||||
int32_t mtmd_gen_state_set_ref(mtmd_context * ctx, mtmd_gen_state * state) {
|
||||
if (!state || state->ctx_clip != ctx->ctx_gen_a) {
|
||||
LOG_ERR("%s: invalid gen state\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
if (ctx->gen_ref_tokens.empty()) {
|
||||
LOG_ERR("%s: the last encode produced no reference data\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
state->ref_tokens = ctx->gen_ref_tokens;
|
||||
state->ref_feat = ctx->gen_ref_feat;
|
||||
state->ref_spk = ctx->gen_ref_spk;
|
||||
return 0;
|
||||
}
|
||||
|
||||
static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, mtmd_gen_state * state, const mtmd_gen_inp * inp, mtmd_gen_out * out) {
|
||||
clip_ctx * ctx_clip = ctx->ctx_gen_a;
|
||||
if (!ctx_clip) {
|
||||
@@ -2202,6 +2664,45 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, mtmd_gen_state *
|
||||
}
|
||||
|
||||
// MTMD_GEN_PROCESS_TYPE_GEN_WAV
|
||||
if (clip_get_projector_type(ctx_clip) == PROJECTOR_TYPE_CHATTERBOX) {
|
||||
if (!inp->codes || inp->n_codes == 0) {
|
||||
LOG_ERR("%s: codes required for gen_wav\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
std::vector<int32_t> in_codes(inp->codes, inp->codes + inp->n_codes);
|
||||
|
||||
// an empty reference selects the precomputed default voice of the model
|
||||
const bool has_ref = !state->ref_tokens.empty();
|
||||
|
||||
// the batch entry is unused, present to satisfy the encode interface
|
||||
clip_image_f32 dummy;
|
||||
dummy.set_size({1, 1}, false, true);
|
||||
dummy.cpy_buf(std::vector<float>(1, 0.0f));
|
||||
clip_image_f32_batch batch;
|
||||
batch.is_audio = true;
|
||||
batch.entries.push_back(std::move(dummy));
|
||||
|
||||
clip_encode_params params;
|
||||
params.imgs = &batch;
|
||||
params.n_threads = ctx->n_threads;
|
||||
params.gen_process = CLIP_GEN_PROCESS_TTS;
|
||||
params.codes = &in_codes;
|
||||
params.out_audio = &state->out_audio;
|
||||
params.ref_tokens = has_ref ? &state->ref_tokens : nullptr;
|
||||
params.ref_feat = has_ref ? &state->ref_feat : nullptr;
|
||||
params.ref_spk = has_ref ? &state->ref_spk : nullptr;
|
||||
params.gen_state = state->clip_state;
|
||||
|
||||
if (!clip_encode(ctx_clip, ¶ms)) {
|
||||
LOG_ERR("%s: clip_encode failed (gen_wav)\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
out->audio = state->out_audio.data();
|
||||
out->n_samples = state->out_audio.size();
|
||||
return 0;
|
||||
}
|
||||
|
||||
const bool has_codes = inp->codes && inp->n_codes > 0;
|
||||
const bool has_feats = inp->feats && inp->n_feats > 0;
|
||||
if (has_codes == has_feats) {
|
||||
|
||||
@@ -335,6 +335,7 @@ MTMD_API int32_t mtmd_encode_chunk(mtmd_context * ctx,
|
||||
MTMD_API float * mtmd_get_output_embd(mtmd_context * ctx);
|
||||
|
||||
|
||||
|
||||
// batch encoding API
|
||||
// chunks are not owned by the batch, they will not be freed by mtmd_batch_free()
|
||||
// batch is valid for a given context, cannot be shared across contexts
|
||||
@@ -420,6 +421,7 @@ enum mtmd_gen_audio_type {
|
||||
MTMD_GEN_AUDIO_TYPE_NONE, // not supported
|
||||
MTMD_GEN_AUDIO_TYPE_QWEN3TTS,
|
||||
MTMD_GEN_AUDIO_TYPE_POCKETTTS,
|
||||
MTMD_GEN_AUDIO_TYPE_CHATTERBOX,
|
||||
};
|
||||
|
||||
struct mtmd_gen_audio_info {
|
||||
@@ -430,6 +432,10 @@ 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 semantic (codes, mel-spectrogram, etc.)
|
||||
@@ -486,6 +492,10 @@ MTMD_API void mtmd_gen_state_free(mtmd_gen_state * state);
|
||||
// clear the decoder state (cold start), keep the RNG stream
|
||||
MTMD_API void mtmd_gen_state_reset(mtmd_gen_state * state);
|
||||
|
||||
// copy the gen-side reference data of the last mtmd_encode_chunk() call into the state (voice cloning)
|
||||
// returns non-zero if the last encode produced no such data
|
||||
MTMD_API int32_t mtmd_gen_state_set_ref(mtmd_context * ctx, mtmd_gen_state * state);
|
||||
|
||||
// defaults tuned for the loaded pipeline, callers override only what they care about
|
||||
MTMD_API struct mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx);
|
||||
|
||||
|
||||
@@ -59,3 +59,35 @@ The [upstream repository](https://huggingface.co/kyutai/pocket-tts) holds one co
|
||||
python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --outfile pocket-tts.gguf
|
||||
python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --mmproj --outfile mmproj-pocket-tts.gguf
|
||||
```
|
||||
|
||||
## Chatterbox
|
||||
|
||||
Two variants are supported: [Chatterbox Turbo](https://huggingface.co/ResembleAI/chatterbox-turbo) (English) and [Chatterbox Multilingual V3](https://huggingface.co/ResembleAI/chatterbox) (23 languages).
|
||||
|
||||
Available params:
|
||||
- `--tts-speaker-file` should point to a speaker reference audio file (wav, mp3). Optional, the default voice of the model is used without it
|
||||
- `--tts-lang` sets the language of the multilingual variant, for example `fr` (default: none). It is not used by Turbo
|
||||
|
||||
Example usage:
|
||||
|
||||
```sh
|
||||
llama-tts -m chatterbox-mtl.gguf \
|
||||
-mm mmproj-chatterbox-mtl.gguf \
|
||||
-p "Bonjour tout le monde" \
|
||||
--tts-lang fr \
|
||||
--tts-speaker-file speaker.mp3 \
|
||||
--output out.wav
|
||||
```
|
||||
|
||||
Note: the multilingual variant runs classifier-free guidance with a second sequence. `llama-tts` sets this up by itself, but `llama-server` does not support it yet.
|
||||
|
||||
**Note for GGUF conversion:**
|
||||
|
||||
The upstream repositories ship raw checkpoints without `config.json`. Download the files below into a directory, then convert it as usual (the variant is detected from the talker file):
|
||||
- Turbo: `t3_turbo_v1.safetensors`, `s3gen_meanflow.safetensors`, `ve.safetensors`, `conds.pt` and the GPT-2 tokenizer files (`vocab.json`, `merges.txt`, `added_tokens.json`, `special_tokens_map.json`, `tokenizer_config.json`)
|
||||
- Multilingual V3: `t3_mtl23ls_v3.safetensors`, `s3gen_v3.safetensors`, `ve.safetensors`, `conds.pt`, `mtl_tokenizer.json`
|
||||
|
||||
```sh
|
||||
python convert_hf_to_gguf.py path/to/chatterbox --outfile chatterbox.gguf
|
||||
python convert_hf_to_gguf.py path/to/chatterbox --mmproj --outfile mmproj-chatterbox.gguf
|
||||
```
|
||||
@@ -68,6 +68,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);
|
||||
|
||||
|
||||
Reference in new issue
Block a user