Kazeia-engine/dist/jni/tts_engine.cpp

830 lines
40 KiB
C++
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#include "tts_engine.h"
#include "llama.h"
#include "ggml-backend.h"
#include "cp_inference.h"
#include "sampler.h"
#include "kazeia_text_tokenizer.h"
#include "decoder.h"
#include "speaker_encoder.h"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <chrono>
#include <fstream>
#include <string>
#include <vector>
#include <sys/wait.h>
#include <unistd.h>
using namespace kazeia::tts;
namespace {
double now_s() {
using clk = std::chrono::steady_clock;
return std::chrono::duration<double>(clk::now().time_since_epoch()).count();
}
std::vector<float> read_f32(const std::string& p, size_t n_expected) {
std::ifstream f(p, std::ios::binary | std::ios::ate);
if (!f) { fprintf(stderr, "tts_engine: open %s\n", p.c_str()); return {}; }
size_t n = (size_t)f.tellg() / sizeof(float);
if (n_expected && n != n_expected) {
fprintf(stderr, "tts_engine: %s: %zu f32, attendu %zu\n", p.c_str(), n, n_expected); return {};
}
f.seekg(0); std::vector<float> v(n); f.read((char*)v.data(), n * sizeof(float)); return v;
}
ggml_backend_dev_t find_htp() {
for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
auto d = ggml_backend_dev_get(i);
if (!strcmp(ggml_backend_dev_name(d), "HTP0")) return d;
}
return nullptr;
}
inline float silu(float x) { return x / (1.0f + std::exp(-x)); }
void linear(const float* W, const float* b, int out_dim, int in_dim, const float* x, float* out) {
for (int m = 0; m < out_dim; ++m) {
float s = b ? b[m] : 0.0f;
const float* wr = W + (size_t)m * in_dim;
for (int k = 0; k < in_dim; ++k) s += wr[k] * x[k];
out[m] = s;
}
}
void text_projection(const float* in, int N,
const float* fc1_w, const float* fc1_b,
const float* fc2_w, const float* fc2_b,
float* mid_buf, float* out,
int text_hidden, int hidden) {
for (int n = 0; n < N; ++n) {
linear(fc1_w, fc1_b, text_hidden, text_hidden, in + n * text_hidden, mid_buf);
for (int d = 0; d < text_hidden; ++d) mid_buf[d] = silu(mid_buf[d]);
linear(fc2_w, fc2_b, hidden, text_hidden, mid_buf, out + n * hidden);
}
}
bool write_wav_pcm16_mono(const char* path, const float* samples, size_t N, int sr = 24000) {
FILE* f = fopen(path, "wb");
if (!f) return false;
const uint32_t data_bytes = (uint32_t)(N * 2);
const uint32_t riff_size = 36 + data_bytes;
auto w16 = [&](uint16_t v){ fwrite(&v, 2, 1, f); };
auto w32 = [&](uint32_t v){ fwrite(&v, 4, 1, f); };
fwrite("RIFF", 1, 4, f); w32(riff_size); fwrite("WAVE", 1, 4, f);
fwrite("fmt ", 1, 4, f); w32(16); w16(1); w16(1);
w32((uint32_t)sr); w32((uint32_t)sr * 2); w16(2); w16(16);
fwrite("data", 1, 4, f); w32(data_bytes);
for (size_t i = 0; i < N; ++i) {
float v = samples[i]; if (v > 1.f) v = 1.f; else if (v < -1.f) v = -1.f;
int16_t pcm = (int16_t)std::lround(v * 32767.0f);
fwrite(&pcm, 2, 1, f);
}
fclose(f); return true;
}
} // namespace
struct TtsEngine {
// --- Constantes manifest ---
int text_vocab = 151936, text_hidden = 2048, hidden = 1024;
int tts_bos = 151672, tts_eos = 151673, tts_pad = 151671;
int codec_bos = 2149, codec_eos = 2150, codec_pad = 2148;
int codec_think = 2154, codec_think_bos = 2156, codec_think_eos = 2157;
int lang_fr = 2061;
// --- Fixtures (load-once) ---
std::vector<float> text_embed;
std::vector<float> tp_fc1_w, tp_fc1_b, tp_fc2_w, tp_fc2_b;
std::vector<float> tok_embd;
std::vector<float> xvector;
// Embeds spéciaux pré-projetés (constants une fois le modèle chargé)
std::vector<float> spec_proj; // [3 * hidden] = {tts_bos_emb, tts_eos_emb, tts_pad_emb}
std::vector<float> codec_input_emb; // [7 * hidden] = {think, think_bos, lang_fr, think_eos, xvec, codec_pad, codec_bos}
std::vector<float> role_proj; // [3 * hidden] = projection de <|im_start|>, assistant, \n (constant car template fixe)
std::vector<float> codec_pad_emb; // [hidden] = tok_embd[codec_pad]
// --- Talker ---
llama_model * talker_m = nullptr;
llama_context * talker_ctx = nullptr;
int n_embd = 0;
int n_vocab = 0;
int npe = 1; // 1 (RoPE) ou 4 (M-RoPE/I-MRoPE)
int n_threads = 6;
// --- CP / Decoder / Tokenizer ---
CPState cp_state;
Decoder decoder;
KzTextTokenizer kz_tok;
bool has_kz_tok = false;
bool cp_use_cache = true;
// --- Speaker encoder (optionnel, clonage vocal embarqué) ---
SpeakerEncoder * spk_enc = nullptr; // nullptr -> pas de clonage à la volée
};
// ===========================================================================
// LOAD
// ===========================================================================
TtsEngine * tts_engine_load(const TtsEngineLoadCfg & cfg) {
if (!cfg.talker_gguf || !cfg.dump_dir) {
fprintf(stderr, "tts_engine_load: talker_gguf et dump_dir requis\n"); return nullptr;
}
auto eng = new TtsEngine();
eng->n_threads = cfg.n_threads;
eng->cp_use_cache = cfg.cp_use_cache;
std::string D = cfg.dump_dir;
if (D.back() != '/') D += '/';
// --- 0) Constantes manifest ---
{
std::ifstream f(D + "manifest_text.txt");
if (!f) { fprintf(stderr, "tts_engine_load: no manifest_text\n"); delete eng; return nullptr; }
std::string line;
auto eq = [&](const char* k){ return line.rfind(k, 0) == 0; };
auto val = [&](size_t off){ return atoi(line.c_str() + off); };
while (std::getline(f, line)) {
if (eq("text_vocab_size:")) eng->text_vocab = val(16);
else if (eq("text_hidden_size:")) eng->text_hidden = val(17);
else if (eq("hidden_size:")) eng->hidden = val(12);
else if (eq("tts_bos_token_id:")) eng->tts_bos = val(17);
else if (eq("tts_eos_token_id:")) eng->tts_eos = val(17);
else if (eq("tts_pad_token_id:")) eng->tts_pad = val(17);
else if (eq("codec_bos_id:")) eng->codec_bos = val(13);
else if (eq("codec_eos_id:")) eng->codec_eos = val(13);
else if (eq("codec_pad_id:")) eng->codec_pad = val(13);
else if (eq("codec_think_id:")) eng->codec_think = val(15);
else if (eq("codec_think_bos_id:")) eng->codec_think_bos = val(19);
else if (eq("codec_think_eos_id:")) eng->codec_think_eos = val(19);
else if (eq("codec_language_french:")) eng->lang_fr = val(22);
}
}
// --- 1) Fixtures ---
const int text_hidden = eng->text_hidden;
const int hidden = eng->hidden;
const int text_vocab = eng->text_vocab;
const double t_load0 = now_s();
eng->text_embed = read_f32(D + "text_embed.bin", (size_t)text_vocab * text_hidden);
eng->tp_fc1_w = read_f32(D + "tp_fc1_w.bin", (size_t)text_hidden * text_hidden);
eng->tp_fc1_b = read_f32(D + "tp_fc1_b.bin", (size_t)text_hidden);
eng->tp_fc2_w = read_f32(D + "tp_fc2_w.bin", (size_t)hidden * text_hidden);
eng->tp_fc2_b = read_f32(D + "tp_fc2_b.bin", (size_t)hidden);
eng->tok_embd = read_f32(D + "talker_tok_embd.bin", (size_t)3072 * hidden);
eng->xvector = read_f32(D + "damien_xvector.bin", (size_t)hidden);
if (eng->text_embed.empty() || eng->tp_fc1_w.empty() || eng->tok_embd.empty() || eng->xvector.empty()) {
fprintf(stderr, "tts_engine_load: fixtures FAIL\n"); delete eng; return nullptr;
}
// --- 2) Backend talker (avant le 1er init backend pour HMX off, etc.) ---
setenv("GGML_HEXAGON_USE_HMX", "0", 1);
llama_backend_init();
// --- 2.5) Chargement dynamique du backend Vulkan (DOIT être AVANT talker load,
// pour que llama.cpp voit le device Adreno lors de mp.devices = {vk_dev, ...}).
if (const char * vk_path = getenv("KZTTS_VULKAN_LIB")) {
if (ggml_backend_load(vk_path)) {
fprintf(stderr, "Vulkan backend chargé : %s\n", vk_path);
} else {
fprintf(stderr, "ggml_backend_load(%s) FAIL\n", vk_path);
}
}
// --- 3) Tokenizer optionnel (vocab_only -> ~50 MB RAM) ---
if (cfg.vocab_gguf && *cfg.vocab_gguf) {
if (!kz_tok_load(eng->kz_tok, cfg.vocab_gguf)) {
fprintf(stderr, "tts_engine_load: vocab load FAIL %s\n", cfg.vocab_gguf);
delete eng; return nullptr;
}
eng->has_kz_tok = true;
}
// --- 4) Talker (libllama) ---
auto mp = llama_model_default_params();
// KZTTS_TALKER_VULKAN=1 -> offload talker sur GPU Adreno via Vulkan device.
const bool talker_vulkan = (getenv("KZTTS_TALKER_VULKAN") && atoi(getenv("KZTTS_TALKER_VULKAN")) != 0);
ggml_backend_dev_t vk_dev = nullptr;
if (talker_vulkan) {
// Adreno UMA -> ggml-vulkan registre le device comme IGPU (type=2, integrated GPU
// utilisant la host memory). HTP s'enregistre comme GPU (type=1). On accepte les deux
// GPU et IGPU mais filtre par nom pour exclure HTP/Hexagon explicitement.
for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
auto d = ggml_backend_dev_get(i);
auto t = ggml_backend_dev_type(d);
if (t != GGML_BACKEND_DEVICE_TYPE_GPU && t != GGML_BACKEND_DEVICE_TYPE_IGPU) continue;
const char * nm = ggml_backend_dev_name(d);
if (strstr(nm, "HTP") || strstr(nm, "Hexagon")) continue;
vk_dev = d;
fprintf(stderr, "talker: GPU candidate found : %s (type=%d)\n", nm, (int)t);
break;
}
}
ggml_backend_dev_t devs[2] = {
vk_dev ? vk_dev : (cfg.use_htp ? find_htp() : nullptr),
nullptr
};
if (devs[0]) {
mp.n_gpu_layers = 99; mp.devices = devs;
fprintf(stderr, "talker: %s (GPU offload)\n", ggml_backend_dev_name(devs[0]));
} else {
mp.n_gpu_layers = 0; fprintf(stderr, "talker: CPU\n");
}
eng->talker_m = llama_model_load_from_file(cfg.talker_gguf, mp);
if (!eng->talker_m) { fprintf(stderr, "talker load FAIL %s\n", cfg.talker_gguf); delete eng; return nullptr; }
auto cp = llama_context_default_params();
// Borne ctx généreuse : prefill max ~50 + max_steps 512 + marge. Pas critique en CPU.
cp.n_ctx = 1024;
cp.n_batch = 1024;
cp.n_threads = cfg.n_threads;
cp.flash_attn_type = (getenv("KZTTS_TALKER_FA") && atoi(getenv("KZTTS_TALKER_FA")) != 0)
? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED;
cp.embeddings = true;
cp.pooling_type = LLAMA_POOLING_TYPE_NONE;
eng->talker_ctx = llama_init_from_model(eng->talker_m, cp);
if (!eng->talker_ctx) { fprintf(stderr, "talker ctx FAIL\n"); delete eng; return nullptr; }
eng->n_embd = llama_model_n_embd(eng->talker_m);
eng->n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(eng->talker_m));
if (eng->n_embd != hidden || eng->n_vocab != 3072) {
fprintf(stderr, "talker dims mismatch (n_embd=%d, n_vocab=%d, attendu %d / 3072)\n",
eng->n_embd, eng->n_vocab, hidden); delete eng; return nullptr;
}
auto rt = llama_model_rope_type(eng->talker_m);
eng->npe = (rt == LLAMA_ROPE_TYPE_MROPE || rt == LLAMA_ROPE_TYPE_IMROPE) ? 4 : 1;
// --- 5) CP ---
// KZTTS_CP_HTP=1 + KZTTS_CP_SCHED=1 active le path sched HMX :
// poids CP sur HTP buffer, forward via ggml_backend_sched (MUL_MAT -> HMX).
// CP est un transformer 5L 1024-hidden, batches n_new in {1,2}, donc en-dessous de la
// limite hexagon nrows(src1) ≤ 1024 qui bloquait BigVGAN. Pattern identique au talker
// option C déjà validé.
eng->cp_state.n_threads = cfg.n_threads;
const bool cp_htp = (getenv("KZTTS_CP_HTP") && atoi(getenv("KZTTS_CP_HTP")) != 0);
bool cp_ok;
if (cp_htp) {
std::vector<ggml_backend_t> cp_backends;
if (auto htp_dev = find_htp()) {
ggml_backend_t htp = ggml_backend_dev_init(htp_dev, nullptr);
if (htp) { cp_backends.push_back(htp); fprintf(stderr, "CP: HTP backend initialisé\n"); }
else { fprintf(stderr, "CP: HTP init FAIL, fallback CPU only\n"); }
}
ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
if (!cpu) { fprintf(stderr, "CP: CPU backend init FAIL\n"); delete eng; return nullptr; }
cp_backends.push_back(cpu); // CPU TOUJOURS en dernier (assertion sched)
cp_ok = cp_load_with_backends(eng->cp_state,
(D + "cp_f16.gguf").c_str(),
(D + "cp_heads.bin").c_str(),
(D + "cp_codec_embs.bin").c_str(),
cp_backends, /*owns=*/true);
} else {
cp_ok = cp_load(eng->cp_state,
(D + "cp_f16.gguf").c_str(),
(D + "cp_heads.bin").c_str(),
(D + "cp_codec_embs.bin").c_str());
}
if (!cp_ok) { fprintf(stderr, "CP load FAIL\n"); delete eng; return nullptr; }
// --- 6) Decoder ---
// KZTTS_DECODER_HTP=1 -> charger via load_with_backends({HTP, CPU}) pour activer
// le sched ggml qui route les MUL_MAT (im2col + matmul issu de ggml_conv_1d) sur HMX.
// IM2COL, CONV_TRANSPOSE_1D, SIN restent CPU (sched fait le split + copies auto).
// Sinon : load() classique = CPU pur via ggml_graph_compute_with_ctx (path actuel).
const bool dec_htp = (getenv("KZTTS_DECODER_HTP") && atoi(getenv("KZTTS_DECODER_HTP")) != 0);
bool dec_ok;
const bool dec_vulkan = (getenv("KZTTS_DECODER_VULKAN") && atoi(getenv("KZTTS_DECODER_VULKAN")) != 0);
if (dec_htp || dec_vulkan) {
std::vector<ggml_backend_t> dec_backends;
// Vulkan prioritaire si demandé (Adreno GPU compute, BW partagée mémoire).
if (dec_vulkan) {
// Note : Hexagon HTP est aussi enregistré comme DEVICE_TYPE_GPU. On filtre par
// nom non-HTP pour trouver le VRAI GPU (Adreno...). Liste tous les devs et choisis.
fprintf(stderr, "Devices enregistrés (%zu):\n", ggml_backend_dev_count());
for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
auto d = ggml_backend_dev_get(i);
fprintf(stderr, " [%zu] %s (type=%d)\n", i,
ggml_backend_dev_name(d), (int)ggml_backend_dev_type(d));
}
for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
auto d = ggml_backend_dev_get(i);
auto t = ggml_backend_dev_type(d);
// Adreno UMA = type IGPU. HTP = type GPU. On accepte les deux et filtre par nom.
if (t != GGML_BACKEND_DEVICE_TYPE_GPU && t != GGML_BACKEND_DEVICE_TYPE_IGPU) continue;
const char * nm = ggml_backend_dev_name(d);
if (strstr(nm, "HTP") || strstr(nm, "Hexagon")) continue;
ggml_backend_t vk = ggml_backend_dev_init(d, nullptr);
if (vk) {
dec_backends.push_back(vk);
fprintf(stderr, "decoder: Vulkan backend initialisé (%s)\n", nm);
}
break;
}
}
// HTP (option backend si pas Vulkan ou ajout)
if (dec_htp) {
if (auto htp_dev = find_htp()) {
ggml_backend_t htp = ggml_backend_dev_init(htp_dev, nullptr);
if (htp) {
dec_backends.push_back(htp);
fprintf(stderr, "decoder: HTP backend initialisé\n");
} else {
fprintf(stderr, "decoder: ggml_backend_dev_init(HTP0) FAIL, fallback CPU\n");
}
}
}
// CPU TOUJOURS en DERNIER (ggml_backend_sched_new assert que backends[end]=CPU,
// c'est le fallback universel pour les ops non-supportées par les autres).
ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
if (!cpu) { fprintf(stderr, "decoder: CPU backend init FAIL\n"); delete eng; return nullptr; }
dec_backends.push_back(cpu);
eng->decoder.owns_backends = true;
dec_ok = eng->decoder.load_with_backends((D + "qwen3tts_decoder.gguf").c_str(), dec_backends);
} else {
// CPU-only : sched sur backend CPU -> stage_bigvgan_sched (gallocr, peak mem bas)
// au lieu du legacy stage_bigvgan (malloc 12 GB -> ggml_aligned_malloc NULL -> crash in-app).
ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
if (!cpu) { fprintf(stderr, "decoder: CPU backend init FAIL\n"); delete eng; return nullptr; }
std::vector<ggml_backend_t> dec_backends{cpu};
eng->decoder.owns_backends = true;
setenv("KZTTS_DECODER_SCHED", "1", 1);
dec_ok = eng->decoder.load_with_backends((D + "qwen3tts_decoder.gguf").c_str(), dec_backends);
}
if (!dec_ok) {
fprintf(stderr, "decoder load FAIL\n"); delete eng; return nullptr;
}
// --- 7) Pre-compute des embeds constants ---
std::vector<float> mid(text_hidden);
// special tokens projetés (tts_bos/eos/pad)
std::vector<float> spec_text(3 * text_hidden);
int sp[3] = { eng->tts_bos, eng->tts_eos, eng->tts_pad };
for (int i = 0; i < 3; ++i)
std::memcpy(spec_text.data() + i * text_hidden,
eng->text_embed.data() + (size_t)sp[i] * text_hidden,
text_hidden * sizeof(float));
eng->spec_proj.resize(3 * hidden);
text_projection(spec_text.data(), 3,
eng->tp_fc1_w.data(), eng->tp_fc1_b.data(),
eng->tp_fc2_w.data(), eng->tp_fc2_b.data(),
mid.data(), eng->spec_proj.data(), text_hidden, hidden);
// codec prefix [think, think_bos, lang_fr, think_eos, x_vec, codec_pad, codec_bos]
int codec_prefill[4] = { eng->codec_think, eng->codec_think_bos, eng->lang_fr, eng->codec_think_eos };
eng->codec_input_emb.resize(7 * hidden);
for (int i = 0; i < 4; ++i)
std::memcpy(eng->codec_input_emb.data() + i * hidden,
eng->tok_embd.data() + (size_t)codec_prefill[i] * hidden,
hidden * sizeof(float));
std::memcpy(eng->codec_input_emb.data() + 4 * hidden, eng->xvector.data(), hidden * sizeof(float));
std::memcpy(eng->codec_input_emb.data() + 5 * hidden,
eng->tok_embd.data() + (size_t)eng->codec_pad * hidden, hidden * sizeof(float));
std::memcpy(eng->codec_input_emb.data() + 6 * hidden,
eng->tok_embd.data() + (size_t)eng->codec_bos * hidden, hidden * sizeof(float));
// role tokens (<|im_start|>, assistant, \n) — projection constante car template fixe.
// Pré-calculé une fois, indépendant du texte d'entrée.
const int role_ids[3] = { 151644, 77091, 198 };
std::vector<float> role_text(3 * text_hidden);
for (int i = 0; i < 3; ++i)
std::memcpy(role_text.data() + i * text_hidden,
eng->text_embed.data() + (size_t)role_ids[i] * text_hidden,
text_hidden * sizeof(float));
eng->role_proj.resize(3 * hidden);
text_projection(role_text.data(), 3,
eng->tp_fc1_w.data(), eng->tp_fc1_b.data(),
eng->tp_fc2_w.data(), eng->tp_fc2_b.data(),
mid.data(), eng->role_proj.data(), text_hidden, hidden);
// codec_pad_emb (utilisé partout dans le body)
eng->codec_pad_emb.assign(eng->tok_embd.data() + (size_t)eng->codec_pad * hidden,
eng->tok_embd.data() + (size_t)(eng->codec_pad + 1) * hidden);
// --- Speaker encoder optionnel (clonage vocal embarqué) ---
if (cfg.speaker_encoder_gguf && cfg.mel_basis_path) {
eng->spk_enc = speaker_encoder_load(cfg.speaker_encoder_gguf, cfg.mel_basis_path);
if (!eng->spk_enc) {
fprintf(stderr, "tts_engine_load: speaker_encoder_load FAIL (continuing without cloning)\n");
}
}
fprintf(stderr, "tts_engine_load: %.2fs (talker_n_embd=%d, npe=%d, kz_tok=%d, cp_cache=%d, spk_enc=%d)\n",
now_s() - t_load0, eng->n_embd, eng->npe, (int)eng->has_kz_tok,
(int)cfg.cp_use_cache, (int)(eng->spk_enc != nullptr));
return eng;
}
// ===========================================================================
// SYNTHESIZE
// ===========================================================================
TtsSynthesizeResult tts_engine_synthesize(TtsEngine * eng, const TtsSynthesizeCfg & cfg) {
TtsSynthesizeResult R{};
if (!eng || !cfg.text || !cfg.out_wav_path) { R.err = -1; return R; }
if (!eng->has_kz_tok) {
fprintf(stderr, "tts_engine_synthesize: pas de tokenizer chargé (vocab_gguf manquant au load)\n");
R.err = -2; return R;
}
const int text_hidden = eng->text_hidden;
const int hidden = eng->hidden;
const int n_embd = eng->n_embd;
const int n_vocab = eng->n_vocab;
// --- 0) Override x_vector (clonage vocal) ---
// codec_input_emb a été pré-calculé au load() avec eng->xvector au slot 4. Si un
// override est fourni pour ce synth, on patche le slot puis on restaure via RAII
// (le synth peut return-early sur erreur, on évite de laisser l'engine sale).
const bool xvec_override = cfg.xvector_override && cfg.xvector_override_len == hidden;
if (cfg.xvector_override && !xvec_override) {
fprintf(stderr, "tts_engine_synthesize: xvector_override de longueur %d != hidden %d, ignoré\n",
cfg.xvector_override_len, hidden);
}
struct XvecRestore {
TtsEngine* e; int h; std::vector<float> backup;
~XvecRestore() {
if (!backup.empty())
std::memcpy(e->codec_input_emb.data() + 4 * h, backup.data(), (size_t)h * sizeof(float));
}
} _xvec_guard{eng, hidden, {}};
if (xvec_override) {
_xvec_guard.backup.assign(eng->codec_input_emb.data() + 4 * hidden,
eng->codec_input_emb.data() + 5 * hidden);
std::memcpy(eng->codec_input_emb.data() + 4 * hidden,
cfg.xvector_override, (size_t)hidden * sizeof(float));
}
// --- 1) Tokenize ---
auto input_ids = kz_tok_encode_tts_prompt(eng->kz_tok, cfg.text);
const int input_ids_len = (int)input_ids.size();
if (input_ids_len < 8) {
fprintf(stderr, "tts_engine_synthesize: tokens trop courts (%d, min=8)\n", input_ids_len);
R.err = -3; return R;
}
const int Nt = input_ids_len - 5 - 3; // role(3) + body(Nt) + trailing(5)
// --- 2) Body projection ---
std::vector<float> mid(text_hidden);
std::vector<float> body_text((size_t)Nt * text_hidden);
for (int i = 0; i < Nt; ++i)
std::memcpy(body_text.data() + i * text_hidden,
eng->text_embed.data() + (size_t)input_ids[3 + i] * text_hidden,
text_hidden * sizeof(float));
std::vector<float> body_proj((size_t)Nt * hidden);
text_projection(body_text.data(), Nt,
eng->tp_fc1_w.data(), eng->tp_fc1_b.data(),
eng->tp_fc2_w.data(), eng->tp_fc2_b.data(),
mid.data(), body_proj.data(), text_hidden, hidden);
// --- 3) Assemble prefill ---
const float* tts_bos_emb = eng->spec_proj.data() + 0 * hidden;
const float* tts_eos_emb = eng->spec_proj.data() + 1 * hidden;
const float* tts_pad_emb = eng->spec_proj.data() + 2 * hidden;
const float* codec_pad_emb = eng->codec_pad_emb.data();
const int T_prefill = 3 + 6 + (Nt + 1) + 1;
std::vector<float> prefill((size_t)T_prefill * hidden, 0.0f);
std::memcpy(prefill.data(), eng->role_proj.data(), 3 * hidden * sizeof(float));
for (int i = 0; i < 5; ++i) {
const float* ce = eng->codec_input_emb.data() + i * hidden;
float* p = prefill.data() + (3 + i) * hidden;
for (int d = 0; d < hidden; ++d) p[d] = tts_pad_emb[d] + ce[d];
}
{ const float* ce = eng->codec_input_emb.data() + 5 * hidden;
float* p = prefill.data() + 8 * hidden;
for (int d = 0; d < hidden; ++d) p[d] = tts_bos_emb[d] + ce[d]; }
for (int i = 0; i < Nt; ++i) {
const float* bp = body_proj.data() + i * hidden;
float* p = prefill.data() + (9 + i) * hidden;
for (int d = 0; d < hidden; ++d) p[d] = bp[d] + codec_pad_emb[d];
}
{ float* p = prefill.data() + (9 + Nt) * hidden;
for (int d = 0; d < hidden; ++d) p[d] = tts_eos_emb[d] + codec_pad_emb[d]; }
{ const float* cbe = eng->tok_embd.data() + (size_t)eng->codec_bos * hidden;
float* p = prefill.data() + (10 + Nt) * hidden;
for (int d = 0; d < hidden; ++d) p[d] = tts_pad_emb[d] + cbe[d]; }
// --- 4) Reset KV talker (chaque synthèse repart de pos=0) ---
llama_memory_clear(llama_get_memory(eng->talker_ctx), true);
// --- 5) Configure les samplers ---
eng->cp_state.sampler.temp = cfg.cp_temp;
eng->cp_state.sampler.top_k = cfg.cp_top_k;
eng->cp_state.sampler.top_p = cfg.cp_top_p;
eng->cp_state.sampler.rep_penalty = cfg.cp_rep_penalty;
eng->cp_state.sampler.rep_window = 16;
sampler_seed(eng->cp_state.sampler, cfg.seed + 1);
Sampler talker_sampler{};
talker_sampler.temp = cfg.talker_temp;
talker_sampler.top_k = cfg.talker_top_k;
talker_sampler.top_p = cfg.talker_top_p;
talker_sampler.rep_penalty = cfg.talker_rep_penalty;
talker_sampler.rep_window = 64;
sampler_seed(talker_sampler, cfg.seed);
// --- 6) Talker prefill ---
const double t_pfill0 = now_s();
{
const int npe = eng->npe;
std::vector<llama_pos> pos(T_prefill * npe, 0);
std::vector<int32_t> nsd(T_prefill, 1);
std::vector<llama_seq_id> sid0(T_prefill, 0);
std::vector<llama_seq_id*> sids(T_prefill);
std::vector<int8_t> lg(T_prefill, 0);
for (int i = 0; i < T_prefill; ++i) {
if (npe == 4) { pos[i] = i; pos[T_prefill + i] = i; pos[2*T_prefill + i] = i; pos[3*T_prefill + i] = 0; }
else { pos[i] = i; }
sids[i] = &sid0[i];
}
lg[T_prefill - 1] = 1;
llama_batch b{};
b.n_tokens = T_prefill; b.embd = prefill.data();
b.pos = pos.data(); b.n_seq_id = nsd.data(); b.seq_id = sids.data(); b.logits = lg.data();
if (llama_decode(eng->talker_ctx, b) != 0) {
fprintf(stderr, "tts_engine_synthesize: prefill FAIL\n"); R.err = -4; return R;
}
}
R.prefill_s = now_s() - t_pfill0;
// --- 7) Decode loop (CB0 -> CP -> next_embed -> talker step) ---
std::vector<float> logits_buf(n_vocab);
{
const float* lp = llama_get_logits_ith(eng->talker_ctx, -1);
memcpy(logits_buf.data(), lp, n_vocab * sizeof(float));
}
// EOS-rank boost dynamique (port du legacy Qwen3TtsEngine.runInterleavedCpuFromEmbeds).
// Sans ça le talker n'émet quasi jamais codec_eos et tourne jusqu'à max_steps :
// audio fixe ~20 s + synth ~5x trop longue. On masque les tokens non-codebook
// (sauf eos), puis dès 2/3 de la durée attendue on amplifie eos quand son rang
// devient bas (le modèle « pense à terminer »). Désactivable via KZTTS_EOS_BOOST=0.
const bool eos_boost_on = !(getenv("KZTTS_EOS_BOOST") && atoi(getenv("KZTTS_EOS_BOOST")) == 0);
// Arming naturel (rang<trig) : workaround de l'ère NaN où le talker n'émettait jamais
// codec_eos. Depuis le fix CP F32, le talker se termine seul correctement (cb0==eos),
// et cet arming ne fait que TRONQUER la fin de phrase. OFF par défaut ; seul le
// hard-fallback reste actif comme garde anti-runaway. KZTTS_EOS_NATURAL=1 le réactive.
const bool eos_natural_on = (getenv("KZTTS_EOS_NATURAL") && atoi(getenv("KZTTS_EOS_NATURAL")) != 0);
const int CB0_VOCAB = 2048; // codes audio valides [0, 2048)
const int expected_steps = Nt * 4;
const int eos_min_step = (expected_steps * 2) / 3;
const int eos_hard_step = eos_min_step * 3;
const int eos_rank_trig = 300;
const float eos_boost_scl = 4.0f;
const float NEG = -1e30f;
int boost_active = 0, consec_low = 0;
// Anti-repeat « seen » (port legacy l.1866) : pénalise les cb0 déjà émis pour
// faire remonter le rang de codec_eos. Sans ça le rang reste >300 et l'arming
// naturel ne se déclenche jamais -> on retombait sur le hard fallback brutal.
std::vector<char> seen(n_vocab, 0);
auto apply_eos = [&](float* lg, int gen_step) -> bool {
for (int j = CB0_VOCAB; j < n_vocab; ++j) if (j != eng->codec_eos) lg[j] = NEG;
for (int j = 0; j < n_vocab; ++j) if (seen[j]) lg[j] = (lg[j] > 0) ? lg[j] / 1.05f : lg[j] * 1.05f;
if (!eos_boost_on) return false;
const float eos_l = lg[eng->codec_eos];
int rank = 0; for (int j = 0; j < n_vocab; ++j) if (lg[j] > eos_l) rank++;
if (boost_active == 0 && gen_step >= eos_min_step) {
if (eos_natural_on && rank < eos_rank_trig) {
if (++consec_low >= 3) { boost_active = 1;
fprintf(stderr, "eos: armed NATURAL at step %d (rank=%d)\n", gen_step, rank); }
} else { consec_low = 0; }
if (boost_active == 0 && gen_step >= eos_hard_step) { boost_active = 1;
fprintf(stderr, "eos: armed HARD-FALLBACK at step %d (rank=%d)\n", gen_step, rank); }
} else if (boost_active > 0) {
boost_active++;
}
if (boost_active > 0) lg[eng->codec_eos] += boost_active * eos_boost_scl;
int am = 0; float amv = lg[0];
for (int j = 1; j < n_vocab; ++j) if (lg[j] > amv) { amv = lg[j]; am = j; }
return am == eng->codec_eos; // eos boosté en tête -> stop immédiat
};
apply_eos(logits_buf.data(), 0);
int cb0 = sampler_sample(talker_sampler, logits_buf.data(), n_vocab);
int last_cb0 = -1, cb0_run = 0;
std::vector<float> hidden_for_cp(n_embd);
{ const float* hh = llama_get_embeddings_ith(eng->talker_ctx, -1);
if (hh) memcpy(hidden_for_cp.data(), hh, n_embd * sizeof(float)); }
std::vector<int32_t> codes_engine;
codes_engine.reserve((size_t)cfg.max_steps * 16);
int N_done = 0;
const double t_loop0 = now_s();
const double t_synth_start = t_pfill0; // pour TTFB depuis le début total
double t_decode_total = 0, t_cp_total = 0;
// Streaming chunked : KZTTS_STREAM_CHUNK=K (>0) -> appeler decoder sur les codes
// accumulés tous les K frames. Le 1er chunk donne le TTFB. Le pipeline causal des 5
// stages garantit que decoder(K) puis decoder(2K) produisent bit-exact le même audio
// pour les frames [0..K]. On ne garde donc que la portion NOUVELLE de chaque appel.
const int stream_chunk = getenv("KZTTS_STREAM_CHUNK")
? atoi(getenv("KZTTS_STREAM_CHUNK")) : 0;
std::vector<float> wav_stream; // audio accumulé chunk par chunk (si streaming)
int last_decoded_N = 0; // nb de frames déjà décodées dans wav_stream
R.ttfb_s = 0;
// Helper : déclenche un decoder.forward sur N_so_far codes accumulés, append
// seulement la portion nouvelle [last_decoded_N * 1920 .. N_so_far * 1920) à wav_stream.
auto run_chunk = [&](int N_so_far) {
if (N_so_far <= last_decoded_N) return;
std::vector<int32_t> codes_dec_chunk(16 * N_so_far);
for (int t = 0; t < N_so_far; ++t)
for (int c = 0; c < 16; ++c)
codes_dec_chunk[c * N_so_far + t] = codes_engine[t * 16 + c];
const double td0 = now_s();
auto wav_full = eng->decoder.forward(codes_dec_chunk, N_so_far);
R.decoder_s += now_s() - td0;
// Append only NEW samples : [last_decoded_N*1920 .. N_so_far*1920)
const size_t start = (size_t)last_decoded_N * 1920;
const size_t end = (size_t)N_so_far * 1920;
wav_stream.insert(wav_stream.end(),
wav_full.begin() + start, wav_full.begin() + end);
if (R.ttfb_s == 0.0) R.ttfb_s = now_s() - t_synth_start;
R.n_chunks++;
last_decoded_N = N_so_far;
};
for (int s = 0; s < cfg.max_steps; ++s) {
if (cb0 == eng->codec_eos) break;
// Garde anti-dégénérescence : même cb0 répété 9x => boucle morte, on stoppe.
if (cb0 == last_cb0) { if (++cb0_run >= 9) { fprintf(stderr, "eos: degeneration break at step %d\n", s); break; } }
else { last_cb0 = cb0; cb0_run = 1; }
seen[cb0] = 1; // mémorise le cb0 courant pour la pénalité anti-repeat
codes_engine.push_back(cb0);
const double tcp0 = now_s();
const float* cb0_emb = eng->tok_embd.data() + (size_t)cb0 * n_embd;
if (getenv("KZTTS_DBG_CP") && s < 4) {
double hn = 0, en = 0; for (int d = 0; d < n_embd; ++d) { hn += hidden_for_cp[d]*hidden_for_cp[d]; en += cb0_emb[d]*cb0_emb[d]; }
fprintf(stderr, "dbg_cp s=%d cb0=%d |hidden|=%.4f |cb0_emb|=%.4f h[0..3]=%.3f,%.3f,%.3f\n",
s, cb0, sqrt(hn), sqrt(en), hidden_for_cp[0], hidden_for_cp[1], hidden_for_cp[2]);
}
int32_t cb15[15];
if (eng->cp_use_cache) cp_predict_cached(eng->cp_state, hidden_for_cp.data(), cb0_emb, cb15);
else cp_predict (eng->cp_state, hidden_for_cp.data(), cb0_emb, cb15);
t_cp_total += now_s() - tcp0;
for (int i = 0; i < 15; ++i) codes_engine.push_back(cb15[i]);
std::vector<float> next_embed(n_embd, 0.0f);
const float* e_cb0 = eng->tok_embd.data() + (size_t)cb0 * n_embd;
for (int d = 0; d < n_embd; ++d) next_embed[d] = e_cb0[d];
for (int i = 1; i < 16; ++i) {
int code = cb15[i - 1];
const float* e = eng->cp_state.codec_embs.data() + ((size_t)(i-1) * 2048 + code) * n_embd;
for (int d = 0; d < n_embd; ++d) next_embed[d] += e[d];
}
for (int d = 0; d < n_embd; ++d) next_embed[d] += tts_pad_emb[d];
const int npe = eng->npe;
llama_pos pos1[4] = {0,0,0,0};
const llama_pos p = T_prefill + s;
if (npe == 4) { pos1[0] = p; pos1[1] = p; pos1[2] = p; pos1[3] = 0; }
else { pos1[0] = p; }
int32_t nn = 1; llama_seq_id sd = 0; llama_seq_id* sp = &sd; int8_t l = 1;
llama_batch b{};
b.n_tokens = 1; b.embd = next_embed.data();
b.pos = pos1; b.n_seq_id = &nn; b.seq_id = &sp; b.logits = &l;
const double td0 = now_s();
if (llama_decode(eng->talker_ctx, b) != 0) { R.err = -5; break; }
t_decode_total += now_s() - td0;
{
const float* lp = llama_get_logits_ith(eng->talker_ctx, -1);
memcpy(logits_buf.data(), lp, n_vocab * sizeof(float));
}
const bool force_eos = apply_eos(logits_buf.data(), s + 1);
cb0 = force_eos ? eng->codec_eos
: sampler_sample(talker_sampler, logits_buf.data(), n_vocab);
const float* hh = llama_get_embeddings_ith(eng->talker_ctx, -1);
if (hh) memcpy(hidden_for_cp.data(), hh, n_embd * sizeof(float));
N_done = s + 1;
// Trigger streaming chunk : si KZTTS_STREAM_CHUNK=K et N_done atteint un multiple
// de K, décoder ce qu'on a accumulé. Le 1er trigger fournit le TTFB.
if (stream_chunk > 0 && N_done % stream_chunk == 0) {
run_chunk(N_done);
}
}
R.talker_loop_s = (now_s() - t_loop0) - t_cp_total; // ne pas double-compter CP
R.cp_loop_s = t_cp_total;
(void)t_decode_total; // déjà compris dans talker_loop_s
R.frames = N_done;
R.audio_s = N_done / 12.0;
// Dump optionnel des codes pour analyse des attracteurs (KZTTS_DUMP_CODES=path).
// Format : int32 binaire [N, 16] time-major (= codes_engine raw).
if (const char * dump = getenv("KZTTS_DUMP_CODES")) {
FILE * f = fopen(dump, "wb");
if (f) {
fwrite(codes_engine.data(), sizeof(int32_t), codes_engine.size(), f);
fclose(f);
fprintf(stderr, "codes dumped to %s (%d frames × 16)\n", dump, N_done);
}
}
// --- 8) Decoder final -> WAV ---
// KZTTS_DECODER_SUBPROC=path/kazeia_decoder_chraac : fork+exec un sous-process décodeur
// statiquement linké contre chraac (-22% sur decoder). Le talker reste qualcomm (qui marche).
// Demande KZTTS_DECODER_GGUF=path/qwen3tts_decoder.gguf pour passer en argv au sous-process.
const char * subproc_bin = getenv("KZTTS_DECODER_SUBPROC");
const char * subproc_model = getenv("KZTTS_DECODER_GGUF");
if (subproc_bin && subproc_model && stream_chunk == 0) {
// Codes time-major [N_done, 16] -> fichier tmp
char codes_tmp[256], wav_tmp[256], n_str[32];
snprintf(codes_tmp, sizeof(codes_tmp), "/data/local/tmp/kazeia_codes_%d.bin", (int)getpid());
snprintf(wav_tmp, sizeof(wav_tmp), "/data/local/tmp/kazeia_wav_%d.wav", (int)getpid());
snprintf(n_str, sizeof(n_str), "%d", N_done);
{
FILE * f = fopen(codes_tmp, "wb");
if (!f) { fprintf(stderr, "subproc: fopen codes FAIL\n"); R.err = -7; return R; }
fwrite(codes_engine.data(), sizeof(int32_t), codes_engine.size(), f);
fclose(f);
}
const double t_dec0 = now_s();
pid_t pid = fork();
if (pid == 0) {
execl(subproc_bin, subproc_bin, subproc_model, codes_tmp, n_str, wav_tmp, (char*)nullptr);
fprintf(stderr, "subproc execl FAIL\n");
_exit(127);
} else if (pid < 0) {
fprintf(stderr, "subproc fork FAIL\n"); R.err = -8; return R;
}
int status = 0;
waitpid(pid, &status, 0);
R.decoder_s = now_s() - t_dec0;
if (!WIFEXITED(status) || WEXITSTATUS(status) != 0) {
fprintf(stderr, "subproc exited code=%d\n", WEXITSTATUS(status)); R.err = -9; return R;
}
if (rename(wav_tmp, cfg.out_wav_path) != 0) {
fprintf(stderr, "subproc rename WAV %s -> %s FAIL\n", wav_tmp, cfg.out_wav_path);
R.err = -10; return R;
}
unlink(codes_tmp);
R.n_chunks = 1;
} else
if (stream_chunk > 0) {
// Mode streaming : flush le dernier chunk s'il reste des frames non décodées.
if (last_decoded_N < N_done) run_chunk(N_done);
// wav_stream contient déjà la concaténation propre des nouveaux samples par chunk.
if (!write_wav_pcm16_mono(cfg.out_wav_path, wav_stream.data(), wav_stream.size(), 24000)) {
fprintf(stderr, "tts_engine_synthesize: WAV write FAIL %s\n", cfg.out_wav_path);
R.err = -6; return R;
}
} else {
// Mode classique : un seul appel decoder sur all codes
std::vector<int32_t> codes_dec(16 * N_done);
for (int t = 0; t < N_done; ++t)
for (int c = 0; c < 16; ++c)
codes_dec[c * N_done + t] = codes_engine[t * 16 + c];
const double t_dec0 = now_s();
auto wav = eng->decoder.forward(codes_dec, N_done);
R.decoder_s = now_s() - t_dec0;
if (!write_wav_pcm16_mono(cfg.out_wav_path, wav.data(), wav.size(), 24000)) {
fprintf(stderr, "tts_engine_synthesize: WAV write FAIL %s\n", cfg.out_wav_path);
R.err = -6; return R;
}
R.n_chunks = 1;
}
R.total_s = R.prefill_s + R.talker_loop_s + R.cp_loop_s + R.decoder_s;
return R;
}
// ===========================================================================
// FREE
// ===========================================================================
void tts_engine_free(TtsEngine * eng) {
if (!eng) return;
cp_free(eng->cp_state);
if (eng->talker_ctx) llama_free(eng->talker_ctx);
if (eng->talker_m) llama_model_free(eng->talker_m);
if (eng->has_kz_tok) kz_tok_free(eng->kz_tok);
if (eng->spk_enc) speaker_encoder_free(eng->spk_enc);
delete eng;
}
// ===========================================================================
// SPEAKER ENCODER (clonage vocal embarqué)
// ===========================================================================
std::vector<float> tts_engine_encode_speaker_wav(TtsEngine * eng, const char * wav_path) {
if (!eng || !eng->spk_enc) {
fprintf(stderr, "tts_engine_encode_speaker_wav: speaker encoder non chargé\n");
return {};
}
return speaker_encoder_encode_wav(eng->spk_enc, wav_path);
}
std::vector<float> tts_engine_encode_speaker_waveform(TtsEngine * eng, const float * wav, int n_samples) {
if (!eng || !eng->spk_enc || !wav || n_samples <= 0) return {};
std::vector<float> w(wav, wav + n_samples);
return speaker_encoder_encode_waveform(eng->spk_enc, w);
}