762 lines
36 KiB
C++
762 lines
36 KiB
C++
#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 = LLAMA_FLASH_ATTN_TYPE_ENABLED;
|
||
cp.embeddings = true;
|
||
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 {
|
||
dec_ok = eng->decoder.load((D + "qwen3tts_decoder.gguf").c_str());
|
||
}
|
||
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));
|
||
}
|
||
int cb0 = sampler_sample(talker_sampler, logits_buf.data(), n_vocab);
|
||
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;
|
||
codes_engine.push_back(cb0);
|
||
|
||
const double tcp0 = now_s();
|
||
const float* cb0_emb = eng->tok_embd.data() + (size_t)cb0 * n_embd;
|
||
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));
|
||
}
|
||
cb0 = 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);
|
||
}
|