349 lines
17 KiB
C++
349 lines
17 KiB
C++
#pragma once
|
|
// speaker-encoder-forward.h: ECAPA-TDNN forward graph in GGML.
|
|
//
|
|
// Mirrors qwen_tts.core.models.modeling_qwen3_tts.Qwen3TTSSpeakerEncoder
|
|
// for the single utterance unbatched path. The forward fuses the mel
|
|
// spectrogram extraction so the whole pipeline lives in one graph :
|
|
//
|
|
// audio [T_pad] f32
|
|
// -> mel [128, T_frames] (audio-mel.h)
|
|
// -> conv0 TDNN k=5 + ReLU [512, T_frames]
|
|
// -> SE-Res2Net dil=2 [512, T_frames]
|
|
// -> SE-Res2Net dil=3 [512, T_frames]
|
|
// -> SE-Res2Net dil=4 [512, T_frames]
|
|
// -> cat blk[1..3] + MFA k=1 + ReLU [1536, T_frames]
|
|
// -> ASP attentive pooling [3072, 1]
|
|
// -> FC k=1 [2048, 1]
|
|
// -> squeeze [2048]
|
|
//
|
|
// Tensor convention: [C, T] inside the graph (ne[0]=C, ne[1]=T) so that
|
|
// ggml_im2col reads each Conv1d along the time axis and ggml_mul_mat
|
|
// contracts over the input channel axis. This matches the layout the
|
|
// upstream PyTorch code uses after its (1, 2) transpose.
|
|
|
|
#include "audio-mel.h"
|
|
#include "ggml-alloc.h"
|
|
#include "ggml-backend.h"
|
|
#include "ggml.h"
|
|
#include "speaker-encoder-weights.h"
|
|
|
|
#include <cmath>
|
|
#include <cstdio>
|
|
#include <vector>
|
|
|
|
// Conv1d k=K with padding="same" mode="reflect" + bias add. The weight
|
|
// tensor lives in upstream layout [K, in_c, out_c]. We implement it with
|
|
// reflect pad + im2col + matmul.
|
|
//
|
|
// x [in_c, T] input
|
|
// w [K, in_c, out_c] weights
|
|
// b [out_c] bias (broadcast over T)
|
|
// Returns [out_c, T]
|
|
//
|
|
// Padding for "same" with kernel K and dilation d is (K - 1) * d / 2 on
|
|
// each side (PyTorch convention, kernel size always odd here so the
|
|
// division is exact).
|
|
static struct ggml_tensor * spk_conv1d_same(struct ggml_context * ctx,
|
|
struct ggml_tensor * x,
|
|
struct ggml_tensor * w,
|
|
struct ggml_tensor * b,
|
|
int dilation) {
|
|
const int K = (int) w->ne[0];
|
|
const int IC = (int) w->ne[1];
|
|
const int OC = (int) w->ne[2];
|
|
const int pad = ((K - 1) * dilation) / 2;
|
|
|
|
// ggml_pad_reflect_1d pads the innermost axis ne[0]. Our temporal
|
|
// axis is ne[1], so we transpose to bring T to ne[0], pad, and keep
|
|
// it that way: the im2col downstream expects ne[0]=T_pad, ne[1]=IC,
|
|
// which is exactly the layout we end up with here.
|
|
struct ggml_tensor * x_t = ggml_cont(ctx, ggml_transpose(ctx, x)); // ne=(T, IC)
|
|
if (pad > 0) {
|
|
x_t = ggml_pad_reflect_1d(ctx, x_t, pad, pad); // ne=(T+2*pad, IC)
|
|
}
|
|
|
|
// Reshape to 4D for ggml_im2col 1D: ne=(T_pad, IC, 1, 1).
|
|
struct ggml_tensor * x4d = ggml_reshape_4d(ctx, x_t, x_t->ne[0], IC, 1, 1);
|
|
|
|
// Dummy F32 kernel with the (K, IC) shape ggml_im2col needs to read
|
|
// off both axes. Borrowing only the ne and not the real weight data
|
|
// avoids the impl's src0 type assert when w is quantized.
|
|
struct ggml_tensor * dummy = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, K, IC, 1, 1);
|
|
ggml_set_name(dummy, "spk.im2col_kernel");
|
|
|
|
// im2col with is_2D=false: the constructor declares ne[0]=IC*K and
|
|
// ne[1]=OW, and the impl writes the buffer in (k inner, ic middle,
|
|
// t outer) order, which matches that ne directly. A reshape_2d to
|
|
// (IC*K, T_out) reads col_2d[ic*K + k, t], lining up with the
|
|
// weight reshape w_2d[ic*K + k, oc] for the mul_mat below.
|
|
struct ggml_tensor * col = ggml_im2col(ctx, dummy, x4d, 1, 1, 0, 0, dilation, 1, false, GGML_TYPE_F32);
|
|
int T_out = (int) col->ne[1];
|
|
col = ggml_reshape_2d(ctx, col, K * IC, T_out);
|
|
|
|
// Weight reshape: [K, IC, OC] -> [K*IC, OC]. mul_mat returns [OC, T_out].
|
|
struct ggml_tensor * w2d = ggml_reshape_2d(ctx, w, K * IC, OC);
|
|
struct ggml_tensor * y = ggml_mul_mat(ctx, w2d, col);
|
|
ggml_mul_mat_set_prec(y, GGML_PREC_F32);
|
|
|
|
// Add bias broadcast over T_out. b is [out_c], reshape [out_c, 1].
|
|
struct ggml_tensor * b2d = ggml_reshape_2d(ctx, b, OC, 1);
|
|
y = ggml_add(ctx, y, b2d);
|
|
return y;
|
|
}
|
|
|
|
// TDNN block: Conv1d(same, reflect) + ReLU. Used both as the conv0
|
|
// frontend (k=5) and inside SE-Res2Net (k=1) and the MFA / ASP TDNNs.
|
|
static struct ggml_tensor * spk_tdnn(struct ggml_context * ctx,
|
|
const SpkEncTDNN & t,
|
|
struct ggml_tensor * x,
|
|
int dilation) {
|
|
struct ggml_tensor * y = spk_conv1d_same(ctx, x, t.weight, t.bias, dilation);
|
|
y = ggml_relu(ctx, y);
|
|
return y;
|
|
}
|
|
|
|
// Res2Net block: split the channel axis in 8 chunks. chunk 0 passes
|
|
// through, chunk 1 goes through TDNN[0], chunks 2..7 mix with the
|
|
// previous chunk output before going through TDNN[i-1]. The 7 TDNN
|
|
// branches share dilation but operate on hidden / 8 channels each.
|
|
//
|
|
// x [C, T]
|
|
// Returns [C, T]
|
|
static struct ggml_tensor * spk_res2net(struct ggml_context * ctx,
|
|
const SpkEncRes2Net & rn,
|
|
struct ggml_tensor * x,
|
|
int dilation,
|
|
int scale) {
|
|
const int C = (int) x->ne[0];
|
|
const int T = (int) x->ne[1];
|
|
const int Cs = C / scale;
|
|
|
|
std::vector<struct ggml_tensor *> outs;
|
|
outs.reserve(scale);
|
|
|
|
// chunk i is the slice along ne[0] of width Cs starting at i * Cs.
|
|
auto chunk = [&](int i) -> struct ggml_tensor * {
|
|
return ggml_view_2d(ctx, x, Cs, T, x->nb[1], (size_t) (i * Cs) * x->nb[0]);
|
|
};
|
|
|
|
struct ggml_tensor * prev = NULL;
|
|
for (int i = 0; i < scale; i++) {
|
|
struct ggml_tensor * c = ggml_cont(ctx, chunk(i));
|
|
if (i == 0) {
|
|
outs.push_back(c);
|
|
continue;
|
|
}
|
|
struct ggml_tensor * inp = c;
|
|
if (i >= 2) {
|
|
inp = ggml_add(ctx, c, prev);
|
|
}
|
|
struct ggml_tensor * y = spk_conv1d_same(ctx, inp, rn.weight[i - 1], rn.bias[i - 1], dilation);
|
|
y = ggml_relu(ctx, y);
|
|
outs.push_back(y);
|
|
prev = y;
|
|
}
|
|
|
|
// Concat along ne[0]. ggml_concat with dim=0 stacks along the
|
|
// fastest axis. Build the concat tree iteratively.
|
|
struct ggml_tensor * acc = outs[0];
|
|
for (int i = 1; i < scale; i++) {
|
|
acc = ggml_concat(ctx, acc, outs[i], 0);
|
|
}
|
|
return acc;
|
|
}
|
|
|
|
// Squeeze and Excitation: compute the temporal mean per channel,
|
|
// project down to se_c with a 1x1 conv + ReLU, project back up to
|
|
// out_c with a 1x1 conv + sigmoid, then scale the input by the gate
|
|
// broadcast over T.
|
|
static struct ggml_tensor * spk_se(struct ggml_context * ctx, const SpkEncSE & se, struct ggml_tensor * x) {
|
|
const int T = (int) x->ne[1];
|
|
|
|
// Mean over T, keep dim. ggml_mean reduces along ne[0], so transpose
|
|
// to put T on ne[0], reduce, transpose back.
|
|
struct ggml_tensor * x_t = ggml_cont(ctx, ggml_transpose(ctx, x)); // [T, C]
|
|
struct ggml_tensor * mean = ggml_mean(ctx, x_t); // [1, C]
|
|
mean = ggml_cont(ctx, ggml_transpose(ctx, mean)); // [C, 1]
|
|
|
|
// conv1 1x1 reduces C -> se_c. dilation 1, padding "same" trivial
|
|
// since k=1.
|
|
struct ggml_tensor * h = spk_conv1d_same(ctx, mean, se.conv1_w, se.conv1_b, 1);
|
|
h = ggml_relu(ctx, h);
|
|
h = spk_conv1d_same(ctx, h, se.conv2_w, se.conv2_b, 1);
|
|
// Sigmoid over [out_c, 1].
|
|
h = ggml_sigmoid(ctx, h);
|
|
|
|
// Scale x by the gate. h is [C, 1], x is [C, T]. ggml_mul broadcasts
|
|
// ne[1]=1 to T.
|
|
struct ggml_tensor * y = ggml_mul(ctx, x, h);
|
|
(void) T;
|
|
return y;
|
|
}
|
|
|
|
// SE-Res2Net block: tdnn1 (1x1) -> Res2Net -> tdnn2 (1x1) -> SE plus
|
|
// a residual add over the whole stack.
|
|
static struct ggml_tensor * spk_block(struct ggml_context * ctx,
|
|
const SpkEncBlock & blk,
|
|
struct ggml_tensor * x,
|
|
int res2net_scale) {
|
|
struct ggml_tensor * residual = x;
|
|
struct ggml_tensor * h = spk_tdnn(ctx, blk.tdnn1, x, 1);
|
|
h = spk_res2net(ctx, blk.res2net, h, blk.dilation, res2net_scale);
|
|
h = spk_tdnn(ctx, blk.tdnn2, h, 1);
|
|
h = spk_se(ctx, blk.se, h);
|
|
return ggml_add(ctx, h, residual);
|
|
}
|
|
|
|
// Attentive Statistical Pooling: compute global mean and std along T,
|
|
// concat with x, run an attention TDNN + tanh + 1x1 conv, softmax along
|
|
// T, recompute weighted mean and std, return the [2C, 1] concat.
|
|
//
|
|
// x [C, T]
|
|
// Returns [2C, 1]
|
|
static struct ggml_tensor * spk_asp(struct ggml_context * ctx, const SpkEncASP & asp, struct ggml_tensor * x) {
|
|
const int C = (int) x->ne[0];
|
|
const int T = (int) x->ne[1];
|
|
|
|
// Mean and std over T axis. The mask reduction is uniform 1/T.
|
|
// mean: [C, 1]
|
|
struct ggml_tensor * x_t = ggml_cont(ctx, ggml_transpose(ctx, x));
|
|
struct ggml_tensor * mean = ggml_mean(ctx, x_t);
|
|
mean = ggml_cont(ctx, ggml_transpose(ctx, mean));
|
|
|
|
// var = mean( (x - mean)^2 ) over T, then std = sqrt(clamp(var, eps)).
|
|
// Broadcasting (x - mean) requires mean repeated to T. ggml_repeat
|
|
// handles this when shapes are compatible.
|
|
struct ggml_tensor * mean_T = ggml_repeat(ctx, mean, x);
|
|
struct ggml_tensor * centered = ggml_sub(ctx, x, mean_T);
|
|
struct ggml_tensor * var_t = ggml_cont(ctx, ggml_transpose(ctx, ggml_sqr(ctx, centered)));
|
|
struct ggml_tensor * var = ggml_mean(ctx, var_t);
|
|
var = ggml_cont(ctx, ggml_transpose(ctx, var));
|
|
var = ggml_scale_bias(ctx, var, 1.0f, 1e-12f);
|
|
struct ggml_tensor * std = ggml_sqrt(ctx, var);
|
|
|
|
// Build [x, mean_repeat, std_repeat] concat along channel axis.
|
|
struct ggml_tensor * std_T = ggml_repeat(ctx, std, x);
|
|
struct ggml_tensor * cat = ggml_concat(ctx, x, mean_T, 0);
|
|
cat = ggml_concat(ctx, cat, std_T, 0); // [3C, T]
|
|
|
|
// Attention TDNN: 3C -> attn_c, ReLU, then tanh, then 1x1 conv
|
|
// attn_c -> C. Upstream applies tanh on the TDNN output before the
|
|
// second conv; the TDNN itself already runs ReLU so the order is
|
|
// ReLU then tanh which is unusual but mirrored faithfully.
|
|
struct ggml_tensor * a = spk_tdnn(ctx, asp.tdnn, cat, 1);
|
|
a = ggml_tanh(ctx, a);
|
|
a = spk_conv1d_same(ctx, a, asp.conv_w, asp.conv_b, 1);
|
|
|
|
// Softmax along T axis. ggml_soft_max reduces ne[0], transpose first.
|
|
struct ggml_tensor * a_t = ggml_cont(ctx, ggml_transpose(ctx, a)); // [T, C]
|
|
struct ggml_tensor * w_t = ggml_soft_max(ctx, a_t);
|
|
struct ggml_tensor * w = ggml_cont(ctx, ggml_transpose(ctx, w_t)); // [C, T]
|
|
|
|
// Weighted mean: sum(w * x) over T, w already sums to 1 over T.
|
|
struct ggml_tensor * wx = ggml_mul(ctx, w, x);
|
|
struct ggml_tensor * wx_t = ggml_cont(ctx, ggml_transpose(ctx, wx));
|
|
// ggml_mean averages over ne[0]=T, giving 1/T scaling. We want the
|
|
// un-normalized sum since w already encodes the soft selection
|
|
// probability, so multiply back by T.
|
|
struct ggml_tensor * w_mean = ggml_mean(ctx, wx_t);
|
|
w_mean = ggml_scale(ctx, w_mean, (float) T);
|
|
w_mean = ggml_cont(ctx, ggml_transpose(ctx, w_mean)); // [C, 1]
|
|
|
|
// Weighted std: sum(w * (x - w_mean)^2) over T.
|
|
struct ggml_tensor * w_mean_T = ggml_repeat(ctx, w_mean, x);
|
|
struct ggml_tensor * dev = ggml_sub(ctx, x, w_mean_T);
|
|
struct ggml_tensor * w_var_in = ggml_mul(ctx, w, ggml_sqr(ctx, dev));
|
|
struct ggml_tensor * w_var_t = ggml_cont(ctx, ggml_transpose(ctx, w_var_in));
|
|
struct ggml_tensor * w_var = ggml_mean(ctx, w_var_t);
|
|
w_var = ggml_scale(ctx, w_var, (float) T);
|
|
w_var = ggml_cont(ctx, ggml_transpose(ctx, w_var));
|
|
w_var = ggml_scale_bias(ctx, w_var, 1.0f, 1e-12f);
|
|
struct ggml_tensor * w_std = ggml_sqrt(ctx, w_var);
|
|
|
|
// Stack [w_mean, w_std] along channel -> [2C, 1]. Time axis already
|
|
// collapsed.
|
|
struct ggml_tensor * stats = ggml_concat(ctx, w_mean, w_std, 0);
|
|
(void) C;
|
|
return stats;
|
|
}
|
|
|
|
// Full speaker encoder forward graph. Assumes the audio waveform has
|
|
// already been resampled to sr=24000 and reflect padded by
|
|
// (n_fft - hop) / 2 on each side. The padded buffer must outlive the
|
|
// graph compute call.
|
|
//
|
|
// Inputs :
|
|
// audio_padded [T_pad] f32, host or backend tensor
|
|
// mel constants hann/dft_real/dft_imag/mel_basis backend tensors
|
|
// mel_out optional out param. When non NULL, receives the post
|
|
// mel_spectrogram tensor [n_mels, T_frames] so the caller
|
|
// can mark it as a graph output and pull its values back.
|
|
// mag_out optional out param. When non NULL, receives the post
|
|
// STFT magnitude tensor [n_freq, T_frames] for debug
|
|
// bisection between the STFT and the mel filtering.
|
|
// frontend_out optional. Post conv0 TDNN k=5 + ReLU output [512, T].
|
|
// block3_out optional. Post third SE-Res2Net block output [512, T].
|
|
// mfa_out optional. Post multi-layer feature aggregation [1536, T].
|
|
// asp_out optional. Post attentive statistical pooling [3072, 1].
|
|
// Output: [enc_dim] f32, the speaker embedding (typically 2048 dims).
|
|
static struct ggml_tensor * speaker_encoder_forward(struct ggml_context * ctx,
|
|
const SpeakerEncoderWeights * sw,
|
|
struct ggml_tensor * audio_padded,
|
|
struct ggml_tensor * hann,
|
|
struct ggml_tensor * dft_real,
|
|
struct ggml_tensor * dft_imag,
|
|
struct ggml_tensor * mel_basis,
|
|
const AudioMelConfig & mel_cfg,
|
|
struct ggml_tensor ** mel_out = NULL,
|
|
struct ggml_tensor ** mag_out = NULL,
|
|
struct ggml_tensor ** frontend_out = NULL,
|
|
struct ggml_tensor ** block3_out = NULL,
|
|
struct ggml_tensor ** mfa_out = NULL,
|
|
struct ggml_tensor ** asp_out = NULL) {
|
|
// Mel: [n_mels=128, T_frames]
|
|
struct ggml_tensor * mel =
|
|
audio_mel_build_graph(ctx, audio_padded, hann, dft_real, dft_imag, mel_basis, mel_cfg, mag_out);
|
|
if (mel_out) {
|
|
*mel_out = mel;
|
|
}
|
|
|
|
// Frontend conv0 TDNN k=5 + ReLU: 128 -> 512, T preserved.
|
|
struct ggml_tensor * h = spk_tdnn(ctx, sw->conv0, mel, 1);
|
|
if (frontend_out) {
|
|
*frontend_out = h;
|
|
}
|
|
|
|
// Three SE-Res2Net blocks at dilations 2, 3, 4.
|
|
struct ggml_tensor * b1 = spk_block(ctx, sw->blocks[0], h, sw->res2net_scale);
|
|
struct ggml_tensor * b2 = spk_block(ctx, sw->blocks[1], b1, sw->res2net_scale);
|
|
struct ggml_tensor * b3 = spk_block(ctx, sw->blocks[2], b2, sw->res2net_scale);
|
|
if (block3_out) {
|
|
*block3_out = b3;
|
|
}
|
|
|
|
// Multi-layer feature aggregation: cat blk1..3 then 1x1 TDNN + ReLU.
|
|
struct ggml_tensor * cat = ggml_concat(ctx, b1, b2, 0);
|
|
cat = ggml_concat(ctx, cat, b3, 0); // [1536, T]
|
|
struct ggml_tensor * mfa = spk_tdnn(ctx, sw->mfa, cat, 1); // [1536, T]
|
|
if (mfa_out) {
|
|
*mfa_out = mfa;
|
|
}
|
|
|
|
// Attentive statistical pooling: [1536, T] -> [3072, 1].
|
|
struct ggml_tensor * stats = spk_asp(ctx, sw->asp, mfa);
|
|
if (asp_out) {
|
|
*asp_out = stats;
|
|
}
|
|
|
|
// Final FC k=1: [3072, 1] -> [enc_dim, 1].
|
|
struct ggml_tensor * emb = spk_conv1d_same(ctx, stats, sw->fc_w, sw->fc_b, 1);
|
|
|
|
// Squeeze T axis, return [enc_dim]. ggml_cont is required so the sched
|
|
// assigns a fresh backend buffer to the graph output rather than
|
|
// forwarding a view of the FC bias add.
|
|
emb = ggml_reshape_1d(ctx, emb, sw->enc_dim);
|
|
emb = ggml_cont(ctx, emb);
|
|
ggml_set_name(emb, "spk.embedding");
|
|
return emb;
|
|
}
|