592 lines
26 KiB
C++
592 lines
26 KiB
C++
// Speaker encoder Qwen3-TTS (ECAPA-TDNN) en ggml C++.
|
|
// API publique : voir speaker_encoder.h.
|
|
//
|
|
// Pipeline : WAV 24kHz mono -> mel 128 -> ECAPA-TDNN -> x_vector[1024].
|
|
// Bit-correct vs torch.nn.functional.conv1d (cos > 0.999 vs Python ref, validé sur damien_5s).
|
|
//
|
|
// Architecture poids (cf /opt/Kazeia-engine/dist/models/speaker_encoder.gguf) :
|
|
// blocks.0 Conv1d(128 -> 512, k=5)
|
|
// blocks.1,2,3 SE-Res2Net (tdnn1 1x1 + 7x Conv k=3 + tdnn2 1x1 + se_block)
|
|
// mfa Conv1d(1536 -> 1536, k=1) // concat des 3 blocks
|
|
// asp.tdnn Conv1d(4608 -> 128, k=1) // Attentive Stats Pooling
|
|
// asp.conv Conv1d(128 -> 1536, k=1)
|
|
// fc Linear(3072 -> 1024) // x_vector final
|
|
//
|
|
// Compile avec -DSPK_STANDALONE pour le binaire CLI :
|
|
// usage : kazeia_speaker_encode <gguf> <mel_basis.bin> <ref.wav> <out.bin>
|
|
#include "speaker_encoder.h"
|
|
#include "ggml.h"
|
|
#include "gguf.h"
|
|
#include "ggml-cpu.h"
|
|
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <cmath>
|
|
#include <fstream>
|
|
#include <string>
|
|
#include <vector>
|
|
#include <unordered_map>
|
|
|
|
// ============================================================================
|
|
// Constantes mel (cf /opt/Kazeia/qnn_venv/.../modeling_qwen3_tts.py mel_spectrogram)
|
|
// ============================================================================
|
|
static const int SAMPLE_RATE = 24000;
|
|
static const int N_FFT = 1024;
|
|
static const int HOP_SIZE = 256;
|
|
static const int WIN_SIZE = 1024;
|
|
static const int N_MELS = 128;
|
|
static const int FFT_BINS = N_FFT / 2 + 1; // 513
|
|
static const float MEL_CLIP_VAL = 1e-5f;
|
|
|
|
// ============================================================================
|
|
// FFT radix-2 (in-place, n doit être puissance de 2). Identique à mel_extractor.cpp.
|
|
// ============================================================================
|
|
static void fft_radix2(float* real, float* imag, int n) {
|
|
// Bit-reversal
|
|
int j = 0;
|
|
for (int i = 1; i < n; i++) {
|
|
int bit = n >> 1;
|
|
while (j & bit) { j ^= bit; bit >>= 1; }
|
|
j ^= bit;
|
|
if (i < j) { std::swap(real[i], real[j]); std::swap(imag[i], imag[j]); }
|
|
}
|
|
for (int len = 2; len <= n; len <<= 1) {
|
|
int half = len / 2;
|
|
double angle = -2.0 * M_PI / len;
|
|
float wR = (float)std::cos(angle), wI = (float)std::sin(angle);
|
|
for (int i = 0; i < n; i += len) {
|
|
float cR = 1.f, cI = 0.f;
|
|
for (int k = 0; k < half; k++) {
|
|
float tR = cR * real[i+k+half] - cI * imag[i+k+half];
|
|
float tI = cR * imag[i+k+half] + cI * real[i+k+half];
|
|
real[i+k+half] = real[i+k] - tR;
|
|
imag[i+k+half] = imag[i+k] - tI;
|
|
real[i+k] += tR;
|
|
imag[i+k] += tI;
|
|
float ncR = cR * wR - cI * wI;
|
|
cI = cR * wI + cI * wR;
|
|
cR = ncR;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Hann window (cosine), torch.hann_window(N) = 0.5*(1 - cos(2π * n / (N-1)))
|
|
static std::vector<float> make_hann(int n) {
|
|
std::vector<float> w(n);
|
|
for (int i = 0; i < n; i++)
|
|
w[i] = 0.5f * (1.0f - std::cos(2.0f * (float)M_PI * (float)i / (float)(n - 1)));
|
|
return w;
|
|
}
|
|
|
|
// Pad reflect : insère padding samples sur chaque côté, miroir interne (sans répéter
|
|
// la 1ère/dernière). Équivalent torch.nn.functional.pad(..., mode="reflect").
|
|
static std::vector<float> pad_reflect(const std::vector<float>& x, int padding) {
|
|
const int n = (int)x.size();
|
|
std::vector<float> out((size_t)n + 2 * padding);
|
|
for (int i = 0; i < padding; i++) out[i] = x[padding - i]; // miroir gauche
|
|
for (int i = 0; i < n; i++) out[padding + i] = x[i]; // milieu
|
|
for (int i = 0; i < padding; i++) out[padding + n + i] = x[n - 2 - i]; // miroir droite
|
|
return out;
|
|
}
|
|
|
|
// Compute mel-spectrogram [n_mels=128, T] depuis waveform mono 24kHz.
|
|
// Réplique exactement qwen_tts.core.models.modeling_qwen3_tts.mel_spectrogram.
|
|
static std::vector<float> mel_spectrogram(const std::vector<float>& wav,
|
|
const std::vector<float>& mel_basis, // [128, 513]
|
|
int& out_T) {
|
|
const int padding = (N_FFT - HOP_SIZE) / 2; // 384
|
|
auto pad = pad_reflect(wav, padding);
|
|
auto hann = make_hann(WIN_SIZE);
|
|
|
|
// T = floor((len_padded - n_fft) / hop) + 1 (center=False)
|
|
const int len_pad = (int)pad.size();
|
|
if (len_pad < N_FFT) { fprintf(stderr, "wav too short (%d < %d)\n", len_pad, N_FFT); return {}; }
|
|
const int T = (len_pad - N_FFT) / HOP_SIZE + 1;
|
|
out_T = T;
|
|
|
|
std::vector<float> spec(FFT_BINS * T, 0.0f); // [513, T] magnitude
|
|
std::vector<float> re(N_FFT), im(N_FFT);
|
|
for (int t = 0; t < T; t++) {
|
|
const int off = t * HOP_SIZE;
|
|
for (int i = 0; i < N_FFT; i++) {
|
|
re[i] = pad[off + i] * hann[i];
|
|
im[i] = 0.0f;
|
|
}
|
|
fft_radix2(re.data(), im.data(), N_FFT);
|
|
// Magnitude sqrt(re² + im² + 1e-9), onesided -> bins 0..N_FFT/2
|
|
for (int k = 0; k < FFT_BINS; k++) {
|
|
float m = re[k]*re[k] + im[k]*im[k] + 1e-9f;
|
|
spec[k * T + t] = std::sqrt(m);
|
|
}
|
|
}
|
|
|
|
// mel = mel_basis [128, 513] @ spec [513, T] -> [128, T]
|
|
std::vector<float> mel(N_MELS * T, 0.0f);
|
|
for (int m = 0; m < N_MELS; m++) {
|
|
for (int t = 0; t < T; t++) {
|
|
float s = 0.0f;
|
|
for (int k = 0; k < FFT_BINS; k++)
|
|
s += mel_basis[m * FFT_BINS + k] * spec[k * T + t];
|
|
mel[m * T + t] = std::log(std::max(s, MEL_CLIP_VAL)); // dynamic_range_compression
|
|
}
|
|
}
|
|
return mel;
|
|
}
|
|
|
|
// ============================================================================
|
|
// Forward ECAPA-TDNN ggml C++ — bit-match qwen_tts.modeling_qwen3_tts
|
|
// ============================================================================
|
|
// Config Qwen3TTSSpeakerEncoderConfig (cf configuration_qwen3_tts.py):
|
|
// mel_dim = 128, enc_dim = 1024
|
|
// enc_channels = [512, 512, 512, 512, 1536]
|
|
// enc_kernel_sizes = [5, 3, 3, 3, 1] (k_init, k_block_1..3, k_mfa)
|
|
// enc_dilations = [1, 2, 3, 4, 1]
|
|
// enc_res2net_scale = 8 (chunks de 64 ch dans la res2net)
|
|
// enc_se_channels = 128
|
|
// enc_attention_channels = 128 (pour ASP)
|
|
static const int K_INIT = 5;
|
|
static const int K_RES2NET = 3;
|
|
static const int DIL[4] = {1, 2, 3, 4}; // dilation par block (block_1..3)
|
|
static const int CHANNELS = 512;
|
|
static const int RES2_SCALE = 8;
|
|
static const int RES2_CHUNK_CH = CHANNELS / RES2_SCALE; // 64
|
|
static const int MFA_CHANNELS = 1536;
|
|
static const int SE_CHANNELS_ = 128;
|
|
static const int ASP_ATT_CH = 128;
|
|
static const int EMBED_DIM = 1024;
|
|
|
|
struct SpeakerEncoder {
|
|
ggml_context * weights_ctx = nullptr;
|
|
std::unordered_map<std::string, ggml_tensor*> tensors;
|
|
std::vector<float> mel_basis; // [128, 513] librosa slaney
|
|
int n_threads = 6;
|
|
};
|
|
|
|
// Charge tous les tenseurs GGUF en F32 dans weights_ctx (17 MB en f32, OK).
|
|
// Note conv1d : ggml_conv_1d exige weights F16 -> cast au call site (ggml_cast).
|
|
static bool spk_load(SpeakerEncoder& s, const char* gguf_path) {
|
|
ggml_context* meta = nullptr;
|
|
gguf_init_params p; p.no_alloc = true; p.ctx = &meta;
|
|
gguf_context* g = gguf_init_from_file(gguf_path, p);
|
|
if (!g) { fprintf(stderr, "gguf open %s FAIL\n", gguf_path); return false; }
|
|
int64_t n = gguf_get_n_tensors(g);
|
|
size_t bytes = 0;
|
|
for (int64_t i = 0; i < n; i++) {
|
|
auto* mt = ggml_get_tensor(meta, gguf_get_tensor_name(g, i));
|
|
bytes += ggml_nelements(mt) * sizeof(float); // tout en F32
|
|
}
|
|
ggml_init_params ip = { bytes + (size_t)n * ggml_tensor_overhead() + (1u << 20), nullptr, false };
|
|
s.weights_ctx = ggml_init(ip);
|
|
|
|
FILE* f = fopen(gguf_path, "rb");
|
|
const size_t off = gguf_get_data_offset(g);
|
|
std::vector<uint8_t> tmp;
|
|
for (int64_t i = 0; i < n; i++) {
|
|
const char* name = gguf_get_tensor_name(g, i);
|
|
ggml_tensor* mt = ggml_get_tensor(meta, name);
|
|
const int ndims = ggml_n_dims(mt);
|
|
// GGUF -> ne ggml. Tout en F32 dans weights_ctx (17 MB f32, ample).
|
|
// Pour conv : on caste vers F16 au call site (ggml_cast in graph).
|
|
ggml_tensor* t = ggml_new_tensor(s.weights_ctx, GGML_TYPE_F32, ndims, mt->ne);
|
|
size_t nb = ggml_nbytes(mt);
|
|
int64_t nel = ggml_nelements(mt);
|
|
tmp.resize(nb);
|
|
fseek(f, off + gguf_get_tensor_offset(g, i), SEEK_SET);
|
|
if (fread(tmp.data(), 1, nb, f) != nb) { fprintf(stderr, "read fail %s\n", name); fclose(f); return false; }
|
|
if (mt->type == GGML_TYPE_F32) memcpy(t->data, tmp.data(), nb);
|
|
else if (mt->type == GGML_TYPE_F16) ggml_fp16_to_fp32_row((const ggml_fp16_t*)tmp.data(), (float*)t->data, nel);
|
|
else { fprintf(stderr, "type %s not supported\n", ggml_type_name(mt->type)); fclose(f); return false; }
|
|
s.tensors[name] = t;
|
|
}
|
|
fclose(f);
|
|
gguf_free(g);
|
|
return true;
|
|
}
|
|
|
|
// Helper : retrouve un tenseur par nom (abort si absent).
|
|
static ggml_tensor* W(SpeakerEncoder& s, const std::string& key) {
|
|
auto it = s.tensors.find(key);
|
|
if (it == s.tensors.end()) { fprintf(stderr, "missing tensor %s\n", key.c_str()); std::abort(); }
|
|
return it->second;
|
|
}
|
|
|
|
// TimeDelayNetBlock = Conv1d k same+reflect + ReLU.
|
|
// x : [T, C_in], w : [K, C_in, C_out], b : [C_out].
|
|
// Pour "same" avec dilation d, pad_left = pad_right = (k-1)*d/2 si k impair.
|
|
// Note conv_1d : im2col exige weight F16. On charge en F32 et caste ici.
|
|
// Padding zero "same" via le param p de ggml_conv_1d (PyTorch default).
|
|
static ggml_tensor* tdnn_block(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, int k, int dilation) {
|
|
const int pad = (k - 1) * dilation / 2;
|
|
ggml_tensor* w16 = ggml_cast(ctx, w, GGML_TYPE_F16);
|
|
ggml_tensor* y = ggml_conv_1d(ctx, w16, x, /*s=*/1, /*p=*/pad, /*d=*/dilation);
|
|
if (b) {
|
|
ggml_tensor* b2 = ggml_reshape_2d(ctx, b, 1, b->ne[0]);
|
|
y = ggml_add(ctx, y, b2);
|
|
}
|
|
return ggml_relu(ctx, y);
|
|
}
|
|
|
|
// Idem mais sans ReLU (linear conv 1d).
|
|
static ggml_tensor* conv1d_linear(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, int k, int dilation) {
|
|
const int pad = (k - 1) * dilation / 2;
|
|
ggml_tensor* w16 = ggml_cast(ctx, w, GGML_TYPE_F16);
|
|
ggml_tensor* y = ggml_conv_1d(ctx, w16, x, 1, pad, dilation);
|
|
if (b) {
|
|
ggml_tensor* b2 = ggml_reshape_2d(ctx, b, 1, b->ne[0]);
|
|
y = ggml_add(ctx, y, b2);
|
|
}
|
|
return y;
|
|
}
|
|
|
|
// Res2NetBlock : chunk x [T, 512] en scale=8 sous-tensors de 64 ch chacun le long
|
|
// de l'axe C (dim 1). Apply :
|
|
// out[0] = chunks[0]
|
|
// out[1] = tdnn_block(chunks[1]) (avec poids res2net.blocks[0])
|
|
// out[i>=2] = tdnn_block(chunks[i] + out[i-1]) (avec res2net.blocks[i-1])
|
|
// Concat outputs along channel dim -> [T, 512].
|
|
//
|
|
// w_list[i]/b_list[i] = res2net.blocks[i].conv.{weight,bias} pour i in 0..6 (7 sub-blocks).
|
|
static ggml_tensor* res2net_block(ggml_context* ctx, ggml_tensor* x,
|
|
const std::vector<ggml_tensor*>& w_list,
|
|
const std::vector<ggml_tensor*>& b_list,
|
|
int dilation) {
|
|
const int T = (int)x->ne[0];
|
|
std::vector<ggml_tensor*> chunks(RES2_SCALE);
|
|
for (int i = 0; i < RES2_SCALE; i++) {
|
|
// ggml_view_2d offset bytes : on prend C_chunk = 64 colonnes en dim 1 à partir de i*64
|
|
chunks[i] = ggml_view_2d(ctx, x, T, RES2_CHUNK_CH,
|
|
x->nb[1], (size_t)i * RES2_CHUNK_CH * x->nb[1]);
|
|
}
|
|
std::vector<ggml_tensor*> outs(RES2_SCALE);
|
|
outs[0] = ggml_cont(ctx, chunks[0]);
|
|
for (int i = 1; i < RES2_SCALE; i++) {
|
|
ggml_tensor* in = chunks[i];
|
|
if (i >= 2) in = ggml_add(ctx, ggml_cont(ctx, chunks[i]), outs[i-1]);
|
|
else in = ggml_cont(ctx, in);
|
|
outs[i] = tdnn_block(ctx, in, w_list[i-1], b_list[i-1], K_RES2NET, dilation);
|
|
}
|
|
// Concat dim 1 (channels). ggml_concat sur dim 1 :
|
|
ggml_tensor* acc = outs[0];
|
|
for (int i = 1; i < RES2_SCALE; i++) {
|
|
acc = ggml_concat(ctx, acc, outs[i], /*dim=*/1);
|
|
}
|
|
return acc;
|
|
}
|
|
|
|
// SqueezeExcitationBlock :
|
|
// m = mean(x, dim=2) shape [1, C] (T pooled to 1)
|
|
// m = relu(conv1(m)) shape [1, SE_CHANNELS]
|
|
// m = sigmoid(conv2(m)) shape [1, C]
|
|
// out = x * m broadcast on T
|
|
//
|
|
// x : [T, C], c1_w : [1, C, SE], c1_b : [SE], c2_w : [1, SE, C], c2_b : [C].
|
|
static ggml_tensor* se_block(ggml_context* ctx, ggml_tensor* x,
|
|
ggml_tensor* c1_w, ggml_tensor* c1_b,
|
|
ggml_tensor* c2_w, ggml_tensor* c2_b) {
|
|
// Mean sur dim T (dim 0 dans notre layout [T, C]).
|
|
// ggml_mean fait la moyenne sur dim 0. x [T, C] -> [1, C].
|
|
ggml_tensor* m = ggml_mean(ctx, x);
|
|
// m a shape [1, C]. Conv1d k=1 (= mul_mat).
|
|
// c1_w shape [1, C, SE] -> reshape [C, SE] pour mul_mat
|
|
ggml_tensor* c1_w2 = ggml_reshape_2d(ctx, c1_w, c1_w->ne[1], c1_w->ne[2]);
|
|
// m shape [1, C], pour mul_mat(W[C, SE], m[C, 1]) -> [SE, 1] on doit reshape m
|
|
ggml_tensor* m_C1 = ggml_reshape_2d(ctx, m, m->ne[1], 1); // [C, 1]
|
|
ggml_tensor* y = ggml_mul_mat(ctx, c1_w2, m_C1); // [SE, 1]
|
|
ggml_tensor* b1_2 = ggml_reshape_2d(ctx, c1_b, 1, c1_b->ne[0]);
|
|
y = ggml_add(ctx, ggml_reshape_2d(ctx, y, 1, y->ne[0]), b1_2); // [1, SE]
|
|
y = ggml_relu(ctx, y);
|
|
// conv2
|
|
ggml_tensor* c2_w2 = ggml_reshape_2d(ctx, c2_w, c2_w->ne[1], c2_w->ne[2]);
|
|
ggml_tensor* y_C1 = ggml_reshape_2d(ctx, y, y->ne[1], 1); // [SE, 1]
|
|
ggml_tensor* z = ggml_mul_mat(ctx, c2_w2, y_C1); // [C, 1]
|
|
ggml_tensor* b2_2 = ggml_reshape_2d(ctx, c2_b, 1, c2_b->ne[0]);
|
|
z = ggml_add(ctx, ggml_reshape_2d(ctx, z, 1, z->ne[0]), b2_2);
|
|
z = ggml_sigmoid(ctx, z); // [1, C]
|
|
// Broadcast z sur T : repeat dim 0
|
|
z = ggml_repeat(ctx, z, x); // [T, C]
|
|
return ggml_mul(ctx, x, z);
|
|
}
|
|
|
|
// SE-Res2Net block : tdnn1 -> res2net -> tdnn2 -> se + residual.
|
|
static ggml_tensor* se_res2net_block(SpeakerEncoder& s, ggml_context* ctx, ggml_tensor* x,
|
|
int block_idx, int dilation) {
|
|
char p[64]; snprintf(p, sizeof(p), "blocks.%d", block_idx);
|
|
std::string b = p;
|
|
|
|
ggml_tensor* residual = x;
|
|
|
|
// tdnn1 : Conv1d k=1 + ReLU
|
|
x = tdnn_block(ctx, x, W(s, b + ".tdnn1.conv.weight"), W(s, b + ".tdnn1.conv.bias"), 1, 1);
|
|
|
|
// res2net
|
|
std::vector<ggml_tensor*> rw, rb;
|
|
for (int i = 0; i < RES2_SCALE - 1; i++) {
|
|
char rp[80]; snprintf(rp, sizeof(rp), "%s.res2net_block.blocks.%d", p, i);
|
|
rw.push_back(W(s, std::string(rp) + ".conv.weight"));
|
|
rb.push_back(W(s, std::string(rp) + ".conv.bias"));
|
|
}
|
|
x = res2net_block(ctx, x, rw, rb, dilation);
|
|
|
|
// tdnn2 : Conv1d k=1 + ReLU
|
|
x = tdnn_block(ctx, x, W(s, b + ".tdnn2.conv.weight"), W(s, b + ".tdnn2.conv.bias"), 1, 1);
|
|
|
|
// se_block
|
|
x = se_block(ctx, x,
|
|
W(s, b + ".se_block.conv1.weight"), W(s, b + ".se_block.conv1.bias"),
|
|
W(s, b + ".se_block.conv2.weight"), W(s, b + ".se_block.conv2.bias"));
|
|
|
|
return ggml_add(ctx, x, residual);
|
|
}
|
|
|
|
// Attentive Statistics Pooling.
|
|
// x : [T, 1536]
|
|
// 1. mean_global = mean(x, dim=0) -> [1, 1536]
|
|
// std_global = sqrt(mean((x - mean_global)^2, dim=0).clamp(eps))
|
|
// 2. attention_in = concat([x, mean_broadcast, std_broadcast], dim=channels) -> [T, 4608]
|
|
// 3. attention = tdnn(attention_in) Conv1d k=1 4608->128 + ReLU -> [T, 128]
|
|
// 4. attention = tanh(attention)
|
|
// 5. attention = conv(attention) Conv1d k=1 128->1536 -> [T, 1536]
|
|
// (pas de mask -> on saute le masked_fill ; tous les frames sont valides)
|
|
// 6. attention = softmax(attention, dim=T)
|
|
// 7. mean = (attention * x).sum(0) -> [1, 1536]
|
|
// std = sqrt((attention*(x-mean.unsqueeze(0))^2).sum(0).clamp(eps))
|
|
// 8. pooled = concat([mean, std], dim=1) -> [1, 3072]
|
|
static ggml_tensor* asp_pool(SpeakerEncoder& s, ggml_context* ctx, ggml_tensor* x, int T) {
|
|
const float ASP_EPS = 1e-12f;
|
|
const int C = MFA_CHANNELS; // 1536
|
|
|
|
// 1. global stats (uniform mask)
|
|
ggml_tensor* mean_g = ggml_mean(ctx, x); // [1, C]
|
|
ggml_tensor* mean_g_rep = ggml_repeat(ctx, mean_g, x); // [T, C]
|
|
ggml_tensor* d = ggml_sub(ctx, x, mean_g_rep); // [T, C]
|
|
d = ggml_sqr(ctx, d);
|
|
ggml_tensor* std_g = ggml_mean(ctx, d); // [1, C]
|
|
// clamp min eps : on utilise max via add + relu hack ? Plus simple : ggml n'a pas de clamp inplace, mais sqrt(eps) si valeur petite est négligeable.
|
|
// On laisse direct sqrt (les valeurs sont positives) :
|
|
std_g = ggml_sqrt(ctx, ggml_scale_bias(ctx, std_g, 1.0f, ASP_EPS));
|
|
|
|
// 2. concat [x, mean_g_broadcast, std_g_broadcast] dim 1
|
|
ggml_tensor* std_g_rep = ggml_repeat(ctx, std_g, x); // [T, C]
|
|
ggml_tensor* att_in = ggml_concat(ctx, x, mean_g_rep, /*dim=*/1); // [T, 2C]
|
|
att_in = ggml_concat(ctx, att_in, std_g_rep, /*dim=*/1); // [T, 3C=4608]
|
|
|
|
// 3. tdnn (Conv1d k=1 + ReLU) : in=4608 -> 128
|
|
ggml_tensor* att = tdnn_block(ctx, att_in,
|
|
W(s, "asp.tdnn.conv.weight"),
|
|
W(s, "asp.tdnn.conv.bias"), 1, 1); // [T, 128]
|
|
// 4. tanh
|
|
att = ggml_tanh(ctx, att);
|
|
// 5. conv (Conv1d k=1 linear) : 128 -> 1536
|
|
att = conv1d_linear(ctx, att, W(s, "asp.conv.weight"), W(s, "asp.conv.bias"), 1, 1); // [T, C]
|
|
|
|
// 6. softmax dim T (dim 0 dans notre layout [T, C]).
|
|
// ggml_soft_max fait softmax sur dim 0 par défaut (le plus rapide).
|
|
att = ggml_soft_max(ctx, att);
|
|
|
|
// 7. weighted mean = sum(att * x, dim=T) -> [1, C]
|
|
// ggml_sum_rows somme sur dim 0 ; ici dim 0 = T (notre layout [T, C]). Donc direct.
|
|
ggml_tensor* wx = ggml_mul(ctx, att, x);
|
|
ggml_tensor* w_mean = ggml_sum_rows(ctx, wx); // [1, C]
|
|
// weighted std = sqrt(sum(att * (x - w_mean_broadcast)^2, dim=T) + eps)
|
|
ggml_tensor* w_mean_rep = ggml_repeat(ctx, w_mean, x); // [T, C] via broadcast
|
|
ggml_tensor* diff = ggml_sub(ctx, x, w_mean_rep);
|
|
diff = ggml_sqr(ctx, diff);
|
|
ggml_tensor* w_var_sum = ggml_sum_rows(ctx, ggml_mul(ctx, att, diff)); // [1, C]
|
|
ggml_tensor* w_std = ggml_sqrt(ctx, ggml_scale_bias(ctx, w_var_sum, 1.0f, ASP_EPS));
|
|
|
|
// 8. concat mean + std -> [1, 2C]
|
|
return ggml_concat(ctx, w_mean, w_std, /*dim=*/1); // [1, 2C=3072]
|
|
}
|
|
|
|
// Forward speaker encoder complet. mel : [T, 128] f32 row-major.
|
|
// Out : x_vector [1024] f32.
|
|
static std::vector<float> spk_forward(SpeakerEncoder& s, const std::vector<float>& mel_TC, int T) {
|
|
size_t mem = 512ULL * 1024 * 1024; // 512 MB pour le forward (3-4s audio, ample)
|
|
ggml_init_params p = { mem, nullptr, false };
|
|
ggml_context* ctx = ggml_init(p);
|
|
|
|
// Input ggml_conv_1d : ne[0]=T (length), ne[1]=C_in (channels). data[c*T + t].
|
|
// Notre mel est exactement [N_MELS, T] = data[c*T + t]. Memcpy direct.
|
|
ggml_tensor* x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, T, N_MELS);
|
|
memcpy(x->data, mel_TC.data(), T * N_MELS * sizeof(float));
|
|
|
|
// block_0 = TimeDelayNetBlock(Conv1d k=5 in=128 out=512) + ReLU
|
|
x = tdnn_block(ctx, x, W(s, "blocks.0.conv.weight"), W(s, "blocks.0.conv.bias"), K_INIT, DIL[0]);
|
|
ggml_tensor* block_0_out = x; // pour dump intermédiaire optionnel
|
|
|
|
// blocks 1, 2, 3 = SE-Res2Net + residual
|
|
ggml_tensor* block_outs[3];
|
|
for (int bi = 0; bi < 3; bi++) {
|
|
x = se_res2net_block(s, ctx, x, bi + 1, DIL[bi + 1]);
|
|
block_outs[bi] = x;
|
|
}
|
|
|
|
// MFA : concat(block_1, block_2, block_3) dim 1 = channels -> [T, 1536]
|
|
ggml_tensor* mfa_in = ggml_concat(ctx, block_outs[0], block_outs[1], /*dim=*/1);
|
|
mfa_in = ggml_concat(ctx, mfa_in, block_outs[2], 1);
|
|
x = tdnn_block(ctx, mfa_in, W(s, "mfa.conv.weight"), W(s, "mfa.conv.bias"), 1, 1);
|
|
|
|
// ASP pooling -> [1, 3072]
|
|
x = asp_pool(s, ctx, x, T);
|
|
|
|
// FC = Conv1d k=1 in=3072 out=1024 (PAS de ReLU)
|
|
x = conv1d_linear(ctx, x, W(s, "fc.weight"), W(s, "fc.bias"), 1, 1); // [1, 1024]
|
|
|
|
ggml_cgraph* gf = ggml_new_graph_custom(ctx, 4096, false);
|
|
ggml_build_forward_expand(gf, x);
|
|
// Forcer la matérialisation des blocs intermédiaires si dump activé
|
|
ggml_build_forward_expand(gf, block_0_out);
|
|
for (int bi = 0; bi < 3; bi++) ggml_build_forward_expand(gf, block_outs[bi]);
|
|
|
|
ggml_graph_compute_with_ctx(ctx, gf, s.n_threads);
|
|
|
|
// Dump intermédiaires (debug, opt-in via env)
|
|
auto dump_t = [](const char* envkey, ggml_tensor* t) {
|
|
const char* p = std::getenv(envkey);
|
|
if (!p) return;
|
|
FILE* f = fopen(p, "wb");
|
|
if (!f) return;
|
|
size_t n = ggml_nelements(t);
|
|
std::vector<float> buf(n);
|
|
if (t->type == GGML_TYPE_F32) memcpy(buf.data(), t->data, n * sizeof(float));
|
|
else if (t->type == GGML_TYPE_F16) ggml_fp16_to_fp32_row((const ggml_fp16_t*)t->data, buf.data(), n);
|
|
fwrite(buf.data(), sizeof(float), n, f);
|
|
fclose(f);
|
|
};
|
|
dump_t("KZTTS_DUMP_B0", block_0_out);
|
|
dump_t("KZTTS_DUMP_B1", block_outs[0]);
|
|
dump_t("KZTTS_DUMP_B2", block_outs[1]);
|
|
dump_t("KZTTS_DUMP_B3", block_outs[2]);
|
|
|
|
std::vector<float> out(EMBED_DIM);
|
|
memcpy(out.data(), x->data, EMBED_DIM * sizeof(float));
|
|
ggml_free(ctx);
|
|
return out;
|
|
}
|
|
|
|
// ============================================================================
|
|
// WAV reader minimal (PCM16 mono, n'importe quel sample rate, on assume 24kHz)
|
|
// ============================================================================
|
|
// RIFF parser propre : trouve les chunks "fmt " et "data" peu importe leur offset
|
|
// (ffmpeg ajoute parfois LIST/INFO/JUNK avant "data").
|
|
static std::vector<float> read_wav_24k_mono(const char* path, int& sr_out) {
|
|
FILE* f = fopen(path, "rb");
|
|
if (!f) { fprintf(stderr, "open %s FAIL\n", path); return {}; }
|
|
char riff[12];
|
|
if (fread(riff, 1, 12, f) != 12 || memcmp(riff, "RIFF", 4) != 0 || memcmp(riff+8, "WAVE", 4) != 0) {
|
|
fprintf(stderr, "not a RIFF/WAVE file\n"); fclose(f); return {};
|
|
}
|
|
uint16_t channels = 0, bps = 0;
|
|
uint32_t sr = 0, data_sz = 0;
|
|
long data_off = -1;
|
|
while (!feof(f)) {
|
|
char cid[4]; uint32_t csz;
|
|
if (fread(cid, 1, 4, f) != 4) break;
|
|
if (fread(&csz, 4, 1, f) != 1) break;
|
|
if (memcmp(cid, "fmt ", 4) == 0) {
|
|
std::vector<uint8_t> buf(csz);
|
|
fread(buf.data(), 1, csz, f);
|
|
channels = *(uint16_t*)&buf[2];
|
|
sr = *(uint32_t*)&buf[4];
|
|
bps = *(uint16_t*)&buf[14];
|
|
} else if (memcmp(cid, "data", 4) == 0) {
|
|
data_sz = csz;
|
|
data_off = ftell(f);
|
|
break;
|
|
} else {
|
|
fseek(f, csz, SEEK_CUR); // skip LIST/JUNK/etc
|
|
}
|
|
}
|
|
if (data_off < 0 || channels != 1 || bps != 16) {
|
|
fprintf(stderr, "WAV must be mono 16-bit (got ch=%d bps=%d data_off=%ld)\n", channels, bps, data_off);
|
|
fclose(f); return {};
|
|
}
|
|
sr_out = (int)sr;
|
|
fseek(f, data_off, SEEK_SET);
|
|
const int n = data_sz / 2;
|
|
std::vector<int16_t> raw(n);
|
|
fread(raw.data(), 2, n, f);
|
|
fclose(f);
|
|
std::vector<float> wav(n);
|
|
for (int i = 0; i < n; i++) wav[i] = (float)raw[i] / 32768.0f;
|
|
return wav;
|
|
}
|
|
|
|
// ============================================================================
|
|
// API publique (cf speaker_encoder.h)
|
|
// ============================================================================
|
|
SpeakerEncoder * speaker_encoder_load(const char * gguf_path, const char * mel_basis_path) {
|
|
auto * spk = new SpeakerEncoder();
|
|
if (!spk_load(*spk, gguf_path)) {
|
|
delete spk; return nullptr;
|
|
}
|
|
spk->mel_basis.assign((size_t)N_MELS * FFT_BINS, 0.0f);
|
|
std::ifstream f(mel_basis_path, std::ios::binary);
|
|
if (!f) {
|
|
fprintf(stderr, "speaker_encoder_load: open %s FAIL\n", mel_basis_path);
|
|
speaker_encoder_free(spk); return nullptr;
|
|
}
|
|
f.read((char*)spk->mel_basis.data(), spk->mel_basis.size() * sizeof(float));
|
|
if (!f) {
|
|
fprintf(stderr, "speaker_encoder_load: read %s FAIL\n", mel_basis_path);
|
|
speaker_encoder_free(spk); return nullptr;
|
|
}
|
|
if (const char* th = std::getenv("KZTTS_SPK_THREADS")) spk->n_threads = atoi(th);
|
|
return spk;
|
|
}
|
|
|
|
std::vector<float> speaker_encoder_encode_waveform(SpeakerEncoder * spk, const std::vector<float> & wav) {
|
|
if (!spk) return {};
|
|
int T_mel = 0;
|
|
auto mel = mel_spectrogram(wav, spk->mel_basis, T_mel);
|
|
if (mel.empty()) return {};
|
|
return spk_forward(*spk, mel, T_mel);
|
|
}
|
|
|
|
std::vector<float> speaker_encoder_encode_wav(SpeakerEncoder * spk, const char * wav_path) {
|
|
int sr = 0;
|
|
auto wav = read_wav_24k_mono(wav_path, sr);
|
|
if (wav.empty()) return {};
|
|
if (sr != SAMPLE_RATE) {
|
|
fprintf(stderr, "speaker_encoder_encode_wav: sample rate %d != 24000\n", sr);
|
|
}
|
|
return speaker_encoder_encode_waveform(spk, wav);
|
|
}
|
|
|
|
void speaker_encoder_free(SpeakerEncoder * spk) {
|
|
if (!spk) return;
|
|
if (spk->weights_ctx) ggml_free(spk->weights_ctx);
|
|
delete spk;
|
|
}
|
|
|
|
// ============================================================================
|
|
// MAIN (binaire standalone -DSPK_STANDALONE)
|
|
// ============================================================================
|
|
#ifdef SPK_STANDALONE
|
|
int main(int argc, char** argv) {
|
|
setvbuf(stderr, nullptr, _IONBF, 0);
|
|
setvbuf(stdout, nullptr, _IONBF, 0);
|
|
if (argc < 5) {
|
|
fprintf(stderr, "usage: %s <speaker_encoder.gguf> <mel_basis.bin> <ref.wav> <out_xvector.bin>\n",
|
|
argv[0]);
|
|
return 1;
|
|
}
|
|
auto * spk = speaker_encoder_load(argv[1], argv[2]);
|
|
if (!spk) return 2;
|
|
auto xvec = speaker_encoder_encode_wav(spk, argv[3]);
|
|
if (xvec.empty()) { speaker_encoder_free(spk); return 3; }
|
|
|
|
FILE* fo = fopen(argv[4], "wb");
|
|
if (!fo) { fprintf(stderr, "open %s FAIL\n", argv[4]); speaker_encoder_free(spk); return 4; }
|
|
fwrite(xvec.data(), sizeof(float), xvec.size(), fo);
|
|
fclose(fo);
|
|
|
|
float sumsq = 0; for (float v : xvec) sumsq += v*v;
|
|
fprintf(stderr, "x_vector norm=%.3f -> %s\n", std::sqrt(sumsq), argv[4]);
|
|
speaker_encoder_free(spk);
|
|
return 0;
|
|
}
|
|
#endif // SPK_STANDALONE
|