Initial release
This commit is contained in:
Pascal
@@ -0,0 +1,312 @@
|
||||
#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 OC = (int) w->ne[2];
|
||||
const int pad = ((K - 1) * dilation) / 2;
|
||||
|
||||
// ggml_pad_reflect_1d pads ne[0]. Our temporal axis is ne[1], so
|
||||
// transpose first, pad, transpose back.
|
||||
struct ggml_tensor * x_t = ggml_cont(ctx, ggml_transpose(ctx, x)); // [T, in_c]
|
||||
if (pad > 0) {
|
||||
x_t = ggml_pad_reflect_1d(ctx, x_t, pad, pad); // [T+2*pad, in_c]
|
||||
}
|
||||
x_t = ggml_cont(ctx, ggml_transpose(ctx, x_t)); // [in_c, T+2*pad]
|
||||
|
||||
// Reshape as [W=T_pad, H=1, IC=in_c, N=1] for ggml_im2col 1D.
|
||||
struct ggml_tensor * x4d = ggml_reshape_4d(ctx, x_t, x_t->ne[1], 1, x_t->ne[0], 1);
|
||||
struct ggml_tensor * dummy = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, K, 1, x_t->ne[0], 1);
|
||||
ggml_set_name(dummy, "spk.im2col_kernel");
|
||||
|
||||
// im2col output shape with is_2D=false : ne = (a.ne[1]*a.ne[0], OW, b.ne[2], 1)
|
||||
// = (K, T_out, in_c, 1) here. To matmul against the [K*in_c, out_c]
|
||||
// weight we need [K*in_c, T_out] which means permuting (K, IC, T)
|
||||
// before flattening. ggml_permute(0, 2, 1, 3) swaps axes 1 and 2.
|
||||
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];
|
||||
int IC = (int) col->ne[2];
|
||||
col = ggml_cont(ctx, ggml_permute(ctx, col, 0, 2, 1, 3)); // [K, IC, T, 1]
|
||||
col = ggml_reshape_2d(ctx, col, K * IC, T_out);
|
||||
|
||||
// weight [K, in_c, out_c] reshape as [K * in_c, out_c]. mul_mat
|
||||
// returns [out_c, T_out].
|
||||
struct ggml_tensor * w2d = ggml_reshape_2d(ctx, w, K * (int) w->ne[1], 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
|
||||
// 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) {
|
||||
// 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);
|
||||
|
||||
// Frontend conv0 TDNN k=5 + ReLU : 128 -> 512, T preserved.
|
||||
struct ggml_tensor * h = spk_tdnn(ctx, sw->conv0, mel, 1);
|
||||
|
||||
// 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);
|
||||
|
||||
// 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]
|
||||
|
||||
// Attentive statistical pooling : [1536, T] -> [3072, 1].
|
||||
struct ggml_tensor * stats = spk_asp(ctx, sw->asp, mfa);
|
||||
|
||||
// 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;
|
||||
}
|
||||
Reference in New Issue
Block a user