From 3bc4cb0a90019e5a793dd696656cf5de66c0cf14 Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 2 Aug 2026 13:11:24 +0200 Subject: [PATCH 01/14] poc chatterbox --- tools/mtmd/CMakeLists.txt | 1 + tools/mtmd/clip-impl.h | 4 + tools/mtmd/clip-model.h | 5 + tools/mtmd/clip.cpp | 465 +++++++++++++++++- tools/mtmd/clip.h | 22 + tools/mtmd/models/chatterbox.cpp | 784 +++++++++++++++++++++++++++++++ tools/mtmd/models/models.h | 18 + tools/mtmd/mtmd-audio.cpp | 459 ++++++++++++++++++ tools/mtmd/mtmd-audio.h | 51 ++ tools/mtmd/mtmd-helper-gen.cpp | 422 +++++++++++++++++ tools/mtmd/mtmd.cpp | 452 +++++++++++++++++- tools/mtmd/mtmd.h | 36 ++ tools/tts/tts.cpp | 28 ++ 13 files changed, 2738 insertions(+), 9 deletions(-) create mode 100644 tools/mtmd/models/chatterbox.cpp diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 4675fb9a97b6..0b35a226efd9 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -27,6 +27,7 @@ add_library(mtmd clip-model.h clip-graph.h models/models.h + models/chatterbox.cpp models/cogvlm.cpp models/conformer.cpp models/dotsocr.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 7222660c7797..f7d03ad81b43 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -455,6 +455,8 @@ enum projector_type { PROJECTOR_TYPE_MIMO_AUDIO, PROJECTOR_TYPE_QWEN3TTS_SPKENC, PROJECTOR_TYPE_QWEN3TTS_GEN, + PROJECTOR_TYPE_CHATTERBOX, + PROJECTOR_TYPE_CHATTERBOX_SPKENC, PROJECTOR_TYPE_UNKNOWN, }; @@ -514,6 +516,8 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_PARAKEET, "parakeet"}, { PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"}, { PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"}, + { PROJECTOR_TYPE_CHATTERBOX, "chatterbox"}, + { PROJECTOR_TYPE_CHATTERBOX_SPKENC, "chatterbox_spkenc"}, }; static projector_type clip_projector_type_from_string(const std::string & str) { diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 30014f1f506c..1779b98d7e2d 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -687,6 +688,10 @@ struct clip_model { // qwen3tts code2wav: RVQ codes -> raw PCM clip_code2wav c2w; + // chatterbox audio stack: tensors are looked up by their source name at + // graph build time, the set is too large and too nested for named fields + std::map cbx_tensors; + // cogvlm ggml_tensor * mm_post_fc_norm_w = nullptr; ggml_tensor * mm_post_fc_norm_b = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 50644bf1e526..74ab0985d013 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1054,6 +1054,35 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + { + builder = std::make_unique(ctx, img); + } break; + case PROJECTOR_TYPE_CHATTERBOX: + { + const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_CODE_GEN; + int n_tokens = params && params->codes ? (int) params->codes->size() : 0; + int n_prompt_mel = 0; + if (n_tokens > 0) { + if (params->ref_tokens) { + n_tokens += (int) params->ref_tokens->size(); + } else { + auto it = ctx->model.cbx_tensors.find("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("cond.gen_prompt_feat"); + GGML_ASSERT(it != ctx->model.cbx_tensors.end()); + n_prompt_mel = (int) it->second->ne[1]; + } + } + const int vnm = params ? params->vocode_n_mel : 0; + const int vns = params ? params->vocode_n_stft : 0; + builder = std::make_unique(ctx, img, gen_process, n_tokens, n_prompt_mel, vnm, vns); + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_CODE_GEN; @@ -1708,6 +1737,15 @@ struct clip_model_loader { hparams.audio_window_len = 1024; hparams.audio_hop_len = 256; } break; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + { + // CAMPPlus x-vector; kaldi fbank front-end (povey + // window, 25 ms / 10 ms framing, 512-point spectrum) + hparams.audio_sample_rate = 16000; + hparams.audio_n_fft = 512; + hparams.audio_window_len = 400; + hparams.audio_hop_len = 160; + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { // discrete-token autoregressive predictor, no mel-frontend needed @@ -1728,6 +1766,11 @@ struct clip_model_loader { // matches the reference decoder's sliding_window (speech_tokenizer/config.json) hparams.wav_tfm_swa = 72; } break; + case PROJECTOR_TYPE_CHATTERBOX: + { + // s3gen output rate; the s3 tokenizer mel front-end runs at 16 kHz + get_u32("chatterbox.sample_rate", hparams.audio_sample_rate); + } break; case PROJECTOR_TYPE_PADDLEOCR: { hparams.n_merge = 2; @@ -2057,7 +2100,9 @@ struct clip_model_loader { const bool has_standard_layers = ( model.proj_type != PROJECTOR_TYPE_GEMMA3NV && - model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC); + model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC && + model.proj_type != PROJECTOR_TYPE_CHATTERBOX && + model.proj_type != PROJECTOR_TYPE_CHATTERBOX_SPKENC); // layers const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0; @@ -2823,6 +2868,26 @@ 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: + { + // the chatterbox audio stack keeps its source tensor names, + // load everything in the file and index by name for the + // graph builders + for (ggml_tensor * t = ggml_get_first_tensor(ctx_meta.get()); t; t = ggml_get_next_tensor(ctx_meta.get(), t)) { + model.cbx_tensors[t->name] = get_tensor(t->name); + } + } break; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + { + // this context only carries the CAMPPlus x-vector body and + // the flow affine that maps its embedding to the s3gen dim + 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("spk.", 0) == 0 || name.rfind("flow.spk_embed_affine_layer.", 0) == 0) { + model.cbx_tensors[name] = get_tensor(name.c_str()); + } + } + } break; case PROJECTOR_TYPE_VOXTRAL: { model.conv1d_1_w = get_tensor(string_format(TN_CONV1D, 1, "weight")); @@ -3688,6 +3753,17 @@ struct clip_init_result clip_init(const char * fname, struct clip_context_params ctx_gen_audio = new clip_ctx(ctx_params); loader.load_hparams(ctx_gen_audio->model, CLIP_MODALITY_GEN_AUDIO); loader.load_tensors(*ctx_gen_audio); + if (ctx_gen_audio->model.proj_type == PROJECTOR_TYPE_CHATTERBOX) { + // the classic cfm solver unrolls up to 10 steps x 2 cfg + // estimator evaluations in one graph + ctx_gen_audio->max_nodes = 65536; + ctx_gen_audio->sched.reset( + ggml_backend_sched_new(ctx_gen_audio->backend_ptrs.data(), ctx_gen_audio->backend_buft.data(), + ctx_gen_audio->backend_ptrs.size(), ctx_gen_audio->max_nodes, false, true)); + if (ctx_params.cb_eval != nullptr) { + ggml_backend_sched_set_eval_callback(ctx_gen_audio->sched.get(), ctx_params.cb_eval, ctx_params.cb_eval_user_data); + } + } // TODO: fix warmup ctx_gen_audio->buf_compute_meta.resize(ctx_gen_audio->max_nodes * ggml_tensor_overhead() + ggml_graph_overhead()); } @@ -4024,6 +4100,12 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { // a single speaker embedding vector, regardless of its length n_patches = 1; } break; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + { + // statistics pooling collapses the whole clip into a single + // speaker embedding vector, regardless of its length + n_patches = 1; + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { // one hidden-state vector fed back to the talker per call @@ -4070,7 +4152,38 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 return clip_encode(ctx, ¶ms); } +size_t clip_cbx_read_tensor(struct clip_ctx * ctx, const char * name, float * out, size_t n_max) { + const auto & tensors = ctx->model.cbx_tensors; + auto it = tensors.find(name); + if (it == tensors.end() || !it->second) { + return 0; + } + ggml_tensor * t = it->second; + const size_t n = (size_t) ggml_nelements(t); + if (!out) { + return n; + } + if (n_max < n) { + return 0; + } + if (t->type == GGML_TYPE_F32) { + ggml_backend_tensor_get(t, out, 0, n * sizeof(float)); + } else if (t->type == GGML_TYPE_F16) { + std::vector tmp(n); + ggml_backend_tensor_get(t, tmp.data(), 0, n * sizeof(ggml_fp16_t)); + for (size_t i = 0; i < n; i++) out[i] = ggml_fp16_to_fp32(tmp[i]); + } else if (t->type == GGML_TYPE_I32) { + std::vector tmp(n); + ggml_backend_tensor_get(t, tmp.data(), 0, n * sizeof(int32_t)); + for (size_t i = 0; i < n; i++) out[i] = (float) tmp[i]; + } else { + return 0; + } + return n; +} + bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { + const clip_image_f32_batch & imgs = *params->imgs; int n_batch_cur = imgs.entries.size(); @@ -4172,7 +4285,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } set_input_f32("inp_raw", inp_raw); - } else if (!(ctx->proj_type() == PROJECTOR_TYPE_QWEN3TTS_GEN && params->gen_process == CLIP_GEN_PROCESS_CODE2WAV)) { + } else if (!(ctx->proj_type() == PROJECTOR_TYPE_QWEN3TTS_GEN && params->gen_process == CLIP_GEN_PROCESS_CODE2WAV) && + !(ctx->proj_type() == PROJECTOR_TYPE_CHATTERBOX && (params->gen_process == CLIP_GEN_PROCESS_TTS || params->gen_process == CLIP_GEN_PROCESS_TTS_VOCODE))) { // audio input (code2wav has no hidden-state/raw input at all, its only input is the "inp_codes" tensor handled in the switch below) GGML_ASSERT(imgs.entries.size() == 1); @@ -4711,6 +4825,116 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { set_input_f32("qwen2_attn_mask", qwen2_mask); } } break; + case PROJECTOR_TYPE_CHATTERBOX: + { + if (params->gen_process == CLIP_GEN_PROCESS_TTS_VOCODE) { + set_input_f32("inp_mel", *params->mel_in); + set_input_f32("inp_sstft", *params->sstft_in); + break; + } + if (params->gen_process == CLIP_GEN_PROCESS_TOKENIZE) { + const int T1 = (imgs.entries[0].nx() - 1) / 2 + 1; + const int T2 = (T1 - 1) / 2 + 1; + std::vector pos(T2); + for (int i = 0; i < T2; i++) { + pos[(size_t) i] = i; + } + set_input_i32("inp_pos", pos); + break; + } + if (params->gen_process != CLIP_GEN_PROCESS_TTS) { + break; + } + const int n_gen = (int) params->codes->size(); + int n_prompt = 0; + std::vector tokens; + if (params->ref_tokens) { + n_prompt = (int) params->ref_tokens->size(); + tokens = *params->ref_tokens; + tokens.resize((size_t) n_prompt + n_gen); + } else { + ggml_tensor * pt = model.cbx_tensors.at("cond.gen_prompt_token"); + n_prompt = (int) pt->ne[0]; + tokens.resize((size_t) n_prompt + n_gen); + ggml_backend_tensor_get(pt, tokens.data(), 0, (size_t) n_prompt * sizeof(int32_t)); + } + 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("cond.gen_prompt_feat"); + std::vector feat(ggml_nelements(pf)); + ggml_backend_tensor_get(pf, feat.data(), 0, ggml_nbytes(pf)); + set_input_f32("inp_prompt_feat", feat); + } + if (params->ref_spk) { + set_input_f32("inp_spk", *params->ref_spk); + } else { + ggml_tensor * sp = model.cbx_tensors.at("cond.gen_spk80"); + std::vector spk(ggml_nelements(sp)); + ggml_backend_tensor_get(sp, spk.data(), 0, ggml_nbytes(sp)); + set_input_f32("inp_spk", spk); + } + + // espnet relative positional encoding, entry k holds the + // sinusoid of relative position (T-1) - k + auto fill_pos = [&](const char * name, int T) { + const int d = 512; + std::vector pos((size_t) (2 * T - 1) * d); + for (int k = 0; k < 2 * T - 1; k++) { + const double rel = (double) (T - 1 - k); + for (int i = 0; i < d / 2; i++) { + const double div = exp(-(double) (2 * i) * log(10000.0) / d); + pos[(size_t) k * d + 2 * i ] = (float) sin(rel * div); + pos[(size_t) k * d + 2 * i + 1] = (float) cos(rel * div); + } + } + set_input_f32(name, pos); + }; + fill_pos("inp_pos1", T1); + fill_pos("inp_pos2", T2); + + // meanflow inputs: gaussian noise and the sinusoidal time + // embeddings of the solver span points, same schedule as the + // graph builder (matcha layout: sines then cosines, scale 1000) + std::vector noise((size_t) 80 * T2); + std::mt19937 rng(42); + std::normal_distribution nd(0.0f, 1.0f); + for (auto & f : noise) f = nd(rng); + set_input_f32("inp_noise", noise); + + const bool meanflow = model.cbx_tensors.count("est.time_embed_mixer.weight") > 0; + const int n_steps = meanflow ? 2 : 10; + std::vector temb((size_t) 320 * (n_steps + 1)); + for (int s = 0; s <= n_steps; s++) { + const double u = (double) s / n_steps; + const double t = meanflow ? u : 1.0 - cos(u * M_PI / 2.0); + for (int i = 0; i < 160; i++) { + const double div = exp(-(double) i * log(10000.0) / 159.0); + temb[(size_t) s * 320 + i ] = (float) sin(1000.0 * t * div); + temb[(size_t) s * 320 + i + 160] = (float) cos(1000.0 * t * div); + } + } + set_input_f32("inp_temb", temb); + + if (getenv("CBX_DUMP")) { + FILE * f = fopen("/tmp/cbx-noise.bin", "wb"); + fwrite(noise.data(), sizeof(float), noise.size(), f); + fclose(f); + } + + if (getenv("CBX_DUMP")) { + FILE * f = fopen("/tmp/cbx-tokens.bin", "wb"); + fwrite(tokens.data(), sizeof(int32_t), tokens.size(), f); + fclose(f); + } + } break; case PROJECTOR_TYPE_GEMMA3: case PROJECTOR_TYPE_GEMMA3NV: case PROJECTOR_TYPE_IDEFICS3: @@ -4733,6 +4957,19 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { { // do nothing } break; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + { + // batchnorm epsilon and the ceil-mode correction of the cam + // seg pooling: pooled sums are divided by the full segment + // length, the last partial segment gets rescaled + set_input_f32("inp_eps", {1e-5f}); + const int T1 = (imgs.entries[0].nx() - 1) / 2 + 1; + const int S = (T1 + 99) / 100; + std::vector segfix((size_t) S, 1.0f); + const int last = T1 - (S - 1) * 100; + segfix[(size_t) S - 1] = 100.0f / (float) last; + set_input_f32("inp_segfix", segfix); + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { if (params->gen_process == CLIP_GEN_PROCESS_CODE2WAV) { @@ -5252,6 +5489,223 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { // for audio gen models // + if (ctx->proj_type() == PROJECTOR_TYPE_CHATTERBOX && params->gen_process == CLIP_GEN_PROCESS_TOKENIZE) { + ggml_tensor * fsq = ggml_graph_get_tensor(gf, "out_fsq"); + GGML_ASSERT(fsq != nullptr && params->out_codes); + const int n_tok = (int) fsq->ne[1]; + std::vector h((size_t) 8 * n_tok); + ggml_backend_tensor_get(fsq, h.data(), 0, ggml_nbytes(fsq)); + + // fsq round to base 3: h in (-1, 1) maps to digits {0, 1, 2} + params->out_codes->resize(n_tok); + for (int t = 0; t < n_tok; t++) { + int32_t code = 0; + for (int i = 7; i >= 0; i--) { + const int32_t d = (int32_t) roundf(h[(size_t) t * 8 + i] * 0.9990000128746033f) + 1; + code = code * 3 + d; + } + (*params->out_codes)[(size_t) t] = code; + } + + if (getenv("CBX_DUMP")) { + FILE * f = fopen("/tmp/cbx-fsq.bin", "wb"); + fwrite(h.data(), sizeof(float), h.size(), f); + fclose(f); + } + 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 * mu = ggml_graph_get_tensor(gf, "out_mu"); + GGML_ASSERT(mu != nullptr); + std::vector mu_data(ggml_nelements(mu)); + ggml_backend_tensor_get(mu, mu_data.data(), 0, ggml_nbytes(mu)); + + if (getenv("CBX_DUMP")) { + FILE * f = fopen("/tmp/cbx-mu.bin", "wb"); + fwrite(mu_data.data(), sizeof(float), mu_data.size(), f); + fclose(f); + LOG_INF("%s: dumped out_mu [%d, %d] to /tmp/cbx-mu.bin\n", __func__, (int) mu->ne[0], (int) mu->ne[1]); + } + + if (getenv("CBX_DUMP")) { + ggml_tensor * dc = ggml_graph_get_tensor(gf, "out_dcond"); + std::vector dc_data(ggml_nelements(dc)); + ggml_backend_tensor_get(dc, dc_data.data(), 0, ggml_nbytes(dc)); + FILE * f = fopen("/tmp/cbx-dcond.bin", "wb"); + fwrite(dc_data.data(), sizeof(float), dc_data.size(), f); + fclose(f); + } + + ggml_tensor * mel = ggml_graph_get_tensor(gf, "out_mel"); + GGML_ASSERT(mel != nullptr); + std::vector mel_data(ggml_nelements(mel)); + ggml_backend_tensor_get(mel, mel_data.data(), 0, ggml_nbytes(mel)); + if (getenv("CBX_DUMP")) { + FILE * f = fopen("/tmp/cbx-mel.bin", "wb"); + fwrite(mel_data.data(), sizeof(float), mel_data.size(), f); + fclose(f); + LOG_INF("%s: dumped out_mel [%d, %d] to /tmp/cbx-mel.bin\n", __func__, (int) mel->ne[0], (int) mel->ne[1]); + } + + // hift bridge: f0 -> harmonic source -> source stft on the host, + // then the vocoder graph, then the istft + ggml_tensor * f0_t = ggml_graph_get_tensor(gf, "out_f0"); + GGML_ASSERT(f0_t != nullptr); + const int n_mel_out = (int) ggml_nelements(f0_t); + std::vector f0(n_mel_out); + ggml_backend_tensor_get(f0_t, f0.data(), 0, ggml_nbytes(f0_t)); + + if (getenv("CBX_DUMP")) { + FILE * f = fopen("/tmp/cbx-f0.bin", "wb"); + fwrite(f0.data(), sizeof(float), f0.size(), f); + fclose(f); + } + + // source: f0 upsampled x480 nearest, 9 harmonics, cumulative phase, + // uv gating and noise as in SineGen, merged by l_linear + tanh + const int ups_total = 480; + const double sr = 24000.0; + const int64_t n_wav = (int64_t) n_mel_out * ups_total; + std::vector lw(9); + float lb = 0.0f; + GGML_ASSERT(clip_cbx_read_tensor(ctx, "hift.m_source.l_linear.weight", lw.data(), lw.size()) == 9); + clip_cbx_read_tensor(ctx, "hift.m_source.l_linear.bias", &lb, 1); + + const bool det = getenv("CBX_DUMP") != nullptr; // deterministic phases and no noise for validation + std::mt19937 srng(1234); + std::uniform_real_distribution ud(-M_PI, M_PI); + std::normal_distribution snd(0.0f, 1.0f); + double phase[9]; + for (int h = 0; h < 9; h++) { + phase[h] = (h == 0 || det) ? 0.0 : ud(srng); + } + std::vector src((size_t) n_wav); + double cum[9] = {0.0}; + for (int64_t t = 0; t < n_wav; t++) { + const float f = f0[(size_t) (t / ups_total)]; + const float uv = f > 10.0f ? 1.0f : 0.0f; // nsf_voiced_threshold + const float namp = uv * 0.003f + (1.0f - uv) * 0.1f / 3.0f; + float merged = lb; + for (int h = 0; h < 9; h++) { + cum[h] += (double) f * (h + 1) / sr; + cum[h] -= floor(cum[h]); + float sine = 0.1f * (float) sin(2.0 * M_PI * cum[h] + phase[h]); + sine = sine * uv + (det ? 0.0f : namp * snd(srng)); + merged += lw[(size_t) h] * sine; + } + src[(size_t) t] = tanhf(merged); + } + + // stft of the source: n_fft 16, hop 4, hann window, centered + const int n_fft = 16, hop = 4, n_bins = 9; + std::vector win(n_fft); + for (int i = 0; i < n_fft; i++) win[(size_t) i] = 0.5f - 0.5f * cosf(2.0f * (float) M_PI * i / n_fft); + const int n_stft = (int) (n_wav / hop) + 1; + std::vector sstft((size_t) n_stft * 18); + for (int fr = 0; fr < n_stft; fr++) { + const int64_t c0 = (int64_t) fr * hop - n_fft / 2; // centered, reflect padded + for (int k = 0; k < n_bins; k++) { + double re = 0.0, im = 0.0; + for (int i = 0; i < n_fft; i++) { + int64_t idx = c0 + i; + if (idx < 0) idx = -idx; + if (idx >= n_wav) idx = 2 * (n_wav - 1) - idx; + const double v = (double) src[(size_t) idx] * win[(size_t) i]; + const double a = 2.0 * M_PI * k * i / n_fft; + re += v * cos(a); + im -= v * sin(a); + } + sstft[(size_t) fr * 18 + k ] = (float) re; + sstft[(size_t) fr * 18 + 9 + k] = (float) im; + } + } + + if (getenv("CBX_DUMP")) { + FILE * f = fopen("/tmp/cbx-src.bin", "wb"); + fwrite(src.data(), sizeof(float), src.size(), f); + fclose(f); + f = fopen("/tmp/cbx-sstft.bin", "wb"); + fwrite(sstft.data(), sizeof(float), sstft.size(), f); + fclose(f); + } + + // vocoder graph on mel + source stft + std::vector spec; + { + clip_encode_params vp = *params; + vp.gen_process = CLIP_GEN_PROCESS_TTS_VOCODE; + vp.mel_in = &mel_data; + vp.sstft_in = &sstft; + vp.vocode_n_mel = n_mel_out; + vp.vocode_n_stft = n_stft; + vp.out_audio = nullptr; + vp.out_spec = &spec; + if (!clip_encode(ctx, &vp)) { + LOG_ERR("%s: vocoder stage failed\n", __func__); + return false; + } + } + + if (getenv("CBX_DUMP")) { + FILE * f = fopen("/tmp/cbx-spec.bin", "wb"); + fwrite(spec.data(), sizeof(float), spec.size(), f); + fclose(f); + } + + // istft: mag = clipped exp, phase = sin, hann overlap-add + const int n_frames_out = (int) (spec.size() / 18); + const int64_t n_out = (int64_t) (n_frames_out - 1) * hop; + std::vector acc((size_t) n_out + n_fft, 0.0); + std::vector wsum((size_t) n_out + n_fft, 0.0); + for (int fr = 0; fr < n_frames_out; fr++) { + double frame[16]; + for (int i = 0; i < n_fft; i++) { + double v = 0.0; + for (int k = 0; k < n_bins; k++) { + const double mag = fmin(exp((double) spec[(size_t) fr * 18 + k]), 1e2); + const double ph = sin((double) spec[(size_t) fr * 18 + 9 + k]); + const double re = mag * cos(ph), im = mag * sin(ph); + const double a = 2.0 * M_PI * k * i / n_fft; + const double w = (k == 0 || k == n_fft / 2) ? 1.0 : 2.0; + v += w * (re * cos(a) - im * sin(a)); + } + frame[i] = v / n_fft; + } + const int64_t o = (int64_t) fr * hop; + for (int i = 0; i < n_fft; i++) { + acc [(size_t) (o + i)] += frame[i] * win[(size_t) i]; + wsum[(size_t) (o + i)] += (double) win[(size_t) i] * win[(size_t) i]; + } + } + GGML_ASSERT(params->out_audio); + auto & out_audio = *params->out_audio; + out_audio.resize((size_t) std::min(n_out - n_fft / 2, n_wav)); + for (size_t i = 0; i < out_audio.size(); i++) { + const size_t j = i + n_fft / 2; // drop the centering pad + const double v = wsum[j] > 1e-11 ? acc[j] / wsum[j] : 0.0; + out_audio[(size_t) i] = (float) fmax(-0.99, fmin(0.99, v)); + } + + // the reference silences the first 20 ms and fades the next 20 ms in + // to hide the onset artifact of the flow prompt boundary (trim_fade) + const size_t n_trim = 24000 / 50; + for (size_t i = 0; i < 2 * n_trim && i < out_audio.size(); i++) { + const double g = i < n_trim ? 0.0 + : (cos(M_PI * (1.0 - (double) (i - n_trim) / n_trim)) + 1.0) / 2.0; + out_audio[(size_t) i] *= (float) g; + } + return true; + } + if (params->out_codes != nullptr) { ggml_tensor * codes = ggml_graph_get_tensor(gf, "out_codes"); if (codes == nullptr) { @@ -5434,6 +5888,13 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_fc_w->ne[2]; case PROJECTOR_TYPE_QWEN3TTS_GEN: return ctx->model.gen_code_out_embd_w->ne[0]; + case PROJECTOR_TYPE_CHATTERBOX: + // gen-only stack, no input projection into the backbone; the mel + // channel count stands in for the interface dimension + return 80; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + // x-vector projected through the s3gen speaker affine + return 80; case PROJECTOR_TYPE_PARAKEET: return ctx->model.mm_1_w->ne[1]; default: diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index e969b2b9b192..6c57ebdd8bd9 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -89,6 +89,9 @@ bool clip_image_batch_encode(struct clip_ctx * ctx, int n_threads, const struct enum clip_gen_process_type { CLIP_GEN_PROCESS_CODE_GEN, // h_state to codes CLIP_GEN_PROCESS_CODE2WAV, // codes to raw PCM audio + CLIP_GEN_PROCESS_TTS, // full utterance of codes to raw PCM audio + CLIP_GEN_PROCESS_TTS_VOCODE, // internal: mel + source stft to istft input + CLIP_GEN_PROCESS_TOKENIZE, // raw PCM audio to semantic speech tokens }; struct clip_encode_params { int n_threads = 1; @@ -112,12 +115,31 @@ struct clip_encode_params { // call (null or wrong size means cold start, state is zero-filled). // state_out receives the state to pass into the next call. const std::vector * codes = nullptr; + // TOKENIZE input + const float * pcm_in = nullptr; + size_t n_pcm = 0; + // TTS reference conditioning overriding the precomputed cond.gen_* + // defaults: speech tokens, mel-rate features and the 80-dim speaker + // vector of the reference clip (null means default) + const std::vector * ref_tokens = nullptr; + const std::vector * ref_feat = nullptr; + const std::vector * ref_spk = nullptr; + // TTS_VOCODE internal stage inputs + const std::vector * mel_in = nullptr; + const std::vector * sstft_in = nullptr; + int vocode_n_mel = 0; + int vocode_n_stft = 0; + std::vector * out_spec = nullptr; std::vector * out_audio = nullptr; const std::vector * state_in = nullptr; std::vector * state_out = nullptr; }; bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params); +// read a chatterbox tensor by source name, converted to F32. returns the +// element count, 0 if not found. out may be null to only query the size. +size_t clip_cbx_read_tensor(struct clip_ctx * ctx, const char * name, float * out, size_t n_max); + bool clip_is_llava(const struct clip_ctx * ctx); // note for contributor: this clip_is_(model) pattern is deprecated // do NOT add new functions like this diff --git a/tools/mtmd/models/chatterbox.cpp b/tools/mtmd/models/chatterbox.cpp new file mode 100644 index 000000000000..31ffe7e312fe --- /dev/null +++ b/tools/mtmd/models/chatterbox.cpp @@ -0,0 +1,784 @@ +#include "models.h" + +// Chatterbox s3gen, stage 1: speech tokens -> mu (flow encoder output). +// Weights come from the source-named tensor map (model.cbx_tensors), the +// estimator and hift stages extend this file. + +static ggml_tensor * cbx_t(const clip_model & model, const std::string & name) { + auto it = model.cbx_tensors.find(name); + if (it == model.cbx_tensors.end()) { + GGML_ABORT("missing chatterbox tensor: %s", name.c_str()); + } + return it->second; +} + +// x [C, T]: y = W x + b with torch Linear weights stored as [in, out] +static ggml_tensor * cbx_linear(ggml_context * ctx0, ggml_tensor * w, ggml_tensor * b, ggml_tensor * x) { + 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 +static ggml_tensor * cbx_conv1d(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * b, ggml_tensor * x, + int stride, int pad_l, int pad_r) { + 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; +} + +// 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] +static ggml_tensor * cbx_enc_layer(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, ggml_tensor * pos, + const std::string & p, 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, cbx_t(model, p + ".norm_mha.weight"), cbx_t(model, p + ".norm_mha.bias"), x, 1e-5f); + + ggml_tensor * q = cbx_linear(ctx0, cbx_t(model, p + ".self_attn.linear_q.weight"), cbx_t(model, p + ".self_attn.linear_q.bias"), cur); + ggml_tensor * k = cbx_linear(ctx0, cbx_t(model, p + ".self_attn.linear_k.weight"), cbx_t(model, p + ".self_attn.linear_k.bias"), cur); + ggml_tensor * v = cbx_linear(ctx0, cbx_t(model, p + ".self_attn.linear_v.weight"), cbx_t(model, p + ".self_attn.linear_v.bias"), cur); + ggml_tensor * pe = cbx_linear(ctx0, cbx_t(model, p + ".self_attn.linear_pos.weight"), 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 = cbx_t(model, p + ".self_attn.pos_bias_u"); // [64, 8] + ggml_tensor * w = cbx_t(model, p + ".self_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(ctx0, cbx_t(model, p + ".self_attn.linear_out.weight"), cbx_t(model, p + ".self_attn.linear_out.bias"), o); + x = ggml_add(ctx0, res, o); + + res = x; + cur = cbx_layer_norm(ctx0, cbx_t(model, p + ".norm_ff.weight"), cbx_t(model, p + ".norm_ff.bias"), x, 1e-5f); + cur = cbx_linear(ctx0, cbx_t(model, p + ".feed_forward.w_1.weight"), cbx_t(model, p + ".feed_forward.w_1.bias"), cur); + cur = ggml_silu(ctx0, cur); // swish + cur = cbx_linear(ctx0, cbx_t(model, p + ".feed_forward.w_2.weight"), cbx_t(model, p + ".feed_forward.w_2.bias"), 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] +static ggml_tensor * cbx_causal_block(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, const std::string & p) { + ggml_tensor * k = cbx_t(model, p + ".block.0.weight"); + x = cbx_conv1d(ctx0, k, cbx_t(model, p + ".block.0.bias"), x, 1, (int) k->ne[0] - 1, 0); + x = cbx_layer_norm(ctx0, cbx_t(model, p + ".block.2.weight"), cbx_t(model, p + ".block.2.bias"), x, 1e-5f); + return cbx_mish(ctx0, x); +} + +// resnet block with time conditioning; x [C, T], temb [1024] +static ggml_tensor * cbx_resnet(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, ggml_tensor * temb, const std::string & p) { + ggml_tensor * h = cbx_causal_block(model, ctx0, x, p + ".block1"); + ggml_tensor * tproj = cbx_linear(ctx0, cbx_t(model, p + ".mlp.1.weight"), cbx_t(model, p + ".mlp.1.bias"), cbx_mish(ctx0, temb)); + h = ggml_add(ctx0, h, tproj); // broadcast [256, 1] over T + h = cbx_causal_block(model, ctx0, h, p + ".block2"); + ggml_tensor * res = cbx_conv1d(ctx0, cbx_t(model, p + ".res_conv.weight"), cbx_t(model, p + ".res_conv.bias"), x, 1, 0, 0); + return ggml_add(ctx0, h, res); +} + +// diffusers-style transformer block, full attention; x [256, T] +static ggml_tensor * cbx_tfm_block(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, const std::string & p) { + 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, cbx_t(model, p + ".norm1.weight"), cbx_t(model, p + ".norm1.bias"), x, 1e-5f); + ggml_tensor * q = ggml_mul_mat(ctx0, cbx_t(model, p + ".attn1.to_q.weight"), cur); + ggml_tensor * k = ggml_mul_mat(ctx0, cbx_t(model, p + ".attn1.to_k.weight"), cur); + ggml_tensor * v = ggml_mul_mat(ctx0, cbx_t(model, p + ".attn1.to_v.weight"), 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(ctx0, cbx_t(model, p + ".attn1.to_out.0.weight"), cbx_t(model, p + ".attn1.to_out.0.bias"), o); + x = ggml_add(ctx0, res, o); + + res = x; + cur = cbx_layer_norm(ctx0, cbx_t(model, p + ".norm3.weight"), cbx_t(model, p + ".norm3.bias"), x, 1e-5f); + cur = cbx_linear(ctx0, cbx_t(model, p + ".ff.net.0.proj.weight"), cbx_t(model, p + ".ff.net.0.proj.bias"), cur); + cur = ggml_gelu_erf(ctx0, cur); + cur = cbx_linear(ctx0, cbx_t(model, p + ".ff.net.2.weight"), cbx_t(model, p + ".ff.net.2.bias"), cur); + return ggml_add(ctx0, res, cur); +} + +// one estimator evaluation; x_noise [80, T], mu [80, T], spks [80], +// cond [80, T], temb [1024] +static ggml_tensor * cbx_estimator(const clip_model & model, ggml_context * ctx0, ggml_tensor * x_noise, ggml_tensor * mu, + ggml_tensor * spks, ggml_tensor * cond, ggml_tensor * temb, int T) { + // 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 = cbx_resnet(model, ctx0, x, temb, "est.down_blocks.0.0"); + for (int j = 0; model.cbx_tensors.count("est.down_blocks.0.1." + std::to_string(j) + ".norm1.weight"); j++) { + x = cbx_tfm_block(model, ctx0, x, "est.down_blocks.0.1." + std::to_string(j)); + } + skip = x; + { + ggml_tensor * k = cbx_t(model, "est.down_blocks.0.2.weight"); + x = cbx_conv1d(ctx0, k, cbx_t(model, "est.down_blocks.0.2.bias"), x, 1, (int) k->ne[0] - 1, 0); + } + + // mid + for (int i = 0; model.cbx_tensors.count("est.mid_blocks." + std::to_string(i) + ".0.block1.block.0.weight"); i++) { + const std::string mp = "est.mid_blocks." + std::to_string(i); + x = cbx_resnet(model, ctx0, x, temb, mp + ".0"); + for (int j = 0; model.cbx_tensors.count(mp + ".1." + std::to_string(j) + ".norm1.weight"); j++) { + x = cbx_tfm_block(model, ctx0, x, mp + ".1." + std::to_string(j)); + } + } + + // up with skip + x = ggml_concat(ctx0, x, skip, 0); // [512, T] + x = cbx_resnet(model, ctx0, x, temb, "est.up_blocks.0.0"); + for (int j = 0; model.cbx_tensors.count("est.up_blocks.0.1." + std::to_string(j) + ".norm1.weight"); j++) { + x = cbx_tfm_block(model, ctx0, x, "est.up_blocks.0.1." + std::to_string(j)); + } + { + ggml_tensor * k = cbx_t(model, "est.up_blocks.0.2.weight"); + x = cbx_conv1d(ctx0, k, cbx_t(model, "est.up_blocks.0.2.bias"), x, 1, (int) k->ne[0] - 1, 0); + } + + x = cbx_causal_block(model, ctx0, x, "est.final_block"); + x = cbx_conv1d(ctx0, cbx_t(model, "est.final_proj.weight"), cbx_t(model, "est.final_proj.bias"), 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. +static ggml_cgraph * cbx_build_s3tok(const clip_model & model, ggml_context * ctx0, ggml_cgraph * gf, int T) { + 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(ctx0, cbx_t(model, "s3tok.encoder.conv1.weight"), cbx_t(model, "s3tok.encoder.conv1.bias"), x, 2, 1, 1); + x = ggml_gelu_erf(ctx0, x); + x = cbx_conv1d(ctx0, cbx_t(model, "s3tok.encoder.conv2.weight"), cbx_t(model, "s3tok.encoder.conv2.bias"), x, 2, 1, 1); + x = ggml_gelu_erf(ctx0, x); // [1280, T2] + + for (int li = 0; model.cbx_tensors.count("s3tok.encoder.blocks." + std::to_string(li) + ".attn_ln.weight"); li++) { + const std::string p = "s3tok.encoder.blocks." + std::to_string(li) + ".attn"; + + ggml_tensor * res = x; + ggml_tensor * cur = cbx_layer_norm(ctx0, cbx_t(model, p + "_ln.weight"), cbx_t(model, p + "_ln.bias"), x, 1e-5f); + ggml_tensor * q = cbx_linear(ctx0, cbx_t(model, p + ".query.weight"), cbx_t(model, p + ".query.bias"), cur); + ggml_tensor * k = ggml_mul_mat(ctx0, cbx_t(model, p + ".key.weight"), cur); + ggml_tensor * v = cbx_linear(ctx0, cbx_t(model, p + ".value.weight"), cbx_t(model, p + ".value.bias"), 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 = cbx_t(model, p + ".fsmn_block.weight"); + 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(ctx0, cbx_t(model, p + ".out.weight"), cbx_t(model, p + ".out.bias"), o); + x = ggml_add(ctx0, res, ggml_add(ctx0, o, fsm)); + + const std::string mp = "s3tok.encoder.blocks." + std::to_string(li) + ".mlp"; + res = x; + cur = cbx_layer_norm(ctx0, cbx_t(model, mp + "_ln.weight"), cbx_t(model, mp + "_ln.bias"), x, 1e-5f); + cur = cbx_linear(ctx0, cbx_t(model, mp + ".0.weight"), cbx_t(model, mp + ".0.bias"), cur); + cur = ggml_gelu_erf(ctx0, cur); + cur = cbx_linear(ctx0, cbx_t(model, mp + ".2.weight"), cbx_t(model, mp + ".2.bias"), cur); + x = ggml_add(ctx0, res, cur); + } + + x = cbx_linear(ctx0, cbx_t(model, "s3tok.quantizer._codebook.project_down.weight"), + cbx_t(model, "s3tok.quantizer._codebook.project_down.bias"), 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; +} + +static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * ctx0, ggml_cgraph * gf, int n_mel, int n_stft); + +ggml_cgraph * clip_graph_chatterbox::build() { + if (gen_process == CLIP_GEN_PROCESS_TTS_VOCODE) { + return cbx_build_vocoder(model, ctx0, gf, vocode_n_mel, vocode_n_stft); + } + if (gen_process == CLIP_GEN_PROCESS_TOKENIZE) { + return cbx_build_s3tok(model, ctx0, gf, 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 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, cbx_t(model, "flow.input_embedding.weight"), inp_tokens); // [512, T1] + + // embed: linear + layer norm, then the espnet xscale + x = cbx_linear(ctx0, cbx_t(model, "fenc.embed.out.0.weight"), cbx_t(model, "fenc.embed.out.0.bias"), x); + x = cbx_layer_norm(ctx0, cbx_t(model, "fenc.embed.out.1.weight"), cbx_t(model, "fenc.embed.out.1.bias"), x, 1e-5f); + x = ggml_scale(ctx0, x, sqrtf(512.0f)); + + // 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(ctx0, cbx_t(model, "fenc.pre_lookahead_layer.conv1.weight"), + cbx_t(model, "fenc.pre_lookahead_layer.conv1.bias"), x, 1, 0, 3); + cur = ggml_leaky_relu(ctx0, cur, 0.01f, false); + cur = cbx_conv1d(ctx0, cbx_t(model, "fenc.pre_lookahead_layer.conv2.weight"), + cbx_t(model, "fenc.pre_lookahead_layer.conv2.bias"), cur, 1, 2, 0); + x = ggml_add(ctx0, res, cur); + } + + for (int i = 0; model.cbx_tensors.count("fenc.encoders." + std::to_string(i) + ".norm_mha.weight"); i++) { + x = cbx_enc_layer(model, ctx0, x, pos1, "fenc.encoders." + std::to_string(i), T1); + } + + // upsample x2: nearest repeat, left pad 4, conv k=5 + { + ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [T1, 512] + xt = ggml_upscale_ext(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, cbx_t(model, "fenc.up_layer.conv.weight"), xt, 1, 0, 1); + x = ggml_cont(ctx0, ggml_transpose(ctx0, y)); // [512, T2] + x = ggml_add(ctx0, x, cbx_t(model, "fenc.up_layer.conv.bias")); + } + + // up embed: linear + layer norm + xscale + x = cbx_linear(ctx0, cbx_t(model, "fenc.up_embed.out.0.weight"), cbx_t(model, "fenc.up_embed.out.0.bias"), x); + x = cbx_layer_norm(ctx0, cbx_t(model, "fenc.up_embed.out.1.weight"), cbx_t(model, "fenc.up_embed.out.1.bias"), x, 1e-5f); + x = ggml_scale(ctx0, x, sqrtf(512.0f)); + + for (int i = 0; model.cbx_tensors.count("fenc.up_encoders." + std::to_string(i) + ".norm_mha.weight"); i++) { + x = cbx_enc_layer(model, ctx0, x, pos2, "fenc.up_encoders." + std::to_string(i), T2); + } + + x = cbx_layer_norm(ctx0, cbx_t(model, "fenc.after_norm.weight"), cbx_t(model, "fenc.after_norm.bias"), x, 1e-5f); + + // encoder projection to the mel channel count + ggml_tensor * mu = cbx_linear(ctx0, cbx_t(model, "flow.encoder_proj.weight"), cbx_t(model, "flow.encoder_proj.bias"), x); // [80, T2] + ggml_set_name(mu, "out_mu"); + ggml_set_output(mu); + ggml_build_forward_expand(gf, mu); + + // 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 = model.cbx_tensors.count("est.time_embed_mixer.weight") > 0; + 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(ctx0, cbx_t(model, "est.time_mlp.linear_1.weight"), cbx_t(model, "est.time_mlp.linear_1.bias"), e); + e = ggml_silu(ctx0, e); + e = cbx_linear(ctx0, cbx_t(model, "est.time_mlp.linear_2.weight"), cbx_t(model, "est.time_mlp.linear_2.bias"), 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, cbx_t(model, "est.time_embed_mixer.weight"), 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] + ggml_tensor * spks = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 80); + ggml_set_name(spks, "inp_spk"); + ggml_set_input(spks); + + 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 = cbx_estimator(model, ctx0, mx, mu, spks, cond, temb, T2); + if (i == 0) { + ggml_set_name(d, "out_dcond"); + ggml_set_output(d); + } + if (!meanflow) { + ggml_tensor * du = cbx_estimator(model, ctx0, 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])); + } + 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 (int i = 0; model.cbx_tensors.count("hift.f0_predictor.condnet." + std::to_string(i) + ".weight"); i += 2) { + const std::string cp = "hift.f0_predictor.condnet." + std::to_string(i); + fx = cbx_conv1d(ctx0, cbx_t(model, cp + ".weight"), cbx_t(model, cp + ".bias"), fx, 1, 1, 1); + fx = ggml_elu(ctx0, fx); + } + fx = cbx_linear(ctx0, cbx_t(model, "hift.f0_predictor.classifier.weight"), cbx_t(model, "hift.f0_predictor.classifier.bias"), 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; +} + + +static ggml_tensor * cbx_hift_resblock(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, const std::string & p); + +// x [C, T] -> symmetric-padded dilated conv -> [OC, T] +static ggml_tensor * cbx_conv1d_dil(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * b, ggml_tensor * x, int pad, int dil) { + 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; +} + +// mel [80, T] + source stft [18, T_stft] -> conv_post output [18, T_stft2] +static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * ctx0, ggml_cgraph * gf, int n_mel, int n_stft) { + 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(ctx0, cbx_t(model, "hift.conv_pre.weight"), cbx_t(model, "hift.conv_pre.bias"), mel, 3, 1); + + for (int i = 0; model.cbx_tensors.count("hift.ups." + std::to_string(i) + ".weight"); i++) { + const std::string is = std::to_string(i); + ggml_tensor * uk = cbx_t(model, "hift.ups." + is + ".weight"); + 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, cbx_t(model, "hift.ups." + is + ".bias")); + + const bool is_last = !model.cbx_tensors.count("hift.ups." + std::to_string(i + 1) + ".weight"); + 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 = cbx_t(model, "hift.source_downs." + is + ".weight"); + 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, cbx_t(model, "hift.source_downs." + is + ".bias")); + } + si = cbx_hift_resblock(model, ctx0, si, "hift.source_resblocks." + is); + // 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 = nullptr; + for (int j = 3 * i; j < 3 * (i + 1); j++) { + ggml_tensor * r = cbx_hift_resblock(model, ctx0, x, "hift.resblocks." + std::to_string(j)); + acc = acc ? ggml_add(ctx0, acc, r) : r; + } + x = ggml_scale(ctx0, acc, 1.0f / 3.0f); + } + + x = ggml_leaky_relu(ctx0, x, 0.01f, false); + x = cbx_conv1d_dil(ctx0, cbx_t(model, "hift.conv_post.weight"), cbx_t(model, "hift.conv_post.bias"), x, 3, 1); + ggml_set_name(x, "out_spec"); + ggml_set_output(x); + ggml_build_forward_expand(gf, x); + return gf; +} + +// snake activation with per-channel alpha: x + sin^2(alpha x) / (alpha + eps) +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 +static ggml_tensor * cbx_hift_resblock(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, const std::string & p) { + static const int dil[3] = {1, 3, 5}; + for (int j = 0; model.cbx_tensors.count(p + ".convs1." + std::to_string(j) + ".weight"); j++) { + const std::string js = std::to_string(j); + ggml_tensor * a1 = cbx_t(model, p + ".activations1." + js + ".alpha"); + ggml_tensor * a2 = cbx_t(model, p + ".activations2." + js + ".alpha"); + ggml_tensor * k1 = cbx_t(model, p + ".convs1." + js + ".weight"); + ggml_tensor * k2 = cbx_t(model, p + ".convs2." + js + ".weight"); + const int d = dil[j % 3]; + const int p1 = (int) (k1->ne[0] - 1) / 2 * d; + const int p2 = (int) (k2->ne[0] - 1) / 2; + ggml_tensor * xt = cbx_snake(ctx0, x, a1); + xt = cbx_conv1d_dil(ctx0, k1, cbx_t(model, p + ".convs1." + js + ".bias"), xt, p1, d); + xt = cbx_snake(ctx0, xt, a2); + xt = cbx_conv1d_dil(ctx0, k2, cbx_t(model, p + ".convs2." + js + ".bias"), xt, p2, 1); + x = ggml_add(ctx0, x, xt); + } + return x; +} + + + +// Chatterbox speaker encoder: CAMPPlus x-vector on kaldi fbank features, +// projected through the s3gen speaker affine. Mirrors s3gen/xvector.py. + +// per-channel batchnorm on x [C, T]; scale = w / sqrt(var + eps), shift folds +// the running mean. pass null w/b for the affine=False variant +static ggml_tensor * cbx_bn1d(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, + const std::string & p, ggml_tensor * eps) { + ggml_tensor * mean = cbx_t(model, p + ".running_mean"); + ggml_tensor * var = cbx_t(model, p + ".running_var"); + ggml_tensor * sd = ggml_sqrt(ctx0, ggml_add(ctx0, var, eps)); + if (!model.cbx_tensors.count(p + ".weight")) { + return ggml_div(ctx0, ggml_sub(ctx0, x, mean), sd); + } + ggml_tensor * a = ggml_div(ctx0, cbx_t(model, p + ".weight"), sd); + ggml_tensor * shift = ggml_sub(ctx0, cbx_t(model, p + ".bias"), ggml_mul(ctx0, mean, a)); + return ggml_add(ctx0, ggml_mul(ctx0, x, a), shift); +} + +// batchnorm + relu on a conv2d activation [W=T, H=F, C, 1], stats on ne2 +static ggml_tensor * cbx_bn2d_relu(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, + const std::string & p, ggml_tensor * eps) { + const int C = (int) x->ne[2]; + ggml_tensor * mean = ggml_reshape_4d(ctx0, cbx_t(model, p + ".running_mean"), 1, 1, C, 1); + ggml_tensor * var = ggml_reshape_4d(ctx0, cbx_t(model, p + ".running_var"), 1, 1, C, 1); + ggml_tensor * w = ggml_reshape_4d(ctx0, cbx_t(model, p + ".weight"), 1, 1, C, 1); + ggml_tensor * b = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bias"), 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_relu(ctx0, ggml_add(ctx0, ggml_mul(ctx0, x, a), shift)); +} + +// fcm residual 2d block, stride on the frequency axis only +static ggml_tensor * cbx_res2d(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, + const std::string & p, int stride, ggml_tensor * eps) { + ggml_tensor * cur = ggml_conv_2d(ctx0, cbx_t(model, p + ".conv1.weight"), x, 1, stride, 1, 1, 1, 1); + cur = cbx_bn2d_relu(model, ctx0, cur, p + ".bn1", eps); + cur = ggml_conv_2d(ctx0, cbx_t(model, p + ".conv2.weight"), cur, 1, 1, 1, 1, 1, 1); + // bn2 without the relu, applied before the residual add + { + const int C = (int) cur->ne[2]; + ggml_tensor * mean = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.running_mean"), 1, 1, C, 1); + ggml_tensor * var = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.running_var"), 1, 1, C, 1); + ggml_tensor * w = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.weight"), 1, 1, C, 1); + ggml_tensor * b = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.bias"), 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)); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, a), shift); + } + ggml_tensor * res = x; + if (model.cbx_tensors.count(p + ".shortcut.0.weight")) { + res = ggml_conv_2d(ctx0, cbx_t(model, p + ".shortcut.0.weight"), x, 1, stride, 0, 0, 1, 1); + const int C = (int) res->ne[2]; + ggml_tensor * mean = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.running_mean"), 1, 1, C, 1); + ggml_tensor * var = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.running_var"), 1, 1, C, 1); + ggml_tensor * w = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.weight"), 1, 1, C, 1); + ggml_tensor * b = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.bias"), 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)); + res = ggml_add(ctx0, ggml_mul(ctx0, res, a), shift); + } + return ggml_relu(ctx0, ggml_add(ctx0, cur, res)); +} + +// cam dense tdnn layer: bottleneck then context-gated conv; x [C_in, T] -> [growth, T] +static ggml_tensor * cbx_cam_layer(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, + const std::string & p, int dil, ggml_tensor * eps, ggml_tensor * segfix) { + ggml_tensor * h = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, p + ".nonlinear1.batchnorm", eps)); + h = cbx_conv1d(ctx0, cbx_t(model, p + ".linear1.weight"), nullptr, h, 1, 0, 0); + h = ggml_relu(ctx0, cbx_bn1d(model, ctx0, h, p + ".nonlinear2.batchnorm", eps)); + + const std::string cp = p + ".cam_layer"; + ggml_tensor * k = cbx_t(model, cp + ".linear_local.weight"); + const int pad = ((int) k->ne[0] - 1) / 2 * dil; + ggml_tensor * y = cbx_conv1d_dil(ctx0, 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_upscale_ext(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(ctx0, cbx_t(model, cp + ".linear1.weight"), cbx_t(model, cp + ".linear1.bias"), context, 1, 0, 0); + context = ggml_relu(ctx0, context); + context = cbx_conv1d(ctx0, cbx_t(model, cp + ".linear2.weight"), cbx_t(model, cp + ".linear2.bias"), context, 1, 0, 0); + ggml_tensor * m = ggml_sigmoid(ctx0, context); + return ggml_mul(ctx0, y, m); +} + +ggml_cgraph * clip_graph_chatterbox_spkenc::build() { + 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, cbx_t(model, "spk.head.conv1.weight"), x, 1, 1, 1, 1, 1, 1); + x = cbx_bn2d_relu(model, ctx0, x, "spk.head.bn1", eps); + x = cbx_res2d(model, ctx0, x, "spk.head.layer1.0", 2, eps); + x = cbx_res2d(model, ctx0, x, "spk.head.layer1.1", 1, eps); + x = cbx_res2d(model, ctx0, x, "spk.head.layer2.0", 2, eps); + x = cbx_res2d(model, ctx0, x, "spk.head.layer2.1", 1, eps); + x = ggml_conv_2d(ctx0, cbx_t(model, "spk.head.conv2.weight"), x, 1, 2, 1, 1, 1, 1); + x = cbx_bn2d_relu(model, ctx0, x, "spk.head.bn2", eps); + x = ggml_reshape_2d(ctx0, x, T, 320); + x = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [320, T] + + // tdnn k5 stride 2 over time, then the three cam dense blocks + x = cbx_conv1d(ctx0, cbx_t(model, "spk.xvector.tdnn.linear.weight"), nullptr, x, 2, 2, 2); // [128, T1] + x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, "spk.xvector.tdnn.nonlinear.batchnorm", eps)); + + static const int block_dil[3] = {1, 2, 2}; + for (int bi = 1; bi <= 3; bi++) { + const std::string bp = "spk.xvector.block" + std::to_string(bi); + for (int li = 1; model.cbx_tensors.count(bp + ".tdnnd" + std::to_string(li) + ".linear1.weight"); li++) { + ggml_tensor * out = cbx_cam_layer(model, ctx0, x, bp + ".tdnnd" + std::to_string(li), + block_dil[bi - 1], eps, segfix); + x = ggml_concat(ctx0, x, out, 0); + } + const std::string tp = "spk.xvector.transit" + std::to_string(bi); + x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, tp + ".nonlinear.batchnorm", eps)); + x = cbx_conv1d(ctx0, cbx_t(model, tp + ".linear.weight"), nullptr, x, 1, 0, 0); + } + x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, "spk.xvector.out_nonlinear.batchnorm", 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] + + // dense 1024 -> 192, batchnorm without affine, into the x-vector + ggml_tensor * dw = ggml_reshape_2d(ctx0, cbx_t(model, "spk.xvector.dense.linear.weight"), 1024, 192); + ggml_tensor * emb = ggml_mul_mat(ctx0, dw, stats); // [192, 1] + emb = cbx_bn1d(model, ctx0, emb, "spk.xvector.dense.nonlinear.batchnorm", eps); + emb = ggml_reshape_1d(ctx0, emb, 192); + ggml_set_name(emb, "out_xvec"); + ggml_set_output(emb); + ggml_build_forward_expand(gf, emb); + + // normalize then the s3gen speaker affine + ggml_tensor * n2 = ggml_sqrt(ctx0, ggml_sum(ctx0, ggml_mul(ctx0, emb, emb))); + ggml_tensor * unit = ggml_div(ctx0, emb, n2); + ggml_tensor * spk80 = cbx_linear(ctx0, cbx_t(model, "flow.spk_embed_affine_layer.weight"), + cbx_t(model, "flow.spk_embed_affine_layer.bias"), + ggml_reshape_2d(ctx0, unit, 192, 1)); + spk80 = ggml_cont(ctx0, ggml_reshape_1d(ctx0, spk80, 80)); + ggml_build_forward_expand(gf, spk80); + return gf; +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 67b10b2a4610..4080885ddd96 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -151,6 +151,24 @@ struct clip_graph_conformer : clip_graph { ggml_cgraph * build() override; }; +struct clip_graph_chatterbox_spkenc : clip_graph { + clip_graph_chatterbox_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; +}; + +struct clip_graph_chatterbox : clip_graph { + clip_gen_process_type gen_process = CLIP_GEN_PROCESS_CODE_GEN; + int n_tokens = 0; + int n_prompt_mel = 0; + int vocode_n_mel = 0; + int vocode_n_stft = 0; + clip_graph_chatterbox(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int n_tokens, + int n_prompt_mel, int vocode_n_mel, int vocode_n_stft) + : clip_graph(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; +}; + struct clip_graph_granite_speech : clip_graph { clip_graph_granite_speech(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index 2811d24df764..b0be418f87cf 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1,6 +1,7 @@ #include "mtmd-audio.h" #define _USE_MATH_DEFINES // for M_PI +#include #include #include #include @@ -855,6 +856,464 @@ bool mtmd_audio_preprocessor_qwen3tts_spk::preprocess(const float * return true; } +// whisper style log-mel of the chatterbox s3 tokenizer, matching the +// reference log_mel_spectrogram (s3tokenizer.py): torch.stft with a periodic +// hann 400 window, hop 160, centered frames reflect-padded at the edges, the +// last frame dropped, power spectrum against the librosa mel filters shipped +// in the gguf, then the whisper normalization. +bool mtmd_audio_s3tok_log_mel(const float * samples, size_t n_samples, + const float * filters, int n_mel, + std::vector & out, int & n_frames) { + const int n_fft = 400; + const int hop = 160; + const int n_bins = n_fft / 2 + 1; + const int half = n_fft / 2; + const int n = (int) n_samples; + + n_frames = n / hop; + if (n_frames <= 0) { + return false; + } + + std::vector window(n_fft); + for (int i = 0; i < n_fft; i++) { + window[(size_t) i] = 0.5 * (1.0 - cos(2.0 * M_PI * i / n_fft)); + } + + std::vector mel((size_t) n_mel * n_frames, 0.0); + std::vector frame(n_fft); + std::vector power(n_bins); + for (int fr = 0; fr < n_frames; fr++) { + for (int i = 0; i < n_fft; i++) { + int idx = fr * hop - half + i; + if (idx < 0) { + idx = -idx; + } + if (idx >= n) { + idx = 2 * n - 2 - idx; + } + frame[(size_t) i] = samples[idx] * window[(size_t) i]; + } + + for (int k = 0; k < n_bins; k++) { + double re = 0.0, im = 0.0; + for (int i = 0; i < n_fft; i++) { + const double a = 2.0 * M_PI * k * i / n_fft; + re += frame[(size_t) i] * cos(a); + im -= frame[(size_t) i] * sin(a); + } + power[(size_t) k] = re * re + im * im; + } + + for (int m = 0; m < n_mel; m++) { + double e = 0.0; + const float * w = filters + (size_t) m * n_bins; + for (int k = 0; k < n_bins; k++) { + e += w[k] * power[(size_t) k]; + } + mel[(size_t) m * n_frames + fr] = log10(std::max(e, 1e-10)); + } + } + + double mx = mel[0]; + for (const double v : mel) { + mx = std::max(mx, v); + } + out.resize(mel.size()); + for (size_t i = 0; i < mel.size(); i++) { + out[i] = (float) ((std::max(mel[i], mx - 8.0) + 4.0) / 4.0); + } + return true; +} + +// rational 3/2 upsampler: every output sample sits at source position +// 2 n / 3, interpolated by a hann windowed sinc cut just under the source +// nyquist. edges are zero extended. +void mtmd_audio_upsample_3_2(const float * samples, size_t n_samples, std::vector & out) { + const int W = 16; // sinc half width in source samples + const double fc = 0.495; // cutoff, normalized to the source rate + + const size_t n_out = n_samples * 3 / 2; + out.assign(n_out, 0.0f); + for (size_t n = 0; n < n_out; n++) { + const double t = (double) (2 * n) / 3.0; + const int k0 = (int) floor(t) - W + 1; + double acc = 0.0; + for (int k = k0; k < k0 + 2 * W; k++) { + if (k < 0 || k >= (int) n_samples) { + continue; + } + const double x = t - k; + const double s = x == 0.0 ? 1.0 : sin(2.0 * M_PI * fc * x) / (M_PI * x); + acc += samples[k] * s * 0.5 * (1.0 + cos(M_PI * x / (W + 1))); + } + out[n] = (float) acc; + } +} + +// matcha style log-mel of the s3gen prompt features, matching the reference +// mel_spectrogram (s3gen/utils/mel.py): (n_fft - hop) / 2 reflect padding, +// torch.stft center false, hann 1920 periodic, magnitude spectrum, slaney +// mel fmin 0 fmax 8000, natural log clamped to 1e-5. +bool mtmd_audio_matcha_log_mel(const float * samples, size_t n_samples, + std::vector & out, int & n_frames) { + const int n_fft = 1920; + const int hop = 480; + const int n_bins = n_fft / 2 + 1; + const int pad = (n_fft - hop) / 2; + const int n_mel = 80; + const int n = (int) n_samples; + + n_frames = 1 + (n + 2 * pad - n_fft) / hop; + if (n <= pad || n_frames <= 0) { + return false; + } + + mtmd_audio_cache cache; + cache.fill_sin_cos_table(n_fft); + cache.fill_hann_window(n_fft, true); + cache.fill_mel_filterbank_matrix(n_mel, n_fft, 24000, 0.0f, 8000.0f); + + std::vector fft_in((size_t) n_fft * 2, 0.0f); + std::vector fft_out((size_t) n_fft * 8); + std::vector mag(n_bins); + out.resize((size_t) n_mel * n_frames); + for (int fr = 0; fr < n_frames; fr++) { + for (int i = 0; i < n_fft; i++) { + int idx = fr * hop - pad + i; + if (idx < 0) { + idx = -idx; + } + if (idx >= n) { + idx = 2 * n - 2 - idx; + } + fft_in[(size_t) i] = samples[idx] * cache.hann_window[(size_t) i]; + } + fft(cache, fft_in.data(), n_fft, fft_out.data()); + + for (int k = 0; k < n_bins; k++) { + const float re = fft_out[2 * k + 0]; + const float im = fft_out[2 * k + 1]; + mag[(size_t) k] = sqrtf(re * re + im * im + 1e-9f); + } + for (int m = 0; m < n_mel; m++) { + float e = 0.0f; + const float * w = cache.filters.data.data() + (size_t) m * n_bins; + for (int k = 0; k < n_bins; k++) { + e += w[k] * mag[(size_t) k]; + } + out[(size_t) fr * n_mel + m] = logf(std::max(e, 1e-5f)); + } + } + return true; +} + +// librosa.effects.trim replica, rms windows 2048/512 centered with zero +// padding, non-silent where the window sits less than top_db under the peak +void mtmd_audio_trim_silence(const float * samples, size_t n_samples, float top_db, + size_t & start, size_t & end) { + const int win = 2048; + const int hop = 512; + const int n = (int) n_samples; + + const int n_fr = 1 + n / hop; + std::vector rms((size_t) n_fr); + double mx = 0.0; + for (int fr = 0; fr < n_fr; fr++) { + double acc = 0.0; + for (int i = 0; i < win; i++) { + const int idx = fr * hop - win / 2 + i; + if (idx >= 0 && idx < n) { + acc += (double) samples[idx] * samples[idx]; + } + } + rms[(size_t) fr] = sqrt(acc / win); + mx = std::max(mx, rms[(size_t) fr]); + } + + const double thr = mx * pow(10.0, -top_db / 20.0); + int first = -1, last = -1; + for (int fr = 0; fr < n_fr; fr++) { + if (rms[(size_t) fr] > thr) { + if (first < 0) { + first = fr; + } + last = fr; + } + } + if (first < 0) { + start = end = 0; + return; + } + start = (size_t) first * hop; + end = std::min((size_t) (last + 1) * hop, n_samples); +} + +// power mel of the voice encoder front-end: centered reflect padded frames, +// hann 400 periodic, hop 160, squared magnitude, slaney mel 40 bins, no log +bool mtmd_audio_ve_mel(const float * samples, size_t n_samples, + std::vector & out, int & n_frames) { + const int n_fft = 400; + const int hop = 160; + const int n_bins = n_fft / 2 + 1; + const int half = n_fft / 2; + const int n_mel = 40; + const int n = (int) n_samples; + + n_frames = 1 + n / hop; + if (n < 2) { + return false; + } + + mtmd_audio_cache cache; + cache.fill_mel_filterbank_matrix(n_mel, n_fft, 16000, 0.0f, 8000.0f); + + std::vector window(n_fft); + for (int i = 0; i < n_fft; i++) { + window[(size_t) i] = 0.5 * (1.0 - cos(2.0 * M_PI * i / n_fft)); + } + + std::vector frame(n_fft); + std::vector power(n_bins); + out.resize((size_t) n_mel * n_frames); + for (int fr = 0; fr < n_frames; fr++) { + for (int i = 0; i < n_fft; i++) { + int idx = fr * hop - half + i; + if (idx < 0) { + idx = -idx; + } + if (idx >= n) { + idx = 2 * n - 2 - idx; + } + frame[(size_t) i] = samples[idx] * window[(size_t) i]; + } + for (int k = 0; k < n_bins; k++) { + double re = 0.0, im = 0.0; + for (int i = 0; i < n_fft; i++) { + const double a = 2.0 * M_PI * k * i / n_fft; + re += frame[(size_t) i] * cos(a); + im -= frame[(size_t) i] * sin(a); + } + power[(size_t) k] = re * re + im * im; + } + for (int m = 0; m < n_mel; m++) { + double e = 0.0; + const float * w = cache.filters.data.data() + (size_t) m * n_bins; + for (int k = 0; k < n_bins; k++) { + e += w[k] * power[(size_t) k]; + } + out[(size_t) fr * n_mel + m] = (float) e; + } + } + return true; +} + +// ITU-R BS.1770 integrated loudness of a mono signal, matching pyloudnorm: +// K-weighting as two RBJ biquads, 400 ms blocks with 75% overlap, absolute +// -70 LUFS gate then a relative -10 LU gate +float mtmd_audio_lufs(const float * samples, size_t n_samples, int sample_rate) { + std::vector y(samples, samples + n_samples); + + auto biquad = [&](double b0, double b1, double b2, double a1, double a2) { + double x1 = 0.0, x2 = 0.0, y1 = 0.0, y2 = 0.0; + for (double & v : y) { + const double x0 = v; + v = b0 * x0 + b1 * x1 + b2 * x2 - a1 * y1 - a2 * y2; + x2 = x1; x1 = x0; + y2 = y1; y1 = v; + } + }; + + // stage 1: high shelf, f0 1681.9744509555319 Hz, +3.99984385397 dB, Q 0.7071752369554196 + { + const double A = pow(10.0, 3.99984385397 / 40.0); + const double w0 = 2.0 * M_PI * 1681.9744509555319 / sample_rate; + const double alpha = sin(w0) / (2.0 * 0.7071752369554196); + const double c = cos(w0); + const double sq = 2.0 * sqrt(A) * alpha; + const double b0 = A * ((A + 1.0) + (A - 1.0) * c + sq); + const double b1 = -2.0 * A * ((A - 1.0) + (A + 1.0) * c); + const double b2 = A * ((A + 1.0) + (A - 1.0) * c - sq); + const double a0 = (A + 1.0) - (A - 1.0) * c + sq; + const double a1 = 2.0 * ((A - 1.0) - (A + 1.0) * c); + const double a2 = (A + 1.0) - (A - 1.0) * c - sq; + biquad(b0 / a0, b1 / a0, b2 / a0, a1 / a0, a2 / a0); + } + // stage 2: high pass, f0 38.13547087602444 Hz, Q 0.5003270373238773 + { + const double w0 = 2.0 * M_PI * 38.13547087602444 / sample_rate; + const double alpha = sin(w0) / (2.0 * 0.5003270373238773); + const double c = cos(w0); + const double b0 = (1.0 + c) / 2.0; + const double b1 = -(1.0 + c); + const double b2 = (1.0 + c) / 2.0; + const double a0 = 1.0 + alpha; + const double a1 = -2.0 * c; + const double a2 = 1.0 - alpha; + biquad(b0 / a0, b1 / a0, b2 / a0, a1 / a0, a2 / a0); + } + + const int win = (int) (0.4 * sample_rate); + const int hop = win / 4; + if ((int) n_samples < win) { + return -HUGE_VALF; + } + const int n_blocks = 1 + ((int) n_samples - win) / hop; + std::vector z((size_t) n_blocks); + for (int b = 0; b < n_blocks; b++) { + double acc = 0.0; + for (int i = 0; i < win; i++) { + acc += y[(size_t) b * hop + i] * y[(size_t) b * hop + i]; + } + z[(size_t) b] = acc / win; + } + + auto gated_mean = [&](double thr_lufs) { + double acc = 0.0; + int n = 0; + for (const double v : z) { + if (-0.691 + 10.0 * log10(std::max(v, 1e-30)) > thr_lufs) { + acc += v; + n++; + } + } + return n > 0 ? acc / n : 0.0; + }; + + const double z_abs = gated_mean(-70.0); + if (z_abs <= 0.0) { + return -HUGE_VALF; + } + const double thr_rel = -0.691 + 10.0 * log10(z_abs) - 10.0; + const double z_rel = gated_mean(thr_rel); + if (z_rel <= 0.0) { + return -HUGE_VALF; + } + return (float) (-0.691 + 10.0 * log10(z_rel)); +} + +// +// mtmd_audio_preprocessor_chatterbox_spk +// +// 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_spk::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_spk::preprocess(const float * samples, + size_t n_samples, + std::vector & output) { + 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) n_samples < frame_len) { + return false; + } + const int n_frames = 1 + ((int) n_samples - frame_len) / hop; + + GGML_ASSERT(!window.empty()); + GGML_ASSERT(!filters.empty()); + + mtmd_audio_mel out; + out.n_len = n_frames; + out.n_len_org = n_frames; + out.n_mel = n_mel; + out.data.assign((size_t) n_mel * n_frames, 0.0f); + + std::vector frame(n_fft); + std::vector power(n_bins); + for (int fr = 0; fr < n_frames; fr++) { + const float * x = samples + (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)); + } + } + + // reference extract_feature subtracts the per-channel mean over time + 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)); + return true; +} + // // mtmd_audio_preprocessor_conformer // diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index b4d6f7259808..91f11704fe0c 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -50,6 +50,47 @@ struct mtmd_audio_cache { ); }; +// whisper style log-mel used by the chatterbox s3 tokenizer front-end: +// hann 400 periodic, hop 160, centered frames with reflect padding, power +// spectrum, caller-supplied mel filters [n_mel x (n_fft / 2 + 1)], log10 +// clamped to 1e-10, global max - 8 dynamic range, (x + 4) / 4 scaling. +// output layout matches the audio batch entries: out[m * n_frames + t]. +bool mtmd_audio_s3tok_log_mel(const float * samples, size_t n_samples, + const float * filters, int n_mel, + std::vector & out, int & n_frames); + +// rational 3/2 upsampler (16 kHz -> 24 kHz), windowed sinc polyphase. +// output length is exactly n_samples * 3 / 2 for even n_samples. +void mtmd_audio_upsample_3_2(const float * samples, size_t n_samples, std::vector & out); + +// matcha style log-mel of the chatterbox s3gen prompt features (utils/mel.py): +// 24 kHz, n_fft 1920, hop 480, hann 1920 periodic, (n_fft - hop) / 2 reflect +// padding with center false, magnitude spectrum, slaney mel 80 bins fmin 0 +// fmax 8000, natural log clamped to 1e-5. +// output layout is frame major: out[t * n_mel + m], as the prompt features +// are consumed row by row at mel rate. +bool mtmd_audio_matcha_log_mel(const float * samples, size_t n_samples, + std::vector & out, int & n_frames); + +// librosa.effects.trim replica: rms over centered 2048/512 windows with zero +// padding, threshold top_db below the loudest window, returns the sample +// span [start, end) of the non-silent region (start == end when all silent). +void mtmd_audio_trim_silence(const float * samples, size_t n_samples, float top_db, + size_t & start, size_t & end); + +// power mel of the chatterbox voice encoder (voice_encoder/melspec.py): +// 16 kHz, hann 400 periodic, hop 160, centered frames with reflect padding, +// squared magnitude against slaney mel 40 bins fmin 0 fmax 8000, no log. +// output layout is frame major: out[t * 40 + m]. +bool mtmd_audio_ve_mel(const float * samples, size_t n_samples, + std::vector & out, int & n_frames); + +// ITU-R BS.1770 integrated loudness (pyloudnorm replica, mono): K-weighting +// (RBJ high shelf 1681.97 Hz +4 dB then high pass 38.14 Hz), 400 ms blocks +// with 75% overlap, absolute -70 then relative -10 gating. +// returns the loudness in LUFS, or -HUGE_VALF when everything is gated out. +float mtmd_audio_lufs(const float * samples, size_t n_samples, int sample_rate); + struct mtmd_audio_preprocessor { const clip_hparams & hparams; @@ -129,6 +170,16 @@ struct mtmd_audio_preprocessor_qwen3tts_spk : mtmd_audio_preprocessor { mtmd_audio_cache cache; }; +struct mtmd_audio_preprocessor_chatterbox_spk : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_chatterbox_spk(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} + void initialize() override; + bool preprocess(const float * samples, size_t n_samples, std::vector & output) override; + + private: + std::vector window; // povey window, frame_length points + std::vector filters; // kaldi mel filterbank, n_mel x (n_fft / 2) dense +}; + struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor { mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { } void initialize() override; diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 3b7762c2145f..6d24505332ed 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -413,10 +413,432 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::vector out_buf; }; + +// Chatterbox: single-track discrete AR (the backbone emits the s3 speech tokens +// directly) into a one-shot flow-matching mel decode and NSF-iSTFT vocoder. +// The prompt is pure embedding concat [spkr, cond speech, text, speech bos], +// positions are handled by the backbone (gpt2 wpe / llama learned pos). +class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { +public: + using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline; + + void reset() override { + pos = 0; + ar_idx = 0; + codes_buf.clear(); + audio_pcm.clear(); + h_state_buf.clear(); + out_buf.clear(); + } + + int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override { + reset(); + + if (!ensure_cache()) { + return 1; + } + + if (inp->speaker_ref) { + // the turbo reference is loudness-normalized before any use + // (speaker encoder included), and its talker conditioning prompt + // is capped at 15 s / 375 tokens where the multilingual variant + // uses 6 s / 150 + mtmd_gen_audio_norm_ref(mctx, inp->speaker_ref); + const size_t t3_cap = (size_t) (t3_cond.empty() ? 15 : 6) * 16000; + + if (!encode_speaker(inp->speaker_ref, spk80)) { + return 1; + } + if (!tokenize_ref(inp->speaker_ref, 10 * 16000, ref_prompt_tokens) || + !tokenize_ref(inp->speaker_ref, t3_cap, ref_t3_tokens)) { + return 1; + } + + // the tts stage derives the mel-rate prompt features from the + // same capped reference clip + { + const float * pcm = (const float *) mtmd_bitmap_get_data(inp->speaker_ref); + const size_t n = mtmd_bitmap_get_n_bytes(inp->speaker_ref) / sizeof(float); + ref_pcm16.assign(pcm, pcm + std::min(n, (size_t) 10 * 16000)); + + // talker conditioning rows from the voice encoder chain; the + // multilingual perceiver consumes the embedding rows of the + // reference speech tokens, built from the fused talker vocab + std::vector pse; + if (!t3_cond.empty()) { + for (size_t i = 0; i < ref_t3_tokens.size(); i++) { + std::vector r(tok_embd.begin() + (size_t) (speech_base + ref_t3_tokens[i]) * n_embd, + tok_embd.begin() + (size_t) (speech_base + ref_t3_tokens[i] + 1) * n_embd); + const float * p = speech_pos.data() + i * (size_t) n_embd; + for (int j = 0; j < n_embd; j++) { + r[(size_t) j] += p[j]; + } + pse.insert(pse.end(), r.begin(), r.end()); + } + } + mtmd_gen_inp gi{}; + gi.type = MTMD_GEN_PROCESS_TYPE_SPEAKER_COND; + gi.pcm = pcm; + gi.n_pcm = n; + gi.ref_speech_embd = pse.empty() ? nullptr : pse.data(); + gi.n_ref_speech_rows = pse.size() / (size_t) n_embd; + mtmd_gen_out go{}; + if (mtmd_gen_audio_process(mctx, &gi, &go) != 0) { + LOG_ERR("mtmd_helper_gen_audio: speaker conditioning failed\n"); + return 1; + } + ref_cond.assign(go.embd, go.embd + go.n_embd); + } + if (getenv("CBX_DUMP")) { + FILE * f = fopen("/tmp/cbx-spk80.bin", "wb"); + fwrite(spk80.data(), sizeof(float), spk80.size(), f); + fclose(f); + f = fopen("/tmp/cbx-s3tok-gen.bin", "wb"); + fwrite(ref_prompt_tokens.data(), sizeof(int32_t), ref_prompt_tokens.size(), f); + fclose(f); + f = fopen("/tmp/cbx-s3tok-t3.bin", "wb"); + fwrite(ref_t3_tokens.data(), sizeof(int32_t), ref_t3_tokens.size(), f); + fclose(f); + f = fopen("/tmp/cbx-ref-cond.bin", "wb"); + fwrite(ref_cond.data(), sizeof(float), ref_cond.size(), f); + fclose(f); + f = fopen("/tmp/cbx-ref-pcm16.bin", "wb"); + fwrite(mtmd_bitmap_get_data(inp->speaker_ref), 1, mtmd_bitmap_get_n_bytes(inp->speaker_ref), f); + fclose(f); + } + } + + const int n_e = n_embd; + auto row = [&](llama_token t) { + return std::vector(tok_embd.begin() + (size_t) t * n_e, + tok_embd.begin() + (size_t) (t + 1) * n_e); + }; + auto add_pos = [&](std::vector & r, const std::vector & tab, int idx) { + const float * p = tab.data() + (size_t) idx * n_e; + for (int j = 0; j < n_e; j++) { + r[(size_t) j] += p[j]; + } + }; + const bool mtl = !t3_cond.empty(); + + std::vector> prompt; + + if (mtl) { + // conditioning: [spkr, perceiver, emotion] block, cloned from the + // reference clip when present, precomputed default otherwise + const auto & cond = ref_cond.empty() ? t3_cond : ref_cond; + for (size_t i = 0; i < cond.size() / (size_t) n_e; i++) { + prompt.emplace_back(cond.begin() + i * (size_t) n_e, cond.begin() + (i + 1) * (size_t) n_e); + } + } else { + // conditioning: projected speaker row, then the speech token prompt + // (already fused ids after the text vocab) + prompt.push_back(ref_cond.empty() ? cond_spkr : ref_cond); + if (ref_cond.empty()) { + for (float f : cond_speech_tokens) { + prompt.push_back(row(speech_base + (llama_token) f)); + } + } else { + for (int32_t t : ref_t3_tokens) { + prompt.push_back(row(speech_base + t)); + } + } + } + + // text, wrapped in the start/stop text tokens of the reference config. + // the multilingual tokenizer expects lowercased text with spaces + // rewritten as the [SPACE] token, language tag left to the caller + std::string txt(inp->prompt, inp->prompt_len); + if (mtl) { + std::string norm; + for (char c : txt) { + if (c == ' ') { + norm += "[SPACE]"; + } else if (c >= 'A' && c <= 'Z') { + norm += (char) (c - 'A' + 'a'); + } else { + norm += c; + } + } + txt = norm; + } + std::vector ids(txt.size() + 16); + int n_ids = llama_tokenize(vocab, txt.c_str(), (int32_t) txt.size(), ids.data(), (int32_t) ids.size(), + false, true); + if (n_ids < 1) { + LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n"); + return 1; + } + ids.resize((size_t) n_ids); + ids.insert(ids.begin(), text_start); + ids.push_back(text_stop); + for (size_t i = 0; i < ids.size(); i++) { + prompt.push_back(row(ids[i])); + if (mtl) { + add_pos(prompt.back(), text_pos, (int) i); + } + } + + // speech bos opens the AR stream + prompt.push_back(row(llama_vocab_bos(vocab))); + if (mtl) { + add_pos(prompt.back(), speech_pos, 0); + } + ar_idx = 1; + + const int n_prompt = (int) prompt.size(); + std::vector embd_buf((size_t) n_prompt * (size_t) n_e); + for (int i = 0; i < n_prompt; i++) { + memcpy(embd_buf.data() + (size_t) i * n_e, prompt[(size_t) i].data(), (size_t) n_e * sizeof(float)); + } + + decode_embd_batch batch_embd(embd_buf.data(), n_prompt, 1, n_e); + batch_embd.set_position_normal(0, 0); + batch_embd.batch.logits[n_prompt - 1] = 1; + + if (getenv("CBX_DUMP")) { + FILE * f = fopen("/tmp/cbx-prompt-embd.bin", "wb"); + fwrite(embd_buf.data(), sizeof(float), embd_buf.size(), f); + fclose(f); + } + + if (llama_decode(lctx, batch_embd.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: prefill decode failed\n"); + return 1; + } + + pos = n_prompt; + out_type = inp->out_type; + return 0; + } + + int32_t step(llama_token sampled, const float * h_state_in, const float ** h_state_out) override { + GGML_UNUSED(h_state_in); + + if (sampled >= speech_base && sampled < speech_base + n_speech) { + codes_buf.push_back(sampled - speech_base); + } + + if (!t3_cond.empty()) { + // multilingual backbone reads embeddings with the learned speech + // position added on top of the token row + std::vector e(tok_embd.begin() + (size_t) sampled * n_embd, + tok_embd.begin() + (size_t) (sampled + 1) * n_embd); + const float * p = speech_pos.data() + (size_t) ar_idx * n_embd; + for (int j = 0; j < n_embd; j++) { + e[(size_t) j] += p[j]; + } + ar_idx++; + decode_embd_batch batch_embd(e.data(), 1, 1, n_embd); + batch_embd.set_position_normal(pos, 0); + batch_embd.batch.logits[0] = 1; + if (llama_decode(lctx, batch_embd.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: step decode failed\n"); + return 1; + } + } else { + llama_batch batch = llama_batch_get_one(&sampled, 1); + if (llama_decode(lctx, batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: step decode failed\n"); + return 1; + } + } + pos++; + + const float * h = llama_get_embeddings_ith(lctx, -1); + h_state_buf.assign(h, h + n_embd); + *h_state_out = h_state_buf.data(); + return 0; + } + + int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override { + if (codes_buf.empty()) { + LOG_ERR("mtmd_helper_gen_audio: no speech tokens generated\n"); + return 1; + } + + mtmd_gen_inp gen_inp{}; + gen_inp.type = MTMD_GEN_PROCESS_TYPE_TTS; + gen_inp.codes = codes_buf.data(); + gen_inp.n_codes = codes_buf.size(); + if (!spk80.empty()) { + gen_inp.ref_spk = spk80.data(); + gen_inp.ref_tokens = ref_prompt_tokens.data(); + gen_inp.n_ref_tokens = ref_prompt_tokens.size(); + gen_inp.ref_pcm = ref_pcm16.data(); + gen_inp.n_ref_pcm = ref_pcm16.size(); + } + mtmd_gen_out gen_out{}; + if (mtmd_gen_audio_process(mctx, &gen_inp, &gen_out) != 0) { + LOG_ERR("mtmd_helper_gen_audio: tts decode failed\n"); + return 1; + } + audio_pcm.assign(gen_out.audio, gen_out.audio + gen_out.n_samples); + + *out_sample_rate = info.sample_rate; + *out_n_samples = (int64_t) audio_pcm.size(); + out_buf.clear(); + if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV) { + write_wav16(out_buf, audio_pcm, info.sample_rate); + } else { + out_buf.resize(audio_pcm.size() * sizeof(float)); + memcpy(out_buf.data(), audio_pcm.data(), out_buf.size()); + } + *out_data = out_buf.data(); + *out_data_len = out_buf.size(); + return 0; + } + +private: + bool ensure_cache() { + if (!tok_embd.empty()) { + return true; + } + // fused vocab layout: [text 0..speech_base) then the speech tokens + speech_base = find_special_token(vocab, "<|speech_0|>"); + if (speech_base == LLAMA_TOKEN_NULL) { + LOG_ERR("mtmd_helper_gen_audio: fused speech tokens not found in vocab\n"); + return false; + } + n_speech = llama_vocab_n_tokens(vocab) - speech_base; + + // reference turbo config: start_text_token = 255, stop_text_token = 0 + text_start = 255; + text_stop = 0; + + const uint32_t n_tok_embd = llama_model_get_tok_embd(model, nullptr); + if (n_tok_embd != (uint32_t) llama_vocab_n_tokens(vocab) * (uint32_t) n_embd) { + LOG_ERR("mtmd_helper_gen_audio: unexpected token embedding size\n"); + return false; + } + tok_embd.resize(n_tok_embd); + llama_model_get_tok_embd(model, tok_embd.data()); + + // multilingual variant: the mmproj ships a precomputed t3 conditioning + // block [spkr, perceiver, emotion] and the learned positional tables + // that the backbone needs added to its input embeddings + size_t n_t3 = mtmd_gen_audio_read_tensor(mctx, "cond.t3_cond", nullptr, 0); + if (n_t3 > 0) { + t3_cond.resize(n_t3); + if (mtmd_gen_audio_read_tensor(mctx, "cond.t3_cond", t3_cond.data(), n_t3) != n_t3 || + n_t3 % (size_t) n_embd != 0) { + LOG_ERR("mtmd_helper_gen_audio: cond.t3_cond read failed\n"); + return false; + } + auto read_table = [&](const char * name, std::vector & dst) { + size_t n = mtmd_gen_audio_read_tensor(mctx, name, nullptr, 0); + dst.resize(n); + if (n == 0 || mtmd_gen_audio_read_tensor(mctx, name, dst.data(), n) != n || + n % (size_t) n_embd != 0) { + LOG_ERR("mtmd_helper_gen_audio: %s read failed\n", name); + return false; + } + return true; + }; + if (!read_table("t3.text_pos_emb", text_pos) || !read_table("t3.speech_pos_emb", speech_pos)) { + return false; + } + return true; + } + + cond_spkr.resize((size_t) n_embd); + if (mtmd_gen_audio_read_tensor(mctx, "cond.spkr_default", cond_spkr.data(), cond_spkr.size()) != (size_t) n_embd) { + LOG_ERR("mtmd_helper_gen_audio: cond.spkr_default missing\n"); + return false; + } + size_t n_ct = mtmd_gen_audio_read_tensor(mctx, "cond.prompt_speech_tokens", nullptr, 0); + cond_speech_tokens.resize(n_ct); + if (n_ct == 0 || mtmd_gen_audio_read_tensor(mctx, "cond.prompt_speech_tokens", cond_speech_tokens.data(), n_ct) != n_ct) { + LOG_ERR("mtmd_helper_gen_audio: cond.prompt_speech_tokens missing\n"); + return false; + } + return true; + } + + // runs the s3 tokenizer on the reference clip, capped to the reference + // conditioning length, into speech tokens for the flow prompt (10 s cap) + // and the talker conditioning (6 s cap) + bool tokenize_ref(mtmd_bitmap * bitmap, size_t n_cap, std::vector & out) { + const float * pcm = (const float *) mtmd_bitmap_get_data(bitmap); + const size_t n = mtmd_bitmap_get_n_bytes(bitmap) / sizeof(float); + + mtmd_gen_inp gi{}; + gi.type = MTMD_GEN_PROCESS_TYPE_TOKENIZE; + gi.pcm = pcm; + gi.n_pcm = std::min(n, n_cap); + mtmd_gen_out go{}; + if (mtmd_gen_audio_process(mctx, &gi, &go) != 0) { + LOG_ERR("mtmd_helper_gen_audio: reference tokenize failed\n"); + return false; + } + out.assign(go.codes, go.codes + go.n_codes); + return true; + } + + // runs the speaker encoder on the reference clip through the standard + // audio chunk path; the CAMPPlus graph outputs the 80-dim s3gen vector + bool encode_speaker(mtmd_bitmap * bitmap, std::vector & out) { + if (!mtmd_support_audio(mctx)) { + LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n"); + return false; + } + const std::string marker = mtmd_default_marker(); + mtmd_input_text text{ marker.c_str(), marker.size(), false, true }; + mtmd_input_chunks * chunks = mtmd_input_chunks_init(); + const mtmd_bitmap * bptr = bitmap; + bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0; + if (ok) { + ok = false; + for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) { + const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i); + if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) { + continue; + } + if (mtmd_encode_chunk(mctx, chunk) != 0) { + LOG_ERR("mtmd_helper_gen_audio: speaker encode failed\n"); + break; + } + const float * embd = mtmd_get_output_embd(mctx); + out.assign(embd, embd + 80); + ok = true; + break; + } + } + mtmd_input_chunks_free(chunks); + return ok; + } + + std::vector tok_embd; + std::vector cond_spkr; + std::vector cond_speech_tokens; + std::vector t3_cond; + std::vector text_pos; + std::vector speech_pos; + std::vector spk80; + std::vector ref_prompt_tokens; + std::vector ref_t3_tokens; + std::vector ref_pcm16; + std::vector ref_cond; + llama_token speech_base = LLAMA_TOKEN_NULL; + int n_speech = 0; + llama_token text_start = LLAMA_TOKEN_NULL; + llama_token text_stop = LLAMA_TOKEN_NULL; + int ar_idx = 0; + + llama_pos pos = 0; + std::vector codes_buf; + std::vector audio_pcm; + std::vector h_state_buf; + std::vector out_buf; + mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; +}; + static std::unique_ptr make_pipeline(llama_context * lctx, mtmd_context * mctx) { switch (mtmd_gen_audio_get_info(mctx).type) { case MTMD_GEN_AUDIO_TYPE_QWEN3TTS: return std::unique_ptr(new qwen3tts_gen_audio_pipeline(lctx, mctx)); + case MTMD_GEN_AUDIO_TYPE_CHATTERBOX: + return std::unique_ptr(new chatterbox_gen_audio_pipeline(lctx, mctx)); default: return nullptr; } diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 361af6dfb51c..8adee93a0478 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -61,6 +61,10 @@ struct mtmd_bitmap { return data; } + std::vector & get_rw_buf() { + return data; + } + bool is_placeholder() const { return data.empty(); } @@ -362,7 +366,7 @@ struct mtmd_context { ctx_v = res.ctx_v; ctx_a = res.ctx_a; ctx_gen_a = res.ctx_gen_a; - if (!ctx_v && !ctx_a) { + if (!ctx_v && !ctx_a && !ctx_gen_a) { throw std::runtime_error(string_format("Failed to load CLIP model from %s\n", mmproj_fname)); } @@ -379,12 +383,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)); + // gen-only mmproj has no input projection, and a speaker encoder + // attached to a gen model outputs a conditioning vector, not tokens + // in the backbone embedding space + if (ctx_v || (ctx_a && !ctx_gen_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) { int n_embd_gen = clip_n_mmproj_embd(ctx_gen_a); @@ -761,6 +770,10 @@ struct mtmd_context { { audio_preproc = std::make_unique(ctx_a); } break; + case PROJECTOR_TYPE_CHATTERBOX_SPKENC: + { + audio_preproc = std::make_unique(ctx_a); + } break; default: throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj)); } @@ -1590,6 +1603,10 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) { info.type = MTMD_GEN_AUDIO_TYPE_QWEN3TTS; info.sample_rate = 24000; break; + case PROJECTOR_TYPE_CHATTERBOX: + info.type = MTMD_GEN_AUDIO_TYPE_CHATTERBOX; + info.sample_rate = 24000; + break; default: info.type = MTMD_GEN_AUDIO_TYPE_NONE; break; @@ -1597,6 +1614,38 @@ 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); +} + +void mtmd_gen_audio_norm_ref(mtmd_context * ctx, mtmd_bitmap * bitmap) { + if (!ctx->ctx_gen_a || !bitmap || !bitmap->is_audio || bitmap->is_placeholder()) { + return; + } + // only the turbo variant normalizes the reference (tts_turbo.py, -27 LUFS); + // the multilingual variant is identified by its perceiver + if (clip_cbx_read_tensor(ctx->ctx_gen_a, "cenc.perceiver.pre_attention_query", nullptr, 0) > 0 || + clip_cbx_read_tensor(ctx->ctx_gen_a, "cenc.spkr_enc.weight", nullptr, 0) == 0) { + return; + } + float * pcm = (float *) bitmap->get_rw_buf().data(); + const size_t n = bitmap->n_bytes() / sizeof(float); + const float lufs = mtmd_audio_lufs(pcm, n, clip_get_hparams(ctx->ctx_a ? ctx->ctx_a : ctx->ctx_gen_a)->audio_sample_rate); + if (lufs == -HUGE_VALF) { + return; + } + const float gain = powf(10.0f, (-27.0f - lufs) / 20.0f); + if (!std::isfinite(gain) || gain <= 0.0f) { + return; + } + for (size_t i = 0; i < n; i++) { + pcm[i] *= gain; + } +} + static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) { clip_ctx * ctx_clip = ctx->ctx_gen_a; if (!ctx_clip) { @@ -1642,6 +1691,395 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in return 0; } + if (inp->type == MTMD_GEN_PROCESS_TYPE_TTS) { + if (!inp->codes || inp->n_codes == 0) { + LOG_ERR("%s: codes required for tts\n", __func__); + return 1; + } + std::vector in_codes(inp->codes, inp->codes + inp->n_codes); + std::vector out_audio; + + std::vector ref_tokens; + std::vector ref_feat; + std::vector ref_spk; + if (inp->ref_tokens) { + ref_tokens.assign(inp->ref_tokens, inp->ref_tokens + inp->n_ref_tokens); + } + if (inp->ref_spk) { + ref_spk.assign(inp->ref_spk, inp->ref_spk + 80); + } + if (inp->ref_pcm) { + // mel-rate prompt features of the reference clip at the 24 kHz + // s3gen rate, padded to the token grid so that the mel length + // stays twice the token length + std::vector pcm16(inp->ref_pcm, inp->ref_pcm + inp->n_ref_pcm); + pcm16.resize((pcm16.size() + 639) / 640 * 640, 0.0f); + std::vector pcm24; + mtmd_audio_upsample_3_2(pcm16.data(), pcm16.size(), pcm24); + int n_feat = 0; + if (!mtmd_audio_matcha_log_mel(pcm24.data(), pcm24.size(), ref_feat, n_feat)) { + LOG_ERR("%s: reference mel failed\n", __func__); + return 1; + } + if ((size_t) n_feat != 2 * ref_tokens.size()) { + LOG_ERR("%s: reference mel length %d does not match %zu tokens\n", + __func__, n_feat, ref_tokens.size()); + return 1; + } + if (getenv("CBX_DUMP")) { + FILE * f = fopen("/tmp/cbx-ref24k.bin", "wb"); + fwrite(pcm24.data(), sizeof(float), pcm24.size(), f); + fclose(f); + f = fopen("/tmp/cbx-ref-feat.bin", "wb"); + fwrite(ref_feat.data(), sizeof(float), ref_feat.size(), f); + fclose(f); + } + } + + // the batch entry is unused, present to satisfy the encode interface + clip_image_f32 dummy; + dummy.set_size({1, 1}, false, true); + dummy.cpy_buf(std::vector(1, 0.0f)); + clip_image_f32_batch batch; + batch.is_audio = true; + batch.entries.push_back(std::move(dummy)); + + clip_encode_params params; + params.imgs = &batch; + params.n_threads = ctx->n_threads; + params.gen_process = CLIP_GEN_PROCESS_TTS; + params.codes = &in_codes; + params.out_audio = &out_audio; + params.ref_tokens = inp->ref_tokens ? &ref_tokens : nullptr; + params.ref_feat = inp->ref_pcm ? &ref_feat : nullptr; + params.ref_spk = inp->ref_spk ? &ref_spk : nullptr; + + if (!clip_encode(ctx_clip, ¶ms)) { + LOG_ERR("%s: clip_encode failed (tts)\n", __func__); + return 1; + } + + ctx->gen_out_audio = std::move(out_audio); + out->audio = ctx->gen_out_audio.data(); + out->n_samples = ctx->gen_out_audio.size(); + return 0; + } + + if (inp->type == MTMD_GEN_PROCESS_TYPE_TOKENIZE) { + if (!inp->pcm || inp->n_pcm == 0) { + LOG_ERR("%s: pcm required for tokenize\n", __func__); + return 1; + } + + // mel filters shipped in the mmproj, [n_mels x (n_fft / 2 + 1)] + const size_t n_filt = clip_cbx_read_tensor(ctx_clip, "s3tok.mel_filters", nullptr, 0); + if (n_filt == 0) { + LOG_ERR("%s: model has no s3 tokenizer\n", __func__); + return 1; + } + std::vector filters(n_filt); + clip_cbx_read_tensor(ctx_clip, "s3tok.mel_filters", filters.data(), n_filt); + const int n_mel = (int) (n_filt / (400 / 2 + 1)); + + // pad to a whole number of 40 ms tokens so that the mel length stays + // twice the token length, as the reference prompt features expect + std::vector pcm(inp->pcm, inp->pcm + inp->n_pcm); + pcm.resize((pcm.size() + 639) / 640 * 640, 0.0f); + + std::vector mel; + int n_frames = 0; + if (!mtmd_audio_s3tok_log_mel(pcm.data(), pcm.size(), filters.data(), n_mel, mel, n_frames)) { + LOG_ERR("%s: log mel failed\n", __func__); + return 1; + } + + clip_image_f32 mel_img; + mel_img.set_size({n_frames, n_mel}, false, true); + mel_img.cpy_buf(std::move(mel)); + + clip_image_f32_batch batch; + batch.is_audio = true; + batch.entries.push_back(std::move(mel_img)); + + std::vector out_codes; + + 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; + + if (!clip_encode(ctx_clip, ¶ms)) { + LOG_ERR("%s: clip_encode failed (tokenize)\n", __func__); + return 1; + } + + ctx->gen_out_codes = std::move(out_codes); + out->codes = ctx->gen_out_codes.data(); + out->n_codes = ctx->gen_out_codes.size(); + return 0; + } + + if (inp->type == MTMD_GEN_PROCESS_TYPE_SPEAKER_COND) { + if (!inp->pcm || inp->n_pcm == 0) { + LOG_ERR("%s: pcm required for speaker cond\n", __func__); + return 1; + } + + auto read_t = [&](const char * name, std::vector & v) -> bool { + const size_t n = clip_cbx_read_tensor(ctx_clip, name, nullptr, 0); + if (n == 0) { + return false; + } + v.resize(n); + return clip_cbx_read_tensor(ctx_clip, name, v.data(), n) == n; + }; + + // voice encoder reference chain (embeds_from_wavs): silence trim, + // 40-bin power mel, overlapping 160-frame partials at rate 1.3, + // 3-layer lstm per partial, projected/relu/normalized embeddings + // averaged into the utterance embedding + size_t t0 = 0, t1 = 0; + mtmd_audio_trim_silence(inp->pcm, inp->n_pcm, 20.0f, t0, t1); + if (t1 <= t0) { + LOG_ERR("%s: reference clip is silent\n", __func__); + return 1; + } + std::vector mel; + int n_frames = 0; + if (!mtmd_audio_ve_mel(inp->pcm + t0, t1 - t0, mel, n_frames)) { + LOG_ERR("%s: voice encoder mel failed\n", __func__); + return 1; + } + + 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 + 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); // zero pad (or trim) to the partial grid + + std::vector w_ih[3], w_hh[3], b_ih[3], b_hh[3]; + std::vector w_proj, b_proj; + for (int l = 0; l < 3; l++) { + const std::string s = std::to_string(l); + if (!read_t(("ve.lstm.weight_ih_l" + s).c_str(), w_ih[l]) || + !read_t(("ve.lstm.weight_hh_l" + s).c_str(), w_hh[l]) || + !read_t(("ve.lstm.bias_ih_l" + s).c_str(), b_ih[l]) || + !read_t(("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("ve.proj.weight", w_proj) || !read_t("ve.proj.bias", b_proj)) { + LOG_ERR("%s: model has no voice encoder projection\n", __func__); + return 1; + } + + const int n_h = 256; + std::vector ve(n_h, 0.0f); + std::vector h((size_t) 3 * n_h), c((size_t) 3 * n_h), x(n_h), g((size_t) 4 * n_h); + for (int p = 0; p < n_wins; p++) { + std::fill(h.begin(), h.end(), 0.0f); + std::fill(c.begin(), c.end(), 0.0f); + for (int t = 0; t < n_partial; t++) { + const float * in = mel.data() + (size_t) (p * step + t) * n_mel; + int n_in = n_mel; + for (int l = 0; l < 3; l++) { + float * hl = h.data() + (size_t) l * n_h; + float * cl = c.data() + (size_t) l * n_h; + for (int j = 0; j < 4 * n_h; j++) { + double acc = b_ih[l][(size_t) j] + b_hh[l][(size_t) j]; + const float * wi = w_ih[l].data() + (size_t) j * n_in; + for (int i = 0; i < n_in; i++) { + acc += (double) wi[i] * in[i]; + } + const float * wh = w_hh[l].data() + (size_t) j * n_h; + for (int i = 0; i < n_h; i++) { + acc += (double) wh[i] * hl[i]; + } + g[(size_t) j] = (float) acc; + } + // torch gate order: input, forget, cell, output + for (int i = 0; i < n_h; i++) { + const float gi = 1.0f / (1.0f + expf(-g[(size_t) i])); + const float gf = 1.0f / (1.0f + expf(-g[(size_t) i + n_h])); + const float gc = tanhf(g[(size_t) i + 2 * n_h]); + const float go = 1.0f / (1.0f + expf(-g[(size_t) i + 3 * n_h])); + cl[i] = gf * cl[i] + gi * gc; + x[(size_t) i] = go * tanhf(cl[i]); + } + memcpy(hl, x.data(), (size_t) n_h * sizeof(float)); + in = hl; + n_in = n_h; + } + } + // projected, relu'd, normalized partial embedding + std::vector e(n_h); + double norm = 0.0; + for (int o = 0; o < n_h; o++) { + double acc = b_proj[(size_t) o]; + const float * w = w_proj.data() + (size_t) o * n_h; + const float * hl = h.data() + (size_t) 2 * n_h; + for (int i = 0; i < n_h; i++) { + acc += (double) w[i] * hl[i]; + } + e[(size_t) o] = (float) std::max(acc, 0.0); + norm += (double) e[(size_t) o] * e[(size_t) o]; + } + norm = sqrt(norm); + for (int o = 0; o < n_h; o++) { + ve[(size_t) o] += (float) (e[(size_t) o] / norm); + } + } + double norm = 0.0; + for (float v : ve) { + norm += (double) v * v; + } + norm = sqrt(norm); + for (float & v : ve) { + v = (float) (v / norm); + } + + // speaker projection row + std::vector w_spkr, b_spkr; + if (!read_t("cenc.spkr_enc.weight", w_spkr) || !read_t("cenc.spkr_enc.bias", b_spkr)) { + LOG_ERR("%s: model has no speaker conditioning projection\n", __func__); + return 1; + } + const int n_e = (int) b_spkr.size(); + std::vector rows((size_t) n_e); + for (int o = 0; o < n_e; o++) { + double acc = b_spkr[(size_t) o]; + const float * w = w_spkr.data() + (size_t) o * n_h; + for (int i = 0; i < n_h; i++) { + acc += (double) w[i] * ve[(size_t) i]; + } + rows[(size_t) o] = (float) acc; + } + + // multilingual variant: [spkr, perceiver x32, emotion] block, the + // perceiver runs its shared attention block as cross then self + // attention over the reference speech embedding rows + std::vector query; + if (read_t("cenc.perceiver.pre_attention_query", query)) { + if (!inp->ref_speech_embd || inp->n_ref_speech_rows == 0) { + LOG_ERR("%s: reference speech embeddings required for the perceiver\n", __func__); + return 1; + } + std::vector ln_w, ln_b, wq, bq, wk, bk, wv, bv, wo, bo, emo; + if (!read_t("cenc.perceiver.attn.norm.weight", ln_w) || !read_t("cenc.perceiver.attn.norm.bias", ln_b) || + !read_t("cenc.perceiver.attn.to_q.weight", wq) || !read_t("cenc.perceiver.attn.to_q.bias", bq) || + !read_t("cenc.perceiver.attn.to_k.weight", wk) || !read_t("cenc.perceiver.attn.to_k.bias", bk) || + !read_t("cenc.perceiver.attn.to_v.weight", wv) || !read_t("cenc.perceiver.attn.to_v.bias", bv) || + !read_t("cenc.perceiver.attn.proj_out.weight", wo) || !read_t("cenc.perceiver.attn.proj_out.bias", bo) || + !read_t("cenc.emotion_adv_fc.weight", emo)) { + LOG_ERR("%s: model has an incomplete perceiver\n", __func__); + return 1; + } + const int n_head = 4; + const int d_head = n_e / n_head; + + auto layer_norm = [&](const float * in, float * out) { + double mean = 0.0, var = 0.0; + for (int i = 0; i < n_e; i++) { + mean += in[i]; + } + mean /= n_e; + for (int i = 0; i < n_e; i++) { + var += (in[i] - mean) * (in[i] - mean); + } + const double sd = sqrt(var / n_e + 1e-5); + for (int i = 0; i < n_e; i++) { + out[i] = (float) ((in[i] - mean) / sd * ln_w[(size_t) i] + ln_b[(size_t) i]); + } + }; + auto linear = [&](const std::vector & w, const std::vector & b, + const std::vector & in, int n_rows, std::vector & out) { + out.resize((size_t) n_rows * n_e); + for (int r = 0; r < n_rows; r++) { + for (int o = 0; o < n_e; o++) { + double acc = b[(size_t) o]; + const float * wr = w.data() + (size_t) o * n_e; + const float * ir = in.data() + (size_t) r * n_e; + for (int i = 0; i < n_e; i++) { + acc += (double) wr[i] * ir[i]; + } + out[(size_t) r * n_e + o] = (float) acc; + } + } + }; + auto attn_block = [&](const std::vector & x1, int n1, + const std::vector & x2, int n2, std::vector & out) { + std::vector nx1((size_t) n1 * n_e), nx2((size_t) n2 * n_e); + for (int r = 0; r < n1; r++) { + layer_norm(x1.data() + (size_t) r * n_e, nx1.data() + (size_t) r * n_e); + } + for (int r = 0; r < n2; r++) { + layer_norm(x2.data() + (size_t) r * n_e, nx2.data() + (size_t) r * n_e); + } + std::vector q, k, v; + linear(wq, bq, nx1, n1, q); + linear(wk, bk, nx2, n2, k); + linear(wv, bv, nx2, n2, v); + + std::vector ctxt((size_t) n1 * n_e); + std::vector sc((size_t) n2); + for (int hd = 0; hd < n_head; hd++) { + const int off = hd * d_head; + for (int t = 0; t < n1; t++) { + double mx = -1e30; + for (int s = 0; s < n2; s++) { + double acc = 0.0; + for (int i = 0; i < d_head; i++) { + acc += (double) q[(size_t) t * n_e + off + i] * k[(size_t) s * n_e + off + i]; + } + sc[(size_t) s] = acc / sqrt((double) d_head); + mx = std::max(mx, sc[(size_t) s]); + } + double sum = 0.0; + for (int s = 0; s < n2; s++) { + sc[(size_t) s] = exp(sc[(size_t) s] - mx); + sum += sc[(size_t) s]; + } + for (int i = 0; i < d_head; i++) { + double acc = 0.0; + for (int s = 0; s < n2; s++) { + acc += sc[(size_t) s] * v[(size_t) s * n_e + off + i]; + } + ctxt[(size_t) t * n_e + off + i] = (float) (acc / sum); + } + } + } + linear(wo, bo, ctxt, n1, out); + for (size_t i = 0; i < out.size(); i++) { + out[i] += x1[i]; + } + }; + + const int n_q = (int) (query.size() / n_e); + std::vector x2(inp->ref_speech_embd, inp->ref_speech_embd + (size_t) inp->n_ref_speech_rows * n_e); + std::vector pre, p32; + attn_block(query, n_q, x2, (int) inp->n_ref_speech_rows, 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); + } + } + + ctx->gen_out_embd = std::move(rows); + out->embd = ctx->gen_out_embd.data(); + out->n_embd = ctx->gen_out_embd.size(); + return 0; + } + // MTMD_GEN_PROCESS_TYPE_CODE2WAV if (!inp->codes || inp->n_codes == 0) { LOG_ERR("%s: codes required for code2wav\n", __func__); diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index f7d1fc65b79c..008548bebba6 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -334,6 +334,7 @@ MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname); enum mtmd_gen_audio_type { MTMD_GEN_AUDIO_TYPE_NONE, // not supported MTMD_GEN_AUDIO_TYPE_QWEN3TTS, + MTMD_GEN_AUDIO_TYPE_CHATTERBOX, }; struct mtmd_gen_audio_info { enum mtmd_gen_audio_type type; @@ -341,9 +342,20 @@ 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); + +// normalize a reference clip in place to the loudness the gen audio model +// expects for voice cloning; no-op when the model does not require it. +MTMD_API void mtmd_gen_audio_norm_ref(mtmd_context * ctx, mtmd_bitmap * bitmap); + enum mtmd_gen_process_type { MTMD_GEN_PROCESS_TYPE_GEN_CODE, // h_state to codes MTMD_GEN_PROCESS_TYPE_CODE2WAV, // codes to raw PCM audio + MTMD_GEN_PROCESS_TYPE_TTS, // full utterance of codes to raw PCM audio + MTMD_GEN_PROCESS_TYPE_TOKENIZE, // raw PCM audio to semantic speech tokens + MTMD_GEN_PROCESS_TYPE_SPEAKER_COND, // raw PCM audio to talker conditioning rows }; struct mtmd_gen_inp { enum mtmd_gen_process_type type; @@ -359,6 +371,27 @@ struct mtmd_gen_inp { size_t n_codes; const char * state_data; size_t state_size; + + // for MTMD_GEN_PROCESS_TYPE_TOKENIZE and MTMD_GEN_PROCESS_TYPE_SPEAKER_COND + const float * pcm; // mono float samples at the audio encoder sample rate + size_t n_pcm; + + // for MTMD_GEN_PROCESS_TYPE_SPEAKER_COND: the speech token embedding rows + // of the reference conditioning prompt (n_text_embd elements each), input + // to the perceiver of the multilingual variant (null on turbo) + const float * ref_speech_embd; + size_t n_ref_speech_rows; + + // for MTMD_GEN_PROCESS_TYPE_TTS: optional reference conditioning of the + // cloned voice, overriding the model's precomputed defaults (null means + // default). ref_spk is the 80-dim speaker vector, ref_tokens the speech + // tokens of the reference clip (from a TOKENIZE call), ref_pcm the same + // clip as mono float samples at the audio encoder sample rate. + const float * ref_spk; + const int32_t * ref_tokens; + size_t n_ref_tokens; + const float * ref_pcm; + size_t n_ref_pcm; }; struct mtmd_gen_out { // note: output memory is allocated by the context, valid until next process() call @@ -368,6 +401,9 @@ struct mtmd_gen_out { size_t n_codes; const float * embd; // the generated hidden state, to be fed back to backbone // it must have n_text_embd elements + // for MTMD_GEN_PROCESS_TYPE_SPEAKER_COND: embd holds the conditioning + // rows and n_embd their total element count + size_t n_embd; // for MTMD_GEN_PROCESS_TYPE_CODE2WAV const float * audio; diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index 23e797dc1255..9fec7b9a1c0a 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -1,3 +1,4 @@ +#include #include "arg.h" #include "common.h" #include "sampling.h" @@ -121,6 +122,27 @@ int main(int argc, char ** argv) { for (llama_token t = 0; t < llama_vocab_n_tokens(vocab); t++) { if (!strcmp(llama_vocab_get_text(vocab, t), "<|codec_eos_token|>")) { codec_eos_tok = t; break; } } + if (codec_eos_tok == LLAMA_TOKEN_NULL) { + // models without a dedicated codec eos (e.g. chatterbox) end the audio + // stream with the regular vocab eos + codec_eos_tok = llama_vocab_eos(vocab); + + // fused text+speech vocab: the reference implementation samples the + // speech head only, mask the text zone out of the sampling chain + llama_token speech_base = LLAMA_TOKEN_NULL; + for (llama_token t = 0; t < llama_vocab_n_tokens(vocab); t++) { + if (!strcmp(llama_vocab_get_text(vocab, t), "<|speech_0|>")) { speech_base = t; break; } + } + if (speech_base != LLAMA_TOKEN_NULL) { + params.sampling.logit_bias.reserve(params.sampling.logit_bias.size() + speech_base); + for (llama_token t = 0; t < speech_base; t++) { + params.sampling.logit_bias.push_back(llama_logit_bias{t, -std::numeric_limits::infinity()}); + } + common_sampler_free(smpl); + smpl = common_sampler_init(model, params.sampling); + if (!smpl) { LOG_ERR("failed to reinit sampler\n"); return 1; } + } + } if (codec_eos_tok == LLAMA_TOKEN_NULL) { LOG_ERR("missing codec eos token in vocab\n"); return 1; @@ -134,6 +156,12 @@ int main(int argc, char ** argv) { const int max_new = params.n_predict > 0 ? params.n_predict : 512; int n_frames = 0; + if (getenv("CBX_DUMP")) { + const float * lg = llama_get_logits_ith(lctx, -1); + FILE * f = fopen("/tmp/cbx-logits0.bin", "wb"); + fwrite(lg, sizeof(float), llama_vocab_n_tokens(vocab), f); + fclose(f); + } llama_token sampled = sample_codec0(); const float * h_state = llama_get_embeddings_ith(lctx, -1); From 4e505862e90014d634e1fa9b873e9b507568d5ab Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 2 Aug 2026 13:24:48 +0200 Subject: [PATCH 02/14] mtmd: fix spurious word at the start of chatterbox turbo cloning The turbo reference feeds the raw tokenizer output to the talker; only the multilingual variant wraps the text in the start/stop text tokens. Token 0 is ! in the turbo BPE vocab, so wrapping vocalized a stray exclamation right before the speech BOS. Also align the turbo path with the reference inference chain: - apply punc_norm to the text prompt - drop SoS/EoS/OOV ids from the collected speech codes - append the short silence tail before vocoding - exempt the chatterbox gen mmproj from the n_embd match check (it emits mel channels, not text embeddings) Prompt validated bit-exact against the reference (438/438 rows). --- tools/mtmd/mtmd-helper-gen.cpp | 66 ++++++++++++++++++++++++++++++---- tools/mtmd/mtmd.cpp | 3 +- 2 files changed, 61 insertions(+), 8 deletions(-) diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 6d24505332ed..516eddec58db 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -545,9 +545,11 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } } - // text, wrapped in the start/stop text tokens of the reference config. - // the multilingual tokenizer expects lowercased text with spaces - // rewritten as the [SPACE] token, language tag left to the caller + // text preprocessing per variant. the multilingual tokenizer expects + // lowercased text with spaces rewritten as the [SPACE] token, language + // tag left to the caller; the turbo reference applies punc_norm: + // capitalized first letter, whitespace runs collapsed, uncommon + // punctuation replaced, and a trailing sentence ender enforced std::string txt(inp->prompt, inp->prompt_len); if (mtl) { std::string norm; @@ -561,6 +563,43 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } } txt = norm; + } else if (!txt.empty()) { + if (txt[0] >= 'a' && txt[0] <= 'z') { + txt[0] = (char) (txt[0] - 'a' + 'A'); + } + std::string norm; + for (char c : txt) { + if (c == ' ' || c == '\t' || c == '\n' || c == '\r') { + if (!norm.empty() && norm.back() != ' ') { + norm += ' '; + } + } else { + norm += c; + } + } + auto replace_all = [&norm](const char * from, const char * to) { + const size_t nf = strlen(from); + const size_t nt = strlen(to); + for (size_t p = 0; (p = norm.find(from, p)) != std::string::npos; p += nt) { + norm.replace(p, nf, to); + } + }; + replace_all("\xE2\x80\xA6", ", "); // ellipsis + replace_all(":", ","); + replace_all("\xE2\x80\x94", "-"); // em dash + replace_all("\xE2\x80\x93", "-"); // en dash + replace_all(" ,", ","); + replace_all("\xE2\x80\x9C", "\""); // curly double quotes + replace_all("\xE2\x80\x9D", "\""); + replace_all("\xE2\x80\x98", "'"); // curly single quotes + replace_all("\xE2\x80\x99", "'"); + while (!norm.empty() && norm.back() == ' ') { + norm.pop_back(); + } + if (!norm.empty() && strchr(".!?-,", norm.back()) == nullptr) { + norm += '.'; + } + txt = norm; } std::vector ids(txt.size() + 16); int n_ids = llama_tokenize(vocab, txt.c_str(), (int32_t) txt.size(), ids.data(), (int32_t) ids.size(), @@ -570,8 +609,12 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { return 1; } ids.resize((size_t) n_ids); - ids.insert(ids.begin(), text_start); - ids.push_back(text_stop); + if (mtl) { + // only the multilingual reference wraps the text in the start/stop + // text tokens; the turbo reference feeds the raw tokenizer output + ids.insert(ids.begin(), text_start); + ids.push_back(text_stop); + } for (size_t i = 0; i < ids.size(); i++) { prompt.push_back(row(ids[i])); if (mtl) { @@ -615,7 +658,9 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { int32_t step(llama_token sampled, const float * h_state_in, const float ** h_state_out) override { GGML_UNUSED(h_state_in); - if (sampled >= speech_base && sampled < speech_base + n_speech) { + // the s3gen speech vocab holds 6561 codes; the fused start/stop + // tokens and any ids beyond are dropped like the reference + if (sampled >= speech_base && sampled - speech_base < 6561) { codes_buf.push_back(sampled - speech_base); } @@ -657,6 +702,12 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { return 1; } + if (t3_cond.empty()) { + // the turbo reference appends a short silence tail (3 tokens of + // the s3gen silence code) before vocoding + codes_buf.insert(codes_buf.end(), 3, 4299); + } + mtmd_gen_inp gen_inp{}; gen_inp.type = MTMD_GEN_PROCESS_TYPE_TTS; gen_inp.codes = codes_buf.data(); @@ -702,7 +753,8 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } n_speech = llama_vocab_n_tokens(vocab) - speech_base; - // reference turbo config: start_text_token = 255, stop_text_token = 0 + // reference config: start_text_token = 255, stop_text_token = 0 + // (only the multilingual prompt wraps the text with them) text_start = 255; text_stop = 0; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 8adee93a0478..96fc3bd8db11 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -395,7 +395,8 @@ struct mtmd_context { 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( From c69f5ac91df450c2ba764efa8bcf8e8993d44701 Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 2 Aug 2026 13:33:59 +0200 Subject: [PATCH 03/14] mtmd: remove chatterbox validation instrumentation Every stage of the port was validated against the reference PyTorch implementation by dumping intermediate tensors (mel, f0, source, spec, speech tokens, prompt embeddings, ...) behind a CBX_DUMP env var and measuring cosine similarity against reference dumps computed on identical inputs. All stages now match at F16 noise level, so the dump points are no longer needed: - remove all CBX_DUMP dump sites across clip.cpp, mtmd.cpp, mtmd-helper-gen.cpp and tts.cpp - remove the deterministic source mode (zeroed phases and noise) that made the NSF vocoder reproducible for comparison; the source is now always stochastic like the reference - remove the out_mu and out_dcond graph output markers and the dump-only out_mu readback Also migrate the two deprecated ggml_upscale_ext calls to ggml_interpolate for a warning-free build. --- tools/mtmd/clip.cpp | 71 +------------------------------- tools/mtmd/models/chatterbox.cpp | 11 +---- tools/mtmd/mtmd-helper-gen.cpp | 23 ----------- tools/mtmd/mtmd.cpp | 8 ---- tools/tts/tts.cpp | 6 --- 5 files changed, 4 insertions(+), 115 deletions(-) diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 74ab0985d013..b333509bc825 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -4922,18 +4922,6 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } } set_input_f32("inp_temb", temb); - - if (getenv("CBX_DUMP")) { - FILE * f = fopen("/tmp/cbx-noise.bin", "wb"); - fwrite(noise.data(), sizeof(float), noise.size(), f); - fclose(f); - } - - if (getenv("CBX_DUMP")) { - FILE * f = fopen("/tmp/cbx-tokens.bin", "wb"); - fwrite(tokens.data(), sizeof(int32_t), tokens.size(), f); - fclose(f); - } } break; case PROJECTOR_TYPE_GEMMA3: case PROJECTOR_TYPE_GEMMA3NV: @@ -5506,12 +5494,6 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } (*params->out_codes)[(size_t) t] = code; } - - if (getenv("CBX_DUMP")) { - FILE * f = fopen("/tmp/cbx-fsq.bin", "wb"); - fwrite(h.data(), sizeof(float), h.size(), f); - fclose(f); - } return true; } @@ -5524,37 +5506,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } if (ctx->proj_type() == PROJECTOR_TYPE_CHATTERBOX && params->gen_process == CLIP_GEN_PROCESS_TTS) { - ggml_tensor * mu = ggml_graph_get_tensor(gf, "out_mu"); - GGML_ASSERT(mu != nullptr); - std::vector mu_data(ggml_nelements(mu)); - ggml_backend_tensor_get(mu, mu_data.data(), 0, ggml_nbytes(mu)); - - if (getenv("CBX_DUMP")) { - FILE * f = fopen("/tmp/cbx-mu.bin", "wb"); - fwrite(mu_data.data(), sizeof(float), mu_data.size(), f); - fclose(f); - LOG_INF("%s: dumped out_mu [%d, %d] to /tmp/cbx-mu.bin\n", __func__, (int) mu->ne[0], (int) mu->ne[1]); - } - - if (getenv("CBX_DUMP")) { - ggml_tensor * dc = ggml_graph_get_tensor(gf, "out_dcond"); - std::vector dc_data(ggml_nelements(dc)); - ggml_backend_tensor_get(dc, dc_data.data(), 0, ggml_nbytes(dc)); - FILE * f = fopen("/tmp/cbx-dcond.bin", "wb"); - fwrite(dc_data.data(), sizeof(float), dc_data.size(), f); - fclose(f); - } - ggml_tensor * mel = ggml_graph_get_tensor(gf, "out_mel"); GGML_ASSERT(mel != nullptr); std::vector mel_data(ggml_nelements(mel)); ggml_backend_tensor_get(mel, mel_data.data(), 0, ggml_nbytes(mel)); - if (getenv("CBX_DUMP")) { - FILE * f = fopen("/tmp/cbx-mel.bin", "wb"); - fwrite(mel_data.data(), sizeof(float), mel_data.size(), f); - fclose(f); - LOG_INF("%s: dumped out_mel [%d, %d] to /tmp/cbx-mel.bin\n", __func__, (int) mel->ne[0], (int) mel->ne[1]); - } // hift bridge: f0 -> harmonic source -> source stft on the host, // then the vocoder graph, then the istft @@ -5564,12 +5519,6 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { std::vector f0(n_mel_out); ggml_backend_tensor_get(f0_t, f0.data(), 0, ggml_nbytes(f0_t)); - if (getenv("CBX_DUMP")) { - FILE * f = fopen("/tmp/cbx-f0.bin", "wb"); - fwrite(f0.data(), sizeof(float), f0.size(), f); - fclose(f); - } - // 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; @@ -5580,13 +5529,12 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { GGML_ASSERT(clip_cbx_read_tensor(ctx, "hift.m_source.l_linear.weight", lw.data(), lw.size()) == 9); clip_cbx_read_tensor(ctx, "hift.m_source.l_linear.bias", &lb, 1); - const bool det = getenv("CBX_DUMP") != nullptr; // deterministic phases and no noise for validation std::mt19937 srng(1234); std::uniform_real_distribution ud(-M_PI, M_PI); std::normal_distribution snd(0.0f, 1.0f); double phase[9]; for (int h = 0; h < 9; h++) { - phase[h] = (h == 0 || det) ? 0.0 : ud(srng); + phase[h] = h == 0 ? 0.0 : ud(srng); } std::vector src((size_t) n_wav); double cum[9] = {0.0}; @@ -5599,7 +5547,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { 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 + (det ? 0.0f : namp * snd(srng)); + sine = sine * uv + namp * snd(srng); merged += lw[(size_t) h] * sine; } src[(size_t) t] = tanhf(merged); @@ -5629,15 +5577,6 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } } - if (getenv("CBX_DUMP")) { - FILE * f = fopen("/tmp/cbx-src.bin", "wb"); - fwrite(src.data(), sizeof(float), src.size(), f); - fclose(f); - f = fopen("/tmp/cbx-sstft.bin", "wb"); - fwrite(sstft.data(), sizeof(float), sstft.size(), f); - fclose(f); - } - // vocoder graph on mel + source stft std::vector spec; { @@ -5655,12 +5594,6 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } } - if (getenv("CBX_DUMP")) { - FILE * f = fopen("/tmp/cbx-spec.bin", "wb"); - fwrite(spec.data(), sizeof(float), spec.size(), f); - fclose(f); - } - // 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; diff --git a/tools/mtmd/models/chatterbox.cpp b/tools/mtmd/models/chatterbox.cpp index 31ffe7e312fe..57de6ab45a5a 100644 --- a/tools/mtmd/models/chatterbox.cpp +++ b/tools/mtmd/models/chatterbox.cpp @@ -367,7 +367,7 @@ ggml_cgraph * clip_graph_chatterbox::build() { // upsample x2: nearest repeat, left pad 4, conv k=5 { ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [T1, 512] - xt = ggml_upscale_ext(ctx0, xt, 2 * T1, 512, 1, 1, GGML_SCALE_MODE_NEAREST); + 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); @@ -389,9 +389,6 @@ ggml_cgraph * clip_graph_chatterbox::build() { // encoder projection to the mel channel count ggml_tensor * mu = cbx_linear(ctx0, cbx_t(model, "flow.encoder_proj.weight"), cbx_t(model, "flow.encoder_proj.bias"), x); // [80, T2] - ggml_set_name(mu, "out_mu"); - ggml_set_output(mu); - ggml_build_forward_expand(gf, mu); // 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 @@ -453,10 +450,6 @@ ggml_cgraph * clip_graph_chatterbox::build() { for (int i = 0; i < n_steps; i++) { ggml_tensor * temb = step_temb(i); ggml_tensor * d = cbx_estimator(model, ctx0, mx, mu, spks, cond, temb, T2); - if (i == 0) { - ggml_set_name(d, "out_dcond"); - ggml_set_output(d); - } if (!meanflow) { ggml_tensor * du = cbx_estimator(model, ctx0, 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)); @@ -695,7 +688,7 @@ static ggml_tensor * cbx_cam_layer(const clip_model & model, ggml_context * ctx0 } 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_upscale_ext(ctx0, pooled, S * 100, C, 1, 1, GGML_SCALE_MODE_NEAREST); + 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] } diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 516eddec58db..677c037e1ba9 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -489,23 +489,6 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } ref_cond.assign(go.embd, go.embd + go.n_embd); } - if (getenv("CBX_DUMP")) { - FILE * f = fopen("/tmp/cbx-spk80.bin", "wb"); - fwrite(spk80.data(), sizeof(float), spk80.size(), f); - fclose(f); - f = fopen("/tmp/cbx-s3tok-gen.bin", "wb"); - fwrite(ref_prompt_tokens.data(), sizeof(int32_t), ref_prompt_tokens.size(), f); - fclose(f); - f = fopen("/tmp/cbx-s3tok-t3.bin", "wb"); - fwrite(ref_t3_tokens.data(), sizeof(int32_t), ref_t3_tokens.size(), f); - fclose(f); - f = fopen("/tmp/cbx-ref-cond.bin", "wb"); - fwrite(ref_cond.data(), sizeof(float), ref_cond.size(), f); - fclose(f); - f = fopen("/tmp/cbx-ref-pcm16.bin", "wb"); - fwrite(mtmd_bitmap_get_data(inp->speaker_ref), 1, mtmd_bitmap_get_n_bytes(inp->speaker_ref), f); - fclose(f); - } } const int n_e = n_embd; @@ -639,12 +622,6 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { batch_embd.set_position_normal(0, 0); batch_embd.batch.logits[n_prompt - 1] = 1; - if (getenv("CBX_DUMP")) { - FILE * f = fopen("/tmp/cbx-prompt-embd.bin", "wb"); - fwrite(embd_buf.data(), sizeof(float), embd_buf.size(), f); - fclose(f); - } - if (llama_decode(lctx, batch_embd.batch) != 0) { LOG_ERR("mtmd_helper_gen_audio: prefill decode failed\n"); return 1; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 96fc3bd8db11..a63fdc14d2c4 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1727,14 +1727,6 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in __func__, n_feat, ref_tokens.size()); return 1; } - if (getenv("CBX_DUMP")) { - FILE * f = fopen("/tmp/cbx-ref24k.bin", "wb"); - fwrite(pcm24.data(), sizeof(float), pcm24.size(), f); - fclose(f); - f = fopen("/tmp/cbx-ref-feat.bin", "wb"); - fwrite(ref_feat.data(), sizeof(float), ref_feat.size(), f); - fclose(f); - } } // the batch entry is unused, present to satisfy the encode interface diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index 9fec7b9a1c0a..44e281331989 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -156,12 +156,6 @@ int main(int argc, char ** argv) { const int max_new = params.n_predict > 0 ? params.n_predict : 512; int n_frames = 0; - if (getenv("CBX_DUMP")) { - const float * lg = llama_get_logits_ith(lctx, -1); - FILE * f = fopen("/tmp/cbx-logits0.bin", "wb"); - fwrite(lg, sizeof(float), llama_vocab_n_tokens(vocab), f); - fclose(f); - } llama_token sampled = sample_codec0(); const float * h_state = llama_get_embeddings_ith(lctx, -1); From 346d25f2f4f008d88ddf65ebf6a08604f9aa8d7b Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 2 Aug 2026 13:47:27 +0200 Subject: [PATCH 04/14] shorten overly long LLM comments --- tools/mtmd/mtmd-audio.cpp | 15 +++------------ tools/mtmd/mtmd-helper-gen.cpp | 22 +++++++--------------- 2 files changed, 10 insertions(+), 27 deletions(-) diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index b0be418f87cf..c98fdaa120f2 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -856,11 +856,7 @@ bool mtmd_audio_preprocessor_qwen3tts_spk::preprocess(const float * return true; } -// whisper style log-mel of the chatterbox s3 tokenizer, matching the -// reference log_mel_spectrogram (s3tokenizer.py): torch.stft with a periodic -// hann 400 window, hop 160, centered frames reflect-padded at the edges, the -// last frame dropped, power spectrum against the librosa mel filters shipped -// in the gguf, then the whisper normalization. +// whisper style log-mel of the chatterbox s3 tokenizer (s3tokenizer.py) bool mtmd_audio_s3tok_log_mel(const float * samples, size_t n_samples, const float * filters, int n_mel, std::vector & out, int & n_frames) { @@ -951,10 +947,7 @@ void mtmd_audio_upsample_3_2(const float * samples, size_t n_samples, std::vecto } } -// matcha style log-mel of the s3gen prompt features, matching the reference -// mel_spectrogram (s3gen/utils/mel.py): (n_fft - hop) / 2 reflect padding, -// torch.stft center false, hann 1920 periodic, magnitude spectrum, slaney -// mel fmin 0 fmax 8000, natural log clamped to 1e-5. +// matcha style log-mel of the s3gen prompt features (s3gen/utils/mel.py) bool mtmd_audio_matcha_log_mel(const float * samples, size_t n_samples, std::vector & out, int & n_frames) { const int n_fft = 1920; @@ -1108,9 +1101,7 @@ bool mtmd_audio_ve_mel(const float * samples, size_t n_samples, return true; } -// ITU-R BS.1770 integrated loudness of a mono signal, matching pyloudnorm: -// K-weighting as two RBJ biquads, 400 ms blocks with 75% overlap, absolute -// -70 LUFS gate then a relative -10 LU gate +// ITU-R BS.1770 integrated loudness of a mono signal, matching pyloudnorm float mtmd_audio_lufs(const float * samples, size_t n_samples, int sample_rate) { std::vector y(samples, samples + n_samples); diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 677c037e1ba9..34acb652acb4 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -439,10 +439,8 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } if (inp->speaker_ref) { - // the turbo reference is loudness-normalized before any use - // (speaker encoder included), and its talker conditioning prompt - // is capped at 15 s / 375 tokens where the multilingual variant - // uses 6 s / 150 + // turbo reference chain: loudness normalize the clip, then cap + // the talker conditioning at 15 s (multilingual: 6 s) mtmd_gen_audio_norm_ref(mctx, inp->speaker_ref); const size_t t3_cap = (size_t) (t3_cond.empty() ? 15 : 6) * 16000; @@ -528,11 +526,8 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } } - // text preprocessing per variant. the multilingual tokenizer expects - // lowercased text with spaces rewritten as the [SPACE] token, language - // tag left to the caller; the turbo reference applies punc_norm: - // capitalized first letter, whitespace runs collapsed, uncommon - // punctuation replaced, and a trailing sentence ender enforced + // text preprocessing per variant: multilingual lowercases with [SPACE] + // tokens (language tag left to the caller), turbo applies punc_norm std::string txt(inp->prompt, inp->prompt_len); if (mtl) { std::string norm; @@ -593,8 +588,7 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } ids.resize((size_t) n_ids); if (mtl) { - // only the multilingual reference wraps the text in the start/stop - // text tokens; the turbo reference feeds the raw tokenizer output + // only the multilingual prompt wraps the text in start/stop tokens ids.insert(ids.begin(), text_start); ids.push_back(text_stop); } @@ -635,8 +629,7 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { int32_t step(llama_token sampled, const float * h_state_in, const float ** h_state_out) override { GGML_UNUSED(h_state_in); - // the s3gen speech vocab holds 6561 codes; the fused start/stop - // tokens and any ids beyond are dropped like the reference + // 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); } @@ -680,8 +673,7 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } if (t3_cond.empty()) { - // the turbo reference appends a short silence tail (3 tokens of - // the s3gen silence code) before vocoding + // turbo appends a short silence tail before vocoding codes_buf.insert(codes_buf.end(), 3, 4299); } From 069ef544f3aa3cb7eb9d70aa1c03bbaefee34c08 Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 2 Aug 2026 13:55:36 +0200 Subject: [PATCH 05/14] mtmd: split chatterbox model into gen and spkenc files Mirror the qwen3tts layout: chatterbox-gen.cpp holds the generation graphs (flow encoder, cfm estimator, s3 tokenizer, hift vocoder) and chatterbox-spkenc.cpp holds the CAMPPlus speaker encoder. The helpers used by both graphs are declared in models.h and defined in the gen file. Also add cb() callbacks at the stage milestones of both graphs (embedding, encoder layers, upsample, mu, solver steps, speaker encoder blocks and pooling) so debug output embeddings work like the other models. --- tools/mtmd/CMakeLists.txt | 3 +- .../{chatterbox.cpp => chatterbox-gen.cpp} | 201 ++---------------- tools/mtmd/models/chatterbox-spkenc.cpp | 187 ++++++++++++++++ tools/mtmd/models/models.h | 9 + 4 files changed, 212 insertions(+), 188 deletions(-) rename tools/mtmd/models/{chatterbox.cpp => chatterbox-gen.cpp} (73%) create mode 100644 tools/mtmd/models/chatterbox-spkenc.cpp diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 0b35a226efd9..516a9b47898b 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -27,7 +27,8 @@ add_library(mtmd clip-model.h clip-graph.h models/models.h - models/chatterbox.cpp + models/chatterbox-gen.cpp + models/chatterbox-spkenc.cpp models/cogvlm.cpp models/conformer.cpp models/dotsocr.cpp diff --git a/tools/mtmd/models/chatterbox.cpp b/tools/mtmd/models/chatterbox-gen.cpp similarity index 73% rename from tools/mtmd/models/chatterbox.cpp rename to tools/mtmd/models/chatterbox-gen.cpp index 57de6ab45a5a..f788d188e46a 100644 --- a/tools/mtmd/models/chatterbox.cpp +++ b/tools/mtmd/models/chatterbox-gen.cpp @@ -1,10 +1,10 @@ #include "models.h" -// Chatterbox s3gen, stage 1: speech tokens -> mu (flow encoder output). -// Weights come from the source-named tensor map (model.cbx_tensors), the -// estimator and hift stages extend this file. +// Chatterbox generation graphs: flow encoder, cfm estimator, s3 tokenizer +// and hift vocoder. Weights come from the source-named tensor map +// (model.cbx_tensors); the speaker encoder lives in chatterbox-spkenc.cpp. -static ggml_tensor * cbx_t(const clip_model & model, const std::string & name) { +ggml_tensor * cbx_t(const clip_model & model, const std::string & name) { auto it = model.cbx_tensors.find(name); if (it == model.cbx_tensors.end()) { GGML_ABORT("missing chatterbox tensor: %s", name.c_str()); @@ -13,7 +13,7 @@ static ggml_tensor * cbx_t(const clip_model & model, const std::string & name) { } // x [C, T]: y = W x + b with torch Linear weights stored as [in, out] -static ggml_tensor * cbx_linear(ggml_context * ctx0, ggml_tensor * w, ggml_tensor * b, ggml_tensor * x) { +ggml_tensor * cbx_linear(ggml_context * ctx0, ggml_tensor * w, ggml_tensor * b, ggml_tensor * x) { ggml_tensor * y = ggml_mul_mat(ctx0, w, x); if (b) { y = ggml_add(ctx0, y, b); @@ -30,7 +30,7 @@ static ggml_tensor * cbx_layer_norm(ggml_context * ctx0, ggml_tensor * w, ggml_t // 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 -static ggml_tensor * cbx_conv1d(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * b, ggml_tensor * x, +ggml_tensor * cbx_conv1d(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * b, ggml_tensor * x, int stride, int pad_l, int pad_r) { ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [T, C] if (pad_l > 0) { @@ -348,6 +348,7 @@ ggml_cgraph * clip_graph_chatterbox::build() { x = cbx_linear(ctx0, cbx_t(model, "fenc.embed.out.0.weight"), cbx_t(model, "fenc.embed.out.0.bias"), x); x = cbx_layer_norm(ctx0, cbx_t(model, "fenc.embed.out.1.weight"), cbx_t(model, "fenc.embed.out.1.bias"), 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 { @@ -358,10 +359,12 @@ ggml_cgraph * clip_graph_chatterbox::build() { cur = cbx_conv1d(ctx0, cbx_t(model, "fenc.pre_lookahead_layer.conv2.weight"), cbx_t(model, "fenc.pre_lookahead_layer.conv2.bias"), cur, 1, 2, 0); x = ggml_add(ctx0, res, cur); + cb(x, "fenc_pre_lookahead", -1); } for (int i = 0; model.cbx_tensors.count("fenc.encoders." + std::to_string(i) + ".norm_mha.weight"); i++) { x = cbx_enc_layer(model, ctx0, x, pos1, "fenc.encoders." + std::to_string(i), T1); + cb(x, "fenc_enc", i); } // upsample x2: nearest repeat, left pad 4, conv k=5 @@ -374,6 +377,7 @@ ggml_cgraph * clip_graph_chatterbox::build() { ggml_tensor * y = ggml_conv_1d(ctx0, cbx_t(model, "fenc.up_layer.conv.weight"), xt, 1, 0, 1); x = ggml_cont(ctx0, ggml_transpose(ctx0, y)); // [512, T2] x = ggml_add(ctx0, x, cbx_t(model, "fenc.up_layer.conv.bias")); + cb(x, "fenc_upsample", -1); } // up embed: linear + layer norm + xscale @@ -383,12 +387,14 @@ ggml_cgraph * clip_graph_chatterbox::build() { for (int i = 0; model.cbx_tensors.count("fenc.up_encoders." + std::to_string(i) + ".norm_mha.weight"); i++) { x = cbx_enc_layer(model, ctx0, x, pos2, "fenc.up_encoders." + std::to_string(i), T2); + cb(x, "fenc_up_enc", i); } x = cbx_layer_norm(ctx0, cbx_t(model, "fenc.after_norm.weight"), cbx_t(model, "fenc.after_norm.bias"), x, 1e-5f); // encoder projection to the mel channel count ggml_tensor * mu = cbx_linear(ctx0, cbx_t(model, "flow.encoder_proj.weight"), cbx_t(model, "flow.encoder_proj.bias"), 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 @@ -455,6 +461,7 @@ ggml_cgraph * clip_graph_chatterbox::build() { 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; @@ -486,7 +493,7 @@ ggml_cgraph * clip_graph_chatterbox::build() { static ggml_tensor * cbx_hift_resblock(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, const std::string & p); // x [C, T] -> symmetric-padded dilated conv -> [OC, T] -static ggml_tensor * cbx_conv1d_dil(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * b, ggml_tensor * x, int pad, int dil) { +ggml_tensor * cbx_conv1d_dil(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * b, ggml_tensor * x, int pad, int dil) { 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)); @@ -595,183 +602,3 @@ static ggml_tensor * cbx_hift_resblock(const clip_model & model, ggml_context * return x; } - - -// Chatterbox speaker encoder: CAMPPlus x-vector on kaldi fbank features, -// projected through the s3gen speaker affine. Mirrors s3gen/xvector.py. - -// per-channel batchnorm on x [C, T]; scale = w / sqrt(var + eps), shift folds -// the running mean. pass null w/b for the affine=False variant -static ggml_tensor * cbx_bn1d(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, - const std::string & p, ggml_tensor * eps) { - ggml_tensor * mean = cbx_t(model, p + ".running_mean"); - ggml_tensor * var = cbx_t(model, p + ".running_var"); - ggml_tensor * sd = ggml_sqrt(ctx0, ggml_add(ctx0, var, eps)); - if (!model.cbx_tensors.count(p + ".weight")) { - return ggml_div(ctx0, ggml_sub(ctx0, x, mean), sd); - } - ggml_tensor * a = ggml_div(ctx0, cbx_t(model, p + ".weight"), sd); - ggml_tensor * shift = ggml_sub(ctx0, cbx_t(model, p + ".bias"), ggml_mul(ctx0, mean, a)); - return ggml_add(ctx0, ggml_mul(ctx0, x, a), shift); -} - -// batchnorm + relu on a conv2d activation [W=T, H=F, C, 1], stats on ne2 -static ggml_tensor * cbx_bn2d_relu(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, - const std::string & p, ggml_tensor * eps) { - const int C = (int) x->ne[2]; - ggml_tensor * mean = ggml_reshape_4d(ctx0, cbx_t(model, p + ".running_mean"), 1, 1, C, 1); - ggml_tensor * var = ggml_reshape_4d(ctx0, cbx_t(model, p + ".running_var"), 1, 1, C, 1); - ggml_tensor * w = ggml_reshape_4d(ctx0, cbx_t(model, p + ".weight"), 1, 1, C, 1); - ggml_tensor * b = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bias"), 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_relu(ctx0, ggml_add(ctx0, ggml_mul(ctx0, x, a), shift)); -} - -// fcm residual 2d block, stride on the frequency axis only -static ggml_tensor * cbx_res2d(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, - const std::string & p, int stride, ggml_tensor * eps) { - ggml_tensor * cur = ggml_conv_2d(ctx0, cbx_t(model, p + ".conv1.weight"), x, 1, stride, 1, 1, 1, 1); - cur = cbx_bn2d_relu(model, ctx0, cur, p + ".bn1", eps); - cur = ggml_conv_2d(ctx0, cbx_t(model, p + ".conv2.weight"), cur, 1, 1, 1, 1, 1, 1); - // bn2 without the relu, applied before the residual add - { - const int C = (int) cur->ne[2]; - ggml_tensor * mean = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.running_mean"), 1, 1, C, 1); - ggml_tensor * var = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.running_var"), 1, 1, C, 1); - ggml_tensor * w = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.weight"), 1, 1, C, 1); - ggml_tensor * b = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.bias"), 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)); - cur = ggml_add(ctx0, ggml_mul(ctx0, cur, a), shift); - } - ggml_tensor * res = x; - if (model.cbx_tensors.count(p + ".shortcut.0.weight")) { - res = ggml_conv_2d(ctx0, cbx_t(model, p + ".shortcut.0.weight"), x, 1, stride, 0, 0, 1, 1); - const int C = (int) res->ne[2]; - ggml_tensor * mean = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.running_mean"), 1, 1, C, 1); - ggml_tensor * var = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.running_var"), 1, 1, C, 1); - ggml_tensor * w = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.weight"), 1, 1, C, 1); - ggml_tensor * b = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.bias"), 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)); - res = ggml_add(ctx0, ggml_mul(ctx0, res, a), shift); - } - return ggml_relu(ctx0, ggml_add(ctx0, cur, res)); -} - -// cam dense tdnn layer: bottleneck then context-gated conv; x [C_in, T] -> [growth, T] -static ggml_tensor * cbx_cam_layer(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, - const std::string & p, int dil, ggml_tensor * eps, ggml_tensor * segfix) { - ggml_tensor * h = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, p + ".nonlinear1.batchnorm", eps)); - h = cbx_conv1d(ctx0, cbx_t(model, p + ".linear1.weight"), nullptr, h, 1, 0, 0); - h = ggml_relu(ctx0, cbx_bn1d(model, ctx0, h, p + ".nonlinear2.batchnorm", eps)); - - const std::string cp = p + ".cam_layer"; - ggml_tensor * k = cbx_t(model, cp + ".linear_local.weight"); - const int pad = ((int) k->ne[0] - 1) / 2 * dil; - ggml_tensor * y = cbx_conv1d_dil(ctx0, 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(ctx0, cbx_t(model, cp + ".linear1.weight"), cbx_t(model, cp + ".linear1.bias"), context, 1, 0, 0); - context = ggml_relu(ctx0, context); - context = cbx_conv1d(ctx0, cbx_t(model, cp + ".linear2.weight"), cbx_t(model, cp + ".linear2.bias"), context, 1, 0, 0); - ggml_tensor * m = ggml_sigmoid(ctx0, context); - return ggml_mul(ctx0, y, m); -} - -ggml_cgraph * clip_graph_chatterbox_spkenc::build() { - 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, cbx_t(model, "spk.head.conv1.weight"), x, 1, 1, 1, 1, 1, 1); - x = cbx_bn2d_relu(model, ctx0, x, "spk.head.bn1", eps); - x = cbx_res2d(model, ctx0, x, "spk.head.layer1.0", 2, eps); - x = cbx_res2d(model, ctx0, x, "spk.head.layer1.1", 1, eps); - x = cbx_res2d(model, ctx0, x, "spk.head.layer2.0", 2, eps); - x = cbx_res2d(model, ctx0, x, "spk.head.layer2.1", 1, eps); - x = ggml_conv_2d(ctx0, cbx_t(model, "spk.head.conv2.weight"), x, 1, 2, 1, 1, 1, 1); - x = cbx_bn2d_relu(model, ctx0, x, "spk.head.bn2", eps); - x = ggml_reshape_2d(ctx0, x, T, 320); - x = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [320, T] - - // tdnn k5 stride 2 over time, then the three cam dense blocks - x = cbx_conv1d(ctx0, cbx_t(model, "spk.xvector.tdnn.linear.weight"), nullptr, x, 2, 2, 2); // [128, T1] - x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, "spk.xvector.tdnn.nonlinear.batchnorm", eps)); - - static const int block_dil[3] = {1, 2, 2}; - for (int bi = 1; bi <= 3; bi++) { - const std::string bp = "spk.xvector.block" + std::to_string(bi); - for (int li = 1; model.cbx_tensors.count(bp + ".tdnnd" + std::to_string(li) + ".linear1.weight"); li++) { - ggml_tensor * out = cbx_cam_layer(model, ctx0, x, bp + ".tdnnd" + std::to_string(li), - block_dil[bi - 1], eps, segfix); - x = ggml_concat(ctx0, x, out, 0); - } - const std::string tp = "spk.xvector.transit" + std::to_string(bi); - x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, tp + ".nonlinear.batchnorm", eps)); - x = cbx_conv1d(ctx0, cbx_t(model, tp + ".linear.weight"), nullptr, x, 1, 0, 0); - } - x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, "spk.xvector.out_nonlinear.batchnorm", 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] - - // dense 1024 -> 192, batchnorm without affine, into the x-vector - ggml_tensor * dw = ggml_reshape_2d(ctx0, cbx_t(model, "spk.xvector.dense.linear.weight"), 1024, 192); - ggml_tensor * emb = ggml_mul_mat(ctx0, dw, stats); // [192, 1] - emb = cbx_bn1d(model, ctx0, emb, "spk.xvector.dense.nonlinear.batchnorm", eps); - emb = ggml_reshape_1d(ctx0, emb, 192); - ggml_set_name(emb, "out_xvec"); - ggml_set_output(emb); - ggml_build_forward_expand(gf, emb); - - // normalize then the s3gen speaker affine - ggml_tensor * n2 = ggml_sqrt(ctx0, ggml_sum(ctx0, ggml_mul(ctx0, emb, emb))); - ggml_tensor * unit = ggml_div(ctx0, emb, n2); - ggml_tensor * spk80 = cbx_linear(ctx0, cbx_t(model, "flow.spk_embed_affine_layer.weight"), - cbx_t(model, "flow.spk_embed_affine_layer.bias"), - ggml_reshape_2d(ctx0, unit, 192, 1)); - spk80 = ggml_cont(ctx0, ggml_reshape_1d(ctx0, spk80, 80)); - ggml_build_forward_expand(gf, spk80); - return gf; -} diff --git a/tools/mtmd/models/chatterbox-spkenc.cpp b/tools/mtmd/models/chatterbox-spkenc.cpp new file mode 100644 index 000000000000..0f63db9c396e --- /dev/null +++ b/tools/mtmd/models/chatterbox-spkenc.cpp @@ -0,0 +1,187 @@ +#include "models.h" + +#include + +// Chatterbox speaker encoder: CAMPPlus x-vector on kaldi fbank features, +// projected through the s3gen speaker affine. Mirrors s3gen/xvector.py. + +// per-channel batchnorm on x [C, T]; scale = w / sqrt(var + eps), shift folds +// the running mean. pass null w/b for the affine=False variant +static ggml_tensor * cbx_bn1d(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, + const std::string & p, ggml_tensor * eps) { + ggml_tensor * mean = cbx_t(model, p + ".running_mean"); + ggml_tensor * var = cbx_t(model, p + ".running_var"); + ggml_tensor * sd = ggml_sqrt(ctx0, ggml_add(ctx0, var, eps)); + if (!model.cbx_tensors.count(p + ".weight")) { + return ggml_div(ctx0, ggml_sub(ctx0, x, mean), sd); + } + ggml_tensor * a = ggml_div(ctx0, cbx_t(model, p + ".weight"), sd); + ggml_tensor * shift = ggml_sub(ctx0, cbx_t(model, p + ".bias"), ggml_mul(ctx0, mean, a)); + return ggml_add(ctx0, ggml_mul(ctx0, x, a), shift); +} + +// batchnorm + relu on a conv2d activation [W=T, H=F, C, 1], stats on ne2 +static ggml_tensor * cbx_bn2d_relu(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, + const std::string & p, ggml_tensor * eps) { + const int C = (int) x->ne[2]; + ggml_tensor * mean = ggml_reshape_4d(ctx0, cbx_t(model, p + ".running_mean"), 1, 1, C, 1); + ggml_tensor * var = ggml_reshape_4d(ctx0, cbx_t(model, p + ".running_var"), 1, 1, C, 1); + ggml_tensor * w = ggml_reshape_4d(ctx0, cbx_t(model, p + ".weight"), 1, 1, C, 1); + ggml_tensor * b = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bias"), 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_relu(ctx0, ggml_add(ctx0, ggml_mul(ctx0, x, a), shift)); +} + +// fcm residual 2d block, stride on the frequency axis only +static ggml_tensor * cbx_res2d(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, + const std::string & p, int stride, ggml_tensor * eps) { + ggml_tensor * cur = ggml_conv_2d(ctx0, cbx_t(model, p + ".conv1.weight"), x, 1, stride, 1, 1, 1, 1); + cur = cbx_bn2d_relu(model, ctx0, cur, p + ".bn1", eps); + cur = ggml_conv_2d(ctx0, cbx_t(model, p + ".conv2.weight"), cur, 1, 1, 1, 1, 1, 1); + // bn2 without the relu, applied before the residual add + { + const int C = (int) cur->ne[2]; + ggml_tensor * mean = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.running_mean"), 1, 1, C, 1); + ggml_tensor * var = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.running_var"), 1, 1, C, 1); + ggml_tensor * w = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.weight"), 1, 1, C, 1); + ggml_tensor * b = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.bias"), 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)); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, a), shift); + } + ggml_tensor * res = x; + if (model.cbx_tensors.count(p + ".shortcut.0.weight")) { + res = ggml_conv_2d(ctx0, cbx_t(model, p + ".shortcut.0.weight"), x, 1, stride, 0, 0, 1, 1); + const int C = (int) res->ne[2]; + ggml_tensor * mean = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.running_mean"), 1, 1, C, 1); + ggml_tensor * var = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.running_var"), 1, 1, C, 1); + ggml_tensor * w = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.weight"), 1, 1, C, 1); + ggml_tensor * b = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.bias"), 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)); + res = ggml_add(ctx0, ggml_mul(ctx0, res, a), shift); + } + return ggml_relu(ctx0, ggml_add(ctx0, cur, res)); +} + +// cam dense tdnn layer: bottleneck then context-gated conv; x [C_in, T] -> [growth, T] +static ggml_tensor * cbx_cam_layer(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, + const std::string & p, int dil, ggml_tensor * eps, ggml_tensor * segfix) { + ggml_tensor * h = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, p + ".nonlinear1.batchnorm", eps)); + h = cbx_conv1d(ctx0, cbx_t(model, p + ".linear1.weight"), nullptr, h, 1, 0, 0); + h = ggml_relu(ctx0, cbx_bn1d(model, ctx0, h, p + ".nonlinear2.batchnorm", eps)); + + const std::string cp = p + ".cam_layer"; + ggml_tensor * k = cbx_t(model, cp + ".linear_local.weight"); + const int pad = ((int) k->ne[0] - 1) / 2 * dil; + ggml_tensor * y = cbx_conv1d_dil(ctx0, 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(ctx0, cbx_t(model, cp + ".linear1.weight"), cbx_t(model, cp + ".linear1.bias"), context, 1, 0, 0); + context = ggml_relu(ctx0, context); + context = cbx_conv1d(ctx0, cbx_t(model, cp + ".linear2.weight"), cbx_t(model, cp + ".linear2.bias"), context, 1, 0, 0); + ggml_tensor * m = ggml_sigmoid(ctx0, context); + return ggml_mul(ctx0, y, m); +} + +ggml_cgraph * clip_graph_chatterbox_spkenc::build() { + 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, cbx_t(model, "spk.head.conv1.weight"), x, 1, 1, 1, 1, 1, 1); + x = cbx_bn2d_relu(model, ctx0, x, "spk.head.bn1", eps); + x = cbx_res2d(model, ctx0, x, "spk.head.layer1.0", 2, eps); + x = cbx_res2d(model, ctx0, x, "spk.head.layer1.1", 1, eps); + x = cbx_res2d(model, ctx0, x, "spk.head.layer2.0", 2, eps); + x = cbx_res2d(model, ctx0, x, "spk.head.layer2.1", 1, eps); + x = ggml_conv_2d(ctx0, cbx_t(model, "spk.head.conv2.weight"), x, 1, 2, 1, 1, 1, 1); + x = cbx_bn2d_relu(model, ctx0, x, "spk.head.bn2", 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(ctx0, cbx_t(model, "spk.xvector.tdnn.linear.weight"), nullptr, x, 2, 2, 2); // [128, T1] + x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, "spk.xvector.tdnn.nonlinear.batchnorm", eps)); + cb(x, "spk_tdnn", -1); + + static const int block_dil[3] = {1, 2, 2}; + for (int bi = 1; bi <= 3; bi++) { + const std::string bp = "spk.xvector.block" + std::to_string(bi); + for (int li = 1; model.cbx_tensors.count(bp + ".tdnnd" + std::to_string(li) + ".linear1.weight"); li++) { + ggml_tensor * out = cbx_cam_layer(model, ctx0, x, bp + ".tdnnd" + std::to_string(li), + block_dil[bi - 1], eps, segfix); + x = ggml_concat(ctx0, x, out, 0); + } + const std::string tp = "spk.xvector.transit" + std::to_string(bi); + x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, tp + ".nonlinear.batchnorm", eps)); + x = cbx_conv1d(ctx0, cbx_t(model, tp + ".linear.weight"), nullptr, x, 1, 0, 0); + cb(x, "spk_block", bi); + } + x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, "spk.xvector.out_nonlinear.batchnorm", 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 x-vector + ggml_tensor * dw = ggml_reshape_2d(ctx0, cbx_t(model, "spk.xvector.dense.linear.weight"), 1024, 192); + ggml_tensor * emb = ggml_mul_mat(ctx0, dw, stats); // [192, 1] + emb = cbx_bn1d(model, ctx0, emb, "spk.xvector.dense.nonlinear.batchnorm", eps); + emb = ggml_reshape_1d(ctx0, emb, 192); + ggml_set_name(emb, "out_xvec"); + ggml_set_output(emb); + ggml_build_forward_expand(gf, emb); + + // normalize then the s3gen speaker affine + ggml_tensor * n2 = ggml_sqrt(ctx0, ggml_sum(ctx0, ggml_mul(ctx0, emb, emb))); + ggml_tensor * unit = ggml_div(ctx0, emb, n2); + ggml_tensor * spk80 = cbx_linear(ctx0, cbx_t(model, "flow.spk_embed_affine_layer.weight"), + cbx_t(model, "flow.spk_embed_affine_layer.bias"), + ggml_reshape_2d(ctx0, unit, 192, 1)); + spk80 = ggml_cont(ctx0, ggml_reshape_1d(ctx0, spk80, 80)); + cb(spk80, "spk_embd", -1); + ggml_build_forward_expand(gf, spk80); + return gf; +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 4080885ddd96..a2934a98e9c6 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -151,6 +151,15 @@ struct clip_graph_conformer : clip_graph { ggml_cgraph * build() override; }; +// chatterbox helpers shared between the gen and spkenc graphs +// (defined in chatterbox-gen.cpp) +ggml_tensor * cbx_t(const clip_model & model, const std::string & name); +ggml_tensor * cbx_linear(ggml_context * ctx0, ggml_tensor * w, ggml_tensor * b, ggml_tensor * x); +ggml_tensor * cbx_conv1d(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * b, ggml_tensor * x, + int stride, int pad_l, int pad_r); +ggml_tensor * cbx_conv1d_dil(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * b, ggml_tensor * x, + int pad, int dil); + struct clip_graph_chatterbox_spkenc : clip_graph { clip_graph_chatterbox_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; From 2b91b88b5811693ed701f19a67bbb57865eccc60 Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 2 Aug 2026 19:44:24 +0200 Subject: [PATCH 06/14] mtmd: address review from @ngxson - in-tree conversion/chatterbox.py: bf16 talkers, a.* / a.gen.* tensor namespace, speech embedding table shipped in the mmproj - tokenize stage returns the embedding rows of its codes, the reference chain consumes them instead of the talker vocab; turbo default conditioning ids resolve through the same table - single SPK_REF stage encodes the reference clip end to end: loudness normalization, speaker vector, speech tokenization, conditioning rows, and an opaque decoder state - TTS folds into CODE2WAV, the reference state is passed back caller-owned - gen API surface reduced: one added enum, pcm-only input, embd and state as the only outputs --- conversion/__init__.py | 2 + conversion/chatterbox.py | 520 +++++++++++++++ gguf-py/gguf/constants.py | 2 + tools/mtmd/clip.cpp | 64 +- tools/mtmd/clip.h | 4 + tools/mtmd/models/chatterbox-gen.cpp | 110 ++-- tools/mtmd/models/chatterbox-spkenc.cpp | 34 +- tools/mtmd/mtmd-helper-gen.cpp | 181 ++---- tools/mtmd/mtmd.cpp | 808 ++++++++++++++---------- tools/mtmd/mtmd.h | 38 +- 10 files changed, 1177 insertions(+), 586 deletions(-) create mode 100644 conversion/chatterbox.py diff --git a/conversion/__init__.py b/conversion/__init__.py index 31da4963cf86..1a5546cc5821 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -37,6 +37,7 @@ "BloomModel": "bloom", "CamembertModel": "bert", "ChameleonForCausalLM": "chameleon", + "ChatterboxModel": "chatterbox", "ChameleonForConditionalGeneration": "chameleon", "ChatGLMForConditionalGeneration": "chatglm", "ChatGLMModel": "chatglm", @@ -260,6 +261,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = { + "ChatterboxModel": "chatterbox", "AudioFlamingo3ForConditionalGeneration": "ultravox", "CogVLMForCausalLM": "cogvlm", "DeepseekOCR2ForCausalLM": "deepseek", diff --git a/conversion/chatterbox.py b/conversion/chatterbox.py new file mode 100644 index 000000000000..f6e8d580660f --- /dev/null +++ b/conversion/chatterbox.py @@ -0,0 +1,520 @@ +# Chatterbox (ResembleAI) conversion: turbo and multilingual v3 variants. +# +# Checkpoint layout is the official repo layout (raw safetensors, no HF weights): +# - turbo (ResembleAI/chatterbox-turbo): t3_turbo_v1.safetensors (GPT-2 medium talker), +# s3gen_meanflow.safetensors, ve.safetensors, conds.pt, GPT-2 BPE tokenizer files +# - multilingual v3 (ResembleAI/chatterbox): t3_mtl23ls_v3.safetensors (Llama 520M talker), +# s3gen_v3.safetensors, ve.safetensors, conds.pt, mtl_tokenizer.json +# The variant is detected by which talker file is present. A minimal config.json with +# architectures ["ChatterboxModel"] routes the directory to these classes. +# +# Talker: the transformer input embeddings (wte / embed_tokens) are dead in the +# reference (inputs_embeds everywhere); the live tables are text_emb and speech_emb, +# fused here into one [text | speech] vocab. Text tokens keep their ids, speech token i +# becomes <|speech_i|> at text_vocab + i. The speech start/stop tokens map to bos/eos. +# +# Mmproj: the whole s3gen sidecar (flow encoder, CFM estimator, HiFT vocoder, CAMPPlus +# speaker encoder, S3 tokenizer), the voice encoder, the talker conditioning encoder, +# the learned position tables, precomputed default-voice conditioning from conds.pt, +# and the talker speech embedding table so that reference speech tokens can be turned +# into talker-space embeddings without a lookup on the text model side. + +from __future__ import annotations + +import json +import math +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Iterable + +import torch +import torch.nn.functional as F + +if TYPE_CHECKING: + from torch import Tensor + +from .base import ModelBase, TextModel, MmprojModel, LazyTorchTensor, gguf, logger + +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) +TURBO_SAMPLING = {"top_k": 1000, "top_p": 0.95, "temp": 0.8, "penalty_repeat": 1.2} +MTL_SAMPLING = {"min_p": 0.05, "top_p": 1.0, "temp": 0.8, "penalty_repeat": 1.2} + + +def _s3tok_mel_filters() -> Tensor: + # slaney-normalized mel filterbank of the s3 tokenizer front end (librosa + # defaults: sr 16000, n_fft 400, 128 bands); not every checkpoint ships + # it, so it is synthesized here for both variants + sr, n_fft, n_mels = 16000, 400, 128 + + def hz_to_mel(f: Tensor) -> Tensor: + lin = f / (200.0 / 3.0) + logstep = math.log(6.4) / 27.0 + return torch.where(f >= 1000.0, 15.0 + torch.log(f.clamp(min=1000.0) / 1000.0) / logstep, lin) + + def mel_to_hz(m: Tensor) -> Tensor: + logstep = math.log(6.4) / 27.0 + return torch.where(m >= 15.0, 1000.0 * torch.exp(logstep * (m - 15.0)), m * (200.0 / 3.0)) + + fftfreqs = torch.arange(n_fft // 2 + 1, dtype=torch.float64) * (sr / n_fft) + bounds = hz_to_mel(torch.tensor([0.0, sr / 2.0], dtype=torch.float64)) + mel_f = mel_to_hz(torch.linspace(bounds[0], bounds[1], n_mels + 2, dtype=torch.float64)) + fdiff = mel_f.diff() + ramps = mel_f[:, None] - fftfreqs[None, :] + lower = -ramps[:n_mels] / fdiff[:n_mels, None] + upper = ramps[2:] / fdiff[1:, None] + weights = torch.minimum(lower, upper).clamp(min=0.0) + weights *= (2.0 / (mel_f[2:] - mel_f[:n_mels]))[:, None] + return weights.float() + + +def _is_turbo(dir_model: Path) -> bool: + if (dir_model / TURBO_TALKER).is_file(): + return True + if (dir_model / MTL_TALKER).is_file(): + return False + raise FileNotFoundError(f"no chatterbox talker checkpoint in {dir_model}") + + +def _index_safetensors(path: Path, lazy: bool, rename: Callable[[str], str | None]) -> dict[str, Callable[[], Tensor]]: + tensors: dict[str, Callable[[], Tensor]] = {} + with gguf.utility.SafetensorsLocal(path) as model_part: + for name in model_part.keys(): + new_name = rename(name) + if new_name is None: + continue + data: gguf.utility.LocalTensor = model_part[name] + if lazy: + data_gen = lambda data=data: LazyTorchTensor.from_local_tensor(data) # noqa: E731 + else: + dtype = LazyTorchTensor._dtype_str_map[data.dtype] + data_gen = lambda data=data, dtype=dtype: torch.from_numpy(data.mmap_bytes()).view(dtype).reshape(data.shape) # noqa: E731 + tensors[new_name] = data_gen + return tensors + + +@ModelBase.register("ChatterboxModel") +class ChatterboxTalkerModel(TextModel): + model_arch = gguf.MODEL_ARCH.LLAMA # multilingual; the turbo constructor switches to GPT2 + + def __init__(self, dir_model: Path, *args, **kwargs): + self.is_turbo = _is_turbo(dir_model) + if self.is_turbo: + self.model_arch = gguf.MODEL_ARCH.GPT2 + super().__init__(dir_model, *args, **kwargs) + self._text_embd: Tensor | None = None + self._speech_embd: Tensor | None = None + self._text_head: Tensor | None = None + self._speech_head: Tensor | None = None + + def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]: + talker = TURBO_TALKER if self.is_turbo else MTL_TALKER + + def rename(name: str) -> str | None: + # transformer input embeddings are dead in the reference (inputs_embeds + # everywhere); the conditioning encoder and position tables go to the mmproj + if name in ("tfmr.wte.weight", "tfmr.embed_tokens.weight"): + return None + if name.startswith(("cond_enc.", "text_pos_emb.", "speech_pos_emb.")): + return None + return name + + return _index_safetensors(self.dir_model / talker, self.lazy, rename) + + def set_vocab(self): + if self.is_turbo: + self._set_vocab_turbo() + else: + self._set_vocab_mtl() + + def _speech_token_names(self, n_speech: int) -> list[str]: + return [f"<|speech_{i}|>" for i in range(n_speech)] + + def _set_vocab_turbo(self): + # stock GPT-2 BPE from the checkpoint dir, extended with the speech tokens + tokens, toktypes, tokpre = self.get_vocab_base() + n_text = len(tokens) + n_speech = self.hparams["speech_vocab_size"] + speech = self._speech_token_names(n_speech) + tokens += speech + toktypes += [gguf.TokenType.CONTROL] * n_speech + + with open(self.dir_model / "merges.txt", "r", encoding="utf-8") as f: + merges = [line.rstrip("\n") for line in f if line.strip() and not line.startswith("#version")] + + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + self.gguf_writer.add_token_merges(merges) + + self.gguf_writer.add_bos_token_id(n_text + SPEECH_BOS) + self.gguf_writer.add_eos_token_id(n_text + SPEECH_EOS) + self.gguf_writer.add_add_bos_token(False) + self.gguf_writer.add_add_eos_token(False) + + def _set_vocab_mtl(self): + # custom multilingual BPE (mtl_tokenizer.json), extended with the speech tokens + with open(self.dir_model / "mtl_tokenizer.json", "r", encoding="utf-8") as f: + tok = json.load(f) + + 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 + for t, i in tok["model"]["vocab"].items(): + tokens[i] = t + toktypes[i] = int(gguf.TokenType.NORMAL) + for entry in tok.get("added_tokens", []): + tokens[entry["id"]] = entry["content"] + toktypes[entry["id"]] = int(gguf.TokenType.CONTROL) + + n_speech = self.hparams["speech_vocab_size"] + tokens += self._speech_token_names(n_speech) + toktypes += [int(gguf.TokenType.CONTROL)] * n_speech + + merges = [" ".join(m) if isinstance(m, list) else m for m in tok["model"].get("merges", [])] + + 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) + + 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 + if "top_k" in sampling: + self.gguf_writer.add_sampling_top_k(sampling["top_k"]) + if "min_p" in sampling: + 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"]) + + 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": + 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):] + # the affine closes the speaker encoding chain, the rest of the + # flow module belongs to the generation stage + if name.startswith("flow.spk_embed_affine_layer."): + return "a." + name[len("flow."):] + if name.startswith("flow."): + return "a.gen." + name # input_embedding, encoder_proj + return None + + def rename_ve(name: str) -> str | None: + if name.startswith("similarity_"): + return None + return "a.ve." + name + + def rename_talker(name: str) -> str | None: + # conditioning encoder, learned position tables and the speech + # embedding table live on the mmproj side + if name.startswith("cond_enc."): + return "a.cenc." + name[len("cond_enc."):] + if name == "text_pos_emb.emb.weight": + return "a.gen.t3.text_pos_emb" + if name == "speech_pos_emb.emb.weight": + return "a.gen.t3.speech_pos_emb" + if name == "speech_emb.weight": + return self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_OUT_EMBD) + return None + + tensors = _index_safetensors(self.dir_model / s3gen, self.lazy, rename_s3gen) + tensors.update(_index_safetensors(self.dir_model / "ve.safetensors", self.lazy, rename_ve)) + tensors.update(_index_safetensors(self.dir_model / talker, self.lazy, rename_talker)) + return tensors + + def set_gguf_parameters(self): + self.gguf_writer.add_file_type(self.ftype) + + # speaker encoder (CAMPPlus, chatterbox_spkenc projector); the DSP front-end + # hparams are fixed by the projector type on the C++ side + self.gguf_writer.add_clip_has_audio_encoder(True) + self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.CHATTERBOX_SPKENC) + self.gguf_writer.add_audio_projection_dim(80) + self.gguf_writer.add_audio_num_mel_bins(80) + self.gguf_writer.add_audio_block_count(0) + self.gguf_writer.add_audio_embedding_length(192) + self.gguf_writer.add_audio_head_count(1) + self.gguf_writer.add_audio_feed_forward_length(192) + self.gguf_writer.add_audio_attention_layernorm_eps(1e-5) + + # audio generator (s3gen, chatterbox projector); the flow encoder shape + self.gguf_writer.add_clip_has_gen_audio_encoder(True) + self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.CHATTERBOX) + self.gguf_writer.add_gen_audio_projection_dim(80) + self.gguf_writer.add_gen_audio_embedding_length(512) + self.gguf_writer.add_gen_audio_feed_forward_length(2048) + self.gguf_writer.add_gen_audio_block_count(6) + self.gguf_writer.add_gen_audio_head_count(8) + self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5) + + self.gguf_writer.add_uint32("chatterbox.n_mels", 80) + self.gguf_writer.add_uint32("chatterbox.sample_rate", 24000) + self.gguf_writer.add_uint32("chatterbox.speech_vocab", self.global_config["text_config"]["speech_vocab_size"]) + self.gguf_writer.add_uint32("chatterbox.meanflow", 1 if self.is_turbo else 0) + + def _fuse_weight_norm(self, name: str, data_torch: Tensor) -> tuple[str, Tensor] | None: + # torch weight_norm parametrization: weight = g * v / |v| over dims 1..n + base = name.split(".parametrizations.weight.original")[0] + if name.endswith("original0"): + self._wnorm_g[base] = data_torch + else: + self._wnorm_v[base] = data_torch + if base in self._wnorm_g and base in self._wnorm_v: + g = self._wnorm_g.pop(base) + v = self._wnorm_v.pop(base) + norm = v.float().norm(dim=tuple(range(1, v.dim())), keepdim=True) + return (base + ".weight", g.float() * v.float() / norm) + return None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if ".parametrizations.weight.original" in name: + fused = self._fuse_weight_norm(name, data_torch) + if fused is not None: + yield fused + return + yield (name, data_torch) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + del name, bid + # tensors read raw on the host (voice encoder, conditioning encoder, + # precomputed conditioning, position tables, mel filterbank, source + # module) must stay F32: the reader handles F32/F16/I32 only + if new_name.startswith(("a.ve.", "a.cenc.", "a.gen.cond.", "a.gen.t3.", "a.gen.hift.m_source.")) or new_name == "a.s3tok.mel_filters": + return gguf.GGMLQuantizationType.F32 + # conv kernels of the graphs (ggml_conv_1d/_2d/_dw and the transposed + # convs of the vocoder) have no BF16 kernels; F16 is the graph-side + # storage type for everything large, F32 for the rest + if n_dims >= 2 and new_name.endswith((".weight", ".weight_v")): + return gguf.GGMLQuantizationType.F16 + return gguf.GGMLQuantizationType.F32 + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + talker = TURBO_TALKER if self.is_turbo else MTL_TALKER + conds = torch.load(self.dir_model / "conds.pt", map_location="cpu", weights_only=False) + t3c = conds["t3"] if isinstance(conds, dict) else conds.t3 + genc = conds["gen"] if isinstance(conds, dict) else conds.gen + t3c = vars(t3c) if not isinstance(t3c, dict) else t3c + genc = vars(genc) if not isinstance(genc, dict) else genc + + def talker_tensor(name: str) -> Tensor: + with gguf.utility.SafetensorsLocal(self.dir_model / talker) as parts: + data = parts[name] + dtype = LazyTorchTensor._dtype_str_map[data.dtype] + return torch.from_numpy(data.mmap_bytes()).view(dtype).reshape(data.shape).clone() + + yield ("a.s3tok.mel_filters", _s3tok_mel_filters()) + + # default voice: precomputed s3gen conditioning from conds.pt + yield ("a.gen.cond.gen_prompt_token", genc["prompt_token"][0].to(torch.int32)) + yield ("a.gen.cond.gen_prompt_feat", genc["prompt_feat"][0].float()) + # the 80-dim flow speaker vector: spk_embed_affine_layer(normalize(campplus)) + with gguf.utility.SafetensorsLocal(self.dir_model / (TURBO_S3GEN if self.is_turbo else MTL_S3GEN)) as parts: + aw = parts["flow.spk_embed_affine_layer.weight"] + ab = parts["flow.spk_embed_affine_layer.bias"] + affine_w = torch.from_numpy(aw.mmap_bytes()).view(LazyTorchTensor._dtype_str_map[aw.dtype]).reshape(aw.shape).float() + affine_b = torch.from_numpy(ab.mmap_bytes()).view(LazyTorchTensor._dtype_str_map[ab.dtype]).reshape(ab.shape).float() + emb = F.normalize(genc["embedding"][0].float(), dim=0) + yield ("a.gen.cond.gen_spk80", affine_w @ emb + affine_b) + + spkr_w = talker_tensor("cond_enc.spkr_enc.weight").float() + spkr_b = talker_tensor("cond_enc.spkr_enc.bias").float() + spkr_row = spkr_w @ t3c["speaker_emb"][0].float() + spkr_b + + if self.is_turbo: + # default talker conditioning: projected speaker row + speech token ids, + # resolved through the speech embedding table at inference time + yield ("a.gen.cond.spkr_default", spkr_row) + yield ("a.gen.cond.prompt_speech_tokens", t3c["cond_prompt_speech_tokens"][0].to(torch.int32)) + return + + # multilingual default talker conditioning: [spkr, perceiver x32, emotion] + # block precomputed by running the reference perceiver over the embedded + # default cond speech tokens (flash attention path of AttentionBlock2) + with torch.no_grad(): + speech_emb = talker_tensor("speech_emb.weight").float() + speech_pos = talker_tensor("speech_pos_emb.emb.weight").float() + cond_tokens = t3c["cond_prompt_speech_tokens"][0] + pse = speech_emb[cond_tokens] + speech_pos[: cond_tokens.shape[0]] + + ln_w = talker_tensor("cond_enc.perceiver.attn.norm.weight").float() + ln_b = talker_tensor("cond_enc.perceiver.attn.norm.bias").float() + wq = talker_tensor("cond_enc.perceiver.attn.to_q.weight").float() + bq = talker_tensor("cond_enc.perceiver.attn.to_q.bias").float() + wk = talker_tensor("cond_enc.perceiver.attn.to_k.weight").float() + bk = talker_tensor("cond_enc.perceiver.attn.to_k.bias").float() + wv = talker_tensor("cond_enc.perceiver.attn.to_v.weight").float() + bv = talker_tensor("cond_enc.perceiver.attn.to_v.bias").float() + wo = talker_tensor("cond_enc.perceiver.attn.proj_out.weight").float() + bo = talker_tensor("cond_enc.perceiver.attn.proj_out.bias").float() + query = talker_tensor("cond_enc.perceiver.pre_attention_query")[0].float() + + n_head = 4 + n_e = query.shape[1] + + def attn_block(x1: Tensor, x2: Tensor) -> Tensor: + nx1 = F.layer_norm(x1, (n_e,), ln_w, ln_b) + nx2 = F.layer_norm(x2, (n_e,), ln_w, ln_b) + q = (nx1 @ wq.T + bq).view(-1, n_head, n_e // n_head).transpose(0, 1) + k = (nx2 @ wk.T + bk).view(-1, n_head, n_e // n_head).transpose(0, 1) + v = (nx2 @ wv.T + bv).view(-1, n_head, n_e // n_head).transpose(0, 1) + ctx = F.scaled_dot_product_attention(q, k, v) + ctx = ctx.transpose(0, 1).reshape(-1, n_e) + return ctx @ wo.T + bo + x1 + + pre = attn_block(query, pse) + p32 = attn_block(pre, pre) + + emo = talker_tensor("cond_enc.emotion_adv_fc.weight").float() + emo_row = emo[:, 0] * t3c["emotion_adv"].reshape(-1)[0] + yield ("a.gen.cond.t3_cond", torch.cat([spkr_row[None], p32, emo_row[None]], dim=0)) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index e72b6564ff02..7ef4844b688d 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -5084,6 +5084,8 @@ class VisionProjectorType: NEMOTRON_V2_VL = "nemotron_v2_vl" QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor + CHATTERBOX_SPKENC = "chatterbox_spkenc" # audio: CAMPPlus speaker encoder + CHATTERBOX = "chatterbox" # audio generation: s3gen flow matching + HiFT vocoder HUNYUANVL = "hunyuanvl" PARAKEET = "parakeet" # audio MINIMAXM3 = "minimax_m3" diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index b333509bc825..461ed6fca29c 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1067,14 +1067,14 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const if (params->ref_tokens) { n_tokens += (int) params->ref_tokens->size(); } else { - auto it = ctx->model.cbx_tensors.find("cond.gen_prompt_token"); + 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("cond.gen_prompt_feat"); + 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]; } @@ -2883,7 +2883,7 @@ struct clip_model_loader { // the flow affine that maps its embedding to the s3gen dim 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("spk.", 0) == 0 || name.rfind("flow.spk_embed_affine_layer.", 0) == 0) { + if (name.rfind("a.spk.", 0) == 0 || name.rfind("a.spk_embed_affine_layer.", 0) == 0) { model.cbx_tensors[name] = get_tensor(name.c_str()); } } @@ -4853,10 +4853,15 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { tokens = *params->ref_tokens; tokens.resize((size_t) n_prompt + n_gen); } else { - ggml_tensor * pt = model.cbx_tensors.at("cond.gen_prompt_token"); - n_prompt = (int) pt->ne[0]; + // 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); - ggml_backend_tensor_get(pt, tokens.data(), 0, (size_t) n_prompt * sizeof(int32_t)); + std::vector ids((size_t) n_prompt); + clip_cbx_read_tensor(ctx, "a.gen.cond.gen_prompt_token", ids.data(), ids.size()); + for (int i = 0; i < n_prompt; i++) { + tokens[(size_t) i] = (int32_t) ids[(size_t) i]; + } } memcpy(tokens.data() + n_prompt, params->codes->data(), (size_t) n_gen * sizeof(int32_t)); const int T1 = n_prompt + n_gen; @@ -4868,7 +4873,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { if (params->ref_feat) { set_input_f32("inp_prompt_feat", *params->ref_feat); } else { - ggml_tensor * pf = model.cbx_tensors.at("cond.gen_prompt_feat"); + ggml_tensor * pf = model.cbx_tensors.at("a.gen.cond.gen_prompt_feat"); std::vector feat(ggml_nelements(pf)); ggml_backend_tensor_get(pf, feat.data(), 0, ggml_nbytes(pf)); set_input_f32("inp_prompt_feat", feat); @@ -4876,7 +4881,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { if (params->ref_spk) { set_input_f32("inp_spk", *params->ref_spk); } else { - ggml_tensor * sp = model.cbx_tensors.at("cond.gen_spk80"); + ggml_tensor * sp = model.cbx_tensors.at("a.gen.cond.gen_spk80"); std::vector spk(ggml_nelements(sp)); ggml_backend_tensor_get(sp, spk.data(), 0, ggml_nbytes(sp)); set_input_f32("inp_spk", spk); @@ -4909,7 +4914,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { for (auto & f : noise) f = nd(rng); set_input_f32("inp_noise", noise); - const bool meanflow = model.cbx_tensors.count("est.time_embed_mixer.weight") > 0; + const bool meanflow = model.cbx_tensors.count("a.gen.est.time_embed_mixer.weight") > 0; const int n_steps = meanflow ? 2 : 10; std::vector temb((size_t) 320 * (n_steps + 1)); for (int s = 0; s <= n_steps; s++) { @@ -5494,6 +5499,43 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } (*params->out_codes)[(size_t) t] = code; } + + // embedding rows of the produced codes, gathered from the talker + // speech table shipped in the mmproj; the multilingual variant adds + // its learned speech positions so the rows feed the conditioning + // perceiver directly + if (params->out_code_embd) { + const auto & tensors = ctx->model.cbx_tensors; + auto tab_it = tensors.find("a.gen.code.out_embd.weight"); + GGML_ASSERT(tab_it != tensors.end()); + ggml_tensor * tab = tab_it->second; + GGML_ASSERT(tab->type == GGML_TYPE_F16 || tab->type == GGML_TYPE_F32); + const int n_e = (int) tab->ne[0]; + auto pos_it = tensors.find("a.gen.t3.speech_pos_emb"); + ggml_tensor * pos = pos_it == tensors.end() ? nullptr : pos_it->second; + + params->out_code_embd->resize((size_t) n_tok * n_e); + std::vector h16(n_e); + std::vector pr(n_e); + for (int t = 0; t < n_tok; t++) { + float * dst = params->out_code_embd->data() + (size_t) t * n_e; + const size_t r = (size_t) (*params->out_codes)[(size_t) t] * n_e; + if (tab->type == GGML_TYPE_F16) { + ggml_backend_tensor_get(tab, h16.data(), r * sizeof(ggml_fp16_t), (size_t) n_e * sizeof(ggml_fp16_t)); + for (int j = 0; j < n_e; j++) { + dst[j] = ggml_fp16_to_fp32(h16[j]); + } + } else { + ggml_backend_tensor_get(tab, dst, r * sizeof(float), (size_t) n_e * sizeof(float)); + } + if (pos) { + ggml_backend_tensor_get(pos, pr.data(), (size_t) t * n_e * sizeof(float), (size_t) n_e * sizeof(float)); + for (int j = 0; j < n_e; j++) { + dst[j] += pr[j]; + } + } + } + } return true; } @@ -5526,8 +5568,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { const int64_t n_wav = (int64_t) n_mel_out * ups_total; std::vector lw(9); float lb = 0.0f; - GGML_ASSERT(clip_cbx_read_tensor(ctx, "hift.m_source.l_linear.weight", lw.data(), lw.size()) == 9); - clip_cbx_read_tensor(ctx, "hift.m_source.l_linear.bias", &lb, 1); + GGML_ASSERT(clip_cbx_read_tensor(ctx, "a.gen.hift.m_source.l_linear.weight", lw.data(), lw.size()) == 9); + clip_cbx_read_tensor(ctx, "a.gen.hift.m_source.l_linear.bias", &lb, 1); std::mt19937 srng(1234); std::uniform_real_distribution ud(-M_PI, M_PI); diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index 6c57ebdd8bd9..ace45fd021e2 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -109,6 +109,10 @@ struct clip_encode_params { int32_t top_k = 50; float top_p = 1.0f; std::vector * out_codes = nullptr; + // TOKENIZE: out_code_embd receives the speech embedding rows of the + // produced codes (kept apart from out_embd, which is reserved for graphs + // whose last node is the embedding tensor) + std::vector * out_code_embd = nullptr; // CODE2WAV: codes holds this frame's 16 RVQ codes, out_audio receives the // decoded PCM samples (F32). state_in is the state from the previous diff --git a/tools/mtmd/models/chatterbox-gen.cpp b/tools/mtmd/models/chatterbox-gen.cpp index f788d188e46a..b9e6e90cd27d 100644 --- a/tools/mtmd/models/chatterbox-gen.cpp +++ b/tools/mtmd/models/chatterbox-gen.cpp @@ -192,19 +192,19 @@ static ggml_tensor * cbx_estimator(const clip_model & model, ggml_context * ctx0 // down ggml_tensor * skip; - x = cbx_resnet(model, ctx0, x, temb, "est.down_blocks.0.0"); - for (int j = 0; model.cbx_tensors.count("est.down_blocks.0.1." + std::to_string(j) + ".norm1.weight"); j++) { - x = cbx_tfm_block(model, ctx0, x, "est.down_blocks.0.1." + std::to_string(j)); + x = cbx_resnet(model, ctx0, x, temb, "a.gen.est.down_blocks.0.0"); + for (int j = 0; model.cbx_tensors.count("a.gen.est.down_blocks.0.1." + std::to_string(j) + ".norm1.weight"); j++) { + x = cbx_tfm_block(model, ctx0, x, "a.gen.est.down_blocks.0.1." + std::to_string(j)); } skip = x; { - ggml_tensor * k = cbx_t(model, "est.down_blocks.0.2.weight"); - x = cbx_conv1d(ctx0, k, cbx_t(model, "est.down_blocks.0.2.bias"), x, 1, (int) k->ne[0] - 1, 0); + ggml_tensor * k = cbx_t(model, "a.gen.est.down_blocks.0.2.weight"); + x = cbx_conv1d(ctx0, k, cbx_t(model, "a.gen.est.down_blocks.0.2.bias"), x, 1, (int) k->ne[0] - 1, 0); } // mid - for (int i = 0; model.cbx_tensors.count("est.mid_blocks." + std::to_string(i) + ".0.block1.block.0.weight"); i++) { - const std::string mp = "est.mid_blocks." + std::to_string(i); + for (int i = 0; model.cbx_tensors.count("a.gen.est.mid_blocks." + std::to_string(i) + ".0.block1.block.0.weight"); i++) { + const std::string mp = "a.gen.est.mid_blocks." + std::to_string(i); x = cbx_resnet(model, ctx0, x, temb, mp + ".0"); for (int j = 0; model.cbx_tensors.count(mp + ".1." + std::to_string(j) + ".norm1.weight"); j++) { x = cbx_tfm_block(model, ctx0, x, mp + ".1." + std::to_string(j)); @@ -213,17 +213,17 @@ static ggml_tensor * cbx_estimator(const clip_model & model, ggml_context * ctx0 // up with skip x = ggml_concat(ctx0, x, skip, 0); // [512, T] - x = cbx_resnet(model, ctx0, x, temb, "est.up_blocks.0.0"); - for (int j = 0; model.cbx_tensors.count("est.up_blocks.0.1." + std::to_string(j) + ".norm1.weight"); j++) { - x = cbx_tfm_block(model, ctx0, x, "est.up_blocks.0.1." + std::to_string(j)); + x = cbx_resnet(model, ctx0, x, temb, "a.gen.est.up_blocks.0.0"); + for (int j = 0; model.cbx_tensors.count("a.gen.est.up_blocks.0.1." + std::to_string(j) + ".norm1.weight"); j++) { + x = cbx_tfm_block(model, ctx0, x, "a.gen.est.up_blocks.0.1." + std::to_string(j)); } { - ggml_tensor * k = cbx_t(model, "est.up_blocks.0.2.weight"); - x = cbx_conv1d(ctx0, k, cbx_t(model, "est.up_blocks.0.2.bias"), x, 1, (int) k->ne[0] - 1, 0); + ggml_tensor * k = cbx_t(model, "a.gen.est.up_blocks.0.2.weight"); + x = cbx_conv1d(ctx0, k, cbx_t(model, "a.gen.est.up_blocks.0.2.bias"), x, 1, (int) k->ne[0] - 1, 0); } - x = cbx_causal_block(model, ctx0, x, "est.final_block"); - x = cbx_conv1d(ctx0, cbx_t(model, "est.final_proj.weight"), cbx_t(model, "est.final_proj.bias"), x, 1, 0, 0); // [80, T] + x = cbx_causal_block(model, ctx0, x, "a.gen.est.final_block"); + x = cbx_conv1d(ctx0, cbx_t(model, "a.gen.est.final_proj.weight"), cbx_t(model, "a.gen.est.final_proj.bias"), x, 1, 0, 0); // [80, T] return x; } @@ -247,13 +247,13 @@ static ggml_cgraph * cbx_build_s3tok(const clip_model & model, ggml_context * ct ggml_set_input(pos); ggml_tensor * x = ggml_cont(ctx0, ggml_transpose(ctx0, inp)); // [128, T] - x = cbx_conv1d(ctx0, cbx_t(model, "s3tok.encoder.conv1.weight"), cbx_t(model, "s3tok.encoder.conv1.bias"), x, 2, 1, 1); + x = cbx_conv1d(ctx0, cbx_t(model, "a.s3tok.encoder.conv1.weight"), cbx_t(model, "a.s3tok.encoder.conv1.bias"), x, 2, 1, 1); x = ggml_gelu_erf(ctx0, x); - x = cbx_conv1d(ctx0, cbx_t(model, "s3tok.encoder.conv2.weight"), cbx_t(model, "s3tok.encoder.conv2.bias"), x, 2, 1, 1); + x = cbx_conv1d(ctx0, cbx_t(model, "a.s3tok.encoder.conv2.weight"), cbx_t(model, "a.s3tok.encoder.conv2.bias"), x, 2, 1, 1); x = ggml_gelu_erf(ctx0, x); // [1280, T2] - for (int li = 0; model.cbx_tensors.count("s3tok.encoder.blocks." + std::to_string(li) + ".attn_ln.weight"); li++) { - const std::string p = "s3tok.encoder.blocks." + std::to_string(li) + ".attn"; + for (int li = 0; model.cbx_tensors.count("a.s3tok.encoder.blocks." + std::to_string(li) + ".attn_ln.weight"); li++) { + const std::string p = "a.s3tok.encoder.blocks." + std::to_string(li) + ".attn"; ggml_tensor * res = x; ggml_tensor * cur = cbx_layer_norm(ctx0, cbx_t(model, p + "_ln.weight"), cbx_t(model, p + "_ln.bias"), x, 1e-5f); @@ -286,7 +286,7 @@ static ggml_cgraph * cbx_build_s3tok(const clip_model & model, ggml_context * ct o = cbx_linear(ctx0, cbx_t(model, p + ".out.weight"), cbx_t(model, p + ".out.bias"), o); x = ggml_add(ctx0, res, ggml_add(ctx0, o, fsm)); - const std::string mp = "s3tok.encoder.blocks." + std::to_string(li) + ".mlp"; + const std::string mp = "a.s3tok.encoder.blocks." + std::to_string(li) + ".mlp"; res = x; cur = cbx_layer_norm(ctx0, cbx_t(model, mp + "_ln.weight"), cbx_t(model, mp + "_ln.bias"), x, 1e-5f); cur = cbx_linear(ctx0, cbx_t(model, mp + ".0.weight"), cbx_t(model, mp + ".0.bias"), cur); @@ -295,8 +295,8 @@ static ggml_cgraph * cbx_build_s3tok(const clip_model & model, ggml_context * ct x = ggml_add(ctx0, res, cur); } - x = cbx_linear(ctx0, cbx_t(model, "s3tok.quantizer._codebook.project_down.weight"), - cbx_t(model, "s3tok.quantizer._codebook.project_down.bias"), x); // [8, T2] + x = cbx_linear(ctx0, cbx_t(model, "a.s3tok.quantizer._codebook.project_down.weight"), + cbx_t(model, "a.s3tok.quantizer._codebook.project_down.bias"), x); // [8, T2] x = ggml_tanh(ctx0, x); ggml_set_name(x, "out_fsq"); @@ -342,28 +342,28 @@ ggml_cgraph * clip_graph_chatterbox::build() { ggml_set_input(pos2); // token embedding - ggml_tensor * x = ggml_get_rows(ctx0, cbx_t(model, "flow.input_embedding.weight"), inp_tokens); // [512, T1] + ggml_tensor * x = ggml_get_rows(ctx0, cbx_t(model, "a.gen.flow.input_embedding.weight"), inp_tokens); // [512, T1] // embed: linear + layer norm, then the espnet xscale - x = cbx_linear(ctx0, cbx_t(model, "fenc.embed.out.0.weight"), cbx_t(model, "fenc.embed.out.0.bias"), x); - x = cbx_layer_norm(ctx0, cbx_t(model, "fenc.embed.out.1.weight"), cbx_t(model, "fenc.embed.out.1.bias"), x, 1e-5f); + x = cbx_linear(ctx0, cbx_t(model, "a.gen.fenc.embed.out.0.weight"), cbx_t(model, "a.gen.fenc.embed.out.0.bias"), x); + x = cbx_layer_norm(ctx0, cbx_t(model, "a.gen.fenc.embed.out.1.weight"), cbx_t(model, "a.gen.fenc.embed.out.1.bias"), 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(ctx0, cbx_t(model, "fenc.pre_lookahead_layer.conv1.weight"), - cbx_t(model, "fenc.pre_lookahead_layer.conv1.bias"), x, 1, 0, 3); + ggml_tensor * cur = cbx_conv1d(ctx0, cbx_t(model, "a.gen.fenc.pre_lookahead_layer.conv1.weight"), + cbx_t(model, "a.gen.fenc.pre_lookahead_layer.conv1.bias"), x, 1, 0, 3); cur = ggml_leaky_relu(ctx0, cur, 0.01f, false); - cur = cbx_conv1d(ctx0, cbx_t(model, "fenc.pre_lookahead_layer.conv2.weight"), - cbx_t(model, "fenc.pre_lookahead_layer.conv2.bias"), cur, 1, 2, 0); + cur = cbx_conv1d(ctx0, cbx_t(model, "a.gen.fenc.pre_lookahead_layer.conv2.weight"), + cbx_t(model, "a.gen.fenc.pre_lookahead_layer.conv2.bias"), cur, 1, 2, 0); x = ggml_add(ctx0, res, cur); cb(x, "fenc_pre_lookahead", -1); } - for (int i = 0; model.cbx_tensors.count("fenc.encoders." + std::to_string(i) + ".norm_mha.weight"); i++) { - x = cbx_enc_layer(model, ctx0, x, pos1, "fenc.encoders." + std::to_string(i), T1); + for (int i = 0; model.cbx_tensors.count("a.gen.fenc.encoders." + std::to_string(i) + ".norm_mha.weight"); i++) { + x = cbx_enc_layer(model, ctx0, x, pos1, "a.gen.fenc.encoders." + std::to_string(i), T1); cb(x, "fenc_enc", i); } @@ -374,32 +374,32 @@ ggml_cgraph * clip_graph_chatterbox::build() { 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, cbx_t(model, "fenc.up_layer.conv.weight"), xt, 1, 0, 1); + ggml_tensor * y = ggml_conv_1d(ctx0, cbx_t(model, "a.gen.fenc.up_layer.conv.weight"), xt, 1, 0, 1); x = ggml_cont(ctx0, ggml_transpose(ctx0, y)); // [512, T2] - x = ggml_add(ctx0, x, cbx_t(model, "fenc.up_layer.conv.bias")); + x = ggml_add(ctx0, x, cbx_t(model, "a.gen.fenc.up_layer.conv.bias")); cb(x, "fenc_upsample", -1); } // up embed: linear + layer norm + xscale - x = cbx_linear(ctx0, cbx_t(model, "fenc.up_embed.out.0.weight"), cbx_t(model, "fenc.up_embed.out.0.bias"), x); - x = cbx_layer_norm(ctx0, cbx_t(model, "fenc.up_embed.out.1.weight"), cbx_t(model, "fenc.up_embed.out.1.bias"), x, 1e-5f); + x = cbx_linear(ctx0, cbx_t(model, "a.gen.fenc.up_embed.out.0.weight"), cbx_t(model, "a.gen.fenc.up_embed.out.0.bias"), x); + x = cbx_layer_norm(ctx0, cbx_t(model, "a.gen.fenc.up_embed.out.1.weight"), cbx_t(model, "a.gen.fenc.up_embed.out.1.bias"), x, 1e-5f); x = ggml_scale(ctx0, x, sqrtf(512.0f)); - for (int i = 0; model.cbx_tensors.count("fenc.up_encoders." + std::to_string(i) + ".norm_mha.weight"); i++) { - x = cbx_enc_layer(model, ctx0, x, pos2, "fenc.up_encoders." + std::to_string(i), T2); + for (int i = 0; model.cbx_tensors.count("a.gen.fenc.up_encoders." + std::to_string(i) + ".norm_mha.weight"); i++) { + x = cbx_enc_layer(model, ctx0, x, pos2, "a.gen.fenc.up_encoders." + std::to_string(i), T2); cb(x, "fenc_up_enc", i); } - x = cbx_layer_norm(ctx0, cbx_t(model, "fenc.after_norm.weight"), cbx_t(model, "fenc.after_norm.bias"), x, 1e-5f); + x = cbx_layer_norm(ctx0, cbx_t(model, "a.gen.fenc.after_norm.weight"), cbx_t(model, "a.gen.fenc.after_norm.bias"), x, 1e-5f); // encoder projection to the mel channel count - ggml_tensor * mu = cbx_linear(ctx0, cbx_t(model, "flow.encoder_proj.weight"), cbx_t(model, "flow.encoder_proj.bias"), x); // [80, T2] + ggml_tensor * mu = cbx_linear(ctx0, cbx_t(model, "a.gen.flow.encoder_proj.weight"), cbx_t(model, "a.gen.flow.encoder_proj.bias"), 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 = model.cbx_tensors.count("est.time_embed_mixer.weight") > 0; + const bool meanflow = model.cbx_tensors.count("a.gen.est.time_embed_mixer.weight") > 0; const int n_steps = meanflow ? 2 : 10; // span points, same schedule as the host side sinusoid fill in clip.cpp @@ -418,9 +418,9 @@ ggml_cgraph * clip_graph_chatterbox::build() { ggml_set_input(temb_sin); auto time_mlp = [&](ggml_tensor * e) { - e = cbx_linear(ctx0, cbx_t(model, "est.time_mlp.linear_1.weight"), cbx_t(model, "est.time_mlp.linear_1.bias"), e); + e = cbx_linear(ctx0, cbx_t(model, "a.gen.est.time_mlp.linear_1.weight"), cbx_t(model, "a.gen.est.time_mlp.linear_1.bias"), e); e = ggml_silu(ctx0, e); - e = cbx_linear(ctx0, cbx_t(model, "est.time_mlp.linear_2.weight"), cbx_t(model, "est.time_mlp.linear_2.bias"), e); + e = cbx_linear(ctx0, cbx_t(model, "a.gen.est.time_mlp.linear_2.weight"), cbx_t(model, "a.gen.est.time_mlp.linear_2.bias"), e); return e; }; auto span_emb = [&](int i) { @@ -431,7 +431,7 @@ ggml_cgraph * clip_graph_chatterbox::build() { 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, cbx_t(model, "est.time_embed_mixer.weight"), e); // [1024, 1] + return ggml_mul_mat(ctx0, cbx_t(model, "a.gen.est.time_embed_mixer.weight"), e); // [1024, 1] }; // mel-rate conditions: prompt features then zeros, and the 80-dim @@ -475,12 +475,12 @@ ggml_cgraph * clip_graph_chatterbox::build() { // f0 predictor on the trimmed mel: 5x (conv k3 same-pad + elu), abs(linear) { ggml_tensor * fx = mel; - for (int i = 0; model.cbx_tensors.count("hift.f0_predictor.condnet." + std::to_string(i) + ".weight"); i += 2) { - const std::string cp = "hift.f0_predictor.condnet." + std::to_string(i); + for (int i = 0; model.cbx_tensors.count("a.gen.hift.f0_predictor.condnet." + std::to_string(i) + ".weight"); i += 2) { + const std::string cp = "a.gen.hift.f0_predictor.condnet." + std::to_string(i); fx = cbx_conv1d(ctx0, cbx_t(model, cp + ".weight"), cbx_t(model, cp + ".bias"), fx, 1, 1, 1); fx = ggml_elu(ctx0, fx); } - fx = cbx_linear(ctx0, cbx_t(model, "hift.f0_predictor.classifier.weight"), cbx_t(model, "hift.f0_predictor.classifier.bias"), fx); + fx = cbx_linear(ctx0, cbx_t(model, "a.gen.hift.f0_predictor.classifier.weight"), cbx_t(model, "a.gen.hift.f0_predictor.classifier.bias"), fx); fx = ggml_abs(ctx0, fx); // [1, T] ggml_set_name(fx, "out_f0"); ggml_set_output(fx); @@ -512,11 +512,11 @@ static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * ggml_set_name(sstft, "inp_sstft"); ggml_set_input(sstft); - ggml_tensor * x = cbx_conv1d_dil(ctx0, cbx_t(model, "hift.conv_pre.weight"), cbx_t(model, "hift.conv_pre.bias"), mel, 3, 1); + ggml_tensor * x = cbx_conv1d_dil(ctx0, cbx_t(model, "a.gen.hift.conv_pre.weight"), cbx_t(model, "a.gen.hift.conv_pre.bias"), mel, 3, 1); - for (int i = 0; model.cbx_tensors.count("hift.ups." + std::to_string(i) + ".weight"); i++) { + for (int i = 0; model.cbx_tensors.count("a.gen.hift.ups." + std::to_string(i) + ".weight"); i++) { const std::string is = std::to_string(i); - ggml_tensor * uk = cbx_t(model, "hift.ups." + is + ".weight"); + ggml_tensor * uk = cbx_t(model, "a.gen.hift.ups." + is + ".weight"); const int K = (int) uk->ne[0]; const int S = K / 2; const int P = (K - S) / 2; @@ -527,9 +527,9 @@ static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * 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, cbx_t(model, "hift.ups." + is + ".bias")); + x = ggml_add(ctx0, x, cbx_t(model, "a.gen.hift.ups." + is + ".bias")); - const bool is_last = !model.cbx_tensors.count("hift.ups." + std::to_string(i + 1) + ".weight"); + const bool is_last = !model.cbx_tensors.count("a.gen.hift.ups." + std::to_string(i + 1) + ".weight"); if (is_last) { ggml_tensor * xr = ggml_cont(ctx0, ggml_transpose(ctx0, x)); xr = ggml_pad_reflect_1d(ctx0, xr, 1, 0); @@ -537,7 +537,7 @@ static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * } // source injection: strided conv on the source stft, one resblock - ggml_tensor * sk = cbx_t(model, "hift.source_downs." + is + ".weight"); + ggml_tensor * sk = cbx_t(model, "a.gen.hift.source_downs." + is + ".weight"); const int SK = (int) sk->ne[0]; const int SS = SK > 1 ? SK / 2 : 1; const int SP = SK > 1 ? SS / 2 : 0; @@ -546,9 +546,9 @@ static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * 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, cbx_t(model, "hift.source_downs." + is + ".bias")); + si = ggml_add(ctx0, si, cbx_t(model, "a.gen.hift.source_downs." + is + ".bias")); } - si = cbx_hift_resblock(model, ctx0, si, "hift.source_resblocks." + is); + si = cbx_hift_resblock(model, ctx0, si, "a.gen.hift.source_resblocks." + is); // 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]); @@ -559,14 +559,14 @@ static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * ggml_tensor * acc = nullptr; for (int j = 3 * i; j < 3 * (i + 1); j++) { - ggml_tensor * r = cbx_hift_resblock(model, ctx0, x, "hift.resblocks." + std::to_string(j)); + ggml_tensor * r = cbx_hift_resblock(model, ctx0, x, "a.gen.hift.resblocks." + std::to_string(j)); acc = acc ? ggml_add(ctx0, acc, r) : r; } x = ggml_scale(ctx0, acc, 1.0f / 3.0f); } x = ggml_leaky_relu(ctx0, x, 0.01f, false); - x = cbx_conv1d_dil(ctx0, cbx_t(model, "hift.conv_post.weight"), cbx_t(model, "hift.conv_post.bias"), x, 3, 1); + x = cbx_conv1d_dil(ctx0, cbx_t(model, "a.gen.hift.conv_post.weight"), cbx_t(model, "a.gen.hift.conv_post.bias"), x, 3, 1); ggml_set_name(x, "out_spec"); ggml_set_output(x); ggml_build_forward_expand(gf, x); diff --git a/tools/mtmd/models/chatterbox-spkenc.cpp b/tools/mtmd/models/chatterbox-spkenc.cpp index 0f63db9c396e..a9e967161f19 100644 --- a/tools/mtmd/models/chatterbox-spkenc.cpp +++ b/tools/mtmd/models/chatterbox-spkenc.cpp @@ -122,37 +122,37 @@ ggml_cgraph * clip_graph_chatterbox_spkenc::build() { // 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, cbx_t(model, "spk.head.conv1.weight"), x, 1, 1, 1, 1, 1, 1); - x = cbx_bn2d_relu(model, ctx0, x, "spk.head.bn1", eps); - x = cbx_res2d(model, ctx0, x, "spk.head.layer1.0", 2, eps); - x = cbx_res2d(model, ctx0, x, "spk.head.layer1.1", 1, eps); - x = cbx_res2d(model, ctx0, x, "spk.head.layer2.0", 2, eps); - x = cbx_res2d(model, ctx0, x, "spk.head.layer2.1", 1, eps); - x = ggml_conv_2d(ctx0, cbx_t(model, "spk.head.conv2.weight"), x, 1, 2, 1, 1, 1, 1); - x = cbx_bn2d_relu(model, ctx0, x, "spk.head.bn2", eps); + x = ggml_conv_2d(ctx0, cbx_t(model, "a.spk.head.conv1.weight"), x, 1, 1, 1, 1, 1, 1); + x = cbx_bn2d_relu(model, ctx0, x, "a.spk.head.bn1", eps); + x = cbx_res2d(model, ctx0, x, "a.spk.head.layer1.0", 2, eps); + x = cbx_res2d(model, ctx0, x, "a.spk.head.layer1.1", 1, eps); + x = cbx_res2d(model, ctx0, x, "a.spk.head.layer2.0", 2, eps); + x = cbx_res2d(model, ctx0, x, "a.spk.head.layer2.1", 1, eps); + x = ggml_conv_2d(ctx0, cbx_t(model, "a.spk.head.conv2.weight"), x, 1, 2, 1, 1, 1, 1); + x = cbx_bn2d_relu(model, ctx0, x, "a.spk.head.bn2", 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(ctx0, cbx_t(model, "spk.xvector.tdnn.linear.weight"), nullptr, x, 2, 2, 2); // [128, T1] - x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, "spk.xvector.tdnn.nonlinear.batchnorm", eps)); + x = cbx_conv1d(ctx0, cbx_t(model, "a.spk.xvector.tdnn.linear.weight"), nullptr, x, 2, 2, 2); // [128, T1] + x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, "a.spk.xvector.tdnn.nonlinear.batchnorm", eps)); cb(x, "spk_tdnn", -1); static const int block_dil[3] = {1, 2, 2}; for (int bi = 1; bi <= 3; bi++) { - const std::string bp = "spk.xvector.block" + std::to_string(bi); + const std::string bp = "a.spk.xvector.block" + std::to_string(bi); for (int li = 1; model.cbx_tensors.count(bp + ".tdnnd" + std::to_string(li) + ".linear1.weight"); li++) { ggml_tensor * out = cbx_cam_layer(model, ctx0, x, bp + ".tdnnd" + std::to_string(li), block_dil[bi - 1], eps, segfix); x = ggml_concat(ctx0, x, out, 0); } - const std::string tp = "spk.xvector.transit" + std::to_string(bi); + const std::string tp = "a.spk.xvector.transit" + std::to_string(bi); x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, tp + ".nonlinear.batchnorm", eps)); x = cbx_conv1d(ctx0, cbx_t(model, tp + ".linear.weight"), nullptr, x, 1, 0, 0); cb(x, "spk_block", bi); } - x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, "spk.xvector.out_nonlinear.batchnorm", eps)); // [512, T1] + x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, "a.spk.xvector.out_nonlinear.batchnorm", 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] @@ -166,9 +166,9 @@ ggml_cgraph * clip_graph_chatterbox_spkenc::build() { cb(stats, "spk_stats_pool", -1); // dense 1024 -> 192, batchnorm without affine, into the x-vector - ggml_tensor * dw = ggml_reshape_2d(ctx0, cbx_t(model, "spk.xvector.dense.linear.weight"), 1024, 192); + ggml_tensor * dw = ggml_reshape_2d(ctx0, cbx_t(model, "a.spk.xvector.dense.linear.weight"), 1024, 192); ggml_tensor * emb = ggml_mul_mat(ctx0, dw, stats); // [192, 1] - emb = cbx_bn1d(model, ctx0, emb, "spk.xvector.dense.nonlinear.batchnorm", eps); + emb = cbx_bn1d(model, ctx0, emb, "a.spk.xvector.dense.nonlinear.batchnorm", eps); emb = ggml_reshape_1d(ctx0, emb, 192); ggml_set_name(emb, "out_xvec"); ggml_set_output(emb); @@ -177,8 +177,8 @@ ggml_cgraph * clip_graph_chatterbox_spkenc::build() { // normalize then the s3gen speaker affine ggml_tensor * n2 = ggml_sqrt(ctx0, ggml_sum(ctx0, ggml_mul(ctx0, emb, emb))); ggml_tensor * unit = ggml_div(ctx0, emb, n2); - ggml_tensor * spk80 = cbx_linear(ctx0, cbx_t(model, "flow.spk_embed_affine_layer.weight"), - cbx_t(model, "flow.spk_embed_affine_layer.bias"), + ggml_tensor * spk80 = cbx_linear(ctx0, cbx_t(model, "a.spk_embed_affine_layer.weight"), + cbx_t(model, "a.spk_embed_affine_layer.bias"), ggml_reshape_2d(ctx0, unit, 192, 1)); spk80 = ggml_cont(ctx0, ggml_reshape_1d(ctx0, spk80, 80)); cb(spk80, "spk_embd", -1); diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 34acb652acb4..539a12e5275a 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -429,6 +429,8 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { audio_pcm.clear(); h_state_buf.clear(); out_buf.clear(); + ref_cond.clear(); + ref_state.clear(); } int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override { @@ -439,54 +441,21 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } if (inp->speaker_ref) { - // turbo reference chain: loudness normalize the clip, then cap - // the talker conditioning at 15 s (multilingual: 6 s) - mtmd_gen_audio_norm_ref(mctx, inp->speaker_ref); - const size_t t3_cap = (size_t) (t3_cond.empty() ? 15 : 6) * 16000; - - if (!encode_speaker(inp->speaker_ref, spk80)) { - return 1; - } - if (!tokenize_ref(inp->speaker_ref, 10 * 16000, ref_prompt_tokens) || - !tokenize_ref(inp->speaker_ref, t3_cap, ref_t3_tokens)) { + // one spk-ref call encodes the reference clip: conditioning rows + // for the talker prompt, opaque state for the flow decoder + const float * pcm = (const float *) mtmd_bitmap_get_data(inp->speaker_ref); + const size_t n = mtmd_bitmap_get_n_bytes(inp->speaker_ref) / sizeof(float); + mtmd_gen_inp gi{}; + gi.type = MTMD_GEN_PROCESS_TYPE_SPK_REF; + gi.pcm = pcm; + gi.n_pcm = n; + mtmd_gen_out go{}; + if (mtmd_gen_audio_process(mctx, &gi, &go) != 0) { + LOG_ERR("mtmd_helper_gen_audio: speaker reference encoding failed\n"); return 1; } - - // the tts stage derives the mel-rate prompt features from the - // same capped reference clip - { - const float * pcm = (const float *) mtmd_bitmap_get_data(inp->speaker_ref); - const size_t n = mtmd_bitmap_get_n_bytes(inp->speaker_ref) / sizeof(float); - ref_pcm16.assign(pcm, pcm + std::min(n, (size_t) 10 * 16000)); - - // talker conditioning rows from the voice encoder chain; the - // multilingual perceiver consumes the embedding rows of the - // reference speech tokens, built from the fused talker vocab - std::vector pse; - if (!t3_cond.empty()) { - for (size_t i = 0; i < ref_t3_tokens.size(); i++) { - std::vector r(tok_embd.begin() + (size_t) (speech_base + ref_t3_tokens[i]) * n_embd, - tok_embd.begin() + (size_t) (speech_base + ref_t3_tokens[i] + 1) * n_embd); - const float * p = speech_pos.data() + i * (size_t) n_embd; - for (int j = 0; j < n_embd; j++) { - r[(size_t) j] += p[j]; - } - pse.insert(pse.end(), r.begin(), r.end()); - } - } - mtmd_gen_inp gi{}; - gi.type = MTMD_GEN_PROCESS_TYPE_SPEAKER_COND; - gi.pcm = pcm; - gi.n_pcm = n; - gi.ref_speech_embd = pse.empty() ? nullptr : pse.data(); - gi.n_ref_speech_rows = pse.size() / (size_t) n_embd; - mtmd_gen_out go{}; - if (mtmd_gen_audio_process(mctx, &gi, &go) != 0) { - LOG_ERR("mtmd_helper_gen_audio: speaker conditioning failed\n"); - return 1; - } - ref_cond.assign(go.embd, go.embd + go.n_embd); - } + ref_cond.assign(go.embd, go.embd + go.n_embd); + ref_state.assign(go.state_data, go.state_data + go.state_size); } const int n_e = n_embd; @@ -512,16 +481,19 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { prompt.emplace_back(cond.begin() + i * (size_t) n_e, cond.begin() + (i + 1) * (size_t) n_e); } } else { - // conditioning: projected speaker row, then the speech token prompt - // (already fused ids after the text vocab) - prompt.push_back(ref_cond.empty() ? cond_spkr : ref_cond); + // conditioning rows: the reference block from the spk-ref stage, + // or the precomputed default speaker row followed by the + // table-resolved prompt ids if (ref_cond.empty()) { - for (float f : cond_speech_tokens) { - prompt.push_back(row(speech_base + (llama_token) f)); + prompt.push_back(cond_spkr); + for (size_t i = 0; i < cond_speech_rows.size() / (size_t) n_e; i++) { + prompt.emplace_back(cond_speech_rows.begin() + i * (size_t) n_e, + cond_speech_rows.begin() + (i + 1) * (size_t) n_e); } } else { - for (int32_t t : ref_t3_tokens) { - prompt.push_back(row(speech_base + t)); + for (size_t i = 0; i < ref_cond.size() / (size_t) n_e; i++) { + prompt.emplace_back(ref_cond.begin() + i * (size_t) n_e, + ref_cond.begin() + (i + 1) * (size_t) n_e); } } } @@ -678,19 +650,16 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } mtmd_gen_inp gen_inp{}; - gen_inp.type = MTMD_GEN_PROCESS_TYPE_TTS; + gen_inp.type = MTMD_GEN_PROCESS_TYPE_CODE2WAV; gen_inp.codes = codes_buf.data(); gen_inp.n_codes = codes_buf.size(); - if (!spk80.empty()) { - gen_inp.ref_spk = spk80.data(); - gen_inp.ref_tokens = ref_prompt_tokens.data(); - gen_inp.n_ref_tokens = ref_prompt_tokens.size(); - gen_inp.ref_pcm = ref_pcm16.data(); - gen_inp.n_ref_pcm = ref_pcm16.size(); + if (!ref_state.empty()) { + gen_inp.state_data = ref_state.data(); + gen_inp.state_size = ref_state.size(); } mtmd_gen_out gen_out{}; if (mtmd_gen_audio_process(mctx, &gen_inp, &gen_out) != 0) { - LOG_ERR("mtmd_helper_gen_audio: tts decode failed\n"); + LOG_ERR("mtmd_helper_gen_audio: code2wav decode failed\n"); return 1; } audio_pcm.assign(gen_out.audio, gen_out.audio + gen_out.n_samples); @@ -738,12 +707,12 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { // multilingual variant: the mmproj ships a precomputed t3 conditioning // block [spkr, perceiver, emotion] and the learned positional tables // that the backbone needs added to its input embeddings - size_t n_t3 = mtmd_gen_audio_read_tensor(mctx, "cond.t3_cond", nullptr, 0); + size_t n_t3 = mtmd_gen_audio_read_tensor(mctx, "a.gen.cond.t3_cond", nullptr, 0); if (n_t3 > 0) { t3_cond.resize(n_t3); - if (mtmd_gen_audio_read_tensor(mctx, "cond.t3_cond", t3_cond.data(), n_t3) != n_t3 || + if (mtmd_gen_audio_read_tensor(mctx, "a.gen.cond.t3_cond", t3_cond.data(), n_t3) != n_t3 || n_t3 % (size_t) n_embd != 0) { - LOG_ERR("mtmd_helper_gen_audio: cond.t3_cond read failed\n"); + LOG_ERR("mtmd_helper_gen_audio: a.gen.cond.t3_cond read failed\n"); return false; } auto read_table = [&](const char * name, std::vector & dst) { @@ -756,90 +725,48 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } return true; }; - if (!read_table("t3.text_pos_emb", text_pos) || !read_table("t3.speech_pos_emb", speech_pos)) { + if (!read_table("a.gen.t3.text_pos_emb", text_pos) || !read_table("a.gen.t3.speech_pos_emb", speech_pos)) { return false; } return true; } cond_spkr.resize((size_t) n_embd); - if (mtmd_gen_audio_read_tensor(mctx, "cond.spkr_default", cond_spkr.data(), cond_spkr.size()) != (size_t) n_embd) { - LOG_ERR("mtmd_helper_gen_audio: cond.spkr_default missing\n"); + if (mtmd_gen_audio_read_tensor(mctx, "a.gen.cond.spkr_default", cond_spkr.data(), cond_spkr.size()) != (size_t) n_embd) { + LOG_ERR("mtmd_helper_gen_audio: a.gen.cond.spkr_default missing\n"); return false; } - size_t n_ct = mtmd_gen_audio_read_tensor(mctx, "cond.prompt_speech_tokens", nullptr, 0); - cond_speech_tokens.resize(n_ct); - if (n_ct == 0 || mtmd_gen_audio_read_tensor(mctx, "cond.prompt_speech_tokens", cond_speech_tokens.data(), n_ct) != n_ct) { - LOG_ERR("mtmd_helper_gen_audio: cond.prompt_speech_tokens missing\n"); + // resolve the precomputed conditioning ids through the speech + // embedding table shipped in the mmproj + size_t n_ct = mtmd_gen_audio_read_tensor(mctx, "a.gen.cond.prompt_speech_tokens", nullptr, 0); + std::vector cond_ids(n_ct); + if (n_ct == 0 || mtmd_gen_audio_read_tensor(mctx, "a.gen.cond.prompt_speech_tokens", cond_ids.data(), n_ct) != n_ct) { + LOG_ERR("mtmd_helper_gen_audio: a.gen.cond.prompt_speech_tokens missing\n"); return false; } - return true; - } - - // runs the s3 tokenizer on the reference clip, capped to the reference - // conditioning length, into speech tokens for the flow prompt (10 s cap) - // and the talker conditioning (6 s cap) - bool tokenize_ref(mtmd_bitmap * bitmap, size_t n_cap, std::vector & out) { - const float * pcm = (const float *) mtmd_bitmap_get_data(bitmap); - const size_t n = mtmd_bitmap_get_n_bytes(bitmap) / sizeof(float); - - mtmd_gen_inp gi{}; - gi.type = MTMD_GEN_PROCESS_TYPE_TOKENIZE; - gi.pcm = pcm; - gi.n_pcm = std::min(n, n_cap); - mtmd_gen_out go{}; - if (mtmd_gen_audio_process(mctx, &gi, &go) != 0) { - LOG_ERR("mtmd_helper_gen_audio: reference tokenize failed\n"); + const size_t n_tab = mtmd_gen_audio_read_tensor(mctx, "a.gen.code.out_embd.weight", nullptr, 0); + std::vector table(n_tab); + if (n_tab == 0 || n_tab % (size_t) n_embd != 0 || + mtmd_gen_audio_read_tensor(mctx, "a.gen.code.out_embd.weight", table.data(), n_tab) != n_tab) { + LOG_ERR("mtmd_helper_gen_audio: a.gen.code.out_embd.weight read failed\n"); return false; } - out.assign(go.codes, go.codes + go.n_codes); - return true; - } - - // runs the speaker encoder on the reference clip through the standard - // audio chunk path; the CAMPPlus graph outputs the 80-dim s3gen vector - bool encode_speaker(mtmd_bitmap * bitmap, std::vector & out) { - if (!mtmd_support_audio(mctx)) { - LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n"); - return false; - } - const std::string marker = mtmd_default_marker(); - mtmd_input_text text{ marker.c_str(), marker.size(), false, true }; - mtmd_input_chunks * chunks = mtmd_input_chunks_init(); - const mtmd_bitmap * bptr = bitmap; - bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0; - if (ok) { - ok = false; - for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) { - const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i); - if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) { - continue; - } - if (mtmd_encode_chunk(mctx, chunk) != 0) { - LOG_ERR("mtmd_helper_gen_audio: speaker encode failed\n"); - break; - } - const float * embd = mtmd_get_output_embd(mctx); - out.assign(embd, embd + 80); - ok = true; - break; - } + cond_speech_rows.resize(n_ct * (size_t) n_embd); + for (size_t i = 0; i < n_ct; i++) { + const size_t r = (size_t) cond_ids[i] * (size_t) n_embd; + memcpy(cond_speech_rows.data() + i * (size_t) n_embd, table.data() + r, (size_t) n_embd * sizeof(float)); } - mtmd_input_chunks_free(chunks); - return ok; + return true; } std::vector tok_embd; std::vector cond_spkr; - std::vector cond_speech_tokens; + std::vector cond_speech_rows; // default conditioning ids resolved through the mmproj speech table std::vector t3_cond; std::vector text_pos; std::vector speech_pos; - std::vector spk80; - std::vector ref_prompt_tokens; - std::vector ref_t3_tokens; - std::vector ref_pcm16; std::vector ref_cond; + std::vector ref_state; llama_token speech_base = LLAMA_TOKEN_NULL; int n_speech = 0; llama_token text_start = LLAMA_TOKEN_NULL; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index a63fdc14d2c4..958661ea9850 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1622,18 +1622,13 @@ size_t mtmd_gen_audio_read_tensor(mtmd_context * ctx, const char * name, float * return clip_cbx_read_tensor(ctx->ctx_gen_a, name, out, n_max); } -void mtmd_gen_audio_norm_ref(mtmd_context * ctx, mtmd_bitmap * bitmap) { - if (!ctx->ctx_gen_a || !bitmap || !bitmap->is_audio || bitmap->is_placeholder()) { +// turbo only: -27 LUFS loudness normalization of the reference clip +// (tts_turbo.py); the multilingual variant is identified by its perceiver +static void cbx_ref_norm(mtmd_context * ctx, float * pcm, size_t n) { + if (clip_cbx_read_tensor(ctx->ctx_gen_a, "a.cenc.perceiver.pre_attention_query", nullptr, 0) > 0 || + clip_cbx_read_tensor(ctx->ctx_gen_a, "a.cenc.spkr_enc.weight", nullptr, 0) == 0) { return; } - // only the turbo variant normalizes the reference (tts_turbo.py, -27 LUFS); - // the multilingual variant is identified by its perceiver - if (clip_cbx_read_tensor(ctx->ctx_gen_a, "cenc.perceiver.pre_attention_query", nullptr, 0) > 0 || - clip_cbx_read_tensor(ctx->ctx_gen_a, "cenc.spkr_enc.weight", nullptr, 0) == 0) { - return; - } - float * pcm = (float *) bitmap->get_rw_buf().data(); - const size_t n = bitmap->n_bytes() / sizeof(float); const float lufs = mtmd_audio_lufs(pcm, n, clip_get_hparams(ctx->ctx_a ? ctx->ctx_a : ctx->ctx_gen_a)->audio_sample_rate); if (lufs == -HUGE_VALF) { return; @@ -1647,6 +1642,350 @@ void mtmd_gen_audio_norm_ref(mtmd_context * ctx, mtmd_bitmap * bitmap) { } } +// 80-dim flow speaker vector of the reference clip: fbank features from the +// audio preprocessor into the CAMPPlus graph of the speaker encoding context +static bool cbx_ref_spk80(mtmd_context * ctx, const float * pcm, size_t n_pcm, std::vector & spk80) { + if (!ctx->ctx_a || !ctx->audio_preproc) { + LOG_ERR("%s: mmproj has no speaker encoder\n", __func__); + return false; + } + std::vector mels; + if (!ctx->audio_preproc->preprocess(pcm, n_pcm, mels) || mels.empty() || mels[0].data.empty()) { + LOG_ERR("%s: speaker features failed\n", __func__); + return false; + } + clip_image_f32 mel_img; + mel_img.set_size({(int) mels[0].n_len, (int) mels[0].n_mel}, false, true); + mel_img.cpy_buf(std::move(mels[0].data)); + clip_image_f32_batch batch; + batch.is_audio = true; + batch.entries.push_back(std::move(mel_img)); + spk80.resize((size_t) clip_n_mmproj_embd(ctx->ctx_a)); + if (!clip_image_batch_encode(ctx->ctx_a, ctx->n_threads, &batch, spk80)) { + LOG_ERR("%s: speaker encoder failed\n", __func__); + return false; + } + return true; +} + +// runs the s3 speech tokenizer on a reference clip; 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 float * ref, size_t n_ref, + std::vector & codes, std::vector * rows) { +clip_ctx * ctx_clip = ctx->ctx_gen_a; + // mel filters shipped in the mmproj, [n_mels x (n_fft / 2 + 1)] + const size_t n_filt = clip_cbx_read_tensor(ctx_clip, "a.s3tok.mel_filters", nullptr, 0); + if (n_filt == 0) { + LOG_ERR("%s: model has no s3 tokenizer\n", __func__); + return 1; + } + std::vector filters(n_filt); + clip_cbx_read_tensor(ctx_clip, "a.s3tok.mel_filters", filters.data(), n_filt); + const int n_mel = (int) (n_filt / (400 / 2 + 1)); + + // pad to a whole number of 40 ms tokens so that the mel length stays + // twice the token length, as the reference prompt features expect + std::vector pcm(ref, ref + n_ref); + pcm.resize((pcm.size() + 639) / 640 * 640, 0.0f); + + std::vector mel; + int n_frames = 0; + if (!mtmd_audio_s3tok_log_mel(pcm.data(), pcm.size(), filters.data(), n_mel, mel, n_frames)) { + LOG_ERR("%s: log mel failed\n", __func__); + return 1; + } + + clip_image_f32 mel_img; + mel_img.set_size({n_frames, n_mel}, false, true); + mel_img.cpy_buf(std::move(mel)); + + clip_image_f32_batch batch; + batch.is_audio = true; + batch.entries.push_back(std::move(mel_img)); + + std::vector out_codes; + std::vector out_embd; + + clip_encode_params params; + params.imgs = &batch; + params.n_threads = ctx->n_threads; + params.gen_process = CLIP_GEN_PROCESS_TOKENIZE; + params.out_codes = &out_codes; + params.out_code_embd = rows ? &out_embd : nullptr; + + if (!clip_encode(ctx_clip, ¶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 float * pcm, size_t n_pcm, + const std::vector & pse, std::vector & out_rows) { +clip_ctx * ctx_clip = ctx->ctx_gen_a; + auto read_t = [&](const char * name, std::vector & v) -> bool { + const size_t n = clip_cbx_read_tensor(ctx_clip, name, nullptr, 0); + if (n == 0) { + return false; + } + v.resize(n); + return clip_cbx_read_tensor(ctx_clip, name, v.data(), n) == n; + }; + + // voice encoder reference chain (embeds_from_wavs): silence trim, + // 40-bin power mel, overlapping 160-frame partials at rate 1.3, + // 3-layer lstm per partial, projected/relu/normalized embeddings + // averaged into the utterance embedding + size_t t0 = 0, t1 = 0; + mtmd_audio_trim_silence(pcm, n_pcm, 20.0f, t0, t1); + if (t1 <= t0) { + LOG_ERR("%s: reference clip is silent\n", __func__); + return 1; + } + std::vector mel; + int n_frames = 0; + if (!mtmd_audio_ve_mel(pcm + t0, t1 - t0, mel, n_frames)) { + LOG_ERR("%s: voice encoder mel failed\n", __func__); + return 1; + } + + 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 + 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); // zero pad (or trim) to the partial grid + + std::vector w_ih[3], w_hh[3], b_ih[3], b_hh[3]; + std::vector w_proj, b_proj; + for (int l = 0; l < 3; l++) { + const std::string s = std::to_string(l); + if (!read_t(("a.ve.lstm.weight_ih_l" + s).c_str(), w_ih[l]) || + !read_t(("a.ve.lstm.weight_hh_l" + s).c_str(), w_hh[l]) || + !read_t(("a.ve.lstm.bias_ih_l" + s).c_str(), b_ih[l]) || + !read_t(("a.ve.lstm.bias_hh_l" + s).c_str(), b_hh[l])) { + LOG_ERR("%s: model has no voice encoder\n", __func__); + return 1; + } + } + if (!read_t("a.ve.proj.weight", w_proj) || !read_t("a.ve.proj.bias", b_proj)) { + LOG_ERR("%s: model has no voice encoder projection\n", __func__); + return 1; + } + + const int n_h = 256; + std::vector ve(n_h, 0.0f); + std::vector h((size_t) 3 * n_h), c((size_t) 3 * n_h), x(n_h), g((size_t) 4 * n_h); + for (int p = 0; p < n_wins; p++) { + std::fill(h.begin(), h.end(), 0.0f); + std::fill(c.begin(), c.end(), 0.0f); + for (int t = 0; t < n_partial; t++) { + const float * in = mel.data() + (size_t) (p * step + t) * n_mel; + int n_in = n_mel; + for (int l = 0; l < 3; l++) { + float * hl = h.data() + (size_t) l * n_h; + float * cl = c.data() + (size_t) l * n_h; + for (int j = 0; j < 4 * n_h; j++) { + double acc = b_ih[l][(size_t) j] + b_hh[l][(size_t) j]; + const float * wi = w_ih[l].data() + (size_t) j * n_in; + for (int i = 0; i < n_in; i++) { + acc += (double) wi[i] * in[i]; + } + const float * wh = w_hh[l].data() + (size_t) j * n_h; + for (int i = 0; i < n_h; i++) { + acc += (double) wh[i] * hl[i]; + } + g[(size_t) j] = (float) acc; + } + // torch gate order: input, forget, cell, output + for (int i = 0; i < n_h; i++) { + const float gi = 1.0f / (1.0f + expf(-g[(size_t) i])); + const float gf = 1.0f / (1.0f + expf(-g[(size_t) i + n_h])); + const float gc = tanhf(g[(size_t) i + 2 * n_h]); + const float go = 1.0f / (1.0f + expf(-g[(size_t) i + 3 * n_h])); + cl[i] = gf * cl[i] + gi * gc; + x[(size_t) i] = go * tanhf(cl[i]); + } + memcpy(hl, x.data(), (size_t) n_h * sizeof(float)); + in = hl; + n_in = n_h; + } + } + // projected, relu'd, normalized partial embedding + std::vector e(n_h); + double norm = 0.0; + for (int o = 0; o < n_h; o++) { + double acc = b_proj[(size_t) o]; + const float * w = w_proj.data() + (size_t) o * n_h; + const float * hl = h.data() + (size_t) 2 * n_h; + for (int i = 0; i < n_h; i++) { + acc += (double) w[i] * hl[i]; + } + e[(size_t) o] = (float) std::max(acc, 0.0); + norm += (double) e[(size_t) o] * e[(size_t) o]; + } + norm = sqrt(norm); + for (int o = 0; o < n_h; o++) { + ve[(size_t) o] += (float) (e[(size_t) o] / norm); + } + } + double norm = 0.0; + for (float v : ve) { + norm += (double) v * v; + } + norm = sqrt(norm); + for (float & v : ve) { + v = (float) (v / norm); + } + + // speaker projection row + std::vector w_spkr, b_spkr; + if (!read_t("a.cenc.spkr_enc.weight", w_spkr) || !read_t("a.cenc.spkr_enc.bias", b_spkr)) { + LOG_ERR("%s: model has no speaker conditioning projection\n", __func__); + return 1; + } + const int n_e = (int) b_spkr.size(); + std::vector rows((size_t) n_e); + for (int o = 0; o < n_e; o++) { + double acc = b_spkr[(size_t) o]; + const float * w = w_spkr.data() + (size_t) o * n_h; + for (int i = 0; i < n_h; i++) { + acc += (double) w[i] * ve[(size_t) i]; + } + rows[(size_t) o] = (float) acc; + } + + // multilingual variant: [spkr, perceiver x32, emotion] block, the + // perceiver runs its shared attention block as cross then self + // attention over the reference speech embedding rows + std::vector query; + if (read_t("a.cenc.perceiver.pre_attention_query", query)) { + if (pse.empty()) { + LOG_ERR("%s: reference speech embeddings required for the perceiver\n", __func__); + return 1; + } + std::vector ln_w, ln_b, wq, bq, wk, bk, wv, bv, wo, bo, emo; + if (!read_t("a.cenc.perceiver.attn.norm.weight", ln_w) || !read_t("a.cenc.perceiver.attn.norm.bias", ln_b) || + !read_t("a.cenc.perceiver.attn.to_q.weight", wq) || !read_t("a.cenc.perceiver.attn.to_q.bias", bq) || + !read_t("a.cenc.perceiver.attn.to_k.weight", wk) || !read_t("a.cenc.perceiver.attn.to_k.bias", bk) || + !read_t("a.cenc.perceiver.attn.to_v.weight", wv) || !read_t("a.cenc.perceiver.attn.to_v.bias", bv) || + !read_t("a.cenc.perceiver.attn.proj_out.weight", wo) || !read_t("a.cenc.perceiver.attn.proj_out.bias", bo) || + !read_t("a.cenc.emotion_adv_fc.weight", emo)) { + LOG_ERR("%s: model has an incomplete perceiver\n", __func__); + return 1; + } + const int n_head = 4; + const int d_head = n_e / n_head; + + auto layer_norm = [&](const float * in, float * out) { + double mean = 0.0, var = 0.0; + for (int i = 0; i < n_e; i++) { + mean += in[i]; + } + mean /= n_e; + for (int i = 0; i < n_e; i++) { + var += (in[i] - mean) * (in[i] - mean); + } + const double sd = sqrt(var / n_e + 1e-5); + for (int i = 0; i < n_e; i++) { + out[i] = (float) ((in[i] - mean) / sd * ln_w[(size_t) i] + ln_b[(size_t) i]); + } + }; + auto linear = [&](const std::vector & w, const std::vector & b, + const std::vector & in, int n_rows, std::vector & out) { + out.resize((size_t) n_rows * n_e); + for (int r = 0; r < n_rows; r++) { + for (int o = 0; o < n_e; o++) { + double acc = b[(size_t) o]; + const float * wr = w.data() + (size_t) o * n_e; + const float * ir = in.data() + (size_t) r * n_e; + for (int i = 0; i < n_e; i++) { + acc += (double) wr[i] * ir[i]; + } + out[(size_t) r * n_e + o] = (float) acc; + } + } + }; + auto attn_block = [&](const std::vector & x1, int n1, + const std::vector & x2, int n2, std::vector & out) { + std::vector nx1((size_t) n1 * n_e), nx2((size_t) n2 * n_e); + for (int r = 0; r < n1; r++) { + layer_norm(x1.data() + (size_t) r * n_e, nx1.data() + (size_t) r * n_e); + } + for (int r = 0; r < n2; r++) { + layer_norm(x2.data() + (size_t) r * n_e, nx2.data() + (size_t) r * n_e); + } + std::vector q, k, v; + linear(wq, bq, nx1, n1, q); + linear(wk, bk, nx2, n2, k); + linear(wv, bv, nx2, n2, v); + + std::vector ctxt((size_t) n1 * n_e); + std::vector sc((size_t) n2); + for (int hd = 0; hd < n_head; hd++) { + const int off = hd * d_head; + for (int t = 0; t < n1; t++) { + double mx = -1e30; + for (int s = 0; s < n2; s++) { + double acc = 0.0; + for (int i = 0; i < d_head; i++) { + acc += (double) q[(size_t) t * n_e + off + i] * k[(size_t) s * n_e + off + i]; + } + sc[(size_t) s] = acc / sqrt((double) d_head); + mx = std::max(mx, sc[(size_t) s]); + } + double sum = 0.0; + for (int s = 0; s < n2; s++) { + sc[(size_t) s] = exp(sc[(size_t) s] - mx); + sum += sc[(size_t) s]; + } + for (int i = 0; i < d_head; i++) { + double acc = 0.0; + for (int s = 0; s < n2; s++) { + acc += sc[(size_t) s] * v[(size_t) s * n_e + off + i]; + } + ctxt[(size_t) t * n_e + off + i] = (float) (acc / sum); + } + } + } + linear(wo, bo, ctxt, n1, out); + for (size_t i = 0; i < out.size(); i++) { + out[i] += x1[i]; + } + }; + + const int n_q = (int) (query.size() / n_e); + const std::vector & x2 = pse; + std::vector pre, p32; + attn_block(query, n_q, x2, (int) (pse.size() / (size_t) n_e), pre); + attn_block(pre, n_q, pre, n_q, p32); + + rows.insert(rows.end(), p32.begin(), p32.end()); + const float exaggeration = 0.5f; // reference default + for (int i = 0; i < n_e; i++) { + rows.push_back(emo[(size_t) i] * exaggeration); + } + } + + out_rows = std::move(rows); + return 0; +} + + static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) { clip_ctx * ctx_clip = ctx->ctx_gen_a; if (!ctx_clip) { @@ -1692,43 +2031,128 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in return 0; } - if (inp->type == MTMD_GEN_PROCESS_TYPE_TTS) { - if (!inp->codes || inp->n_codes == 0) { - LOG_ERR("%s: codes required for tts\n", __func__); + if (inp->type == MTMD_GEN_PROCESS_TYPE_SPK_REF) { + if (!inp->pcm || inp->n_pcm == 0) { + LOG_ERR("%s: pcm required for spk ref\n", __func__); return 1; } - std::vector in_codes(inp->codes, inp->codes + inp->n_codes); - std::vector out_audio; - std::vector ref_tokens; - std::vector ref_feat; - std::vector ref_spk; - if (inp->ref_tokens) { - ref_tokens.assign(inp->ref_tokens, inp->ref_tokens + inp->n_ref_tokens); + // the whole reference encoding chain runs inside this stage: + // loudness normalization, speaker vector, speech tokenization at the + // flow and talker caps, conditioning rows, mel-rate prompt features + std::vector pcm(inp->pcm, inp->pcm + inp->n_pcm); + cbx_ref_norm(ctx, pcm.data(), pcm.size()); + + const bool mtl = clip_cbx_read_tensor(ctx_clip, "a.cenc.perceiver.pre_attention_query", nullptr, 0) > 0; + const size_t gen_cap = (size_t) 10 * 16000; + const size_t t3_cap = (size_t) (mtl ? 6 : 15) * 16000; + + std::vector spk80; + if (!cbx_ref_spk80(ctx, pcm.data(), pcm.size(), spk80)) { + return 1; + } + + std::vector flow_tokens, t3_tokens; + std::vector t3_rows; + if (cbx_ref_tokenize(ctx, pcm.data(), std::min(pcm.size(), gen_cap), flow_tokens, nullptr) != 0 || + cbx_ref_tokenize(ctx, pcm.data(), std::min(pcm.size(), t3_cap), t3_tokens, &t3_rows) != 0) { + return 1; + } + + // conditioning rows: [spkr] then, multilingual, the perceiver block + // over the reference rows, or, turbo, the raw reference rows + std::vector pse; + if (mtl) { + pse = std::move(t3_rows); + } + std::vector rows; + if (cbx_ref_cond(ctx, pcm.data(), pcm.size(), pse, rows) != 0) { + return 1; } - if (inp->ref_spk) { - ref_spk.assign(inp->ref_spk, inp->ref_spk + 80); + if (!mtl) { + rows.insert(rows.end(), t3_rows.begin(), t3_rows.end()); } - if (inp->ref_pcm) { - // mel-rate prompt features of the reference clip at the 24 kHz - // s3gen rate, padded to the token grid so that the mel length - // stays twice the token length - std::vector pcm16(inp->ref_pcm, inp->ref_pcm + inp->n_ref_pcm); + + // mel-rate prompt features of the flow reference at the 24 kHz s3gen + // rate, padded to the token grid so that the mel length stays twice + // the token length + std::vector feat; + { + std::vector pcm16(pcm.begin(), pcm.begin() + std::min(pcm.size(), gen_cap)); pcm16.resize((pcm16.size() + 639) / 640 * 640, 0.0f); std::vector pcm24; mtmd_audio_upsample_3_2(pcm16.data(), pcm16.size(), pcm24); int n_feat = 0; - if (!mtmd_audio_matcha_log_mel(pcm24.data(), pcm24.size(), ref_feat, n_feat)) { + if (!mtmd_audio_matcha_log_mel(pcm24.data(), pcm24.size(), feat, n_feat)) { LOG_ERR("%s: reference mel failed\n", __func__); return 1; } - if ((size_t) n_feat != 2 * ref_tokens.size()) { + if ((size_t) n_feat != 2 * flow_tokens.size()) { LOG_ERR("%s: reference mel length %d does not match %zu tokens\n", - __func__, n_feat, ref_tokens.size()); + __func__, n_feat, flow_tokens.size()); return 1; } } + // decoder reference state, opaque to the caller: + // [n_tokens, n_feat] i32 header, tokens, features, 80-dim speaker vector + std::vector blob(2 * sizeof(int32_t) + + flow_tokens.size() * sizeof(int32_t) + + (feat.size() + spk80.size()) * sizeof(float)); + { + uint8_t * q = blob.data(); + const int32_t hdr[2] = { (int32_t) flow_tokens.size(), (int32_t) feat.size() }; + memcpy(q, hdr, sizeof(hdr)); q += sizeof(hdr); + memcpy(q, flow_tokens.data(), flow_tokens.size() * sizeof(int32_t)); q += flow_tokens.size() * sizeof(int32_t); + memcpy(q, feat.data(), feat.size() * sizeof(float)); q += feat.size() * sizeof(float); + memcpy(q, spk80.data(), spk80.size() * sizeof(float)); + } + + ctx->gen_out_embd = std::move(rows); + ctx->gen_out_state = std::move(blob); + out->embd = ctx->gen_out_embd.data(); + out->n_embd = ctx->gen_out_embd.size(); + out->state_data = (const char *) ctx->gen_out_state.data(); + out->state_size = ctx->gen_out_state.size(); + return 0; + } + + // MTMD_GEN_PROCESS_TYPE_CODE2WAV + if (clip_get_projector_type(ctx_clip) == PROJECTOR_TYPE_CHATTERBOX) { + if (!inp->codes || inp->n_codes == 0) { + LOG_ERR("%s: codes required for code2wav\n", __func__); + return 1; + } + std::vector in_codes(inp->codes, inp->codes + inp->n_codes); + std::vector out_audio; + + // optional reference state from the spk-ref stage: flow prompt + // tokens, mel-rate prompt features, 80-dim speaker vector; a null + // state selects the model's precomputed default voice + std::vector ref_tokens; + std::vector ref_feat; + std::vector ref_spk; + if (inp->state_data) { + int32_t hdr[2]; + if (inp->state_size < sizeof(hdr)) { + LOG_ERR("%s: malformed reference state\n", __func__); + return 1; + } + memcpy(hdr, inp->state_data, sizeof(hdr)); + const size_t n_tok = (size_t) hdr[0], n_feat = (size_t) hdr[1]; + if (inp->state_size != sizeof(hdr) + n_tok * sizeof(int32_t) + (n_feat + 80) * sizeof(float)) { + LOG_ERR("%s: malformed reference state\n", __func__); + return 1; + } + const char * q = inp->state_data + sizeof(hdr); + ref_tokens.resize(n_tok); + memcpy(ref_tokens.data(), q, n_tok * sizeof(int32_t)); q += n_tok * sizeof(int32_t); + ref_feat.resize(n_feat); + memcpy(ref_feat.data(), q, n_feat * sizeof(float)); q += n_feat * sizeof(float); + ref_spk.resize(80); + memcpy(ref_spk.data(), q, 80 * sizeof(float)); + } + // the batch entry is unused, present to satisfy the encode interface clip_image_f32 dummy; dummy.set_size({1, 1}, false, true); @@ -1743,12 +2167,12 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in params.gen_process = CLIP_GEN_PROCESS_TTS; params.codes = &in_codes; params.out_audio = &out_audio; - params.ref_tokens = inp->ref_tokens ? &ref_tokens : nullptr; - params.ref_feat = inp->ref_pcm ? &ref_feat : nullptr; - params.ref_spk = inp->ref_spk ? &ref_spk : nullptr; + params.ref_tokens = inp->state_data ? &ref_tokens : nullptr; + params.ref_feat = inp->state_data ? &ref_feat : nullptr; + params.ref_spk = inp->state_data ? &ref_spk : nullptr; if (!clip_encode(ctx_clip, ¶ms)) { - LOG_ERR("%s: clip_encode failed (tts)\n", __func__); + LOG_ERR("%s: clip_encode failed (code2wav)\n", __func__); return 1; } @@ -1758,322 +2182,6 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in return 0; } - if (inp->type == MTMD_GEN_PROCESS_TYPE_TOKENIZE) { - if (!inp->pcm || inp->n_pcm == 0) { - LOG_ERR("%s: pcm required for tokenize\n", __func__); - return 1; - } - - // mel filters shipped in the mmproj, [n_mels x (n_fft / 2 + 1)] - const size_t n_filt = clip_cbx_read_tensor(ctx_clip, "s3tok.mel_filters", nullptr, 0); - if (n_filt == 0) { - LOG_ERR("%s: model has no s3 tokenizer\n", __func__); - return 1; - } - std::vector filters(n_filt); - clip_cbx_read_tensor(ctx_clip, "s3tok.mel_filters", filters.data(), n_filt); - const int n_mel = (int) (n_filt / (400 / 2 + 1)); - - // pad to a whole number of 40 ms tokens so that the mel length stays - // twice the token length, as the reference prompt features expect - std::vector pcm(inp->pcm, inp->pcm + inp->n_pcm); - pcm.resize((pcm.size() + 639) / 640 * 640, 0.0f); - - std::vector mel; - int n_frames = 0; - if (!mtmd_audio_s3tok_log_mel(pcm.data(), pcm.size(), filters.data(), n_mel, mel, n_frames)) { - LOG_ERR("%s: log mel failed\n", __func__); - return 1; - } - - clip_image_f32 mel_img; - mel_img.set_size({n_frames, n_mel}, false, true); - mel_img.cpy_buf(std::move(mel)); - - clip_image_f32_batch batch; - batch.is_audio = true; - batch.entries.push_back(std::move(mel_img)); - - std::vector out_codes; - - 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; - - if (!clip_encode(ctx_clip, ¶ms)) { - LOG_ERR("%s: clip_encode failed (tokenize)\n", __func__); - return 1; - } - - ctx->gen_out_codes = std::move(out_codes); - out->codes = ctx->gen_out_codes.data(); - out->n_codes = ctx->gen_out_codes.size(); - return 0; - } - - if (inp->type == MTMD_GEN_PROCESS_TYPE_SPEAKER_COND) { - if (!inp->pcm || inp->n_pcm == 0) { - LOG_ERR("%s: pcm required for speaker cond\n", __func__); - return 1; - } - - auto read_t = [&](const char * name, std::vector & v) -> bool { - const size_t n = clip_cbx_read_tensor(ctx_clip, name, nullptr, 0); - if (n == 0) { - return false; - } - v.resize(n); - return clip_cbx_read_tensor(ctx_clip, name, v.data(), n) == n; - }; - - // voice encoder reference chain (embeds_from_wavs): silence trim, - // 40-bin power mel, overlapping 160-frame partials at rate 1.3, - // 3-layer lstm per partial, projected/relu/normalized embeddings - // averaged into the utterance embedding - size_t t0 = 0, t1 = 0; - mtmd_audio_trim_silence(inp->pcm, inp->n_pcm, 20.0f, t0, t1); - if (t1 <= t0) { - LOG_ERR("%s: reference clip is silent\n", __func__); - return 1; - } - std::vector mel; - int n_frames = 0; - if (!mtmd_audio_ve_mel(inp->pcm + t0, t1 - t0, mel, n_frames)) { - LOG_ERR("%s: voice encoder mel failed\n", __func__); - return 1; - } - - 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 - 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); // zero pad (or trim) to the partial grid - - std::vector w_ih[3], w_hh[3], b_ih[3], b_hh[3]; - std::vector w_proj, b_proj; - for (int l = 0; l < 3; l++) { - const std::string s = std::to_string(l); - if (!read_t(("ve.lstm.weight_ih_l" + s).c_str(), w_ih[l]) || - !read_t(("ve.lstm.weight_hh_l" + s).c_str(), w_hh[l]) || - !read_t(("ve.lstm.bias_ih_l" + s).c_str(), b_ih[l]) || - !read_t(("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("ve.proj.weight", w_proj) || !read_t("ve.proj.bias", b_proj)) { - LOG_ERR("%s: model has no voice encoder projection\n", __func__); - return 1; - } - - const int n_h = 256; - std::vector ve(n_h, 0.0f); - std::vector h((size_t) 3 * n_h), c((size_t) 3 * n_h), x(n_h), g((size_t) 4 * n_h); - for (int p = 0; p < n_wins; p++) { - std::fill(h.begin(), h.end(), 0.0f); - std::fill(c.begin(), c.end(), 0.0f); - for (int t = 0; t < n_partial; t++) { - const float * in = mel.data() + (size_t) (p * step + t) * n_mel; - int n_in = n_mel; - for (int l = 0; l < 3; l++) { - float * hl = h.data() + (size_t) l * n_h; - float * cl = c.data() + (size_t) l * n_h; - for (int j = 0; j < 4 * n_h; j++) { - double acc = b_ih[l][(size_t) j] + b_hh[l][(size_t) j]; - const float * wi = w_ih[l].data() + (size_t) j * n_in; - for (int i = 0; i < n_in; i++) { - acc += (double) wi[i] * in[i]; - } - const float * wh = w_hh[l].data() + (size_t) j * n_h; - for (int i = 0; i < n_h; i++) { - acc += (double) wh[i] * hl[i]; - } - g[(size_t) j] = (float) acc; - } - // torch gate order: input, forget, cell, output - for (int i = 0; i < n_h; i++) { - const float gi = 1.0f / (1.0f + expf(-g[(size_t) i])); - const float gf = 1.0f / (1.0f + expf(-g[(size_t) i + n_h])); - const float gc = tanhf(g[(size_t) i + 2 * n_h]); - const float go = 1.0f / (1.0f + expf(-g[(size_t) i + 3 * n_h])); - cl[i] = gf * cl[i] + gi * gc; - x[(size_t) i] = go * tanhf(cl[i]); - } - memcpy(hl, x.data(), (size_t) n_h * sizeof(float)); - in = hl; - n_in = n_h; - } - } - // projected, relu'd, normalized partial embedding - std::vector e(n_h); - double norm = 0.0; - for (int o = 0; o < n_h; o++) { - double acc = b_proj[(size_t) o]; - const float * w = w_proj.data() + (size_t) o * n_h; - const float * hl = h.data() + (size_t) 2 * n_h; - for (int i = 0; i < n_h; i++) { - acc += (double) w[i] * hl[i]; - } - e[(size_t) o] = (float) std::max(acc, 0.0); - norm += (double) e[(size_t) o] * e[(size_t) o]; - } - norm = sqrt(norm); - for (int o = 0; o < n_h; o++) { - ve[(size_t) o] += (float) (e[(size_t) o] / norm); - } - } - double norm = 0.0; - for (float v : ve) { - norm += (double) v * v; - } - norm = sqrt(norm); - for (float & v : ve) { - v = (float) (v / norm); - } - - // speaker projection row - std::vector w_spkr, b_spkr; - if (!read_t("cenc.spkr_enc.weight", w_spkr) || !read_t("cenc.spkr_enc.bias", b_spkr)) { - LOG_ERR("%s: model has no speaker conditioning projection\n", __func__); - return 1; - } - const int n_e = (int) b_spkr.size(); - std::vector rows((size_t) n_e); - for (int o = 0; o < n_e; o++) { - double acc = b_spkr[(size_t) o]; - const float * w = w_spkr.data() + (size_t) o * n_h; - for (int i = 0; i < n_h; i++) { - acc += (double) w[i] * ve[(size_t) i]; - } - rows[(size_t) o] = (float) acc; - } - - // multilingual variant: [spkr, perceiver x32, emotion] block, the - // perceiver runs its shared attention block as cross then self - // attention over the reference speech embedding rows - std::vector query; - if (read_t("cenc.perceiver.pre_attention_query", query)) { - if (!inp->ref_speech_embd || inp->n_ref_speech_rows == 0) { - LOG_ERR("%s: reference speech embeddings required for the perceiver\n", __func__); - return 1; - } - std::vector ln_w, ln_b, wq, bq, wk, bk, wv, bv, wo, bo, emo; - if (!read_t("cenc.perceiver.attn.norm.weight", ln_w) || !read_t("cenc.perceiver.attn.norm.bias", ln_b) || - !read_t("cenc.perceiver.attn.to_q.weight", wq) || !read_t("cenc.perceiver.attn.to_q.bias", bq) || - !read_t("cenc.perceiver.attn.to_k.weight", wk) || !read_t("cenc.perceiver.attn.to_k.bias", bk) || - !read_t("cenc.perceiver.attn.to_v.weight", wv) || !read_t("cenc.perceiver.attn.to_v.bias", bv) || - !read_t("cenc.perceiver.attn.proj_out.weight", wo) || !read_t("cenc.perceiver.attn.proj_out.bias", bo) || - !read_t("cenc.emotion_adv_fc.weight", emo)) { - LOG_ERR("%s: model has an incomplete perceiver\n", __func__); - return 1; - } - const int n_head = 4; - const int d_head = n_e / n_head; - - auto layer_norm = [&](const float * in, float * out) { - double mean = 0.0, var = 0.0; - for (int i = 0; i < n_e; i++) { - mean += in[i]; - } - mean /= n_e; - for (int i = 0; i < n_e; i++) { - var += (in[i] - mean) * (in[i] - mean); - } - const double sd = sqrt(var / n_e + 1e-5); - for (int i = 0; i < n_e; i++) { - out[i] = (float) ((in[i] - mean) / sd * ln_w[(size_t) i] + ln_b[(size_t) i]); - } - }; - auto linear = [&](const std::vector & w, const std::vector & b, - const std::vector & in, int n_rows, std::vector & out) { - out.resize((size_t) n_rows * n_e); - for (int r = 0; r < n_rows; r++) { - for (int o = 0; o < n_e; o++) { - double acc = b[(size_t) o]; - const float * wr = w.data() + (size_t) o * n_e; - const float * ir = in.data() + (size_t) r * n_e; - for (int i = 0; i < n_e; i++) { - acc += (double) wr[i] * ir[i]; - } - out[(size_t) r * n_e + o] = (float) acc; - } - } - }; - auto attn_block = [&](const std::vector & x1, int n1, - const std::vector & x2, int n2, std::vector & out) { - std::vector nx1((size_t) n1 * n_e), nx2((size_t) n2 * n_e); - for (int r = 0; r < n1; r++) { - layer_norm(x1.data() + (size_t) r * n_e, nx1.data() + (size_t) r * n_e); - } - for (int r = 0; r < n2; r++) { - layer_norm(x2.data() + (size_t) r * n_e, nx2.data() + (size_t) r * n_e); - } - std::vector q, k, v; - linear(wq, bq, nx1, n1, q); - linear(wk, bk, nx2, n2, k); - linear(wv, bv, nx2, n2, v); - - std::vector ctxt((size_t) n1 * n_e); - std::vector sc((size_t) n2); - for (int hd = 0; hd < n_head; hd++) { - const int off = hd * d_head; - for (int t = 0; t < n1; t++) { - double mx = -1e30; - for (int s = 0; s < n2; s++) { - double acc = 0.0; - for (int i = 0; i < d_head; i++) { - acc += (double) q[(size_t) t * n_e + off + i] * k[(size_t) s * n_e + off + i]; - } - sc[(size_t) s] = acc / sqrt((double) d_head); - mx = std::max(mx, sc[(size_t) s]); - } - double sum = 0.0; - for (int s = 0; s < n2; s++) { - sc[(size_t) s] = exp(sc[(size_t) s] - mx); - sum += sc[(size_t) s]; - } - for (int i = 0; i < d_head; i++) { - double acc = 0.0; - for (int s = 0; s < n2; s++) { - acc += sc[(size_t) s] * v[(size_t) s * n_e + off + i]; - } - ctxt[(size_t) t * n_e + off + i] = (float) (acc / sum); - } - } - } - linear(wo, bo, ctxt, n1, out); - for (size_t i = 0; i < out.size(); i++) { - out[i] += x1[i]; - } - }; - - const int n_q = (int) (query.size() / n_e); - std::vector x2(inp->ref_speech_embd, inp->ref_speech_embd + (size_t) inp->n_ref_speech_rows * n_e); - std::vector pre, p32; - attn_block(query, n_q, x2, (int) inp->n_ref_speech_rows, 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); - } - } - - ctx->gen_out_embd = std::move(rows); - out->embd = ctx->gen_out_embd.data(); - out->n_embd = ctx->gen_out_embd.size(); - return 0; - } - - // MTMD_GEN_PROCESS_TYPE_CODE2WAV if (!inp->codes || inp->n_codes == 0) { LOG_ERR("%s: codes required for code2wav\n", __func__); return 1; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 008548bebba6..53c6650385a0 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -346,16 +346,13 @@ MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * // 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); -// normalize a reference clip in place to the loudness the gen audio model -// expects for voice cloning; no-op when the model does not require it. -MTMD_API void mtmd_gen_audio_norm_ref(mtmd_context * ctx, mtmd_bitmap * bitmap); + enum mtmd_gen_process_type { MTMD_GEN_PROCESS_TYPE_GEN_CODE, // h_state to codes MTMD_GEN_PROCESS_TYPE_CODE2WAV, // codes to raw PCM audio - MTMD_GEN_PROCESS_TYPE_TTS, // full utterance of codes to raw PCM audio - MTMD_GEN_PROCESS_TYPE_TOKENIZE, // raw PCM audio to semantic speech tokens - MTMD_GEN_PROCESS_TYPE_SPEAKER_COND, // raw PCM audio to talker conditioning rows + MTMD_GEN_PROCESS_TYPE_SPK_REF, // raw PCM audio to talker conditioning rows + // and the decoder reference state }; struct mtmd_gen_inp { enum mtmd_gen_process_type type; @@ -369,29 +366,15 @@ struct mtmd_gen_inp { // for MTMD_GEN_PROCESS_TYPE_CODE2WAV int32_t * codes; size_t n_codes; + // opaque state: the decoder carry-over between calls, or the reference + // state returned by a MTMD_GEN_PROCESS_TYPE_SPK_REF call (null means the + // model's precomputed default voice) const char * state_data; size_t state_size; - // for MTMD_GEN_PROCESS_TYPE_TOKENIZE and MTMD_GEN_PROCESS_TYPE_SPEAKER_COND + // for MTMD_GEN_PROCESS_TYPE_SPK_REF const float * pcm; // mono float samples at the audio encoder sample rate size_t n_pcm; - - // for MTMD_GEN_PROCESS_TYPE_SPEAKER_COND: the speech token embedding rows - // of the reference conditioning prompt (n_text_embd elements each), input - // to the perceiver of the multilingual variant (null on turbo) - const float * ref_speech_embd; - size_t n_ref_speech_rows; - - // for MTMD_GEN_PROCESS_TYPE_TTS: optional reference conditioning of the - // cloned voice, overriding the model's precomputed defaults (null means - // default). ref_spk is the 80-dim speaker vector, ref_tokens the speech - // tokens of the reference clip (from a TOKENIZE call), ref_pcm the same - // clip as mono float samples at the audio encoder sample rate. - const float * ref_spk; - const int32_t * ref_tokens; - size_t n_ref_tokens; - const float * ref_pcm; - size_t n_ref_pcm; }; struct mtmd_gen_out { // note: output memory is allocated by the context, valid until next process() call @@ -401,13 +384,16 @@ struct mtmd_gen_out { size_t n_codes; const float * embd; // the generated hidden state, to be fed back to backbone // it must have n_text_embd elements - // for MTMD_GEN_PROCESS_TYPE_SPEAKER_COND: embd holds the conditioning - // rows and n_embd their total element count + // for MTMD_GEN_PROCESS_TYPE_SPK_REF: embd holds the talker conditioning + // rows of the reference clip and n_embd their total element count size_t n_embd; // for MTMD_GEN_PROCESS_TYPE_CODE2WAV const float * audio; size_t n_samples; + // opaque state: the decoder carry-over to pass into the next CODE2WAV + // call, or, from MTMD_GEN_PROCESS_TYPE_SPK_REF, the encoded reference + // state of the cloned voice const char * state_data; size_t state_size; }; From 61752e5a4aab27697d21e76a2ae9f2faf29e4ede Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 2 Aug 2026 19:54:38 +0200 Subject: [PATCH 07/14] mtmd: chatterbox sampling KVs carry explicit disable values The general.sampling.* KVs are applied as model defaults by common, so a key absent from the GGUF falls back to the common default instead of the reference behavior. Write all five keys per variant, with the samplers the reference does not use at their disable value (turbo: min_p 0.0, multilingual: top_k 0). --- conversion/chatterbox.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/conversion/chatterbox.py b/conversion/chatterbox.py index f6e8d580660f..6a0b00a1980e 100644 --- a/conversion/chatterbox.py +++ b/conversion/chatterbox.py @@ -43,9 +43,11 @@ SPEECH_BOS = 6561 SPEECH_EOS = 6562 -# reference sampling defaults (tts_turbo.py generate / mtl tts.py) -TURBO_SAMPLING = {"top_k": 1000, "top_p": 0.95, "temp": 0.8, "penalty_repeat": 1.2} -MTL_SAMPLING = {"min_p": 0.05, "top_p": 1.0, "temp": 0.8, "penalty_repeat": 1.2} +# 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} def _s3tok_mel_filters() -> Tensor: @@ -213,10 +215,8 @@ def set_gguf_parameters(self): self.gguf_writer.add_file_type(self.ftype) sampling = TURBO_SAMPLING if self.is_turbo else MTL_SAMPLING - if "top_k" in sampling: - self.gguf_writer.add_sampling_top_k(sampling["top_k"]) - if "min_p" in sampling: - self.gguf_writer.add_sampling_min_p(sampling["min_p"]) + 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"]) From f324a336f68cc780a3ea7bb9fe6f8b49c66a2d5e Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 2 Aug 2026 20:16:09 +0200 Subject: [PATCH 08/14] conversion: remove unused import in chatterbox converter --- conversion/chatterbox.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/conversion/chatterbox.py b/conversion/chatterbox.py index 6a0b00a1980e..2da021e9b2fa 100644 --- a/conversion/chatterbox.py +++ b/conversion/chatterbox.py @@ -32,7 +32,7 @@ if TYPE_CHECKING: from torch import Tensor -from .base import ModelBase, TextModel, MmprojModel, LazyTorchTensor, gguf, logger +from .base import ModelBase, TextModel, MmprojModel, LazyTorchTensor, gguf TURBO_TALKER = "t3_turbo_v1.safetensors" MTL_TALKER = "t3_mtl23ls_v3.safetensors" From 9f0db106dd14c6392adadf4ec6c3455e4321b5cd Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 2 Aug 2026 20:43:44 +0200 Subject: [PATCH 09/14] mtmd: fix multilingual chatterbox tokenization of non-ascii text The reference multilingual tokenizer is a char-level BPE while the gpt2 tokenizer of llama.cpp is byte-level, so every accented char fell back to byte tokens the model never saw. The converter re-encodes the vocab and merges through the gpt2 byte-to-unicode map and prepends synthetic merges that rebuild each multi-byte char from its bytes, so the byte-level closure reproduces the char-level tokenization exactly. The helper lowercase now also covers the latin-1 accented uppercase, matching the reference .lower(). --- conversion/chatterbox.py | 26 +++++++++++++++++++++++--- tools/mtmd/mtmd-helper-gen.cpp | 15 +++++++++++++-- 2 files changed, 36 insertions(+), 5 deletions(-) diff --git a/conversion/chatterbox.py b/conversion/chatterbox.py index 2da021e9b2fa..657162bdcfce 100644 --- a/conversion/chatterbox.py +++ b/conversion/chatterbox.py @@ -163,16 +163,31 @@ def _set_vocab_turbo(self): self.gguf_writer.add_add_eos_token(False) def _set_vocab_mtl(self): - # custom multilingual BPE (mtl_tokenizer.json), extended with the speech tokens + # 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 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(): - tokens[i] = t + 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) @@ -181,7 +196,12 @@ def _set_vocab_mtl(self): tokens += self._speech_token_names(n_speech) toktypes += [int(gguf.TokenType.CONTROL)] * n_speech - merges = [" ".join(m) if isinstance(m, list) else m for m in tok["model"].get("merges", [])] + # 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") diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 539a12e5275a..caf81ab996e5 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -502,14 +502,25 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { // tokens (language tag left to the caller), turbo applies punc_norm std::string txt(inp->prompt, inp->prompt_len); if (mtl) { + // lowercase matches the reference .lower() on the latin-1 range: + // ascii letters plus the accented uppercase (utf-8 c3 80..c3 9e + // maps to c3 a0..c3 be, the multiplication sign c3 97 excepted) std::string norm; - for (char c : txt) { + for (size_t i = 0; i < txt.size(); i++) { + const unsigned char c = (unsigned char) txt[i]; if (c == ' ') { norm += "[SPACE]"; } else if (c >= 'A' && c <= 'Z') { norm += (char) (c - 'A' + 'a'); + } else if (c == 0xC3 && i + 1 < txt.size() + && (unsigned char) txt[i + 1] >= 0x80 + && (unsigned char) txt[i + 1] <= 0x9E + && (unsigned char) txt[i + 1] != 0x97) { + norm += (char) 0xC3; + norm += (char) ((unsigned char) txt[i + 1] + 0x20); + i++; } else { - norm += c; + norm += (char) c; } } txt = norm; From d83af018c06eeb885bd2940e3039d4eaec564eda Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 2 Aug 2026 21:17:36 +0200 Subject: [PATCH 10/14] mtmd: add cfg to the chatterbox multilingual talker The reference decodes a conditional and an unconditional branch and mixes the logits as cond + 0.5 * (cond - uncond). The helper prefills a second sequence where the text rows keep only their learned positions, feeds each sampled token to both branches and writes the mixed logits into the row the tool samples from, so the tool and the public API stay unchanged. llama-tts provisions the second sequence while keeping the per-sequence context window at the requested size. The repetition penalty now covers the whole generated history through the model sampling defaults, matching the reference for both variants. --- conversion/chatterbox.py | 3 +++ tools/mtmd/mtmd-helper-gen.cpp | 48 ++++++++++++++++++++++++++++++++++ tools/tts/tts.cpp | 7 +++++ 3 files changed, 58 insertions(+) diff --git a/conversion/chatterbox.py b/conversion/chatterbox.py index 657162bdcfce..96324584c42a 100644 --- a/conversion/chatterbox.py +++ b/conversion/chatterbox.py @@ -48,6 +48,8 @@ # 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: @@ -240,6 +242,7 @@ def set_gguf_parameters(self): 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: diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index caf81ab996e5..16c73f2f77e9 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -575,6 +575,7 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { ids.insert(ids.begin(), text_start); ids.push_back(text_stop); } + const size_t text0 = prompt.size(); for (size_t i = 0; i < ids.size(); i++) { prompt.push_back(row(ids[i])); if (mtl) { @@ -604,6 +605,27 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { return 1; } + if (mtl) { + // cfg second sequence: the same prompt with the text embeddings + // zeroed, keeping their learned positions (reference prepares the + // uncond branch before the position add) + std::vector cond_logits; + cfg_read(cond_logits); + for (size_t i = 0; i < ids.size(); i++) { + std::vector u((size_t) n_e, 0.0f); + add_pos(u, text_pos, (int) i); + memcpy(embd_buf.data() + (text0 + i) * (size_t) n_e, u.data(), (size_t) n_e * sizeof(float)); + } + decode_embd_batch batch_uncond(embd_buf.data(), n_prompt, 1, n_e); + batch_uncond.set_position_normal(0, 1); + batch_uncond.batch.logits[n_prompt - 1] = 1; + if (llama_decode(lctx, batch_uncond.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: cfg prefill decode failed\n"); + return 1; + } + cfg_apply(cond_logits); + } + pos = n_prompt; out_type = inp->out_type; return 0; @@ -634,6 +656,17 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { LOG_ERR("mtmd_helper_gen_audio: step decode failed\n"); return 1; } + // cfg second sequence: the sampled token feeds both branches + std::vector cond_logits; + cfg_read(cond_logits); + decode_embd_batch batch_uncond(e.data(), 1, 1, n_embd); + batch_uncond.set_position_normal(pos, 1); + batch_uncond.batch.logits[0] = 1; + if (llama_decode(lctx, batch_uncond.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: cfg step decode failed\n"); + return 1; + } + cfg_apply(cond_logits); } else { llama_batch batch = llama_batch_get_one(&sampled, 1); if (llama_decode(lctx, batch) != 0) { @@ -690,6 +723,21 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } private: + // multilingual cfg, reference t3 combine: logits = cond + w * (cond - uncond) + // with the reference default weight 0.5. the cond row is saved after the + // first decode, the mix lands in the row the tool samples from (the last + // one with logits enabled, which the second decode produced) + void cfg_read(std::vector & cond_logits) { + const float * c = llama_get_logits_ith(lctx, -1); + cond_logits.assign(c, c + llama_vocab_n_tokens(vocab)); + } + void cfg_apply(const std::vector & cond_logits) { + float * u = llama_get_logits_ith(lctx, -1); + for (size_t i = 0; i < cond_logits.size(); i++) { + u[i] = cond_logits[i] + 0.5f * (cond_logits[i] - u[i]); + } + } + bool ensure_cache() { if (!tok_embd.empty()) { return true; diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index 44e281331989..f84fbfc51cf2 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -66,6 +66,13 @@ 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 a second sequence for pipelines that decode a cfg pair, + // scaling the context so the per-sequence window keeps the requested size + if (params.n_parallel < 2) { + params.n_parallel = 2; + params.n_ctx *= 2; + } + llama_backend_init(); llama_numa_init(params.numa); From 80be614db645a4067827fde6255121e8edb87eb0 Mon Sep 17 00:00:00 2001 From: Pascal Date: Mon, 3 Aug 2026 04:59:33 +0200 Subject: [PATCH 11/14] mtmd: move chatterbox weights into nested model structs The gen and spkenc graphs read typed fields through a shared clip_graph_chatterbox_base instead of a name to tensor map; the map now only carries host-read side data. Variant probes move to tensors the reduced map still holds. --- conversion/chatterbox.py | 8 + tools/mtmd/clip-model.h | 240 ++++++++++++++- tools/mtmd/clip.cpp | 278 ++++++++++++++++- tools/mtmd/models/chatterbox-gen.cpp | 321 +++++++++----------- tools/mtmd/models/chatterbox-spkenc.cpp | 137 ++++----- tools/mtmd/models/models.h | 44 ++- tools/mtmd/mtmd-audio.cpp | 215 ++++++++++---- tools/mtmd/mtmd-audio.h | 61 ++-- tools/mtmd/mtmd-helper-gen.cpp | 104 ++++--- tools/mtmd/mtmd.cpp | 379 +++++++++++------------- tools/mtmd/mtmd.h | 36 ++- tools/tts/tts.cpp | 36 +-- 12 files changed, 1183 insertions(+), 676 deletions(-) diff --git a/conversion/chatterbox.py b/conversion/chatterbox.py index 96324584c42a..a7e580caedd8 100644 --- a/conversion/chatterbox.py +++ b/conversion/chatterbox.py @@ -164,6 +164,10 @@ def _set_vocab_turbo(self): 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 @@ -216,6 +220,10 @@ def enc(s: str) -> str: 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_gguf_parameters(self): if self.is_turbo: self.gguf_writer.add_block_count(self.hparams["n_layer"]) diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 1779b98d7e2d..27ff5eef3ec7 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -457,6 +457,237 @@ struct clip_code2wav { ggml_tensor * dac_post_conv_b = nullptr; }; +struct clip_chatterbox { + // espnet rel-pos conformer layer of the flow encoder + struct enc_layer { + ggml_tensor * norm_mha_w = nullptr; + ggml_tensor * norm_mha_b = nullptr; + ggml_tensor * attn_q_w = nullptr; + ggml_tensor * attn_q_b = nullptr; + ggml_tensor * attn_k_w = nullptr; + ggml_tensor * attn_k_b = nullptr; + ggml_tensor * attn_v_w = nullptr; + ggml_tensor * attn_v_b = nullptr; + ggml_tensor * attn_pos_w = nullptr; // linear_pos, no bias + ggml_tensor * attn_pos_bias_u = nullptr; + ggml_tensor * attn_pos_bias_v = nullptr; + ggml_tensor * attn_out_w = nullptr; + ggml_tensor * attn_out_b = nullptr; + ggml_tensor * norm_ff_w = nullptr; + ggml_tensor * norm_ff_b = nullptr; + ggml_tensor * ffn_1_w = nullptr; + ggml_tensor * ffn_1_b = nullptr; + ggml_tensor * ffn_2_w = nullptr; + ggml_tensor * ffn_2_b = nullptr; + }; + + // causal conv block of the estimator (conv k3 left padded + layer norm) + struct causal_block { + ggml_tensor * conv_w = nullptr; + ggml_tensor * conv_b = nullptr; + ggml_tensor * norm_w = nullptr; + ggml_tensor * norm_b = nullptr; + }; + + // time conditioned resnet of the estimator + struct resnet { + causal_block block1; + causal_block block2; + ggml_tensor * mlp_w = nullptr; // time projection (mlp.1) + ggml_tensor * mlp_b = nullptr; + ggml_tensor * res_conv_w = nullptr; + ggml_tensor * res_conv_b = nullptr; + }; + + // diffusers style transformer block of the estimator + struct tfm_block { + ggml_tensor * norm1_w = nullptr; + ggml_tensor * norm1_b = nullptr; + ggml_tensor * attn_q_w = nullptr; // no bias on qkv + ggml_tensor * attn_k_w = nullptr; + ggml_tensor * attn_v_w = nullptr; + ggml_tensor * attn_out_w = nullptr; // to_out.0 + ggml_tensor * attn_out_b = nullptr; + ggml_tensor * norm3_w = nullptr; + ggml_tensor * norm3_b = nullptr; + ggml_tensor * ff_in_w = nullptr; // ff.net.0.proj + ggml_tensor * ff_in_b = nullptr; + ggml_tensor * ff_out_w = nullptr; // ff.net.2 + ggml_tensor * ff_out_b = nullptr; + }; + + // one estimator stage: resnet, transformer stack, boundary causal conv + // (the boundary conv stays null on mid stages) + struct est_stage { + resnet res; + std::vector tfm; + ggml_tensor * conv_w = nullptr; + ggml_tensor * conv_b = nullptr; + }; + + // snake resblock unit of the hift vocoder + struct hift_res_unit { + ggml_tensor * act1_alpha = nullptr; + ggml_tensor * act2_alpha = nullptr; + ggml_tensor * conv1_w = nullptr; + ggml_tensor * conv1_b = nullptr; + ggml_tensor * conv2_w = nullptr; + ggml_tensor * conv2_b = nullptr; + }; + struct hift_res { + std::vector units; // dilations 1/3/5 on conv1 + }; + + // one hift upsample stage: conv transpose, source injection, 3 resblocks + struct hift_up { + ggml_tensor * up_w = nullptr; + ggml_tensor * up_b = nullptr; + ggml_tensor * source_down_w = nullptr; + ggml_tensor * source_down_b = nullptr; + hift_res source_res; + hift_res res_0; + hift_res res_1; + hift_res res_2; + }; + + // f0 predictor conv of the hift vocoder + struct f0_conv { + ggml_tensor * w = nullptr; + ggml_tensor * b = nullptr; + }; + + // s3 tokenizer attention block (neox rope on q/k, fsmn memory on v) + struct s3tok_block { + ggml_tensor * attn_ln_w = nullptr; + ggml_tensor * attn_ln_b = nullptr; + ggml_tensor * attn_q_w = nullptr; + ggml_tensor * attn_q_b = nullptr; + ggml_tensor * attn_k_w = nullptr; // no bias + ggml_tensor * attn_v_w = nullptr; + ggml_tensor * attn_v_b = nullptr; + ggml_tensor * fsmn_w = nullptr; + ggml_tensor * attn_out_w = nullptr; + ggml_tensor * attn_out_b = nullptr; + ggml_tensor * mlp_ln_w = nullptr; + ggml_tensor * mlp_ln_b = nullptr; + ggml_tensor * mlp_in_w = nullptr; + ggml_tensor * mlp_in_b = nullptr; + ggml_tensor * mlp_out_w = nullptr; + ggml_tensor * mlp_out_b = nullptr; + }; + + // folded batchnorm (w/b stay null on the affine free variant) + struct bn { + ggml_tensor * mean = nullptr; + ggml_tensor * var = nullptr; + ggml_tensor * w = nullptr; + ggml_tensor * b = nullptr; + }; + + // fcm residual 2d block of the speaker encoder + // (shortcut stays null on the stride 1 blocks) + struct spk_res2d { + ggml_tensor * conv1_w = nullptr; + bn bn1; + ggml_tensor * conv2_w = nullptr; + bn bn2; + ggml_tensor * shortcut_w = nullptr; + bn shortcut_bn; + }; + + // cam dense tdnn layer of the speaker encoder + struct spk_cam_layer { + bn nl1_bn; + ggml_tensor * linear1_w = nullptr; + bn nl2_bn; + ggml_tensor * local_w = nullptr; // cam_layer.linear_local + ggml_tensor * ctx1_w = nullptr; // cam_layer.linear1 + ggml_tensor * ctx1_b = nullptr; + ggml_tensor * ctx2_w = nullptr; // cam_layer.linear2 + ggml_tensor * ctx2_b = nullptr; + }; + struct spk_cam_block { + std::vector layers; + bn transit_bn; + ggml_tensor * transit_w = nullptr; + }; + + // flow encoder + ggml_tensor * input_embedding_w = nullptr; // flow.input_embedding + ggml_tensor * embed_linear_w = nullptr; // fenc.embed.out.0 + ggml_tensor * embed_linear_b = nullptr; + ggml_tensor * embed_norm_w = nullptr; // fenc.embed.out.1 + ggml_tensor * embed_norm_b = nullptr; + ggml_tensor * pre_conv1_w = nullptr; // pre_lookahead_layer + ggml_tensor * pre_conv1_b = nullptr; + ggml_tensor * pre_conv2_w = nullptr; + ggml_tensor * pre_conv2_b = nullptr; + std::vector enc; + ggml_tensor * up_conv_w = nullptr; // fenc.up_layer.conv + ggml_tensor * up_conv_b = nullptr; + ggml_tensor * up_embed_linear_w = nullptr; // fenc.up_embed.out.0 + ggml_tensor * up_embed_linear_b = nullptr; + ggml_tensor * up_embed_norm_w = nullptr; // fenc.up_embed.out.1 + ggml_tensor * up_embed_norm_b = nullptr; + std::vector up_enc; + ggml_tensor * after_norm_w = nullptr; + ggml_tensor * after_norm_b = nullptr; + ggml_tensor * encoder_proj_w = nullptr; // flow.encoder_proj + ggml_tensor * encoder_proj_b = nullptr; + + // cfm estimator + ggml_tensor * time_mlp_1_w = nullptr; + ggml_tensor * time_mlp_1_b = nullptr; + ggml_tensor * time_mlp_2_w = nullptr; + ggml_tensor * time_mlp_2_b = nullptr; + ggml_tensor * time_embed_mixer_w = nullptr; // meanflow variant only + est_stage est_down; + std::vector est_mid; + est_stage est_up; + causal_block est_final_block; + ggml_tensor * est_final_proj_w = nullptr; + ggml_tensor * est_final_proj_b = nullptr; + + // hift vocoder + ggml_tensor * hift_pre_w = nullptr; + ggml_tensor * hift_pre_b = nullptr; + std::vector hift_ups; + ggml_tensor * hift_post_w = nullptr; + ggml_tensor * hift_post_b = nullptr; + std::vector f0_condnet; + ggml_tensor * f0_classifier_w = nullptr; + ggml_tensor * f0_classifier_b = nullptr; + + // s3 tokenizer + ggml_tensor * s3tok_conv1_w = nullptr; + ggml_tensor * s3tok_conv1_b = nullptr; + ggml_tensor * s3tok_conv2_w = nullptr; + ggml_tensor * s3tok_conv2_b = nullptr; + std::vector s3tok_blocks; + ggml_tensor * s3tok_down_w = nullptr; // quantizer._codebook.project_down + ggml_tensor * s3tok_down_b = nullptr; + + // CAMPPlus speaker encoder + ggml_tensor * spk_conv1_w = nullptr; + bn spk_bn1; + spk_res2d spk_layer1_0; + spk_res2d spk_layer1_1; + spk_res2d spk_layer2_0; + spk_res2d spk_layer2_1; + ggml_tensor * spk_conv2_w = nullptr; + bn spk_bn2; + ggml_tensor * spk_tdnn_w = nullptr; + bn spk_tdnn_bn; + spk_cam_block spk_block1; + spk_cam_block spk_block2; + spk_cam_block spk_block3; + bn spk_out_bn; + ggml_tensor * spk_dense_w = nullptr; + bn spk_dense_bn; + ggml_tensor * spk_affine_w = nullptr; + ggml_tensor * spk_affine_b = nullptr; +}; + struct clip_model { clip_modality modality = CLIP_MODALITY_VISION; projector_type proj_type = PROJECTOR_TYPE_MLP; @@ -688,8 +919,13 @@ struct clip_model { // qwen3tts code2wav: RVQ codes -> raw PCM clip_code2wav c2w; - // chatterbox audio stack: tensors are looked up by their source name at - // graph build time, the set is too large and too nested for named fields + // chatterbox flow encoder, cfm estimator, hift vocoder, s3 tokenizer + // and CAMPPlus speaker encoder + clip_chatterbox cbx; + + // chatterbox host-read side data (conditioning defaults, embedding + // tables, filterbanks, voice/conditioning encoders run on the host), + // accessed by name through clip_cbx_read_tensor std::map cbx_tensors; // cogvlm diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 461ed6fca29c..4bbd2c0232ea 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -2870,23 +2870,274 @@ struct clip_model_loader { } break; case PROJECTOR_TYPE_CHATTERBOX: { - // the chatterbox audio stack keeps its source tensor names, - // load everything in the file and index by name for the - // graph builders + 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.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)) { - model.cbx_tensors[t->name] = get_tensor(t->name); + 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: { - // this context only carries the CAMPPlus x-vector body and - // the flow affine that maps its embedding to the s3gen dim + 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); + c.spk_affine_w = get_tensor("a.spk_embed_affine_layer.weight"); + c.spk_affine_b = get_tensor("a.spk_embed_affine_layer.bias"); + + // host-read side data, accessed by name through + // clip_cbx_read_tensor: voice encoder lstm, conditioning + // encoder, speaker affine interface dim 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.spk.", 0) == 0 || name.rfind("a.spk_embed_affine_layer.", 0) == 0) { - model.cbx_tensors[name] = get_tensor(name.c_str()); + if (name.rfind("a.ve.", 0) == 0 || name.rfind("a.cenc.", 0) == 0) { + model.cbx_tensors[name] = get_tensor(name); } } + // the affine bias doubles as the host-side presence and + // dimension probe of the speaker encoder + model.cbx_tensors["a.spk_embed_affine_layer.bias"] = c.spk_affine_b; } break; case PROJECTOR_TYPE_VOXTRAL: { @@ -4914,7 +5165,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { for (auto & f : noise) f = nd(rng); set_input_f32("inp_noise", noise); - const bool meanflow = model.cbx_tensors.count("a.gen.est.time_embed_mixer.weight") > 0; + const bool meanflow = model.cbx.time_embed_mixer_w != nullptr; const int n_steps = meanflow ? 2 : 10; std::vector temb((size_t) 320 * (n_steps + 1)); for (int s = 0; s <= n_steps; s++) { @@ -5868,8 +6119,13 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { // channel count stands in for the interface dimension return 80; case PROJECTOR_TYPE_CHATTERBOX_SPKENC: - // x-vector projected through the s3gen speaker affine - return 80; + { + // the encoder emits talker conditioning rows in the backbone + // embedding space, whose dim the speaker projection carries + auto it = ctx->model.cbx_tensors.find("a.cenc.spkr_enc.weight"); + GGML_ASSERT(it != ctx->model.cbx_tensors.end()); + return it->second->ne[1]; + } case PROJECTOR_TYPE_PARAKEET: return ctx->model.mm_1_w->ne[1]; default: diff --git a/tools/mtmd/models/chatterbox-gen.cpp b/tools/mtmd/models/chatterbox-gen.cpp index b9e6e90cd27d..0beda00fe3ae 100644 --- a/tools/mtmd/models/chatterbox-gen.cpp +++ b/tools/mtmd/models/chatterbox-gen.cpp @@ -1,19 +1,11 @@ #include "models.h" // Chatterbox generation graphs: flow encoder, cfm estimator, s3 tokenizer -// and hift vocoder. Weights come from the source-named tensor map -// (model.cbx_tensors); the speaker encoder lives in chatterbox-spkenc.cpp. - -ggml_tensor * cbx_t(const clip_model & model, const std::string & name) { - auto it = model.cbx_tensors.find(name); - if (it == model.cbx_tensors.end()) { - GGML_ABORT("missing chatterbox tensor: %s", name.c_str()); - } - return it->second; -} +// 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 * cbx_linear(ggml_context * ctx0, ggml_tensor * w, ggml_tensor * b, ggml_tensor * x) { +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); @@ -30,8 +22,8 @@ static ggml_tensor * cbx_layer_norm(ggml_context * ctx0, ggml_tensor * w, ggml_t // 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 * cbx_conv1d(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * b, ggml_tensor * x, - int stride, int pad_l, int 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]); @@ -51,6 +43,18 @@ ggml_tensor * cbx_conv1d(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * 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 @@ -69,27 +73,27 @@ static ggml_tensor * cbx_rel_shift(ggml_context * ctx0, ggml_tensor * bd, int T) } // espnet rel-pos self attention block, pre-norm, x [512, T], pos [512, 2T-1] -static ggml_tensor * cbx_enc_layer(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, ggml_tensor * pos, - const std::string & p, int T) { +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, cbx_t(model, p + ".norm_mha.weight"), cbx_t(model, p + ".norm_mha.bias"), x, 1e-5f); + ggml_tensor * cur = cbx_layer_norm(ctx0, l.norm_mha_w, l.norm_mha_b, x, 1e-5f); - ggml_tensor * q = cbx_linear(ctx0, cbx_t(model, p + ".self_attn.linear_q.weight"), cbx_t(model, p + ".self_attn.linear_q.bias"), cur); - ggml_tensor * k = cbx_linear(ctx0, cbx_t(model, p + ".self_attn.linear_k.weight"), cbx_t(model, p + ".self_attn.linear_k.bias"), cur); - ggml_tensor * v = cbx_linear(ctx0, cbx_t(model, p + ".self_attn.linear_v.weight"), cbx_t(model, p + ".self_attn.linear_v.bias"), cur); - ggml_tensor * pe = cbx_linear(ctx0, cbx_t(model, p + ".self_attn.linear_pos.weight"), nullptr, pos); // [512, 2T-1] + ggml_tensor * q = cbx_linear(l.attn_q_w, l.attn_q_b, cur); + ggml_tensor * k = cbx_linear(l.attn_k_w, l.attn_k_b, cur); + ggml_tensor * v = cbx_linear(l.attn_v_w, l.attn_v_b, cur); + ggml_tensor * pe = cbx_linear(l.attn_pos_w, nullptr, pos); // [512, 2T-1] q = ggml_reshape_3d(ctx0, q, d_head, n_head, T); k = ggml_reshape_3d(ctx0, k, d_head, n_head, T); v = ggml_reshape_3d(ctx0, v, d_head, n_head, T); pe = ggml_reshape_3d(ctx0, pe, d_head, n_head, 2 * T - 1); - ggml_tensor * u = cbx_t(model, p + ".self_attn.pos_bias_u"); // [64, 8] - ggml_tensor * w = cbx_t(model, p + ".self_attn.pos_bias_v"); + 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)); @@ -111,14 +115,14 @@ static ggml_tensor * cbx_enc_layer(const clip_model & model, ggml_context * ctx0 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(ctx0, cbx_t(model, p + ".self_attn.linear_out.weight"), cbx_t(model, p + ".self_attn.linear_out.bias"), o); + 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, cbx_t(model, p + ".norm_ff.weight"), cbx_t(model, p + ".norm_ff.bias"), x, 1e-5f); - cur = cbx_linear(ctx0, cbx_t(model, p + ".feed_forward.w_1.weight"), cbx_t(model, p + ".feed_forward.w_1.bias"), cur); + 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(ctx0, cbx_t(model, p + ".feed_forward.w_2.weight"), cbx_t(model, p + ".feed_forward.w_2.bias"), cur); + cur = cbx_linear(l.ffn_2_w, l.ffn_2_b, cur); x = ggml_add(ctx0, res, cur); return x; } @@ -130,34 +134,33 @@ static ggml_tensor * cbx_mish(ggml_context * ctx0, ggml_tensor * x) { } // causal block: conv k3 left-padded, layer norm over channels, mish; x [C, T] -static ggml_tensor * cbx_causal_block(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, const std::string & p) { - ggml_tensor * k = cbx_t(model, p + ".block.0.weight"); - x = cbx_conv1d(ctx0, k, cbx_t(model, p + ".block.0.bias"), x, 1, (int) k->ne[0] - 1, 0); - x = cbx_layer_norm(ctx0, cbx_t(model, p + ".block.2.weight"), cbx_t(model, p + ".block.2.bias"), x, 1e-5f); +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] -static ggml_tensor * cbx_resnet(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, ggml_tensor * temb, const std::string & p) { - ggml_tensor * h = cbx_causal_block(model, ctx0, x, p + ".block1"); - ggml_tensor * tproj = cbx_linear(ctx0, cbx_t(model, p + ".mlp.1.weight"), cbx_t(model, p + ".mlp.1.bias"), cbx_mish(ctx0, temb)); +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 = cbx_causal_block(model, ctx0, h, p + ".block2"); - ggml_tensor * res = cbx_conv1d(ctx0, cbx_t(model, p + ".res_conv.weight"), cbx_t(model, p + ".res_conv.bias"), x, 1, 0, 0); + 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] -static ggml_tensor * cbx_tfm_block(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, const std::string & p) { +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, cbx_t(model, p + ".norm1.weight"), cbx_t(model, p + ".norm1.bias"), x, 1e-5f); - ggml_tensor * q = ggml_mul_mat(ctx0, cbx_t(model, p + ".attn1.to_q.weight"), cur); - ggml_tensor * k = ggml_mul_mat(ctx0, cbx_t(model, p + ".attn1.to_k.weight"), cur); - ggml_tensor * v = ggml_mul_mat(ctx0, cbx_t(model, p + ".attn1.to_v.weight"), cur); + 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)); @@ -166,21 +169,23 @@ static ggml_tensor * cbx_tfm_block(const clip_model & model, ggml_context * ctx0 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(ctx0, cbx_t(model, p + ".attn1.to_out.0.weight"), cbx_t(model, p + ".attn1.to_out.0.bias"), o); + 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, cbx_t(model, p + ".norm3.weight"), cbx_t(model, p + ".norm3.bias"), x, 1e-5f); - cur = cbx_linear(ctx0, cbx_t(model, p + ".ff.net.0.proj.weight"), cbx_t(model, p + ".ff.net.0.proj.bias"), cur); + 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(ctx0, cbx_t(model, p + ".ff.net.2.weight"), cbx_t(model, p + ".ff.net.2.bias"), 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] -static ggml_tensor * cbx_estimator(const clip_model & model, ggml_context * ctx0, ggml_tensor * x_noise, ggml_tensor * mu, - ggml_tensor * spks, ggml_tensor * cond, ggml_tensor * temb, int T) { +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)); @@ -192,38 +197,31 @@ static ggml_tensor * cbx_estimator(const clip_model & model, ggml_context * ctx0 // down ggml_tensor * skip; - x = cbx_resnet(model, ctx0, x, temb, "a.gen.est.down_blocks.0.0"); - for (int j = 0; model.cbx_tensors.count("a.gen.est.down_blocks.0.1." + std::to_string(j) + ".norm1.weight"); j++) { - x = cbx_tfm_block(model, ctx0, x, "a.gen.est.down_blocks.0.1." + std::to_string(j)); + x = resnet(c.est_down.res, x, temb); + for (const auto & b : c.est_down.tfm) { + x = tfm_block(b, x); } skip = x; - { - ggml_tensor * k = cbx_t(model, "a.gen.est.down_blocks.0.2.weight"); - x = cbx_conv1d(ctx0, k, cbx_t(model, "a.gen.est.down_blocks.0.2.bias"), x, 1, (int) k->ne[0] - 1, 0); - } + 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 (int i = 0; model.cbx_tensors.count("a.gen.est.mid_blocks." + std::to_string(i) + ".0.block1.block.0.weight"); i++) { - const std::string mp = "a.gen.est.mid_blocks." + std::to_string(i); - x = cbx_resnet(model, ctx0, x, temb, mp + ".0"); - for (int j = 0; model.cbx_tensors.count(mp + ".1." + std::to_string(j) + ".norm1.weight"); j++) { - x = cbx_tfm_block(model, ctx0, x, mp + ".1." + std::to_string(j)); + 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 = cbx_resnet(model, ctx0, x, temb, "a.gen.est.up_blocks.0.0"); - for (int j = 0; model.cbx_tensors.count("a.gen.est.up_blocks.0.1." + std::to_string(j) + ".norm1.weight"); j++) { - x = cbx_tfm_block(model, ctx0, x, "a.gen.est.up_blocks.0.1." + std::to_string(j)); - } - { - ggml_tensor * k = cbx_t(model, "a.gen.est.up_blocks.0.2.weight"); - x = cbx_conv1d(ctx0, k, cbx_t(model, "a.gen.est.up_blocks.0.2.bias"), x, 1, (int) k->ne[0] - 1, 0); + 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 = cbx_causal_block(model, ctx0, x, "a.gen.est.final_block"); - x = cbx_conv1d(ctx0, cbx_t(model, "a.gen.est.final_proj.weight"), cbx_t(model, "a.gen.est.final_proj.bias"), x, 1, 0, 0); // [80, T] + 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; } @@ -232,7 +230,8 @@ static ggml_tensor * cbx_estimator(const clip_model & model, ggml_context * ctx0 // 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. -static ggml_cgraph * cbx_build_s3tok(const clip_model & model, ggml_context * ctx0, ggml_cgraph * gf, int T) { +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; @@ -247,25 +246,23 @@ static ggml_cgraph * cbx_build_s3tok(const clip_model & model, ggml_context * ct ggml_set_input(pos); ggml_tensor * x = ggml_cont(ctx0, ggml_transpose(ctx0, inp)); // [128, T] - x = cbx_conv1d(ctx0, cbx_t(model, "a.s3tok.encoder.conv1.weight"), cbx_t(model, "a.s3tok.encoder.conv1.bias"), x, 2, 1, 1); + x = cbx_conv1d(c.s3tok_conv1_w, c.s3tok_conv1_b, x, 2, 1, 1); x = ggml_gelu_erf(ctx0, x); - x = cbx_conv1d(ctx0, cbx_t(model, "a.s3tok.encoder.conv2.weight"), cbx_t(model, "a.s3tok.encoder.conv2.bias"), x, 2, 1, 1); + x = cbx_conv1d(c.s3tok_conv2_w, c.s3tok_conv2_b, x, 2, 1, 1); x = ggml_gelu_erf(ctx0, x); // [1280, T2] - for (int li = 0; model.cbx_tensors.count("a.s3tok.encoder.blocks." + std::to_string(li) + ".attn_ln.weight"); li++) { - const std::string p = "a.s3tok.encoder.blocks." + std::to_string(li) + ".attn"; - + for (const auto & blk : c.s3tok_blocks) { ggml_tensor * res = x; - ggml_tensor * cur = cbx_layer_norm(ctx0, cbx_t(model, p + "_ln.weight"), cbx_t(model, p + "_ln.bias"), x, 1e-5f); - ggml_tensor * q = cbx_linear(ctx0, cbx_t(model, p + ".query.weight"), cbx_t(model, p + ".query.bias"), cur); - ggml_tensor * k = ggml_mul_mat(ctx0, cbx_t(model, p + ".key.weight"), cur); - ggml_tensor * v = cbx_linear(ctx0, cbx_t(model, p + ".value.weight"), cbx_t(model, p + ".value.bias"), cur); + 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 = cbx_t(model, p + ".fsmn_block.weight"); + 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); } @@ -283,20 +280,18 @@ static ggml_cgraph * cbx_build_s3tok(const clip_model & model, ggml_context * ct 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(ctx0, cbx_t(model, p + ".out.weight"), cbx_t(model, p + ".out.bias"), o); + o = cbx_linear(blk.attn_out_w, blk.attn_out_b, o); x = ggml_add(ctx0, res, ggml_add(ctx0, o, fsm)); - const std::string mp = "a.s3tok.encoder.blocks." + std::to_string(li) + ".mlp"; res = x; - cur = cbx_layer_norm(ctx0, cbx_t(model, mp + "_ln.weight"), cbx_t(model, mp + "_ln.bias"), x, 1e-5f); - cur = cbx_linear(ctx0, cbx_t(model, mp + ".0.weight"), cbx_t(model, mp + ".0.bias"), cur); + 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(ctx0, cbx_t(model, mp + ".2.weight"), cbx_t(model, mp + ".2.bias"), cur); + cur = cbx_linear(blk.mlp_out_w, blk.mlp_out_b, cur); x = ggml_add(ctx0, res, cur); } - x = cbx_linear(ctx0, cbx_t(model, "a.s3tok.quantizer._codebook.project_down.weight"), - cbx_t(model, "a.s3tok.quantizer._codebook.project_down.bias"), x); // [8, T2] + 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"); @@ -305,14 +300,12 @@ static ggml_cgraph * cbx_build_s3tok(const clip_model & model, ggml_context * ct return gf; } -static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * ctx0, ggml_cgraph * gf, int n_mel, int n_stft); - ggml_cgraph * clip_graph_chatterbox::build() { if (gen_process == CLIP_GEN_PROCESS_TTS_VOCODE) { - return cbx_build_vocoder(model, ctx0, gf, vocode_n_mel, vocode_n_stft); + return build_vocoder(vocode_n_mel, vocode_n_stft); } if (gen_process == CLIP_GEN_PROCESS_TOKENIZE) { - return cbx_build_s3tok(model, ctx0, gf, img.nx()); + return build_s3tok(img.nx()); } if (gen_process != CLIP_GEN_PROCESS_TTS) { // load-time buffer sizing path @@ -326,6 +319,8 @@ ggml_cgraph * clip_graph_chatterbox::build() { 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 @@ -342,29 +337,27 @@ ggml_cgraph * clip_graph_chatterbox::build() { ggml_set_input(pos2); // token embedding - ggml_tensor * x = ggml_get_rows(ctx0, cbx_t(model, "a.gen.flow.input_embedding.weight"), inp_tokens); // [512, T1] + 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(ctx0, cbx_t(model, "a.gen.fenc.embed.out.0.weight"), cbx_t(model, "a.gen.fenc.embed.out.0.bias"), x); - x = cbx_layer_norm(ctx0, cbx_t(model, "a.gen.fenc.embed.out.1.weight"), cbx_t(model, "a.gen.fenc.embed.out.1.bias"), x, 1e-5f); + 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(ctx0, cbx_t(model, "a.gen.fenc.pre_lookahead_layer.conv1.weight"), - cbx_t(model, "a.gen.fenc.pre_lookahead_layer.conv1.bias"), x, 1, 0, 3); + 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(ctx0, cbx_t(model, "a.gen.fenc.pre_lookahead_layer.conv2.weight"), - cbx_t(model, "a.gen.fenc.pre_lookahead_layer.conv2.bias"), cur, 1, 2, 0); + 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 (int i = 0; model.cbx_tensors.count("a.gen.fenc.encoders." + std::to_string(i) + ".norm_mha.weight"); i++) { - x = cbx_enc_layer(model, ctx0, x, pos1, "a.gen.fenc.encoders." + std::to_string(i), T1); - cb(x, "fenc_enc", i); + 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 @@ -374,32 +367,32 @@ ggml_cgraph * clip_graph_chatterbox::build() { 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, cbx_t(model, "a.gen.fenc.up_layer.conv.weight"), xt, 1, 0, 1); + 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, cbx_t(model, "a.gen.fenc.up_layer.conv.bias")); + x = ggml_add(ctx0, x, c.up_conv_b); cb(x, "fenc_upsample", -1); } // up embed: linear + layer norm + xscale - x = cbx_linear(ctx0, cbx_t(model, "a.gen.fenc.up_embed.out.0.weight"), cbx_t(model, "a.gen.fenc.up_embed.out.0.bias"), x); - x = cbx_layer_norm(ctx0, cbx_t(model, "a.gen.fenc.up_embed.out.1.weight"), cbx_t(model, "a.gen.fenc.up_embed.out.1.bias"), x, 1e-5f); + 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 (int i = 0; model.cbx_tensors.count("a.gen.fenc.up_encoders." + std::to_string(i) + ".norm_mha.weight"); i++) { - x = cbx_enc_layer(model, ctx0, x, pos2, "a.gen.fenc.up_encoders." + std::to_string(i), T2); - cb(x, "fenc_up_enc", i); + 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, cbx_t(model, "a.gen.fenc.after_norm.weight"), cbx_t(model, "a.gen.fenc.after_norm.bias"), x, 1e-5f); + 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(ctx0, cbx_t(model, "a.gen.flow.encoder_proj.weight"), cbx_t(model, "a.gen.flow.encoder_proj.bias"), x); // [80, T2] + 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 = model.cbx_tensors.count("a.gen.est.time_embed_mixer.weight") > 0; + 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 @@ -418,9 +411,9 @@ ggml_cgraph * clip_graph_chatterbox::build() { ggml_set_input(temb_sin); auto time_mlp = [&](ggml_tensor * e) { - e = cbx_linear(ctx0, cbx_t(model, "a.gen.est.time_mlp.linear_1.weight"), cbx_t(model, "a.gen.est.time_mlp.linear_1.bias"), e); + e = cbx_linear(c.time_mlp_1_w, c.time_mlp_1_b, e); e = ggml_silu(ctx0, e); - e = cbx_linear(ctx0, cbx_t(model, "a.gen.est.time_mlp.linear_2.weight"), cbx_t(model, "a.gen.est.time_mlp.linear_2.bias"), e); + e = cbx_linear(c.time_mlp_2_w, c.time_mlp_2_b, e); return e; }; auto span_emb = [&](int i) { @@ -431,7 +424,7 @@ ggml_cgraph * clip_graph_chatterbox::build() { 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, cbx_t(model, "a.gen.est.time_embed_mixer.weight"), e); // [1024, 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 @@ -455,9 +448,9 @@ ggml_cgraph * clip_graph_chatterbox::build() { ggml_tensor * mx = noise; for (int i = 0; i < n_steps; i++) { ggml_tensor * temb = step_temb(i); - ggml_tensor * d = cbx_estimator(model, ctx0, mx, mu, spks, cond, temb, T2); + ggml_tensor * d = estimator(mx, mu, spks, cond, temb, T2); if (!meanflow) { - ggml_tensor * du = cbx_estimator(model, ctx0, mx, mu_zero, spks_zero, cond_zero, temb, T2); + 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])); @@ -475,12 +468,11 @@ ggml_cgraph * clip_graph_chatterbox::build() { // f0 predictor on the trimmed mel: 5x (conv k3 same-pad + elu), abs(linear) { ggml_tensor * fx = mel; - for (int i = 0; model.cbx_tensors.count("a.gen.hift.f0_predictor.condnet." + std::to_string(i) + ".weight"); i += 2) { - const std::string cp = "a.gen.hift.f0_predictor.condnet." + std::to_string(i); - fx = cbx_conv1d(ctx0, cbx_t(model, cp + ".weight"), cbx_t(model, cp + ".bias"), fx, 1, 1, 1); + 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(ctx0, cbx_t(model, "a.gen.hift.f0_predictor.classifier.weight"), cbx_t(model, "a.gen.hift.f0_predictor.classifier.bias"), 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); @@ -489,22 +481,35 @@ ggml_cgraph * clip_graph_chatterbox::build() { 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); +} -static ggml_tensor * cbx_hift_resblock(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, const std::string & p); - -// x [C, T] -> symmetric-padded dilated conv -> [OC, T] -ggml_tensor * cbx_conv1d_dil(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * b, ggml_tensor * x, int pad, int dil) { - 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); +// 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 y; + return x; } // mel [80, T] + source stft [18, T_stft] -> conv_post output [18, T_stft2] -static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * ctx0, ggml_cgraph * gf, int n_mel, int n_stft) { +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); @@ -512,11 +517,11 @@ static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * ggml_set_name(sstft, "inp_sstft"); ggml_set_input(sstft); - ggml_tensor * x = cbx_conv1d_dil(ctx0, cbx_t(model, "a.gen.hift.conv_pre.weight"), cbx_t(model, "a.gen.hift.conv_pre.bias"), mel, 3, 1); + ggml_tensor * x = cbx_conv1d_dil(c.hift_pre_w, c.hift_pre_b, mel, 3, 1); - for (int i = 0; model.cbx_tensors.count("a.gen.hift.ups." + std::to_string(i) + ".weight"); i++) { - const std::string is = std::to_string(i); - ggml_tensor * uk = cbx_t(model, "a.gen.hift.ups." + is + ".weight"); + 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; @@ -527,9 +532,9 @@ static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * 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, cbx_t(model, "a.gen.hift.ups." + is + ".bias")); + x = ggml_add(ctx0, x, up.up_b); - const bool is_last = !model.cbx_tensors.count("a.gen.hift.ups." + std::to_string(i + 1) + ".weight"); + 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); @@ -537,7 +542,7 @@ static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * } // source injection: strided conv on the source stft, one resblock - ggml_tensor * sk = cbx_t(model, "a.gen.hift.source_downs." + is + ".weight"); + 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; @@ -546,9 +551,9 @@ static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * 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, cbx_t(model, "a.gen.hift.source_downs." + is + ".bias")); + si = ggml_add(ctx0, si, up.source_down_b); } - si = cbx_hift_resblock(model, ctx0, si, "a.gen.hift.source_resblocks." + is); + 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]); @@ -557,48 +562,16 @@ static ggml_cgraph * cbx_build_vocoder(const clip_model & model, ggml_context * } x = ggml_add(ctx0, x, si); - ggml_tensor * acc = nullptr; - for (int j = 3 * i; j < 3 * (i + 1); j++) { - ggml_tensor * r = cbx_hift_resblock(model, ctx0, x, "a.gen.hift.resblocks." + std::to_string(j)); - acc = acc ? ggml_add(ctx0, acc, r) : r; - } + 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(ctx0, cbx_t(model, "a.gen.hift.conv_post.weight"), cbx_t(model, "a.gen.hift.conv_post.bias"), x, 3, 1); + 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; } - -// snake activation with per-channel alpha: x + sin^2(alpha x) / (alpha + eps) -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 -static ggml_tensor * cbx_hift_resblock(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, const std::string & p) { - static const int dil[3] = {1, 3, 5}; - for (int j = 0; model.cbx_tensors.count(p + ".convs1." + std::to_string(j) + ".weight"); j++) { - const std::string js = std::to_string(j); - ggml_tensor * a1 = cbx_t(model, p + ".activations1." + js + ".alpha"); - ggml_tensor * a2 = cbx_t(model, p + ".activations2." + js + ".alpha"); - ggml_tensor * k1 = cbx_t(model, p + ".convs1." + js + ".weight"); - ggml_tensor * k2 = cbx_t(model, p + ".convs2." + js + ".weight"); - const int d = dil[j % 3]; - const int p1 = (int) (k1->ne[0] - 1) / 2 * d; - const int p2 = (int) (k2->ne[0] - 1) / 2; - ggml_tensor * xt = cbx_snake(ctx0, x, a1); - xt = cbx_conv1d_dil(ctx0, k1, cbx_t(model, p + ".convs1." + js + ".bias"), xt, p1, d); - xt = cbx_snake(ctx0, xt, a2); - xt = cbx_conv1d_dil(ctx0, k2, cbx_t(model, p + ".convs2." + js + ".bias"), xt, p2, 1); - x = ggml_add(ctx0, x, xt); - } - return x; -} - diff --git a/tools/mtmd/models/chatterbox-spkenc.cpp b/tools/mtmd/models/chatterbox-spkenc.cpp index a9e967161f19..20a4729c234c 100644 --- a/tools/mtmd/models/chatterbox-spkenc.cpp +++ b/tools/mtmd/models/chatterbox-spkenc.cpp @@ -1,81 +1,62 @@ #include "models.h" -#include - // Chatterbox speaker encoder: CAMPPlus x-vector on kaldi fbank features, // projected through the s3gen speaker affine. Mirrors s3gen/xvector.py. // per-channel batchnorm on x [C, T]; scale = w / sqrt(var + eps), shift folds -// the running mean. pass null w/b for the affine=False variant -static ggml_tensor * cbx_bn1d(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, - const std::string & p, ggml_tensor * eps) { - ggml_tensor * mean = cbx_t(model, p + ".running_mean"); - ggml_tensor * var = cbx_t(model, p + ".running_var"); - ggml_tensor * sd = ggml_sqrt(ctx0, ggml_add(ctx0, var, eps)); - if (!model.cbx_tensors.count(p + ".weight")) { - return ggml_div(ctx0, ggml_sub(ctx0, x, mean), sd); +// 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, cbx_t(model, p + ".weight"), sd); - ggml_tensor * shift = ggml_sub(ctx0, cbx_t(model, p + ".bias"), ggml_mul(ctx0, mean, a)); + 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 + relu on a conv2d activation [W=T, H=F, C, 1], stats on ne2 -static ggml_tensor * cbx_bn2d_relu(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, - const std::string & p, ggml_tensor * eps) { +// 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, cbx_t(model, p + ".running_mean"), 1, 1, C, 1); - ggml_tensor * var = ggml_reshape_4d(ctx0, cbx_t(model, p + ".running_var"), 1, 1, C, 1); - ggml_tensor * w = ggml_reshape_4d(ctx0, cbx_t(model, p + ".weight"), 1, 1, C, 1); - ggml_tensor * b = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bias"), 1, 1, C, 1); + 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_relu(ctx0, ggml_add(ctx0, ggml_mul(ctx0, x, a), shift)); + 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 -static ggml_tensor * cbx_res2d(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, - const std::string & p, int stride, ggml_tensor * eps) { - ggml_tensor * cur = ggml_conv_2d(ctx0, cbx_t(model, p + ".conv1.weight"), x, 1, stride, 1, 1, 1, 1); - cur = cbx_bn2d_relu(model, ctx0, cur, p + ".bn1", eps); - cur = ggml_conv_2d(ctx0, cbx_t(model, p + ".conv2.weight"), cur, 1, 1, 1, 1, 1, 1); +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 - { - const int C = (int) cur->ne[2]; - ggml_tensor * mean = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.running_mean"), 1, 1, C, 1); - ggml_tensor * var = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.running_var"), 1, 1, C, 1); - ggml_tensor * w = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.weight"), 1, 1, C, 1); - ggml_tensor * b = ggml_reshape_4d(ctx0, cbx_t(model, p + ".bn2.bias"), 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)); - cur = ggml_add(ctx0, ggml_mul(ctx0, cur, a), shift); - } + cur = cbx_bn2d(ctx0, r.bn2, cur, eps); ggml_tensor * res = x; - if (model.cbx_tensors.count(p + ".shortcut.0.weight")) { - res = ggml_conv_2d(ctx0, cbx_t(model, p + ".shortcut.0.weight"), x, 1, stride, 0, 0, 1, 1); - const int C = (int) res->ne[2]; - ggml_tensor * mean = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.running_mean"), 1, 1, C, 1); - ggml_tensor * var = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.running_var"), 1, 1, C, 1); - ggml_tensor * w = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.weight"), 1, 1, C, 1); - ggml_tensor * b = ggml_reshape_4d(ctx0, cbx_t(model, p + ".shortcut.1.bias"), 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)); - res = ggml_add(ctx0, ggml_mul(ctx0, res, a), shift); + 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] -static ggml_tensor * cbx_cam_layer(const clip_model & model, ggml_context * ctx0, ggml_tensor * x, - const std::string & p, int dil, ggml_tensor * eps, ggml_tensor * segfix) { - ggml_tensor * h = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, p + ".nonlinear1.batchnorm", eps)); - h = cbx_conv1d(ctx0, cbx_t(model, p + ".linear1.weight"), nullptr, h, 1, 0, 0); - h = ggml_relu(ctx0, cbx_bn1d(model, ctx0, h, p + ".nonlinear2.batchnorm", eps)); - - const std::string cp = p + ".cam_layer"; - ggml_tensor * k = cbx_t(model, cp + ".linear_local.weight"); +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(ctx0, k, nullptr, h, pad, 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]; @@ -98,14 +79,16 @@ static ggml_tensor * cbx_cam_layer(const clip_model & model, ggml_context * ctx0 seg = ggml_cont(ctx0, ggml_transpose(ctx0, exp)); // [C, T] } ggml_tensor * context = ggml_add(ctx0, seg, gmean); - context = cbx_conv1d(ctx0, cbx_t(model, cp + ".linear1.weight"), cbx_t(model, cp + ".linear1.bias"), context, 1, 0, 0); + context = cbx_conv1d(l.ctx1_w, l.ctx1_b, context, 1, 0, 0); context = ggml_relu(ctx0, context); - context = cbx_conv1d(ctx0, cbx_t(model, cp + ".linear2.weight"), cbx_t(model, cp + ".linear2.bias"), context, 1, 0, 0); + 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); @@ -122,37 +105,36 @@ ggml_cgraph * clip_graph_chatterbox_spkenc::build() { // 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, cbx_t(model, "a.spk.head.conv1.weight"), x, 1, 1, 1, 1, 1, 1); - x = cbx_bn2d_relu(model, ctx0, x, "a.spk.head.bn1", eps); - x = cbx_res2d(model, ctx0, x, "a.spk.head.layer1.0", 2, eps); - x = cbx_res2d(model, ctx0, x, "a.spk.head.layer1.1", 1, eps); - x = cbx_res2d(model, ctx0, x, "a.spk.head.layer2.0", 2, eps); - x = cbx_res2d(model, ctx0, x, "a.spk.head.layer2.1", 1, eps); - x = ggml_conv_2d(ctx0, cbx_t(model, "a.spk.head.conv2.weight"), x, 1, 2, 1, 1, 1, 1); - x = cbx_bn2d_relu(model, ctx0, x, "a.spk.head.bn2", eps); + 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(ctx0, cbx_t(model, "a.spk.xvector.tdnn.linear.weight"), nullptr, x, 2, 2, 2); // [128, T1] - x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, "a.spk.xvector.tdnn.nonlinear.batchnorm", eps)); + 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 std::string bp = "a.spk.xvector.block" + std::to_string(bi); - for (int li = 1; model.cbx_tensors.count(bp + ".tdnnd" + std::to_string(li) + ".linear1.weight"); li++) { - ggml_tensor * out = cbx_cam_layer(model, ctx0, x, bp + ".tdnnd" + std::to_string(li), - block_dil[bi - 1], eps, segfix); + 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); } - const std::string tp = "a.spk.xvector.transit" + std::to_string(bi); - x = ggml_relu(ctx0, cbx_bn1d(model, ctx0, x, tp + ".nonlinear.batchnorm", eps)); - x = cbx_conv1d(ctx0, cbx_t(model, tp + ".linear.weight"), nullptr, x, 1, 0, 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, cbx_bn1d(model, ctx0, x, "a.spk.xvector.out_nonlinear.batchnorm", eps)); // [512, T1] + 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] @@ -166,9 +148,9 @@ ggml_cgraph * clip_graph_chatterbox_spkenc::build() { cb(stats, "spk_stats_pool", -1); // dense 1024 -> 192, batchnorm without affine, into the x-vector - ggml_tensor * dw = ggml_reshape_2d(ctx0, cbx_t(model, "a.spk.xvector.dense.linear.weight"), 1024, 192); + 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 = cbx_bn1d(model, ctx0, emb, "a.spk.xvector.dense.nonlinear.batchnorm", eps); + emb = bn1d(c.spk_dense_bn, emb, eps); emb = ggml_reshape_1d(ctx0, emb, 192); ggml_set_name(emb, "out_xvec"); ggml_set_output(emb); @@ -177,8 +159,7 @@ ggml_cgraph * clip_graph_chatterbox_spkenc::build() { // normalize then the s3gen speaker affine ggml_tensor * n2 = ggml_sqrt(ctx0, ggml_sum(ctx0, ggml_mul(ctx0, emb, emb))); ggml_tensor * unit = ggml_div(ctx0, emb, n2); - ggml_tensor * spk80 = cbx_linear(ctx0, cbx_t(model, "a.spk_embed_affine_layer.weight"), - cbx_t(model, "a.spk_embed_affine_layer.bias"), + ggml_tensor * spk80 = cbx_linear(c.spk_affine_w, c.spk_affine_b, ggml_reshape_2d(ctx0, unit, 192, 1)); spk80 = ggml_cont(ctx0, ggml_reshape_1d(ctx0, spk80, 80)); cb(spk80, "spk_embd", -1); diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index a2934a98e9c6..b89e34dd213e 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -151,21 +151,32 @@ struct clip_graph_conformer : clip_graph { ggml_cgraph * build() override; }; -// chatterbox helpers shared between the gen and spkenc graphs +// linear/conv builders shared between the chatterbox gen and spkenc graphs // (defined in chatterbox-gen.cpp) -ggml_tensor * cbx_t(const clip_model & model, const std::string & name); -ggml_tensor * cbx_linear(ggml_context * ctx0, ggml_tensor * w, ggml_tensor * b, ggml_tensor * x); -ggml_tensor * cbx_conv1d(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * b, ggml_tensor * x, - int stride, int pad_l, int pad_r); -ggml_tensor * cbx_conv1d_dil(ggml_context * ctx0, ggml_tensor * k, ggml_tensor * b, ggml_tensor * x, - int pad, int dil); +struct clip_graph_chatterbox_base : clip_graph { + using clip_graph::clip_graph; -struct clip_graph_chatterbox_spkenc : clip_graph { - clip_graph_chatterbox_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + 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 { +struct clip_graph_chatterbox : clip_graph_chatterbox_base { clip_gen_process_type gen_process = CLIP_GEN_PROCESS_CODE_GEN; int n_tokens = 0; int n_prompt_mel = 0; @@ -173,9 +184,20 @@ struct clip_graph_chatterbox : clip_graph { 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(ctx, img), gen_process(gen_process), n_tokens(n_tokens), n_prompt_mel(n_prompt_mel), + : 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 { diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index c98fdaa120f2..10e36fc1c62b 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -857,7 +857,7 @@ bool mtmd_audio_preprocessor_qwen3tts_spk::preprocess(const float * } // whisper style log-mel of the chatterbox s3 tokenizer (s3tokenizer.py) -bool mtmd_audio_s3tok_log_mel(const float * samples, size_t n_samples, +static bool mtmd_audio_s3tok_log_mel(const float * samples, size_t n_samples, const float * filters, int n_mel, std::vector & out, int & n_frames) { const int n_fft = 400; @@ -925,7 +925,7 @@ bool mtmd_audio_s3tok_log_mel(const float * samples, size_t n_samples, // 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. -void mtmd_audio_upsample_3_2(const float * samples, size_t n_samples, std::vector & out) { +static void mtmd_audio_upsample_3_2(const float * samples, size_t n_samples, std::vector & out) { const int W = 16; // sinc half width in source samples const double fc = 0.495; // cutoff, normalized to the source rate @@ -948,7 +948,7 @@ void mtmd_audio_upsample_3_2(const float * samples, size_t n_samples, std::vecto } // matcha style log-mel of the s3gen prompt features (s3gen/utils/mel.py) -bool mtmd_audio_matcha_log_mel(const float * samples, size_t n_samples, +static bool mtmd_audio_matcha_log_mel(const float * samples, size_t n_samples, std::vector & out, int & n_frames) { const int n_fft = 1920; const int hop = 480; @@ -1003,7 +1003,7 @@ bool mtmd_audio_matcha_log_mel(const float * samples, size_t n_samples, // 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 -void mtmd_audio_trim_silence(const float * samples, size_t n_samples, float top_db, +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; @@ -1044,7 +1044,7 @@ void mtmd_audio_trim_silence(const float * samples, size_t n_samples, float top_ // 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 -bool mtmd_audio_ve_mel(const float * samples, size_t n_samples, +static bool mtmd_audio_ve_mel(const float * samples, size_t n_samples, std::vector & out, int & n_frames) { const int n_fft = 400; const int hop = 160; @@ -1102,7 +1102,7 @@ bool mtmd_audio_ve_mel(const float * samples, size_t n_samples, } // ITU-R BS.1770 integrated loudness of a mono signal, matching pyloudnorm -float mtmd_audio_lufs(const float * samples, size_t n_samples, int sample_rate) { +static float mtmd_audio_lufs(const float * samples, size_t n_samples, int sample_rate) { std::vector y(samples, samples + n_samples); auto biquad = [&](double b0, double b1, double b2, double a1, double a2) { @@ -1184,7 +1184,7 @@ float mtmd_audio_lufs(const float * samples, size_t n_samples, int sample_rate) } // -// mtmd_audio_preprocessor_chatterbox_spk +// 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): @@ -1194,7 +1194,7 @@ float mtmd_audio_lufs(const float * samples, size_t n_samples, int sample_rate) // reference's own cepstral mean subtraction over time. // -void mtmd_audio_preprocessor_chatterbox_spk::initialize() { +void mtmd_audio_preprocessor_chatterbox_ref::initialize() { const int frame_len = 400; const int n_fft = 512; const int n_bins = n_fft / 2; @@ -1228,80 +1228,175 @@ void mtmd_audio_preprocessor_chatterbox_spk::initialize() { } } -bool mtmd_audio_preprocessor_chatterbox_spk::preprocess(const float * samples, +bool mtmd_audio_preprocessor_chatterbox_ref::preprocess(const float * samples, size_t n_samples, std::vector & output) { - 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) n_samples < frame_len) { - return false; + output.clear(); + const int sr = (int) hparams.audio_sample_rate; + + // the turbo variant loudness-normalizes the whole clip before any feature + std::vector pcm(samples, samples + n_samples); + if (!is_mtl) { + const float lufs = mtmd_audio_lufs(pcm.data(), pcm.size(), sr); + if (lufs != -HUGE_VALF) { + const float gain = powf(10.0f, (-27.0f - lufs) / 20.0f); + if (std::isfinite(gain) && gain > 0.0f) { + for (float & v : pcm) { + v *= gain; + } + } + } } - const int n_frames = 1 + ((int) n_samples - frame_len) / hop; - GGML_ASSERT(!window.empty()); - GGML_ASSERT(!filters.empty()); + // 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; - 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); + if ((int) pcm.size() < frame_len) { + return false; + } + const int n_frames = 1 + ((int) pcm.size() - frame_len) / hop; - std::vector frame(n_fft); - std::vector power(n_bins); - for (int fr = 0; fr < n_frames; fr++) { - const float * x = samples + (size_t) fr * hop; + GGML_ASSERT(!window.empty()); + GGML_ASSERT(!filters.empty()); - double mean = 0.0; - for (int i = 0; i < frame_len; i++) { - mean += x[i]; - } - mean /= frame_len; + 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); - 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); + std::vector frame(n_fft); + std::vector power(n_bins); + for (int fr = 0; fr < n_frames; fr++) { + const float * x = pcm.data() + (size_t) fr * hop; - for (int k = 0; k < n_bins; k++) { - double re = 0.0, im = 0.0; + double mean = 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); + 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)); } - 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]; + 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; } - out.data[(size_t) m * n_frames + fr] = (float) log(std::max(e, (double) FLT_EPSILON)); } + + output.push_back(std::move(out)); } - // reference extract_feature subtracts the per-channel mean over time - 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]; + // entries 1 and 2: s3 tokenizer log-mels at the flow and t3 caps, each + // clip padded to whole 40 ms tokens so the mel stays twice the token grid + const int n_mel_s3 = (int) (s3tok_filters.size() / (400 / 2 + 1)); + auto s3tok_entry = [&](size_t cap) -> bool { + std::vector clip(pcm.begin(), pcm.begin() + std::min(pcm.size(), cap)); + clip.resize((clip.size() + 639) / 640 * 640, 0.0f); + std::vector mel; + int n_frames = 0; + if (!mtmd_audio_s3tok_log_mel(clip.data(), clip.size(), s3tok_filters.data(), n_mel_s3, mel, n_frames)) { + return false; } - mean /= n_frames; - for (int fr = 0; fr < n_frames; fr++) { - row[fr] -= (float) mean; + mtmd_audio_mel out; + out.n_len = n_frames; + out.n_len_org = n_frames; + out.n_mel = n_mel_s3; + out.data = std::move(mel); + output.push_back(std::move(out)); + return true; + }; + const size_t gen_cap = (size_t) 10 * sr; + const size_t t3_cap = (size_t) (is_mtl ? 6 : 15) * sr; + if (n_mel_s3 == 0 || !s3tok_entry(gen_cap) || !s3tok_entry(t3_cap)) { + return false; + } + + // entry 3: s3gen prompt features, the flow-capped clip upsampled to the + // 24 kHz decoder rate then through the matcha log-mel + { + std::vector clip(pcm.begin(), pcm.begin() + std::min(pcm.size(), gen_cap)); + clip.resize((clip.size() + 639) / 640 * 640, 0.0f); + std::vector pcm24; + mtmd_audio_upsample_3_2(clip.data(), clip.size(), pcm24); + std::vector feat; + int n_frames = 0; + if (!mtmd_audio_matcha_log_mel(pcm24.data(), pcm24.size(), feat, n_frames)) { + return false; } + mtmd_audio_mel out; + out.n_len = n_frames; + out.n_len_org = n_frames; + out.n_mel = 80; + out.data = std::move(feat); + output.push_back(std::move(out)); } - output.push_back(std::move(out)); + // entry 4: voice encoder power mel of the silence-trimmed clip, padded + // (or trimmed) to the 160-frame partial grid at the reference 1.3 rate + { + size_t t0 = 0, t1 = 0; + mtmd_audio_trim_silence(pcm.data(), pcm.size(), 20.0f, t0, t1); + if (t1 <= t0) { + return false; + } + std::vector mel; + int n_frames = 0; + if (!mtmd_audio_ve_mel(pcm.data() + t0, t1 - t0, mel, n_frames)) { + return false; + } + const int n_mel = 40; + const int n_partial = 160; + const int step = (int) lround((16000.0 / 1.3) / n_partial); + int n_wins = std::max(n_frames - n_partial + step, 0) / step; + const int rem = std::max(n_frames - n_partial + step, 0) % step; + if (n_wins == 0 || (double) (rem + n_partial - step) / n_partial >= 0.8) { + n_wins++; + } + const int target = n_partial + step * (n_wins - 1); + mel.resize((size_t) target * n_mel, 0.0f); + mtmd_audio_mel out; + out.n_len = target; + out.n_len_org = target; + out.n_mel = n_mel; + out.data = std::move(mel); + output.push_back(std::move(out)); + } return true; } diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index 91f11704fe0c..acf7c0605425 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -50,46 +50,6 @@ struct mtmd_audio_cache { ); }; -// whisper style log-mel used by the chatterbox s3 tokenizer front-end: -// hann 400 periodic, hop 160, centered frames with reflect padding, power -// spectrum, caller-supplied mel filters [n_mel x (n_fft / 2 + 1)], log10 -// clamped to 1e-10, global max - 8 dynamic range, (x + 4) / 4 scaling. -// output layout matches the audio batch entries: out[m * n_frames + t]. -bool mtmd_audio_s3tok_log_mel(const float * samples, size_t n_samples, - const float * filters, int n_mel, - std::vector & out, int & n_frames); - -// rational 3/2 upsampler (16 kHz -> 24 kHz), windowed sinc polyphase. -// output length is exactly n_samples * 3 / 2 for even n_samples. -void mtmd_audio_upsample_3_2(const float * samples, size_t n_samples, std::vector & out); - -// matcha style log-mel of the chatterbox s3gen prompt features (utils/mel.py): -// 24 kHz, n_fft 1920, hop 480, hann 1920 periodic, (n_fft - hop) / 2 reflect -// padding with center false, magnitude spectrum, slaney mel 80 bins fmin 0 -// fmax 8000, natural log clamped to 1e-5. -// output layout is frame major: out[t * n_mel + m], as the prompt features -// are consumed row by row at mel rate. -bool mtmd_audio_matcha_log_mel(const float * samples, size_t n_samples, - std::vector & out, int & n_frames); - -// librosa.effects.trim replica: rms over centered 2048/512 windows with zero -// padding, threshold top_db below the loudest window, returns the sample -// span [start, end) of the non-silent region (start == end when all silent). -void mtmd_audio_trim_silence(const float * samples, size_t n_samples, float top_db, - size_t & start, size_t & end); - -// power mel of the chatterbox voice encoder (voice_encoder/melspec.py): -// 16 kHz, hann 400 periodic, hop 160, centered frames with reflect padding, -// squared magnitude against slaney mel 40 bins fmin 0 fmax 8000, no log. -// output layout is frame major: out[t * 40 + m]. -bool mtmd_audio_ve_mel(const float * samples, size_t n_samples, - std::vector & out, int & n_frames); - -// ITU-R BS.1770 integrated loudness (pyloudnorm replica, mono): K-weighting -// (RBJ high shelf 1681.97 Hz +4 dB then high pass 38.14 Hz), 400 ms blocks -// with 75% overlap, absolute -70 then relative -10 gating. -// returns the loudness in LUFS, or -HUGE_VALF when everything is gated out. -float mtmd_audio_lufs(const float * samples, size_t n_samples, int sample_rate); struct mtmd_audio_preprocessor { const clip_hparams & hparams; @@ -170,14 +130,27 @@ struct mtmd_audio_preprocessor_qwen3tts_spk : mtmd_audio_preprocessor { mtmd_audio_cache cache; }; -struct mtmd_audio_preprocessor_chatterbox_spk : mtmd_audio_preprocessor { - mtmd_audio_preprocessor_chatterbox_spk(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} +// full reference chain front-end of the chatterbox speaker encoder: one +// preprocess call emits the host DSP products of a reference clip as five +// entries, consumed positionally by the encoder orchestration in mtmd.cpp +// 0: CAMPPlus kaldi fbank [80 x n_frames] (mel major) +// 1: s3 tokenizer log-mel, flow cap [n_mel x n_frames] (mel major) +// 2: s3 tokenizer log-mel, t3 cap [n_mel x n_frames] (mel major) +// 3: s3gen prompt features [n_frames x 80] (frame major, 24 kHz mel rate) +// 4: voice encoder power mel [n_frames x 40] (frame major, trimmed, partial grid) +// the turbo variant loudness-normalizes the clip to -27 LUFS first and caps +// the t3 clip at 15 s instead of 6 s +struct mtmd_audio_preprocessor_chatterbox_ref : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_chatterbox_ref(const clip_ctx * ctx, bool is_mtl, std::vector s3tok_filters) + : mtmd_audio_preprocessor(ctx), is_mtl(is_mtl), s3tok_filters(std::move(s3tok_filters)) {} void initialize() override; bool preprocess(const float * samples, size_t n_samples, std::vector & output) override; private: - std::vector window; // povey window, frame_length points - std::vector filters; // kaldi mel filterbank, n_mel x (n_fft / 2) dense + bool is_mtl; + std::vector s3tok_filters; // s3 tokenizer filterbank [n_mel x (400 / 2 + 1)], from the mmproj + std::vector window; // povey window of the fbank front-end, frame_length points + std::vector filters; // kaldi mel filterbank of the fbank front-end, n_mel x (n_fft / 2) dense }; struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor { diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 16c73f2f77e9..bc16662b41c3 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -91,6 +91,40 @@ class mtmd_gen_audio_pipeline { virtual int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) = 0; protected: + // encodes a speaker reference clip through the standard audio path: + // bitmap to chunks to chunk encode, out receives the encoder embeddings + bool encode_speaker(mtmd_bitmap * bitmap, std::vector & out) { + if (!mtmd_support_audio(mctx)) { + LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n"); + return false; + } + const std::string marker = mtmd_default_marker(); + mtmd_input_text text{ marker.c_str(), marker.size(), false, true }; + mtmd_input_chunks * chunks = mtmd_input_chunks_init(); + const mtmd_bitmap * bptr = bitmap; + bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0; + if (ok) { + ok = false; + for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) { + const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i); + if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) { + continue; + } + if (mtmd_encode_chunk(mctx, chunk) != 0) { + LOG_ERR("mtmd_helper_gen_audio: speaker encode failed\n"); + break; + } + const float * embd = mtmd_get_output_embd(mctx); + const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk); + out.assign(embd, embd + n); + ok = true; + break; + } + } + mtmd_input_chunks_free(chunks); + return ok; + } + llama_context * lctx; mtmd_context * mctx; const llama_model * model; @@ -325,38 +359,6 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { // encodes a reference wav (already loaded as a bitmap) through the mmproj's // speaker encoder, returning the single x-vector embedding row it produces - bool encode_speaker(mtmd_bitmap * bitmap, std::vector & out) { - if (!mtmd_support_audio(mctx)) { - LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n"); - return false; - } - const std::string marker = mtmd_default_marker(); - mtmd_input_text text{ marker.c_str(), marker.size(), false, true }; - mtmd_input_chunks * chunks = mtmd_input_chunks_init(); - const mtmd_bitmap * bptr = bitmap; - bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0; - if (ok) { - ok = false; - for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) { - const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i); - if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) { - continue; - } - if (mtmd_encode_chunk(mctx, chunk) != 0) { - LOG_ERR("mtmd_helper_gen_audio: speaker encode failed\n"); - break; - } - const float * embd = mtmd_get_output_embd(mctx); - const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk); - out.assign(embd, embd + n); - ok = true; - break; - } - } - mtmd_input_chunks_free(chunks); - return ok; - } - // runs one CODE2WAV process() call on whatever is currently buffered, carrying // the persisted state (KV cache + conv left-context) across batches bool flush_c2w() { @@ -430,7 +432,9 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { h_state_buf.clear(); out_buf.clear(); ref_cond.clear(); - ref_state.clear(); + ref_codes.clear(); + ref_feat.clear(); + ref_spk.clear(); } int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override { @@ -441,21 +445,20 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } if (inp->speaker_ref) { - // one spk-ref call encodes the reference clip: conditioning rows - // for the talker prompt, opaque state for the flow decoder - const float * pcm = (const float *) mtmd_bitmap_get_data(inp->speaker_ref); - const size_t n = mtmd_bitmap_get_n_bytes(inp->speaker_ref) / sizeof(float); - mtmd_gen_inp gi{}; - gi.type = MTMD_GEN_PROCESS_TYPE_SPK_REF; - gi.pcm = pcm; - gi.n_pcm = n; - mtmd_gen_out go{}; - if (mtmd_gen_audio_process(mctx, &gi, &go) != 0) { + // the reference clip goes through the standard audio encode path: + // the chunk encodes to the talker conditioning rows, the flow + // decoder reference comes back through the typed side outputs + if (!encode_speaker(inp->speaker_ref, ref_cond)) { LOG_ERR("mtmd_helper_gen_audio: speaker reference encoding failed\n"); return 1; } - ref_cond.assign(go.embd, go.embd + go.n_embd); - ref_state.assign(go.state_data, go.state_data + go.state_size); + size_t n = 0; + const float * p = mtmd_get_output_typed_embd(mctx, MTMD_EMBD_OUT_TYPE_REF_CODES, &n); + ref_codes.assign(p, p + n); + p = mtmd_get_output_typed_embd(mctx, MTMD_EMBD_OUT_TYPE_REF_FEAT, &n); + ref_feat.assign(p, p + n); + p = mtmd_get_output_typed_embd(mctx, MTMD_EMBD_OUT_TYPE_REF_SPK, &n); + ref_spk.assign(p, p + n); } const int n_e = n_embd; @@ -697,9 +700,10 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { gen_inp.type = MTMD_GEN_PROCESS_TYPE_CODE2WAV; gen_inp.codes = codes_buf.data(); gen_inp.n_codes = codes_buf.size(); - if (!ref_state.empty()) { - gen_inp.state_data = ref_state.data(); - gen_inp.state_size = ref_state.size(); + if (!ref_codes.empty()) { + gen_inp.ref_codes = ref_codes.data(); gen_inp.n_ref_codes = ref_codes.size(); + gen_inp.ref_feat = ref_feat.data(); gen_inp.n_ref_feat = ref_feat.size(); + gen_inp.ref_spk = ref_spk.data(); gen_inp.n_ref_spk = ref_spk.size(); } mtmd_gen_out gen_out{}; if (mtmd_gen_audio_process(mctx, &gen_inp, &gen_out) != 0) { @@ -825,7 +829,9 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::vector text_pos; std::vector speech_pos; std::vector ref_cond; - std::vector ref_state; + std::vector ref_codes; // flow reference codes carried as exact float values + std::vector ref_feat; + std::vector ref_spk; llama_token speech_base = LLAMA_TOKEN_NULL; int n_speech = 0; llama_token text_start = LLAMA_TOKEN_NULL; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 958661ea9850..bce7a2f3ff0d 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -273,6 +273,13 @@ struct mtmd_context { std::vector gen_out_audio; // decoded PCM samples for the current frame (CODE2WAV) std::vector gen_out_state; // state to feed into the next CODE2WAV call + // typed side outputs of the last encoded reference chunk (chatterbox): + // flow codes carried as exact float values, mel-rate features, speaker + // vector; read back through mtmd_get_output_typed_embd + std::vector gen_out_ref_codes; + std::vector gen_out_ref_feat; + std::vector gen_out_ref_spk; + bool print_timings; int n_threads; std::string media_marker; @@ -383,10 +390,7 @@ struct mtmd_context { // since we already validate n_embd of vision and audio mmproj, // we can safely assume that they are the same - // gen-only mmproj has no input projection, and a speaker encoder - // attached to a gen model outputs a conditioning vector, not tokens - // in the backbone embedding space - if (ctx_v || (ctx_a && !ctx_gen_a)) { + if (ctx_v || ctx_a) { int n_embd_clip = clip_n_mmproj_embd(ctx_v ? ctx_v : ctx_a); if (n_embd_text > 0 && n_embd_text != n_embd_clip) { throw std::runtime_error(string_format( @@ -773,7 +777,22 @@ struct mtmd_context { } break; case PROJECTOR_TYPE_CHATTERBOX_SPKENC: { - audio_preproc = std::make_unique(ctx_a); + { + // the ref-chain front-end needs the s3 tokenizer + // filterbank of the generation context and the + // variant to pick its caps and loudness handling + std::vector filt; + bool is_mtl = false; + if (ctx_gen_a) { + const size_t n_filt = clip_cbx_read_tensor(ctx_gen_a, "a.s3tok.mel_filters", nullptr, 0); + if (n_filt > 0) { + filt.resize(n_filt); + clip_cbx_read_tensor(ctx_gen_a, "a.s3tok.mel_filters", filt.data(), n_filt); + } + is_mtl = clip_cbx_read_tensor(ctx_gen_a, "a.gen.t3.speech_pos_emb", nullptr, 0) > 0; + } + audio_preproc = std::make_unique(ctx_a, is_mtl, std::move(filt)); + } } break; default: throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj)); @@ -1351,6 +1370,41 @@ struct mtmd_tokenizer { } } + // the chatterbox reference encoder consumes the five DSP entries + // of one clip as a single chunk; its token count is the number of + // talker conditioning rows the clip encodes to (fixed on the + // multilingual variant, spkr row plus the t3 token grid on turbo) + if (clip_get_projector_type(ctx->ctx_a) == PROJECTOR_TYPE_CHATTERBOX_SPKENC) { + GGML_ASSERT(mel_spec_chunks.size() == 5); + const bool mtl = ctx->ctx_gen_a && + clip_cbx_read_tensor(ctx->ctx_gen_a, "a.gen.t3.speech_pos_emb", nullptr, 0) > 0; + const size_t n_tokens = mtl ? 34 : 1 + (size_t) mel_spec_chunks[2].n_len / 4; + + clip_image_f32_batch batch_f32; + batch_f32.is_audio = true; + for (auto & mel_spec : mel_spec_chunks) { + GGML_ASSERT(mel_spec.n_len <= INT32_MAX && mel_spec.n_len >= 0); + GGML_ASSERT(mel_spec.n_mel <= INT32_MAX && mel_spec.n_mel >= 0); + clip_image_f32 mel_f32; + mel_f32.set_size({(int) mel_spec.n_len, (int) mel_spec.n_mel}, + mel_spec.data.empty(), /* is_audio */ true); + mel_f32.cpy_buf(mel_spec.data); + batch_f32.entries.push_back(std::move(mel_f32)); + } + + mtmd_audio_tokens_ptr audio_tokens(new mtmd_audio_tokens); + audio_tokens->n_tokens = (uint32_t) n_tokens; + audio_tokens->batch_f32 = std::move(batch_f32); + audio_tokens->id = bitmap->id; + + mtmd_input_chunk chunk{ + MTMD_INPUT_CHUNK_TYPE_AUDIO, + {}, // text tokens + nullptr, // image tokens + std::move(audio_tokens), + }; + cur.entries.emplace_back(std::move(chunk)); + } else // consider each mel_spec as a separate audio chunk // TODO: maybe support batching, but this may come with memory cost for (auto & mel_spec : mel_spec_chunks) { @@ -1521,6 +1575,10 @@ static int32_t mtmd_encode_impl(mtmd_context * ctx, const mtmd_image_tokens * im return ok ? 0 : 1; } +// defined with the reference chain below; encodes a chatterbox reference +// chunk into talker conditioning rows and the typed decoder reference outputs +static int32_t cbx_ref_encode(mtmd_context * ctx, const mtmd_audio_tokens * audio_tokens, std::vector & out_embd); + static int32_t mtmd_encode_chunk_impl(mtmd_context * ctx, const mtmd_input_chunk * chunk, std::vector & out_embd) { if (chunk->type == MTMD_INPUT_CHUNK_TYPE_TEXT) { LOG_WRN("mtmd_encode_chunk has no effect for text chunks\n"); @@ -1552,6 +1610,9 @@ static int32_t mtmd_encode_chunk_impl(mtmd_context * ctx, const mtmd_input_chunk LOG_ERR("%s: audio tokens batch is placeholder\n", __func__); return 1; } + if (clip_get_projector_type(ctx->ctx_a) == PROJECTOR_TYPE_CHATTERBOX_SPKENC) { + return cbx_ref_encode(ctx, chunk->tokens_audio.get(), out_embd); + } int n_mmproj_embd = ctx->n_embd_out(); out_embd.resize((size_t)chunk->tokens_audio->n_tokens * n_mmproj_embd); bool ok = clip_image_batch_encode( @@ -1589,6 +1650,23 @@ float * mtmd_get_output_embd(mtmd_context * ctx) { return ctx->out_embd.data(); } +const float * mtmd_get_output_typed_embd(mtmd_context * ctx, + enum mtmd_embd_out_type type, + size_t * n_elements) { + const std::vector * buf = nullptr; + switch (type) { + case MTMD_EMBD_OUT_TYPE_REF_CODES: buf = &ctx->gen_out_ref_codes; break; + case MTMD_EMBD_OUT_TYPE_REF_FEAT: buf = &ctx->gen_out_ref_feat; break; + case MTMD_EMBD_OUT_TYPE_REF_SPK: buf = &ctx->gen_out_ref_spk; break; + } + if (!buf || buf->empty()) { + *n_elements = 0; + return nullptr; + } + *n_elements = buf->size(); + return buf->data(); +} + // // audio generation // @@ -1622,45 +1700,21 @@ size_t mtmd_gen_audio_read_tensor(mtmd_context * ctx, const char * name, float * return clip_cbx_read_tensor(ctx->ctx_gen_a, name, out, n_max); } -// turbo only: -27 LUFS loudness normalization of the reference clip -// (tts_turbo.py); the multilingual variant is identified by its perceiver -static void cbx_ref_norm(mtmd_context * ctx, float * pcm, size_t n) { - if (clip_cbx_read_tensor(ctx->ctx_gen_a, "a.cenc.perceiver.pre_attention_query", nullptr, 0) > 0 || - clip_cbx_read_tensor(ctx->ctx_gen_a, "a.cenc.spkr_enc.weight", nullptr, 0) == 0) { - return; - } - const float lufs = mtmd_audio_lufs(pcm, n, clip_get_hparams(ctx->ctx_a ? ctx->ctx_a : ctx->ctx_gen_a)->audio_sample_rate); - if (lufs == -HUGE_VALF) { - return; - } - const float gain = powf(10.0f, (-27.0f - lufs) / 20.0f); - if (!std::isfinite(gain) || gain <= 0.0f) { - return; - } - for (size_t i = 0; i < n; i++) { - pcm[i] *= gain; - } -} - -// 80-dim flow speaker vector of the reference clip: fbank features from the -// audio preprocessor into the CAMPPlus graph of the speaker encoding context -static bool cbx_ref_spk80(mtmd_context * ctx, const float * pcm, size_t n_pcm, std::vector & spk80) { - if (!ctx->ctx_a || !ctx->audio_preproc) { +// 80-dim flow speaker vector of the reference clip: the precomputed fbank +// entry through the CAMPPlus graph of the speaker encoding context +static bool cbx_ref_spk80(mtmd_context * ctx, const clip_image_f32 & fbank, std::vector & spk80) { + const size_t n = clip_cbx_read_tensor(ctx->ctx_a, "a.spk_embed_affine_layer.bias", nullptr, 0); + if (n == 0) { LOG_ERR("%s: mmproj has no speaker encoder\n", __func__); return false; } - std::vector mels; - if (!ctx->audio_preproc->preprocess(pcm, n_pcm, mels) || mels.empty() || mels[0].data.empty()) { - LOG_ERR("%s: speaker features failed\n", __func__); - return false; - } clip_image_f32 mel_img; - mel_img.set_size({(int) mels[0].n_len, (int) mels[0].n_mel}, false, true); - mel_img.cpy_buf(std::move(mels[0].data)); + 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)); - spk80.resize((size_t) clip_n_mmproj_embd(ctx->ctx_a)); + spk80.resize(n); if (!clip_image_batch_encode(ctx->ctx_a, ctx->n_threads, &batch, spk80)) { LOG_ERR("%s: speaker encoder failed\n", __func__); return false; @@ -1668,37 +1722,14 @@ static bool cbx_ref_spk80(mtmd_context * ctx, const float * pcm, size_t n_pcm, s return true; } -// runs the s3 speech tokenizer on a reference clip; 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 float * ref, size_t n_ref, +// s3 speech tokens of a precomputed reference log-mel entry; rows, when +// requested, receives the speech embedding rows of the codes (with the +// learned speech positions added on the multilingual variant) +static int32_t cbx_ref_tokenize(mtmd_context * ctx, const clip_image_f32 & mel, std::vector & codes, std::vector * rows) { -clip_ctx * ctx_clip = ctx->ctx_gen_a; - // mel filters shipped in the mmproj, [n_mels x (n_fft / 2 + 1)] - const size_t n_filt = clip_cbx_read_tensor(ctx_clip, "a.s3tok.mel_filters", nullptr, 0); - if (n_filt == 0) { - LOG_ERR("%s: model has no s3 tokenizer\n", __func__); - return 1; - } - std::vector filters(n_filt); - clip_cbx_read_tensor(ctx_clip, "a.s3tok.mel_filters", filters.data(), n_filt); - const int n_mel = (int) (n_filt / (400 / 2 + 1)); - - // pad to a whole number of 40 ms tokens so that the mel length stays - // twice the token length, as the reference prompt features expect - std::vector pcm(ref, ref + n_ref); - pcm.resize((pcm.size() + 639) / 640 * 640, 0.0f); - - std::vector mel; - int n_frames = 0; - if (!mtmd_audio_s3tok_log_mel(pcm.data(), pcm.size(), filters.data(), n_mel, mel, n_frames)) { - LOG_ERR("%s: log mel failed\n", __func__); - return 1; - } - clip_image_f32 mel_img; - mel_img.set_size({n_frames, n_mel}, false, true); - mel_img.cpy_buf(std::move(mel)); + 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; @@ -1714,7 +1745,7 @@ clip_ctx * ctx_clip = ctx->ctx_gen_a; params.out_codes = &out_codes; params.out_code_embd = rows ? &out_embd : nullptr; - if (!clip_encode(ctx_clip, ¶ms)) { + if (!clip_encode(ctx->ctx_gen_a, ¶ms)) { LOG_ERR("%s: clip_encode failed (tokenize)\n", __func__); return 1; } @@ -1730,9 +1761,9 @@ clip_ctx * ctx_clip = ctx->ctx_gen_a; // 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 float * pcm, size_t n_pcm, +static int32_t cbx_ref_cond(mtmd_context * ctx, const clip_image_f32 & ve_mel, const std::vector & pse, std::vector & out_rows) { -clip_ctx * ctx_clip = ctx->ctx_gen_a; + clip_ctx * ctx_clip = ctx->ctx_a; auto read_t = [&](const char * name, std::vector & v) -> bool { const size_t n = clip_cbx_read_tensor(ctx_clip, name, nullptr, 0); if (n == 0) { @@ -1742,33 +1773,15 @@ clip_ctx * ctx_clip = ctx->ctx_gen_a; return clip_cbx_read_tensor(ctx_clip, name, v.data(), n) == n; }; - // voice encoder reference chain (embeds_from_wavs): silence trim, - // 40-bin power mel, overlapping 160-frame partials at rate 1.3, - // 3-layer lstm per partial, projected/relu/normalized embeddings - // averaged into the utterance embedding - size_t t0 = 0, t1 = 0; - mtmd_audio_trim_silence(pcm, n_pcm, 20.0f, t0, t1); - if (t1 <= t0) { - LOG_ERR("%s: reference clip is silent\n", __func__); - return 1; - } - std::vector mel; - int n_frames = 0; - if (!mtmd_audio_ve_mel(pcm + t0, t1 - t0, mel, n_frames)) { - LOG_ERR("%s: voice encoder mel failed\n", __func__); - return 1; - } - + // voice encoder reference chain (embeds_from_wavs) over the precomputed + // power mel entry, already on the 160-frame partial grid: 3-layer lstm + // per partial, projected/relu/normalized embeddings averaged into the + // utterance embedding + const std::vector & mel = ve_mel.get_ro_buf(); const int n_mel = 40; const int n_partial = 160; const int step = (int) lround((16000.0 / 1.3) / n_partial); // reference rate 1.3 - 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); // zero pad (or trim) to the partial grid + const int n_wins = (ve_mel.nx() - n_partial) / step + 1; std::vector w_ih[3], w_hh[3], b_ih[3], b_hh[3]; std::vector w_proj, b_proj; @@ -1985,6 +1998,63 @@ clip_ctx * ctx_clip = ctx->ctx_gen_a; return 0; } +// encodes the five DSP entries of a chatterbox reference chunk (fbank, s3 +// log-mels at the flow and t3 caps, s3gen prompt features, voice encoder +// mel): the primary output is the talker conditioning rows in the backbone +// embedding space, the flow decoder reference (codes, mel-rate features, +// speaker vector) lands in the typed side outputs +static int32_t cbx_ref_encode(mtmd_context * ctx, const mtmd_audio_tokens * audio_tokens, std::vector & out_embd) { + if (!ctx->ctx_gen_a) { + LOG_ERR("%s: reference encoding needs the audio generation context\n", __func__); + return 1; + } + const auto & entries = audio_tokens->batch_f32.entries; + GGML_ASSERT(entries.size() == 5); + + std::vector spk80; + if (!cbx_ref_spk80(ctx, entries[0], spk80)) { + return 1; + } + + std::vector flow_codes, t3_codes; + std::vector t3_rows; + if (cbx_ref_tokenize(ctx, entries[1], flow_codes, nullptr) != 0 || + cbx_ref_tokenize(ctx, entries[2], t3_codes, &t3_rows) != 0) { + return 1; + } + + const size_t n_feat = (size_t) entries[3].nx() * 80; + if ((size_t) entries[3].nx() != 2 * flow_codes.size()) { + LOG_ERR("%s: reference mel length %d does not match %zu codes\n", + __func__, entries[3].nx(), flow_codes.size()); + return 1; + } + + // conditioning rows: [spkr] then, multilingual, the perceiver block over + // the reference rows, or, turbo, the raw reference rows + const bool mtl = clip_cbx_read_tensor(ctx->ctx_a, "a.cenc.perceiver.pre_attention_query", nullptr, 0) > 0; + std::vector pse; + if (mtl) { + pse = std::move(t3_rows); + } + std::vector rows; + if (cbx_ref_cond(ctx, entries[4], pse, rows) != 0) { + return 1; + } + if (!mtl) { + rows.insert(rows.end(), t3_rows.begin(), t3_rows.end()); + } + GGML_ASSERT(rows.size() == (size_t) audio_tokens->n_tokens * (size_t) ctx->n_embd_out()); + + ctx->gen_out_ref_codes.assign(flow_codes.begin(), flow_codes.end()); + ctx->gen_out_ref_feat.assign(entries[3].get_ro_buf().begin(), + entries[3].get_ro_buf().begin() + n_feat); + ctx->gen_out_ref_spk = std::move(spk80); + + out_embd = std::move(rows); + return 0; +} + static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) { clip_ctx * ctx_clip = ctx->ctx_gen_a; @@ -2031,92 +2101,6 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in return 0; } - if (inp->type == MTMD_GEN_PROCESS_TYPE_SPK_REF) { - if (!inp->pcm || inp->n_pcm == 0) { - LOG_ERR("%s: pcm required for spk ref\n", __func__); - return 1; - } - - // the whole reference encoding chain runs inside this stage: - // loudness normalization, speaker vector, speech tokenization at the - // flow and talker caps, conditioning rows, mel-rate prompt features - std::vector pcm(inp->pcm, inp->pcm + inp->n_pcm); - cbx_ref_norm(ctx, pcm.data(), pcm.size()); - - const bool mtl = clip_cbx_read_tensor(ctx_clip, "a.cenc.perceiver.pre_attention_query", nullptr, 0) > 0; - const size_t gen_cap = (size_t) 10 * 16000; - const size_t t3_cap = (size_t) (mtl ? 6 : 15) * 16000; - - std::vector spk80; - if (!cbx_ref_spk80(ctx, pcm.data(), pcm.size(), spk80)) { - return 1; - } - - std::vector flow_tokens, t3_tokens; - std::vector t3_rows; - if (cbx_ref_tokenize(ctx, pcm.data(), std::min(pcm.size(), gen_cap), flow_tokens, nullptr) != 0 || - cbx_ref_tokenize(ctx, pcm.data(), std::min(pcm.size(), t3_cap), t3_tokens, &t3_rows) != 0) { - return 1; - } - - // conditioning rows: [spkr] then, multilingual, the perceiver block - // over the reference rows, or, turbo, the raw reference rows - std::vector pse; - if (mtl) { - pse = std::move(t3_rows); - } - std::vector rows; - if (cbx_ref_cond(ctx, pcm.data(), pcm.size(), pse, rows) != 0) { - return 1; - } - if (!mtl) { - rows.insert(rows.end(), t3_rows.begin(), t3_rows.end()); - } - - // mel-rate prompt features of the flow reference at the 24 kHz s3gen - // rate, padded to the token grid so that the mel length stays twice - // the token length - std::vector feat; - { - std::vector pcm16(pcm.begin(), pcm.begin() + std::min(pcm.size(), gen_cap)); - pcm16.resize((pcm16.size() + 639) / 640 * 640, 0.0f); - std::vector pcm24; - mtmd_audio_upsample_3_2(pcm16.data(), pcm16.size(), pcm24); - int n_feat = 0; - if (!mtmd_audio_matcha_log_mel(pcm24.data(), pcm24.size(), feat, n_feat)) { - LOG_ERR("%s: reference mel failed\n", __func__); - return 1; - } - if ((size_t) n_feat != 2 * flow_tokens.size()) { - LOG_ERR("%s: reference mel length %d does not match %zu tokens\n", - __func__, n_feat, flow_tokens.size()); - return 1; - } - } - - // decoder reference state, opaque to the caller: - // [n_tokens, n_feat] i32 header, tokens, features, 80-dim speaker vector - std::vector blob(2 * sizeof(int32_t) - + flow_tokens.size() * sizeof(int32_t) - + (feat.size() + spk80.size()) * sizeof(float)); - { - uint8_t * q = blob.data(); - const int32_t hdr[2] = { (int32_t) flow_tokens.size(), (int32_t) feat.size() }; - memcpy(q, hdr, sizeof(hdr)); q += sizeof(hdr); - memcpy(q, flow_tokens.data(), flow_tokens.size() * sizeof(int32_t)); q += flow_tokens.size() * sizeof(int32_t); - memcpy(q, feat.data(), feat.size() * sizeof(float)); q += feat.size() * sizeof(float); - memcpy(q, spk80.data(), spk80.size() * sizeof(float)); - } - - ctx->gen_out_embd = std::move(rows); - ctx->gen_out_state = std::move(blob); - out->embd = ctx->gen_out_embd.data(); - out->n_embd = ctx->gen_out_embd.size(); - out->state_data = (const char *) ctx->gen_out_state.data(); - out->state_size = ctx->gen_out_state.size(); - return 0; - } - // MTMD_GEN_PROCESS_TYPE_CODE2WAV if (clip_get_projector_type(ctx_clip) == PROJECTOR_TYPE_CHATTERBOX) { if (!inp->codes || inp->n_codes == 0) { @@ -2126,31 +2110,24 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in std::vector in_codes(inp->codes, inp->codes + inp->n_codes); std::vector out_audio; - // optional reference state from the spk-ref stage: flow prompt - // tokens, mel-rate prompt features, 80-dim speaker vector; a null - // state selects the model's precomputed default voice + // optional voice cloning reference from the encoded reference chunk: + // flow prompt codes, mel-rate prompt features, 80-dim speaker vector; + // all null selects the model's precomputed default voice std::vector ref_tokens; std::vector ref_feat; std::vector ref_spk; - if (inp->state_data) { - int32_t hdr[2]; - if (inp->state_size < sizeof(hdr)) { - LOG_ERR("%s: malformed reference state\n", __func__); + const bool has_ref = inp->ref_codes && inp->n_ref_codes > 0; + if (has_ref) { + if (!inp->ref_feat || inp->n_ref_feat == 0 || !inp->ref_spk || inp->n_ref_spk == 0) { + LOG_ERR("%s: incomplete cloning reference\n", __func__); return 1; } - memcpy(hdr, inp->state_data, sizeof(hdr)); - const size_t n_tok = (size_t) hdr[0], n_feat = (size_t) hdr[1]; - if (inp->state_size != sizeof(hdr) + n_tok * sizeof(int32_t) + (n_feat + 80) * sizeof(float)) { - LOG_ERR("%s: malformed reference state\n", __func__); - return 1; + ref_tokens.resize(inp->n_ref_codes); + for (size_t i = 0; i < inp->n_ref_codes; i++) { + ref_tokens[i] = (int32_t) lroundf(inp->ref_codes[i]); } - const char * q = inp->state_data + sizeof(hdr); - ref_tokens.resize(n_tok); - memcpy(ref_tokens.data(), q, n_tok * sizeof(int32_t)); q += n_tok * sizeof(int32_t); - ref_feat.resize(n_feat); - memcpy(ref_feat.data(), q, n_feat * sizeof(float)); q += n_feat * sizeof(float); - ref_spk.resize(80); - memcpy(ref_spk.data(), q, 80 * sizeof(float)); + ref_feat.assign(inp->ref_feat, inp->ref_feat + inp->n_ref_feat); + ref_spk.assign(inp->ref_spk, inp->ref_spk + inp->n_ref_spk); } // the batch entry is unused, present to satisfy the encode interface @@ -2167,9 +2144,9 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in params.gen_process = CLIP_GEN_PROCESS_TTS; params.codes = &in_codes; params.out_audio = &out_audio; - params.ref_tokens = inp->state_data ? &ref_tokens : nullptr; - params.ref_feat = inp->state_data ? &ref_feat : nullptr; - params.ref_spk = inp->state_data ? &ref_spk : nullptr; + params.ref_tokens = has_ref ? &ref_tokens : nullptr; + params.ref_feat = has_ref ? &ref_feat : nullptr; + params.ref_spk = has_ref ? &ref_spk : nullptr; if (!clip_encode(ctx_clip, ¶ms)) { LOG_ERR("%s: clip_encode failed (code2wav)\n", __func__); diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 53c6650385a0..498ac36bfd24 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -295,6 +295,20 @@ MTMD_API int32_t mtmd_encode_chunk(mtmd_context * ctx, // llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk) * sizeof(float) MTMD_API float * mtmd_get_output_embd(mtmd_context * ctx); +// typed side outputs of the last encoded chunk, for encoders that produce +// more than the backbone-space embeddings; entries are model-defined and +// valid until the next encode call. n_elements receives the element count, +// the return is null when the encoder has no output of the requested type. +// integer outputs (audio codes) are carried as exact float values. +enum mtmd_embd_out_type { + MTMD_EMBD_OUT_TYPE_REF_CODES, // speech codes of the reference clip + MTMD_EMBD_OUT_TYPE_REF_FEAT, // mel-rate features of the reference clip + MTMD_EMBD_OUT_TYPE_REF_SPK, // speaker vector of the reference clip +}; +MTMD_API const float * mtmd_get_output_typed_embd(mtmd_context * ctx, + enum mtmd_embd_out_type type, + size_t * n_elements); + // batch encoding API // chunks are not owned by the batch, they will not be freed by mtmd_batch_free() @@ -351,8 +365,6 @@ MTMD_API size_t mtmd_gen_audio_read_tensor(mtmd_context * ctx, const char * name enum mtmd_gen_process_type { MTMD_GEN_PROCESS_TYPE_GEN_CODE, // h_state to codes MTMD_GEN_PROCESS_TYPE_CODE2WAV, // codes to raw PCM audio - MTMD_GEN_PROCESS_TYPE_SPK_REF, // raw PCM audio to talker conditioning rows - // and the decoder reference state }; struct mtmd_gen_inp { enum mtmd_gen_process_type type; @@ -366,15 +378,15 @@ struct mtmd_gen_inp { // for MTMD_GEN_PROCESS_TYPE_CODE2WAV int32_t * codes; size_t n_codes; - // opaque state: the decoder carry-over between calls, or the reference - // state returned by a MTMD_GEN_PROCESS_TYPE_SPK_REF call (null means the - // model's precomputed default voice) + // opaque state: the decoder carry-over between calls const char * state_data; size_t state_size; - - // for MTMD_GEN_PROCESS_TYPE_SPK_REF - const float * pcm; // mono float samples at the audio encoder sample rate - size_t n_pcm; + // voice cloning reference from mtmd_get_output_typed_embd after encoding + // the reference clip; all null selects the model's precomputed default + // voice + const float * ref_codes; size_t n_ref_codes; + const float * ref_feat; size_t n_ref_feat; + const float * ref_spk; size_t n_ref_spk; }; struct mtmd_gen_out { // note: output memory is allocated by the context, valid until next process() call @@ -384,16 +396,12 @@ struct mtmd_gen_out { size_t n_codes; const float * embd; // the generated hidden state, to be fed back to backbone // it must have n_text_embd elements - // for MTMD_GEN_PROCESS_TYPE_SPK_REF: embd holds the talker conditioning - // rows of the reference clip and n_embd their total element count size_t n_embd; // for MTMD_GEN_PROCESS_TYPE_CODE2WAV const float * audio; size_t n_samples; - // opaque state: the decoder carry-over to pass into the next CODE2WAV - // call, or, from MTMD_GEN_PROCESS_TYPE_SPK_REF, the encoded reference - // state of the cloned voice + // opaque state: the decoder carry-over to pass into the next CODE2WAV call const char * state_data; size_t state_size; }; diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index f84fbfc51cf2..07afbdcf33e1 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -122,38 +122,10 @@ int main(int argc, char ** argv) { return 1; } - // codec_0 (backbone) EOS token: ordinary LLM sampling concern, kept out of the - // model-agnostic audio-generation helper + // codec_0 (backbone) generation ends on the model's eog tokens; sampling is + // restricted to the audio zone by the tokenizer.ggml.suppress_tokens GGUF + // metadata, merged into the sampling chain by common_sampler_init const llama_vocab * vocab = llama_model_get_vocab(model); - llama_token codec_eos_tok = LLAMA_TOKEN_NULL; - for (llama_token t = 0; t < llama_vocab_n_tokens(vocab); t++) { - if (!strcmp(llama_vocab_get_text(vocab, t), "<|codec_eos_token|>")) { codec_eos_tok = t; break; } - } - if (codec_eos_tok == LLAMA_TOKEN_NULL) { - // models without a dedicated codec eos (e.g. chatterbox) end the audio - // stream with the regular vocab eos - codec_eos_tok = llama_vocab_eos(vocab); - - // fused text+speech vocab: the reference implementation samples the - // speech head only, mask the text zone out of the sampling chain - llama_token speech_base = LLAMA_TOKEN_NULL; - for (llama_token t = 0; t < llama_vocab_n_tokens(vocab); t++) { - if (!strcmp(llama_vocab_get_text(vocab, t), "<|speech_0|>")) { speech_base = t; break; } - } - if (speech_base != LLAMA_TOKEN_NULL) { - params.sampling.logit_bias.reserve(params.sampling.logit_bias.size() + speech_base); - for (llama_token t = 0; t < speech_base; t++) { - params.sampling.logit_bias.push_back(llama_logit_bias{t, -std::numeric_limits::infinity()}); - } - common_sampler_free(smpl); - smpl = common_sampler_init(model, params.sampling); - if (!smpl) { LOG_ERR("failed to reinit sampler\n"); return 1; } - } - } - if (codec_eos_tok == LLAMA_TOKEN_NULL) { - LOG_ERR("missing codec eos token in vocab\n"); - return 1; - } auto sample_codec0 = [&]() -> llama_token { llama_token t = common_sampler_sample(smpl, lctx, -1); @@ -169,7 +141,7 @@ int main(int argc, char ** argv) { tts_timings timings; const int64_t t_gen_start_us = ggml_time_us(); - for (; n_frames < max_new && sampled != codec_eos_tok; n_frames++) { + for (; n_frames < max_new && !llama_vocab_is_eog(vocab, sampled); n_frames++) { const float * h_next = nullptr; if (gen.step(sampled, h_state, &h_next) != 0) { LOG_ERR("step failed at frame %d\n", n_frames); From a78262f665788e2666dea1ce8b7a240cb7f1235c Mon Sep 17 00:00:00 2001 From: Pascal Date: Mon, 3 Aug 2026 08:22:29 +0200 Subject: [PATCH 12/14] mtmd: move chatterbox text preprocessing into the conversion The multilingual vocab bakes the reference [SPACE] substitution in: the [SPACE] slot takes the byte-level space as content and the dead raw-space entry becomes unused, so raw spaces tokenize to the exact reference ids. Case folding moves to the embedding table, uppercase rows are copies of their lowercase rows. The turbo punc_norm is dropped. The raw prompt now goes straight to llama_tokenize with no runtime text transformation in the helper. --- conversion/chatterbox.py | 40 +++++++++++++++++++- tools/mtmd/mtmd-helper-gen.cpp | 68 ++-------------------------------- 2 files changed, 43 insertions(+), 65 deletions(-) diff --git a/conversion/chatterbox.py b/conversion/chatterbox.py index a7e580caedd8..4059db2ed47d 100644 --- a/conversion/chatterbox.py +++ b/conversion/chatterbox.py @@ -174,7 +174,11 @@ def _set_vocab_mtl(self): # 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 + # 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) @@ -188,6 +192,10 @@ def enc(s: str) -> str: 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: @@ -198,6 +206,13 @@ def enc(s: str) -> str: 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 @@ -224,6 +239,21 @@ def enc(s: str) -> str: # 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"]) @@ -301,6 +331,14 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter 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 diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index bc16662b41c3..88e9d643ba78 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -501,70 +501,10 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } } - // text preprocessing per variant: multilingual lowercases with [SPACE] - // tokens (language tag left to the caller), turbo applies punc_norm - std::string txt(inp->prompt, inp->prompt_len); - if (mtl) { - // lowercase matches the reference .lower() on the latin-1 range: - // ascii letters plus the accented uppercase (utf-8 c3 80..c3 9e - // maps to c3 a0..c3 be, the multiplication sign c3 97 excepted) - std::string norm; - for (size_t i = 0; i < txt.size(); i++) { - const unsigned char c = (unsigned char) txt[i]; - if (c == ' ') { - norm += "[SPACE]"; - } else if (c >= 'A' && c <= 'Z') { - norm += (char) (c - 'A' + 'a'); - } else if (c == 0xC3 && i + 1 < txt.size() - && (unsigned char) txt[i + 1] >= 0x80 - && (unsigned char) txt[i + 1] <= 0x9E - && (unsigned char) txt[i + 1] != 0x97) { - norm += (char) 0xC3; - norm += (char) ((unsigned char) txt[i + 1] + 0x20); - i++; - } else { - norm += (char) c; - } - } - txt = norm; - } else if (!txt.empty()) { - if (txt[0] >= 'a' && txt[0] <= 'z') { - txt[0] = (char) (txt[0] - 'a' + 'A'); - } - std::string norm; - for (char c : txt) { - if (c == ' ' || c == '\t' || c == '\n' || c == '\r') { - if (!norm.empty() && norm.back() != ' ') { - norm += ' '; - } - } else { - norm += c; - } - } - auto replace_all = [&norm](const char * from, const char * to) { - const size_t nf = strlen(from); - const size_t nt = strlen(to); - for (size_t p = 0; (p = norm.find(from, p)) != std::string::npos; p += nt) { - norm.replace(p, nf, to); - } - }; - replace_all("\xE2\x80\xA6", ", "); // ellipsis - replace_all(":", ","); - replace_all("\xE2\x80\x94", "-"); // em dash - replace_all("\xE2\x80\x93", "-"); // en dash - replace_all(" ,", ","); - replace_all("\xE2\x80\x9C", "\""); // curly double quotes - replace_all("\xE2\x80\x9D", "\""); - replace_all("\xE2\x80\x98", "'"); // curly single quotes - replace_all("\xE2\x80\x99", "'"); - while (!norm.empty() && norm.back() == ' ') { - norm.pop_back(); - } - if (!norm.empty() && strchr(".!?-,", norm.back()) == nullptr) { - norm += '.'; - } - txt = norm; - } + // the raw prompt goes straight to the tokenizer: the multilingual + // vocab carries the [SPACE] substitution and the case folding lives + // in the embedding table, both baked in at conversion + const std::string txt(inp->prompt, inp->prompt_len); std::vector ids(txt.size() + 16); int n_ids = llama_tokenize(vocab, txt.c_str(), (int32_t) txt.size(), ids.data(), (int32_t) ids.size(), false, true); From a687fcfd9db4ed6f6285a7965ed619d563aec30c Mon Sep 17 00:00:00 2001 From: Pascal Date: Mon, 3 Aug 2026 09:40:53 +0200 Subject: [PATCH 13/14] mtmd: chatterbox speaker encoder outputs the raw x-vector The CAMPPlus graph ends at the 192-dim x-vector; normalization and the flow speaker affine run inside the code2wav graph, matching the reference layering where both belong to the flow module. The affine follows it to the a.gen.flow namespace and the default voice ships as the raw stored x-vector instead of a precomputed projection. --- conversion/chatterbox.py | 18 +++++------------- tools/mtmd/clip.cpp | 15 ++++++--------- tools/mtmd/models/chatterbox-gen.cpp | 13 ++++++++++--- tools/mtmd/models/chatterbox-spkenc.cpp | 17 +++++------------ tools/mtmd/mtmd.cpp | 19 +++++++++---------- tools/mtmd/mtmd.h | 2 +- 6 files changed, 36 insertions(+), 48 deletions(-) diff --git a/conversion/chatterbox.py b/conversion/chatterbox.py index 4059db2ed47d..54a2bf6008db 100644 --- a/conversion/chatterbox.py +++ b/conversion/chatterbox.py @@ -411,12 +411,9 @@ def rename_s3gen(name: str) -> str | None: ): if name.startswith(src): return dst + name[len(src):] - # the affine closes the speaker encoding chain, the rest of the - # flow module belongs to the generation stage - if name.startswith("flow.spk_embed_affine_layer."): - return "a." + name[len("flow."):] if name.startswith("flow."): - return "a.gen." + name # input_embedding, encoder_proj + # input_embedding, encoder_proj, spk_embed_affine_layer + return "a.gen." + name return None def rename_ve(name: str) -> str | None: @@ -527,14 +524,9 @@ def talker_tensor(name: str) -> Tensor: # 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()) - # the 80-dim flow speaker vector: spk_embed_affine_layer(normalize(campplus)) - with gguf.utility.SafetensorsLocal(self.dir_model / (TURBO_S3GEN if self.is_turbo else MTL_S3GEN)) as parts: - aw = parts["flow.spk_embed_affine_layer.weight"] - ab = parts["flow.spk_embed_affine_layer.bias"] - affine_w = torch.from_numpy(aw.mmap_bytes()).view(LazyTorchTensor._dtype_str_map[aw.dtype]).reshape(aw.shape).float() - affine_b = torch.from_numpy(ab.mmap_bytes()).view(LazyTorchTensor._dtype_str_map[ab.dtype]).reshape(ab.shape).float() - emb = F.normalize(genc["embedding"][0].float(), dim=0) - yield ("a.gen.cond.gen_spk80", affine_w @ emb + affine_b) + # 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() diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 4bbd2c0232ea..b575e60278a4 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -2952,6 +2952,8 @@ struct clip_model_loader { // 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"); @@ -3123,21 +3125,16 @@ struct clip_model_loader { 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); - c.spk_affine_w = get_tensor("a.spk_embed_affine_layer.weight"); - c.spk_affine_b = get_tensor("a.spk_embed_affine_layer.bias"); // host-read side data, accessed by name through // clip_cbx_read_tensor: voice encoder lstm, conditioning - // encoder, speaker affine interface dim + // 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); } } - // the affine bias doubles as the host-side presence and - // dimension probe of the speaker encoder - model.cbx_tensors["a.spk_embed_affine_layer.bias"] = c.spk_affine_b; } break; case PROJECTOR_TYPE_VOXTRAL: { @@ -5130,12 +5127,12 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { set_input_f32("inp_prompt_feat", feat); } if (params->ref_spk) { - set_input_f32("inp_spk", *params->ref_spk); + set_input_f32("inp_xvec", *params->ref_spk); } else { - ggml_tensor * sp = model.cbx_tensors.at("a.gen.cond.gen_spk80"); + ggml_tensor * sp = model.cbx_tensors.at("a.gen.cond.gen_embedding"); std::vector spk(ggml_nelements(sp)); ggml_backend_tensor_get(sp, spk.data(), 0, ggml_nbytes(sp)); - set_input_f32("inp_spk", spk); + set_input_f32("inp_xvec", spk); } // espnet relative positional encoding, entry k holds the diff --git a/tools/mtmd/models/chatterbox-gen.cpp b/tools/mtmd/models/chatterbox-gen.cpp index 0beda00fe3ae..a552d9835917 100644 --- a/tools/mtmd/models/chatterbox-gen.cpp +++ b/tools/mtmd/models/chatterbox-gen.cpp @@ -436,9 +436,16 @@ ggml_cgraph * clip_graph_chatterbox::build() { 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] - ggml_tensor * spks = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 80); - ggml_set_name(spks, "inp_spk"); - ggml_set_input(spks); + // 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); diff --git a/tools/mtmd/models/chatterbox-spkenc.cpp b/tools/mtmd/models/chatterbox-spkenc.cpp index 20a4729c234c..7220319b08e3 100644 --- a/tools/mtmd/models/chatterbox-spkenc.cpp +++ b/tools/mtmd/models/chatterbox-spkenc.cpp @@ -1,7 +1,7 @@ #include "models.h" // Chatterbox speaker encoder: CAMPPlus x-vector on kaldi fbank features, -// projected through the s3gen speaker affine. Mirrors s3gen/xvector.py. +// 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 @@ -147,22 +147,15 @@ ggml_cgraph * clip_graph_chatterbox_spkenc::build() { ggml_cont(ctx0, ggml_transpose(ctx0, sd)), 0); // [1024, 1] cb(stats, "spk_stats_pool", -1); - // dense 1024 -> 192, batchnorm without affine, into the x-vector + // 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_reshape_1d(ctx0, emb, 192); + 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); - - // normalize then the s3gen speaker affine - ggml_tensor * n2 = ggml_sqrt(ctx0, ggml_sum(ctx0, ggml_mul(ctx0, emb, emb))); - ggml_tensor * unit = ggml_div(ctx0, emb, n2); - ggml_tensor * spk80 = cbx_linear(c.spk_affine_w, c.spk_affine_b, - ggml_reshape_2d(ctx0, unit, 192, 1)); - spk80 = ggml_cont(ctx0, ggml_reshape_1d(ctx0, spk80, 80)); - cb(spk80, "spk_embd", -1); - ggml_build_forward_expand(gf, spk80); return gf; } diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index bce7a2f3ff0d..af329c0c413f 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1700,11 +1700,10 @@ size_t mtmd_gen_audio_read_tensor(mtmd_context * ctx, const char * name, float * return clip_cbx_read_tensor(ctx->ctx_gen_a, name, out, n_max); } -// 80-dim flow speaker vector of the reference clip: the precomputed fbank +// 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_spk80(mtmd_context * ctx, const clip_image_f32 & fbank, std::vector & spk80) { - const size_t n = clip_cbx_read_tensor(ctx->ctx_a, "a.spk_embed_affine_layer.bias", nullptr, 0); - if (n == 0) { +static bool cbx_ref_xvec(mtmd_context * ctx, const clip_image_f32 & fbank, std::vector & xvec) { + if (!ctx->ctx_a || clip_get_projector_type(ctx->ctx_a) != PROJECTOR_TYPE_CHATTERBOX_SPKENC) { LOG_ERR("%s: mmproj has no speaker encoder\n", __func__); return false; } @@ -1714,8 +1713,8 @@ static bool cbx_ref_spk80(mtmd_context * ctx, const clip_image_f32 & fbank, std: clip_image_f32_batch batch; batch.is_audio = true; batch.entries.push_back(std::move(mel_img)); - spk80.resize(n); - if (!clip_image_batch_encode(ctx->ctx_a, ctx->n_threads, &batch, spk80)) { + 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; } @@ -2011,8 +2010,8 @@ static int32_t cbx_ref_encode(mtmd_context * ctx, const mtmd_audio_tokens * audi const auto & entries = audio_tokens->batch_f32.entries; GGML_ASSERT(entries.size() == 5); - std::vector spk80; - if (!cbx_ref_spk80(ctx, entries[0], spk80)) { + std::vector xvec; + if (!cbx_ref_xvec(ctx, entries[0], xvec)) { return 1; } @@ -2049,7 +2048,7 @@ static int32_t cbx_ref_encode(mtmd_context * ctx, const mtmd_audio_tokens * audi ctx->gen_out_ref_codes.assign(flow_codes.begin(), flow_codes.end()); ctx->gen_out_ref_feat.assign(entries[3].get_ro_buf().begin(), entries[3].get_ro_buf().begin() + n_feat); - ctx->gen_out_ref_spk = std::move(spk80); + ctx->gen_out_ref_spk = std::move(xvec); out_embd = std::move(rows); return 0; @@ -2111,7 +2110,7 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in std::vector out_audio; // optional voice cloning reference from the encoded reference chunk: - // flow prompt codes, mel-rate prompt features, 80-dim speaker vector; + // flow prompt codes, mel-rate prompt features, raw x-vector; // all null selects the model's precomputed default voice std::vector ref_tokens; std::vector ref_feat; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 498ac36bfd24..d228b65c3013 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -303,7 +303,7 @@ MTMD_API float * mtmd_get_output_embd(mtmd_context * ctx); enum mtmd_embd_out_type { MTMD_EMBD_OUT_TYPE_REF_CODES, // speech codes of the reference clip MTMD_EMBD_OUT_TYPE_REF_FEAT, // mel-rate features of the reference clip - MTMD_EMBD_OUT_TYPE_REF_SPK, // speaker vector of the reference clip + MTMD_EMBD_OUT_TYPE_REF_SPK, // raw x-vector of the reference clip }; MTMD_API const float * mtmd_get_output_typed_embd(mtmd_context * ctx, enum mtmd_embd_out_type type, From 748a5e30be44336b2d46e84cfd77043a81314591 Mon Sep 17 00:00:00 2001 From: Pascal Date: Mon, 3 Aug 2026 09:51:29 +0200 Subject: [PATCH 14/14] tts: pair cfg sequences under a unified kv cache Parallelism doubles so each slot i pairs with the uncond sequence n_seq_max / 2 + i, and the unified kv cache shares the window across sequences, keeping the per-sequence context at the requested size. The helper derives the uncond sequence from the pairing instead of a fixed second sequence. --- tools/mtmd/mtmd-helper-gen.cpp | 8 ++++++-- tools/tts/tts.cpp | 11 +++++------ 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 88e9d643ba78..2a792ee98a7b 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -560,7 +560,10 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { memcpy(embd_buf.data() + (text0 + i) * (size_t) n_e, u.data(), (size_t) n_e * sizeof(float)); } decode_embd_batch batch_uncond(embd_buf.data(), n_prompt, 1, n_e); - batch_uncond.set_position_normal(0, 1); + // the uncond branch runs in the paired sequence half a seq space + // away from the cond one (slot i pairs with n_seq_max / 2 + i) + const llama_seq_id seq_uncond = (llama_seq_id) (llama_n_seq_max(lctx) / 2); + batch_uncond.set_position_normal(0, seq_uncond); batch_uncond.batch.logits[n_prompt - 1] = 1; if (llama_decode(lctx, batch_uncond.batch) != 0) { LOG_ERR("mtmd_helper_gen_audio: cfg prefill decode failed\n"); @@ -603,7 +606,8 @@ class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::vector cond_logits; cfg_read(cond_logits); decode_embd_batch batch_uncond(e.data(), 1, 1, n_embd); - batch_uncond.set_position_normal(pos, 1); + const llama_seq_id seq_uncond = (llama_seq_id) (llama_n_seq_max(lctx) / 2); + batch_uncond.set_position_normal(pos, seq_uncond); batch_uncond.batch.logits[0] = 1; if (llama_decode(lctx, batch_uncond.batch) != 0) { LOG_ERR("mtmd_helper_gen_audio: cfg step decode failed\n"); diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index 07afbdcf33e1..7392c8c4be35 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -66,12 +66,11 @@ int main(int argc, char ** argv) { // always enable embd, so that we can pass hidden states to the audio generation helper params.embedding = true; - // provision a second sequence for pipelines that decode a cfg pair, - // scaling the context so the per-sequence window keeps the requested size - if (params.n_parallel < 2) { - params.n_parallel = 2; - params.n_ctx *= 2; - } + // 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);